MLIR  19.0.0git
LowerVectorTransfer.cpp
Go to the documentation of this file.
1 //===- VectorTransferPermutationMapRewritePatterns.cpp - Xfer map rewrite -===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 //
9 // This file implements rewrite patterns for the permutation_map attribute of
10 // vector.transfer operations.
11 //
12 //===----------------------------------------------------------------------===//
13 
19 
20 using namespace mlir;
21 using namespace mlir::vector;
22 
23 /// Transpose a vector transfer op's `in_bounds` attribute by applying reverse
24 /// permutation based on the given indices.
25 static ArrayAttr
26 inverseTransposeInBoundsAttr(OpBuilder &builder, ArrayAttr attr,
27  const SmallVector<unsigned> &permutation) {
28  SmallVector<bool> newInBoundsValues(permutation.size());
29  size_t index = 0;
30  for (unsigned pos : permutation)
31  newInBoundsValues[pos] =
32  cast<BoolAttr>(attr.getValue()[index++]).getValue();
33  return builder.getBoolArrayAttr(newInBoundsValues);
34 }
35 
36 /// Extend the rank of a vector Value by `addedRanks` by adding outer unit
37 /// dimensions.
38 static Value extendVectorRank(OpBuilder &builder, Location loc, Value vec,
39  int64_t addedRank) {
40  auto originalVecType = cast<VectorType>(vec.getType());
41  SmallVector<int64_t> newShape(addedRank, 1);
42  newShape.append(originalVecType.getShape().begin(),
43  originalVecType.getShape().end());
44 
45  SmallVector<bool> newScalableDims(addedRank, false);
46  newScalableDims.append(originalVecType.getScalableDims().begin(),
47  originalVecType.getScalableDims().end());
48  VectorType newVecType = VectorType::get(
49  newShape, originalVecType.getElementType(), newScalableDims);
50  return builder.create<vector::BroadcastOp>(loc, newVecType, vec);
51 }
52 
53 /// Extend the rank of a vector Value by `addedRanks` by adding inner unit
54 /// dimensions.
55 static Value extendMaskRank(OpBuilder &builder, Location loc, Value vec,
56  int64_t addedRank) {
57  Value broadcasted = extendVectorRank(builder, loc, vec, addedRank);
58  SmallVector<int64_t> permutation;
59  for (int64_t i = addedRank,
60  e = cast<VectorType>(broadcasted.getType()).getRank();
61  i < e; ++i)
62  permutation.push_back(i);
63  for (int64_t i = 0; i < addedRank; ++i)
64  permutation.push_back(i);
65  return builder.create<vector::TransposeOp>(loc, broadcasted, permutation);
66 }
67 
68 //===----------------------------------------------------------------------===//
69 // populateVectorTransferPermutationMapLoweringPatterns
70 //===----------------------------------------------------------------------===//
71 
72 namespace {
73 /// Lower transfer_read op with permutation into a transfer_read with a
74 /// permutation map composed of leading zeros followed by a minor identiy +
75 /// vector.transpose op.
76 /// Ex:
77 /// vector.transfer_read ...
78 /// permutation_map: (d0, d1, d2) -> (0, d1)
79 /// into:
80 /// %v = vector.transfer_read ...
81 /// permutation_map: (d0, d1, d2) -> (d1, 0)
82 /// vector.transpose %v, [1, 0]
83 ///
84 /// vector.transfer_read ...
85 /// permutation_map: (d0, d1, d2, d3) -> (0, 0, 0, d1, d3)
86 /// into:
87 /// %v = vector.transfer_read ...
88 /// permutation_map: (d0, d1, d2, d3) -> (0, 0, d1, 0, d3)
89 /// vector.transpose %v, [0, 1, 3, 2, 4]
90 /// Note that an alternative is to transform it to linalg.transpose +
91 /// vector.transfer_read to do the transpose in memory instead.
92 struct TransferReadPermutationLowering
93  : public MaskableOpRewritePattern<vector::TransferReadOp> {
94  using MaskableOpRewritePattern::MaskableOpRewritePattern;
95 
97  matchAndRewriteMaskableOp(vector::TransferReadOp op,
98  MaskingOpInterface maskOp,
99  PatternRewriter &rewriter) const override {
100  // TODO: support 0-d corner case.
101  if (op.getTransferRank() == 0)
102  return rewriter.notifyMatchFailure(op, "0-d corner case not supported");
103  // TODO: Support transfer_read inside MaskOp case.
104  if (maskOp)
105  return rewriter.notifyMatchFailure(op, "Masked case not supported");
106 
107  SmallVector<unsigned> permutation;
108  AffineMap map = op.getPermutationMap();
109  if (map.getNumResults() == 0)
110  return rewriter.notifyMatchFailure(op, "0 result permutation map");
111  if (!map.isPermutationOfMinorIdentityWithBroadcasting(permutation)) {
112  return rewriter.notifyMatchFailure(
113  op, "map is not permutable to minor identity, apply another pattern");
114  }
115  AffineMap permutationMap =
116  map.getPermutationMap(permutation, op.getContext());
117  if (permutationMap.isIdentity())
118  return rewriter.notifyMatchFailure(op, "map is not identity");
119 
120  permutationMap = map.getPermutationMap(permutation, op.getContext());
121  // Caluclate the map of the new read by applying the inverse permutation.
122  permutationMap = inversePermutation(permutationMap);
123  AffineMap newMap = permutationMap.compose(map);
124  // Apply the reverse transpose to deduce the type of the transfer_read.
125  ArrayRef<int64_t> originalShape = op.getVectorType().getShape();
126  SmallVector<int64_t> newVectorShape(originalShape.size());
127  ArrayRef<bool> originalScalableDims = op.getVectorType().getScalableDims();
128  SmallVector<bool> newScalableDims(originalShape.size());
129  for (const auto &pos : llvm::enumerate(permutation)) {
130  newVectorShape[pos.value()] = originalShape[pos.index()];
131  newScalableDims[pos.value()] = originalScalableDims[pos.index()];
132  }
133 
134  // Transpose in_bounds attribute.
135  ArrayAttr newInBoundsAttr =
136  op.getInBounds() ? inverseTransposeInBoundsAttr(
137  rewriter, op.getInBounds().value(), permutation)
138  : ArrayAttr();
139 
140  // Generate new transfer_read operation.
141  VectorType newReadType = VectorType::get(
142  newVectorShape, op.getVectorType().getElementType(), newScalableDims);
143  Value newRead = rewriter.create<vector::TransferReadOp>(
144  op.getLoc(), newReadType, op.getSource(), op.getIndices(),
145  AffineMapAttr::get(newMap), op.getPadding(), op.getMask(),
146  newInBoundsAttr);
147 
148  // Transpose result of transfer_read.
149  SmallVector<int64_t> transposePerm(permutation.begin(), permutation.end());
150  return rewriter
151  .create<vector::TransposeOp>(op.getLoc(), newRead, transposePerm)
152  .getResult();
153  }
154 };
155 
156 /// Lower transfer_write op with permutation into a transfer_write with a
157 /// minor identity permutation map. (transfer_write ops cannot have broadcasts.)
158 /// Ex:
159 /// vector.transfer_write %v ...
160 /// permutation_map: (d0, d1, d2) -> (d2, d0, d1)
161 /// into:
162 /// %tmp = vector.transpose %v, [2, 0, 1]
163 /// vector.transfer_write %tmp ...
164 /// permutation_map: (d0, d1, d2) -> (d0, d1, d2)
165 ///
166 /// vector.transfer_write %v ...
167 /// permutation_map: (d0, d1, d2, d3) -> (d3, d2)
168 /// into:
169 /// %tmp = vector.transpose %v, [1, 0]
170 /// %v = vector.transfer_write %tmp ...
171 /// permutation_map: (d0, d1, d2, d3) -> (d2, d3)
172 struct TransferWritePermutationLowering
173  : public MaskableOpRewritePattern<vector::TransferWriteOp> {
174  using MaskableOpRewritePattern::MaskableOpRewritePattern;
175 
177  matchAndRewriteMaskableOp(vector::TransferWriteOp op,
178  MaskingOpInterface maskOp,
179  PatternRewriter &rewriter) const override {
180  // TODO: support 0-d corner case.
181  if (op.getTransferRank() == 0)
182  return rewriter.notifyMatchFailure(op, "0-d corner case not supported");
183  // TODO: Support transfer_write inside MaskOp case.
184  if (maskOp)
185  return rewriter.notifyMatchFailure(op, "Masked case not supported");
186 
187  SmallVector<unsigned> permutation;
188  AffineMap map = op.getPermutationMap();
189  if (map.isMinorIdentity())
190  return rewriter.notifyMatchFailure(op, "map is already minor identity");
191 
192  if (!map.isPermutationOfMinorIdentityWithBroadcasting(permutation)) {
193  return rewriter.notifyMatchFailure(
194  op, "map is not permutable to minor identity, apply another pattern");
195  }
196 
197  // Remove unused dims from the permutation map. E.g.:
198  // E.g.: (d0, d1, d2, d3, d4, d5) -> (d5, d3, d4)
199  // comp = (d0, d1, d2) -> (d2, d0, d1)
200  auto comp = compressUnusedDims(map);
201  AffineMap permutationMap = inversePermutation(comp);
202  // Get positions of remaining result dims.
203  SmallVector<int64_t> indices;
204  llvm::transform(permutationMap.getResults(), std::back_inserter(indices),
205  [](AffineExpr expr) {
206  return dyn_cast<AffineDimExpr>(expr).getPosition();
207  });
208 
209  // Transpose in_bounds attribute.
210  ArrayAttr newInBoundsAttr =
211  op.getInBounds() ? inverseTransposeInBoundsAttr(
212  rewriter, op.getInBounds().value(), permutation)
213  : ArrayAttr();
214 
215  // Generate new transfer_write operation.
216  Value newVec = rewriter.create<vector::TransposeOp>(
217  op.getLoc(), op.getVector(), indices);
218  auto newMap = AffineMap::getMinorIdentityMap(
219  map.getNumDims(), map.getNumResults(), rewriter.getContext());
220  auto newWrite = rewriter.create<vector::TransferWriteOp>(
221  op.getLoc(), newVec, op.getSource(), op.getIndices(),
222  AffineMapAttr::get(newMap), op.getMask(), newInBoundsAttr);
223  if (newWrite.hasPureTensorSemantics())
224  return newWrite.getResult();
225  // In the memref case there's no return value. Use empty value to signal
226  // success.
227  return Value();
228  }
229 };
230 
231 /// Convert a transfer.write op with a map which isn't the permutation of a
232 /// minor identity into a vector.broadcast + transfer_write with permutation of
233 /// minor identity map by adding unit dim on inner dimension. Ex:
234 /// ```
235 /// vector.transfer_write %v
236 /// {permutation_map = affine_map<(d0, d1, d2, d3) -> (d1, d2)>} :
237 /// vector<8x16xf32>
238 /// ```
239 /// into:
240 /// ```
241 /// %v1 = vector.broadcast %v : vector<8x16xf32> to vector<1x8x16xf32>
242 /// vector.transfer_write %v1
243 /// {permutation_map = affine_map<(d0, d1, d2, d3) -> (d3, d1, d2)>} :
244 /// vector<1x8x16xf32>
245 /// ```
246 struct TransferWriteNonPermutationLowering
247  : public MaskableOpRewritePattern<vector::TransferWriteOp> {
248  using MaskableOpRewritePattern::MaskableOpRewritePattern;
249 
251  matchAndRewriteMaskableOp(vector::TransferWriteOp op,
252  MaskingOpInterface maskOp,
253  PatternRewriter &rewriter) const override {
254  // TODO: support 0-d corner case.
255  if (op.getTransferRank() == 0)
256  return rewriter.notifyMatchFailure(op, "0-d corner case not supported");
257  // TODO: Support transfer_write inside MaskOp case.
258  if (maskOp)
259  return rewriter.notifyMatchFailure(op, "Masked case not supported");
260 
261  SmallVector<unsigned> permutation;
262  AffineMap map = op.getPermutationMap();
263  if (map.isPermutationOfMinorIdentityWithBroadcasting(permutation)) {
264  return rewriter.notifyMatchFailure(
265  op,
266  "map is already permutable to minor identity, apply another pattern");
267  }
268 
269  // Missing outer dimensions are allowed, find the most outer existing
270  // dimension then deduce the missing inner dimensions.
271  SmallVector<bool> foundDim(map.getNumDims(), false);
272  for (AffineExpr exp : map.getResults())
273  foundDim[cast<AffineDimExpr>(exp).getPosition()] = true;
275  bool foundFirstDim = false;
276  SmallVector<int64_t> missingInnerDim;
277  for (size_t i = 0; i < foundDim.size(); i++) {
278  if (foundDim[i]) {
279  foundFirstDim = true;
280  continue;
281  }
282  if (!foundFirstDim)
283  continue;
284  // Once we found one outer dimension existing in the map keep track of all
285  // the missing dimensions after that.
286  missingInnerDim.push_back(i);
287  exprs.push_back(rewriter.getAffineDimExpr(i));
288  }
289  // Vector: add unit dims at the beginning of the shape.
290  Value newVec = extendVectorRank(rewriter, op.getLoc(), op.getVector(),
291  missingInnerDim.size());
292  // Mask: add unit dims at the end of the shape.
293  Value newMask;
294  if (op.getMask())
295  newMask = extendMaskRank(rewriter, op.getLoc(), op.getMask(),
296  missingInnerDim.size());
297  exprs.append(map.getResults().begin(), map.getResults().end());
298  AffineMap newMap =
299  AffineMap::get(map.getNumDims(), 0, exprs, op.getContext());
300  // All the new dimensions added are inbound.
301  SmallVector<bool> newInBoundsValues(missingInnerDim.size(), true);
302  for (int64_t i = 0, e = op.getVectorType().getRank(); i < e; ++i) {
303  newInBoundsValues.push_back(op.isDimInBounds(i));
304  }
305  ArrayAttr newInBoundsAttr = rewriter.getBoolArrayAttr(newInBoundsValues);
306  auto newWrite = rewriter.create<vector::TransferWriteOp>(
307  op.getLoc(), newVec, op.getSource(), op.getIndices(),
308  AffineMapAttr::get(newMap), newMask, newInBoundsAttr);
309  if (newWrite.hasPureTensorSemantics())
310  return newWrite.getResult();
311  // In the memref case there's no return value. Use empty value to signal
312  // success.
313  return Value();
314  }
315 };
316 
317 /// Lower transfer_read op with broadcast in the leading dimensions into
318 /// transfer_read of lower rank + vector.broadcast.
319 /// Ex: vector.transfer_read ...
320 /// permutation_map: (d0, d1, d2, d3) -> (0, d1, 0, d3)
321 /// into:
322 /// %v = vector.transfer_read ...
323 /// permutation_map: (d0, d1, d2, d3) -> (d1, 0, d3)
324 /// vector.broadcast %v
325 struct TransferOpReduceRank : public OpRewritePattern<vector::TransferReadOp> {
327 
328  LogicalResult matchAndRewrite(vector::TransferReadOp op,
329  PatternRewriter &rewriter) const override {
330  // TODO: support 0-d corner case.
331  if (op.getTransferRank() == 0)
332  return rewriter.notifyMatchFailure(op, "0-d corner case not supported");
333 
334  AffineMap map = op.getPermutationMap();
335  unsigned numLeadingBroadcast = 0;
336  for (auto expr : map.getResults()) {
337  auto dimExpr = dyn_cast<AffineConstantExpr>(expr);
338  if (!dimExpr || dimExpr.getValue() != 0)
339  break;
340  numLeadingBroadcast++;
341  }
342  // If there are no leading zeros in the map there is nothing to do.
343  if (numLeadingBroadcast == 0)
344  return rewriter.notifyMatchFailure(op, "no leading broadcasts in map");
345 
346  VectorType originalVecType = op.getVectorType();
347  unsigned reducedShapeRank = originalVecType.getRank() - numLeadingBroadcast;
348  // Calculate new map, vector type and masks without the leading zeros.
349  AffineMap newMap = AffineMap::get(
350  map.getNumDims(), 0, map.getResults().take_back(reducedShapeRank),
351  op.getContext());
352  // Only remove the leading zeros if the rest of the map is a minor identity
353  // with broadasting. Otherwise we first want to permute the map.
354  if (!newMap.isMinorIdentityWithBroadcasting()) {
355  return rewriter.notifyMatchFailure(
356  op, "map is not a minor identity with broadcasting");
357  }
358 
359  // TODO: support zero-dimension vectors natively. See:
360  // https://llvm.discourse.group/t/should-we-have-0-d-vectors/3097.
361  // In the meantime, lower these to a scalar load when they pop up.
362  if (reducedShapeRank == 0) {
363  Value newRead;
364  if (isa<TensorType>(op.getShapedType())) {
365  newRead = rewriter.create<tensor::ExtractOp>(
366  op.getLoc(), op.getSource(), op.getIndices());
367  } else {
368  newRead = rewriter.create<memref::LoadOp>(
369  op.getLoc(), originalVecType.getElementType(), op.getSource(),
370  op.getIndices());
371  }
372  rewriter.replaceOpWithNewOp<vector::BroadcastOp>(op, originalVecType,
373  newRead);
374  return success();
375  }
376 
377  SmallVector<int64_t> newShape(
378  originalVecType.getShape().take_back(reducedShapeRank));
379  SmallVector<bool> newScalableDims(
380  originalVecType.getScalableDims().take_back(reducedShapeRank));
381  // Vector rank cannot be zero. Handled by TransferReadToVectorLoadLowering.
382  if (newShape.empty())
383  return rewriter.notifyMatchFailure(op, "rank-reduced vector is 0-d");
384 
385  VectorType newReadType = VectorType::get(
386  newShape, originalVecType.getElementType(), newScalableDims);
387  ArrayAttr newInBoundsAttr =
388  op.getInBounds()
389  ? rewriter.getArrayAttr(
390  op.getInBoundsAttr().getValue().take_back(reducedShapeRank))
391  : ArrayAttr();
392  Value newRead = rewriter.create<vector::TransferReadOp>(
393  op.getLoc(), newReadType, op.getSource(), op.getIndices(),
394  AffineMapAttr::get(newMap), op.getPadding(), op.getMask(),
395  newInBoundsAttr);
396  rewriter.replaceOpWithNewOp<vector::BroadcastOp>(op, originalVecType,
397  newRead);
398  return success();
399  }
400 };
401 
402 } // namespace
403 
405  RewritePatternSet &patterns, PatternBenefit benefit) {
406  patterns
407  .add<TransferReadPermutationLowering, TransferWritePermutationLowering,
408  TransferOpReduceRank, TransferWriteNonPermutationLowering>(
409  patterns.getContext(), benefit);
410 }
411 
412 //===----------------------------------------------------------------------===//
413 // populateVectorTransferLoweringPatterns
414 //===----------------------------------------------------------------------===//
415 
416 namespace {
417 /// Progressive lowering of transfer_read. This pattern supports lowering of
418 /// `vector.transfer_read` to a combination of `vector.load` and
419 /// `vector.broadcast` if all of the following hold:
420 /// - Stride of most minor memref dimension must be 1.
421 /// - Out-of-bounds masking is not required.
422 /// - If the memref's element type is a vector type then it coincides with the
423 /// result type.
424 /// - The permutation map doesn't perform permutation (broadcasting is allowed).
425 struct TransferReadToVectorLoadLowering
426  : public OpRewritePattern<vector::TransferReadOp> {
427  TransferReadToVectorLoadLowering(MLIRContext *context,
428  std::optional<unsigned> maxRank,
429  PatternBenefit benefit = 1)
430  : OpRewritePattern<vector::TransferReadOp>(context, benefit),
431  maxTransferRank(maxRank) {}
432 
433  LogicalResult matchAndRewrite(vector::TransferReadOp read,
434  PatternRewriter &rewriter) const override {
435  if (maxTransferRank && read.getVectorType().getRank() > *maxTransferRank) {
436  return rewriter.notifyMatchFailure(
437  read, "vector type is greater than max transfer rank");
438  }
439 
440  SmallVector<unsigned> broadcastedDims;
441  // Permutations are handled by VectorToSCF or
442  // populateVectorTransferPermutationMapLoweringPatterns.
443  // We let the 0-d corner case pass-through as it is supported.
444  if (!read.getPermutationMap().isMinorIdentityWithBroadcasting(
445  &broadcastedDims))
446  return rewriter.notifyMatchFailure(read, "not minor identity + bcast");
447 
448  auto memRefType = dyn_cast<MemRefType>(read.getShapedType());
449  if (!memRefType)
450  return rewriter.notifyMatchFailure(read, "not a memref source");
451 
452  // Non-unit strides are handled by VectorToSCF.
453  if (!isLastMemrefDimUnitStride(memRefType))
454  return rewriter.notifyMatchFailure(read, "!= 1 stride needs VectorToSCF");
455 
456  // If there is broadcasting involved then we first load the unbroadcasted
457  // vector, and then broadcast it with `vector.broadcast`.
458  ArrayRef<int64_t> vectorShape = read.getVectorType().getShape();
459  SmallVector<int64_t> unbroadcastedVectorShape(vectorShape.begin(),
460  vectorShape.end());
461  for (unsigned i : broadcastedDims)
462  unbroadcastedVectorShape[i] = 1;
463  VectorType unbroadcastedVectorType = read.getVectorType().cloneWith(
464  unbroadcastedVectorShape, read.getVectorType().getElementType());
465 
466  // `vector.load` supports vector types as memref's elements only when the
467  // resulting vector type is the same as the element type.
468  auto memrefElTy = memRefType.getElementType();
469  if (isa<VectorType>(memrefElTy) && memrefElTy != unbroadcastedVectorType)
470  return rewriter.notifyMatchFailure(read, "incompatible element type");
471 
472  // Otherwise, element types of the memref and the vector must match.
473  if (!isa<VectorType>(memrefElTy) &&
474  memrefElTy != read.getVectorType().getElementType())
475  return rewriter.notifyMatchFailure(read, "non-matching element type");
476 
477  // Out-of-bounds dims are handled by MaterializeTransferMask.
478  if (read.hasOutOfBoundsDim())
479  return rewriter.notifyMatchFailure(read, "out-of-bounds needs mask");
480 
481  // Create vector load op.
482  Operation *loadOp;
483  if (read.getMask()) {
484  if (read.getVectorType().getRank() != 1)
485  // vector.maskedload operates on 1-D vectors.
486  return rewriter.notifyMatchFailure(
487  read, "vector type is not rank 1, can't create masked load, needs "
488  "VectorToSCF");
489 
490  Value fill = rewriter.create<vector::SplatOp>(
491  read.getLoc(), unbroadcastedVectorType, read.getPadding());
492  loadOp = rewriter.create<vector::MaskedLoadOp>(
493  read.getLoc(), unbroadcastedVectorType, read.getSource(),
494  read.getIndices(), read.getMask(), fill);
495  } else {
496  loadOp = rewriter.create<vector::LoadOp>(
497  read.getLoc(), unbroadcastedVectorType, read.getSource(),
498  read.getIndices());
499  }
500 
501  // Insert a broadcasting op if required.
502  if (!broadcastedDims.empty()) {
503  rewriter.replaceOpWithNewOp<vector::BroadcastOp>(
504  read, read.getVectorType(), loadOp->getResult(0));
505  } else {
506  rewriter.replaceOp(read, loadOp->getResult(0));
507  }
508 
509  return success();
510  }
511 
512  std::optional<unsigned> maxTransferRank;
513 };
514 
515 /// Replace a 0-d vector.load with a memref.load + vector.broadcast.
516 // TODO: we shouldn't cross the vector/scalar domains just for this
517 // but atm we lack the infra to avoid it. Possible solutions include:
518 // - go directly to LLVM + bitcast
519 // - introduce a bitcast op and likely a new pointer dialect
520 // - let memref.load/store additionally support the 0-d vector case
521 // There are still deeper data layout issues lingering even in this
522 // trivial case (for architectures for which this matters).
523 struct VectorLoadToMemrefLoadLowering
524  : public OpRewritePattern<vector::LoadOp> {
526 
527  LogicalResult matchAndRewrite(vector::LoadOp loadOp,
528  PatternRewriter &rewriter) const override {
529  auto vecType = loadOp.getVectorType();
530  if (vecType.getNumElements() != 1)
531  return rewriter.notifyMatchFailure(loadOp, "not a single element vector");
532 
533  auto memrefLoad = rewriter.create<memref::LoadOp>(
534  loadOp.getLoc(), loadOp.getBase(), loadOp.getIndices());
535  rewriter.replaceOpWithNewOp<vector::BroadcastOp>(loadOp, vecType,
536  memrefLoad);
537  return success();
538  }
539 };
540 
541 /// Replace a 0-d vector.store with a vector.extractelement + memref.store.
542 struct VectorStoreToMemrefStoreLowering
543  : public OpRewritePattern<vector::StoreOp> {
545 
546  LogicalResult matchAndRewrite(vector::StoreOp storeOp,
547  PatternRewriter &rewriter) const override {
548  auto vecType = storeOp.getVectorType();
549  if (vecType.getNumElements() != 1)
550  return rewriter.notifyMatchFailure(storeOp, "not single element vector");
551 
552  Value extracted;
553  if (vecType.getRank() == 0) {
554  // TODO: Unifiy once ExtractOp supports 0-d vectors.
555  extracted = rewriter.create<vector::ExtractElementOp>(
556  storeOp.getLoc(), storeOp.getValueToStore());
557  } else {
558  SmallVector<int64_t> indices(vecType.getRank(), 0);
559  extracted = rewriter.create<vector::ExtractOp>(
560  storeOp.getLoc(), storeOp.getValueToStore(), indices);
561  }
562 
563  rewriter.replaceOpWithNewOp<memref::StoreOp>(
564  storeOp, extracted, storeOp.getBase(), storeOp.getIndices());
565  return success();
566  }
567 };
568 
569 /// Progressive lowering of transfer_write. This pattern supports lowering of
570 /// `vector.transfer_write` to `vector.store` if all of the following hold:
571 /// - Stride of most minor memref dimension must be 1.
572 /// - Out-of-bounds masking is not required.
573 /// - If the memref's element type is a vector type then it coincides with the
574 /// type of the written value.
575 /// - The permutation map is the minor identity map (neither permutation nor
576 /// broadcasting is allowed).
577 struct TransferWriteToVectorStoreLowering
578  : public OpRewritePattern<vector::TransferWriteOp> {
579  TransferWriteToVectorStoreLowering(MLIRContext *context,
580  std::optional<unsigned> maxRank,
581  PatternBenefit benefit = 1)
582  : OpRewritePattern<vector::TransferWriteOp>(context, benefit),
583  maxTransferRank(maxRank) {}
584 
585  LogicalResult matchAndRewrite(vector::TransferWriteOp write,
586  PatternRewriter &rewriter) const override {
587  if (maxTransferRank && write.getVectorType().getRank() > *maxTransferRank) {
588  return rewriter.notifyMatchFailure(
589  write, "vector type is greater than max transfer rank");
590  }
591 
592  // Permutations are handled by VectorToSCF or
593  // populateVectorTransferPermutationMapLoweringPatterns.
594  if ( // pass-through for the 0-d corner case.
595  !write.getPermutationMap().isMinorIdentity())
596  return rewriter.notifyMatchFailure(write.getLoc(), [=](Diagnostic &diag) {
597  diag << "permutation map is not minor identity: " << write;
598  });
599 
600  auto memRefType = dyn_cast<MemRefType>(write.getShapedType());
601  if (!memRefType)
602  return rewriter.notifyMatchFailure(write.getLoc(), [=](Diagnostic &diag) {
603  diag << "not a memref type: " << write;
604  });
605 
606  // Non-unit strides are handled by VectorToSCF.
607  if (!isLastMemrefDimUnitStride(memRefType))
608  return rewriter.notifyMatchFailure(write.getLoc(), [=](Diagnostic &diag) {
609  diag << "most minor stride is not 1: " << write;
610  });
611 
612  // `vector.store` supports vector types as memref's elements only when the
613  // type of the vector value being written is the same as the element type.
614  auto memrefElTy = memRefType.getElementType();
615  if (isa<VectorType>(memrefElTy) && memrefElTy != write.getVectorType())
616  return rewriter.notifyMatchFailure(write.getLoc(), [=](Diagnostic &diag) {
617  diag << "elemental type mismatch: " << write;
618  });
619 
620  // Otherwise, element types of the memref and the vector must match.
621  if (!isa<VectorType>(memrefElTy) &&
622  memrefElTy != write.getVectorType().getElementType())
623  return rewriter.notifyMatchFailure(write.getLoc(), [=](Diagnostic &diag) {
624  diag << "elemental type mismatch: " << write;
625  });
626 
627  // Out-of-bounds dims are handled by MaterializeTransferMask.
628  if (write.hasOutOfBoundsDim())
629  return rewriter.notifyMatchFailure(write.getLoc(), [=](Diagnostic &diag) {
630  diag << "out of bounds dim: " << write;
631  });
632  if (write.getMask()) {
633  if (write.getVectorType().getRank() != 1)
634  // vector.maskedstore operates on 1-D vectors.
635  return rewriter.notifyMatchFailure(
636  write.getLoc(), [=](Diagnostic &diag) {
637  diag << "vector type is not rank 1, can't create masked store, "
638  "needs VectorToSCF: "
639  << write;
640  });
641 
642  rewriter.replaceOpWithNewOp<vector::MaskedStoreOp>(
643  write, write.getSource(), write.getIndices(), write.getMask(),
644  write.getVector());
645  } else {
646  rewriter.replaceOpWithNewOp<vector::StoreOp>(
647  write, write.getVector(), write.getSource(), write.getIndices());
648  }
649  return success();
650  }
651 
652  std::optional<unsigned> maxTransferRank;
653 };
654 } // namespace
655 
657  RewritePatternSet &patterns, std::optional<unsigned> maxTransferRank,
658  PatternBenefit benefit) {
659  patterns.add<TransferReadToVectorLoadLowering,
660  TransferWriteToVectorStoreLowering>(patterns.getContext(),
661  maxTransferRank, benefit);
662  patterns
663  .add<VectorLoadToMemrefLoadLowering, VectorStoreToMemrefStoreLowering>(
664  patterns.getContext(), benefit);
665 }
static ArrayAttr inverseTransposeInBoundsAttr(OpBuilder &builder, ArrayAttr attr, const SmallVector< unsigned > &permutation)
Transpose a vector transfer op's in_bounds attribute by applying reverse permutation based on the giv...
static Value extendMaskRank(OpBuilder &builder, Location loc, Value vec, int64_t addedRank)
Extend the rank of a vector Value by addedRanks by adding inner unit dimensions.
static Value extendVectorRank(OpBuilder &builder, Location loc, Value vec, int64_t addedRank)
Extend the rank of a vector Value by addedRanks by adding outer unit dimensions.
static std::string diag(const llvm::Value &value)
static VectorShape vectorShape(Type type)
Base type for affine expression.
Definition: AffineExpr.h:69
A multi-dimensional affine map Affine map's are immutable like Type's, and they are uniqued.
Definition: AffineMap.h:47
static AffineMap getMinorIdentityMap(unsigned dims, unsigned results, MLIRContext *context)
Returns an identity affine map (d0, ..., dn) -> (dp, ..., dn) on the most minor dimensions.
Definition: AffineMap.cpp:132
bool isMinorIdentity() const
Returns true if this affine map is a minor identity, i.e.
Definition: AffineMap.cpp:152
static AffineMap get(MLIRContext *context)
Returns a zero result affine map with no dimensions or symbols: () -> ().
bool isMinorIdentityWithBroadcasting(SmallVectorImpl< unsigned > *broadcastedDims=nullptr) const
Returns true if this affine map is a minor identity up to broadcasted dimensions which are indicated ...
Definition: AffineMap.cpp:160
unsigned getNumDims() const
Definition: AffineMap.cpp:378
ArrayRef< AffineExpr > getResults() const
Definition: AffineMap.cpp:391
bool isPermutationOfMinorIdentityWithBroadcasting(SmallVectorImpl< unsigned > &permutedDims) const
Return true if this affine map can be converted to a minor identity with broadcast by doing a permute...
Definition: AffineMap.cpp:200
unsigned getNumResults() const
Definition: AffineMap.cpp:386
static AffineMap getPermutationMap(ArrayRef< unsigned > permutation, MLIRContext *context)
Returns an AffineMap representing a permutation.
Definition: AffineMap.cpp:248
AffineMap compose(AffineMap map) const
Returns the AffineMap resulting from composing this with map.
Definition: AffineMap.cpp:540
bool isIdentity() const
Returns true if this affine map is an identity affine map.
Definition: AffineMap.cpp:329
AffineExpr getAffineDimExpr(unsigned position)
Definition: Builders.cpp:371
MLIRContext * getContext() const
Definition: Builders.h:55
ArrayAttr getArrayAttr(ArrayRef< Attribute > value)
Definition: Builders.cpp:273
ArrayAttr getBoolArrayAttr(ArrayRef< bool > values)
Definition: Builders.cpp:277
This class contains all of the information necessary to report a diagnostic to the DiagnosticEngine.
Definition: Diagnostics.h:156
This class provides support for representing a failure result, or a valid value of type T.
Definition: LogicalResult.h:78
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
Definition: Location.h:63
MLIRContext is the top-level object for a collection of MLIR operations.
Definition: MLIRContext.h:60
This class helps build Operations.
Definition: Builders.h:209
Operation * create(const OperationState &state)
Creates an operation given the fields represented as an OperationState.
Definition: Builders.cpp:464
Operation is the basic unit of execution within MLIR.
Definition: Operation.h:88
OpResult getResult(unsigned idx)
Get the 'idx'th result of this operation.
Definition: Operation.h:402
MLIRContext * getContext()
Return the context this operation is associated with.
Definition: Operation.h:216
Location getLoc()
The source location the operation was defined or derived from.
Definition: Operation.h:223
This class represents the benefit of a pattern match in a unitless scheme that ranges from 0 (very li...
Definition: PatternMatch.h:34
A special type of RewriterBase that coordinates the application of a rewrite pattern on the current I...
Definition: PatternMatch.h:785
MLIRContext * getContext() const
Definition: PatternMatch.h:822
RewritePatternSet & add(ConstructorArg &&arg, ConstructorArgs &&...args)
Add an instance of each of the pattern types 'Ts' to the pattern list with the given arguments.
Definition: PatternMatch.h:846
std::enable_if_t<!std::is_convertible< CallbackT, Twine >::value, LogicalResult > notifyMatchFailure(Location loc, CallbackT &&reasonCallback)
Used to notify the listener that the IR failed to be rewritten because of a match failure,...
Definition: PatternMatch.h:718
virtual void replaceOp(Operation *op, ValueRange newValues)
Replace the results of the given (original) operation with the specified list of values (replacements...
OpTy replaceOpWithNewOp(Operation *op, Args &&...args)
Replace the results of the given (original) op with a new op that is created without verification (re...
Definition: PatternMatch.h:536
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Definition: Value.h:96
Type getType() const
Return the type of this value.
Definition: Value.h:129
constexpr void enumerate(std::tuple< Tys... > &tuple, CallbackT &&callback)
Definition: Matchers.h:285
void populateVectorTransferPermutationMapLoweringPatterns(RewritePatternSet &patterns, PatternBenefit benefit=1)
Collect a set of transfer read/write lowering patterns that simplify the permutation map (e....
void populateVectorTransferLoweringPatterns(RewritePatternSet &patterns, std::optional< unsigned > maxTransferRank=std::nullopt, PatternBenefit benefit=1)
Populate the pattern set with the following patterns:
Include the generated interface declarations.
bool isLastMemrefDimUnitStride(MemRefType type)
Return "true" if the last dimension of the given type has a static unit stride.
AffineMap inversePermutation(AffineMap map)
Returns a map of codomain to domain dimensions such that the first codomain dimension for a particula...
Definition: AffineMap.cpp:753
LogicalResult success(bool isSuccess=true)
Utility function to generate a LogicalResult.
Definition: LogicalResult.h:56
AffineMap compressUnusedDims(AffineMap map)
Drop the dims that are not used.
Definition: AffineMap.cpp:683
auto get(MLIRContext *context, Ts &&...params)
Helper method that injects context only if needed, this helps unify some of the attribute constructio...
This class represents an efficient way to signal success or failure.
Definition: LogicalResult.h:26
OpRewritePattern is a wrapper around RewritePattern that allows for matching and rewriting against an...
Definition: PatternMatch.h:358
OpRewritePattern(MLIRContext *context, PatternBenefit benefit=1, ArrayRef< StringRef > generatedNames={})
Patterns must specify the root operation name they match against, and can also specify the benefit of...
Definition: PatternMatch.h:362
A pattern for ops that implement MaskableOpInterface and that might be masked (i.e.
Definition: VectorUtils.h:133