MLIR 24.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
16
17using namespace mlir;
18using namespace mlir::vector;
19
20/// Transpose a vector transfer op's `in_bounds` attribute by applying reverse
21/// permutation based on the given indices.
22static ArrayAttr
24 const SmallVector<unsigned> &permutation) {
25 SmallVector<bool> newInBoundsValues(permutation.size());
26 size_t index = 0;
27 for (unsigned pos : permutation)
28 newInBoundsValues[pos] =
29 cast<BoolAttr>(attr.getValue()[index++]).getValue();
30 return builder.getBoolArrayAttr(newInBoundsValues);
31}
32
33/// Extend the rank of a vector Value by `addedRanks` by adding outer unit
34/// dimensions.
35static Value extendVectorRank(OpBuilder &builder, Location loc, Value vec,
36 int64_t addedRank) {
37 auto originalVecType = cast<VectorType>(vec.getType());
38 SmallVector<int64_t> newShape(addedRank, 1);
39 newShape.append(originalVecType.getShape().begin(),
40 originalVecType.getShape().end());
41
42 SmallVector<bool> newScalableDims(addedRank, false);
43 newScalableDims.append(originalVecType.getScalableDims().begin(),
44 originalVecType.getScalableDims().end());
45 VectorType newVecType = VectorType::get(
46 newShape, originalVecType.getElementType(), newScalableDims);
47 return vector::BroadcastOp::create(builder, loc, newVecType, vec);
48}
49
50/// Extend the rank of a vector Value by `addedRanks` by adding inner unit
51/// dimensions.
52static Value extendMaskRank(OpBuilder &builder, Location loc, Value vec,
53 int64_t addedRank) {
54 Value broadcasted = extendVectorRank(builder, loc, vec, addedRank);
55 SmallVector<int64_t> permutation;
56 for (int64_t i = addedRank,
57 e = cast<VectorType>(broadcasted.getType()).getRank();
58 i < e; ++i)
59 permutation.push_back(i);
60 for (int64_t i = 0; i < addedRank; ++i)
61 permutation.push_back(i);
62 return vector::TransposeOp::create(builder, loc, broadcasted, permutation);
63}
64
65//===----------------------------------------------------------------------===//
66// populateVectorTransferPermutationMapLoweringPatterns
67//===----------------------------------------------------------------------===//
68
69namespace {
70/// Lower transfer_read op with permutation into a transfer_read with a
71/// permutation map composed of leading zeros followed by a minor identiy +
72/// vector.transpose op.
73/// Ex:
74/// vector.transfer_read ...
75/// permutation_map: (d0, d1, d2) -> (0, d1)
76/// into:
77/// %v = vector.transfer_read ...
78/// permutation_map: (d0, d1, d2) -> (d1, 0)
79/// vector.transpose %v, [1, 0]
80///
81/// vector.transfer_read ...
82/// permutation_map: (d0, d1, d2, d3) -> (0, 0, 0, d1, d3)
83/// into:
84/// %v = vector.transfer_read ...
85/// permutation_map: (d0, d1, d2, d3) -> (0, 0, d1, 0, d3)
86/// vector.transpose %v, [0, 1, 3, 2, 4]
87/// Note that an alternative is to transform it to linalg.transpose +
88/// vector.transfer_read to do the transpose in memory instead.
89struct TransferReadPermutationLowering
90 : public MaskableOpRewritePattern<vector::TransferReadOp> {
91 using MaskableOpRewritePattern::MaskableOpRewritePattern;
92
93 FailureOr<mlir::Value>
94 matchAndRewriteMaskableOp(vector::TransferReadOp op,
95 MaskingOpInterface maskOp,
96 PatternRewriter &rewriter) const override {
97 // TODO: support 0-d corner case.
98 if (op.getTransferRank() == 0)
99 return rewriter.notifyMatchFailure(op, "0-d corner case not supported");
100
101 SmallVector<unsigned> permutation;
102 AffineMap map = op.getPermutationMap();
103 if (map.getNumResults() == 0)
104 return rewriter.notifyMatchFailure(op, "0 result permutation map");
105 if (!map.isPermutationOfMinorIdentityWithBroadcasting(permutation)) {
106 return rewriter.notifyMatchFailure(
107 op, "map is not permutable to minor identity, apply another pattern");
108 }
109 AffineMap permutationMap =
110 map.getPermutationMap(permutation, op.getContext());
111 if (permutationMap.isIdentity())
112 return rewriter.notifyMatchFailure(op, "map is not identity");
113
114 permutationMap = map.getPermutationMap(permutation, op.getContext());
115 // Caluclate the map of the new read by applying the inverse permutation.
116 permutationMap = inversePermutation(permutationMap);
117 AffineMap newMap = permutationMap.compose(map);
118 // Apply the reverse transpose to deduce the type of the transfer_read.
119 ArrayRef<int64_t> originalShape = op.getVectorType().getShape();
120 SmallVector<int64_t> newVectorShape(originalShape.size());
121 ArrayRef<bool> originalScalableDims = op.getVectorType().getScalableDims();
122 SmallVector<bool> newScalableDims(originalShape.size());
123 for (const auto &pos : llvm::enumerate(permutation)) {
124 newVectorShape[pos.value()] = originalShape[pos.index()];
125 newScalableDims[pos.value()] = originalScalableDims[pos.index()];
126 }
127
128 // Transpose in_bounds attribute.
129 ArrayAttr newInBoundsAttr =
130 inverseTransposeInBoundsAttr(rewriter, op.getInBounds(), permutation);
131
132 // Generate new transfer_read operation.
133 VectorType newReadType = VectorType::get(
134 newVectorShape, op.getVectorType().getElementType(), newScalableDims);
135 Operation *newRead = vector::TransferReadOp::create(
136 rewriter, op.getLoc(), newReadType, op.getBase(), op.getIndices(),
137 AffineMapAttr::get(newMap), op.getPadding(), op.getMask(),
138 newInBoundsAttr);
139
140 SmallVector<int64_t> transposePerm(permutation.begin(), permutation.end());
141
142 // Re-apply an enclosing vector.mask. Its mask is indexed in memory order,
143 // so only the passthru has to be transposed.
144 if (maskOp) {
145 Value passthru = maskOp.getPassthru();
146 if (passthru)
147 passthru =
148 vector::TransposeOp::create(rewriter, op.getLoc(), passthru,
149 invertPermutationVector(transposePerm));
150 newRead =
151 vector::maskOperation(rewriter, newRead, maskOp.getMask(), passthru);
152 }
153
154 // Transpose result of transfer_read.
155 return vector::TransposeOp::create(rewriter, op.getLoc(),
156 newRead->getResult(0), transposePerm)
157 .getResult();
158 }
159};
160
161/// Lower transfer_write op with permutation into a transfer_write with a
162/// minor identity permutation map. (transfer_write ops cannot have broadcasts.)
163/// Ex:
164/// vector.transfer_write %v ...
165/// permutation_map: (d0, d1, d2) -> (d2, d0, d1)
166/// into:
167/// %tmp = vector.transpose %v, [2, 0, 1]
168/// vector.transfer_write %tmp ...
169/// permutation_map: (d0, d1, d2) -> (d0, d1, d2)
170///
171/// vector.transfer_write %v ...
172/// permutation_map: (d0, d1, d2, d3) -> (d3, d2)
173/// into:
174/// %tmp = vector.transpose %v, [1, 0]
175/// %v = vector.transfer_write %tmp ...
176/// permutation_map: (d0, d1, d2, d3) -> (d2, d3)
177struct TransferWritePermutationLowering
178 : public MaskableOpRewritePattern<vector::TransferWriteOp> {
179 using MaskableOpRewritePattern::MaskableOpRewritePattern;
180
181 FailureOr<mlir::Value>
182 matchAndRewriteMaskableOp(vector::TransferWriteOp op,
183 MaskingOpInterface maskOp,
184 PatternRewriter &rewriter) const override {
185 // TODO: support 0-d corner case.
186 if (op.getTransferRank() == 0)
187 return rewriter.notifyMatchFailure(op, "0-d corner case not supported");
188
189 SmallVector<unsigned> permutation;
190 AffineMap map = op.getPermutationMap();
191 if (map.isMinorIdentity())
192 return rewriter.notifyMatchFailure(op, "map is already minor identity");
193
194 if (!map.isPermutationOfMinorIdentityWithBroadcasting(permutation)) {
195 return rewriter.notifyMatchFailure(
196 op, "map is not permutable to minor identity, apply another pattern");
197 }
198
199 // Remove unused dims from the permutation map. E.g.:
200 // E.g.: (d0, d1, d2, d3, d4, d5) -> (d5, d3, d4)
201 // comp = (d0, d1, d2) -> (d2, d0, d1)
202 auto comp = compressUnusedDims(map);
203 AffineMap permutationMap = inversePermutation(comp);
204 // Get positions of remaining result dims.
205 SmallVector<int64_t> indices;
206 llvm::transform(permutationMap.getResults(), std::back_inserter(indices),
207 [](AffineExpr expr) {
208 return dyn_cast<AffineDimExpr>(expr).getPosition();
209 });
210
211 // Transpose in_bounds attribute.
212 ArrayAttr newInBoundsAttr =
213 inverseTransposeInBoundsAttr(rewriter, op.getInBounds(), permutation);
214
215 // Generate new transfer_write operation.
216 Value newVec = vector::TransposeOp::create(rewriter, op.getLoc(),
217 op.getVector(), indices);
218 auto newMap = AffineMap::getMinorIdentityMap(
219 map.getNumDims(), map.getNumResults(), rewriter.getContext());
220 auto newWrite = vector::TransferWriteOp::create(
221 rewriter, op.getLoc(), newVec, op.getBase(), op.getIndices(),
222 AffineMapAttr::get(newMap), op.getMask(), newInBoundsAttr);
223
224 // Re-apply an enclosing vector.mask. Its mask is indexed in memory order,
225 // so transposing the written vector leaves it unchanged.
226 Operation *rewritten = newWrite;
227 if (maskOp)
228 rewritten = vector::maskOperation(rewriter, newWrite, maskOp.getMask());
229
230 if (newWrite.hasPureTensorSemantics())
231 return rewritten->getResult(0);
232 // In the memref case there's no return value. Use empty value to signal
233 // success.
234 return Value();
235 }
236};
237
238/// Convert a transfer.write op with a map which isn't the permutation of a
239/// minor identity into a vector.broadcast + transfer_write with permutation of
240/// minor identity map by adding unit dim on inner dimension. Ex:
241/// ```
242/// vector.transfer_write %v
243/// {permutation_map = affine_map<(d0, d1, d2, d3) -> (d1, d2)>} :
244/// vector<8x16xf32>
245/// ```
246/// into:
247/// ```
248/// %v1 = vector.broadcast %v : vector<8x16xf32> to vector<1x8x16xf32>
249/// vector.transfer_write %v1
250/// {permutation_map = affine_map<(d0, d1, d2, d3) -> (d3, d1, d2)>} :
251/// vector<1x8x16xf32>
252/// ```
253struct TransferWriteNonPermutationLowering
254 : public MaskableOpRewritePattern<vector::TransferWriteOp> {
255 using MaskableOpRewritePattern::MaskableOpRewritePattern;
256
257 FailureOr<mlir::Value>
258 matchAndRewriteMaskableOp(vector::TransferWriteOp op,
259 MaskingOpInterface maskOp,
260 PatternRewriter &rewriter) const override {
261 // TODO: support 0-d corner case.
262 if (op.getTransferRank() == 0)
263 return rewriter.notifyMatchFailure(op, "0-d corner case not supported");
264 // TODO: Support transfer_write inside MaskOp case.
265 if (maskOp)
266 return rewriter.notifyMatchFailure(op, "Masked case not supported");
267
268 SmallVector<unsigned> permutation;
269 AffineMap map = op.getPermutationMap();
271 return rewriter.notifyMatchFailure(
272 op,
273 "map is already permutable to minor identity, apply another pattern");
274 }
275
276 // Missing outer dimensions are allowed, find the most outer existing
277 // dimension then deduce the missing inner dimensions.
278 SmallVector<bool> foundDim(map.getNumDims(), false);
279 for (AffineExpr exp : map.getResults())
280 foundDim[cast<AffineDimExpr>(exp).getPosition()] = true;
281 SmallVector<AffineExpr> exprs;
282 bool foundFirstDim = false;
283 SmallVector<int64_t> missingInnerDim;
284 for (size_t i = 0; i < foundDim.size(); i++) {
285 if (foundDim[i]) {
286 foundFirstDim = true;
287 continue;
288 }
289 if (!foundFirstDim)
290 continue;
291 // Once we found one outer dimension existing in the map keep track of all
292 // the missing dimensions after that.
293 missingInnerDim.push_back(i);
294 exprs.push_back(rewriter.getAffineDimExpr(i));
295 }
296 // Vector: add unit dims at the beginning of the shape.
297 Value newVec = extendVectorRank(rewriter, op.getLoc(), op.getVector(),
298 missingInnerDim.size());
299 // Mask: add unit dims at the end of the shape.
300 Value newMask;
301 if (op.getMask())
302 newMask = extendMaskRank(rewriter, op.getLoc(), op.getMask(),
303 missingInnerDim.size());
304 exprs.append(map.getResults().begin(), map.getResults().end());
305 AffineMap newMap =
306 AffineMap::get(map.getNumDims(), 0, exprs, op.getContext());
307 // All the new dimensions added are inbound.
308 SmallVector<bool> newInBoundsValues(missingInnerDim.size(), true);
309 for (int64_t i = 0, e = op.getVectorType().getRank(); i < e; ++i) {
310 newInBoundsValues.push_back(op.isDimInBounds(i));
311 }
312 ArrayAttr newInBoundsAttr = rewriter.getBoolArrayAttr(newInBoundsValues);
313 auto newWrite = vector::TransferWriteOp::create(
314 rewriter, op.getLoc(), newVec, op.getBase(), op.getIndices(),
315 AffineMapAttr::get(newMap), newMask, newInBoundsAttr);
316 if (newWrite.hasPureTensorSemantics())
317 return newWrite.getResult();
318 // In the memref case there's no return value. Use empty value to signal
319 // success.
320 return Value();
321 }
322};
323
324/// Lower transfer_read op with broadcast in the leading dimensions into
325/// transfer_read of lower rank + vector.broadcast.
326/// Ex: vector.transfer_read ...
327/// permutation_map: (d0, d1, d2, d3) -> (0, d1, 0, d3)
328/// into:
329/// %v = vector.transfer_read ...
330/// permutation_map: (d0, d1, d2, d3) -> (d1, 0, d3)
331/// vector.broadcast %v
332struct TransferOpReduceRank
333 : public MaskableOpRewritePattern<vector::TransferReadOp> {
334 using MaskableOpRewritePattern::MaskableOpRewritePattern;
335
336 FailureOr<mlir::Value>
337 matchAndRewriteMaskableOp(vector::TransferReadOp op,
338 MaskingOpInterface maskOp,
339 PatternRewriter &rewriter) const override {
340 // TODO: support 0-d corner case.
341 if (op.getTransferRank() == 0)
342 return rewriter.notifyMatchFailure(op, "0-d corner case not supported");
343 // TODO: support masked case.
344 if (maskOp)
345 return rewriter.notifyMatchFailure(op, "Masked case not supported");
346
347 AffineMap map = op.getPermutationMap();
348 unsigned numLeadingBroadcast = 0;
349 for (auto expr : map.getResults()) {
350 auto dimExpr = dyn_cast<AffineConstantExpr>(expr);
351 if (!dimExpr || dimExpr.getValue() != 0)
352 break;
353 numLeadingBroadcast++;
354 }
355 // If there are no leading zeros in the map there is nothing to do.
356 if (numLeadingBroadcast == 0)
357 return rewriter.notifyMatchFailure(op, "no leading broadcasts in map");
358
359 VectorType originalVecType = op.getVectorType();
360 unsigned reducedShapeRank = originalVecType.getRank() - numLeadingBroadcast;
361 // Calculate new map, vector type and masks without the leading zeros.
362 AffineMap newMap = AffineMap::get(
363 map.getNumDims(), 0, map.getResults().take_back(reducedShapeRank),
364 op.getContext());
365 // Only remove the leading zeros if the rest of the map is a minor identity
366 // with broadasting. Otherwise we first want to permute the map.
367 if (!newMap.isMinorIdentityWithBroadcasting()) {
368 return rewriter.notifyMatchFailure(
369 op, "map is not a minor identity with broadcasting");
370 }
371
372 SmallVector<int64_t> newShape(
373 originalVecType.getShape().take_back(reducedShapeRank));
374 SmallVector<bool> newScalableDims(
375 originalVecType.getScalableDims().take_back(reducedShapeRank));
376
377 VectorType newReadType = VectorType::get(
378 newShape, originalVecType.getElementType(), newScalableDims);
379 ArrayAttr newInBoundsAttr =
380 op.getInBounds()
381 ? rewriter.getArrayAttr(
382 op.getInBoundsAttr().getValue().take_back(reducedShapeRank))
383 : ArrayAttr();
384 Value newRead = vector::TransferReadOp::create(
385 rewriter, op.getLoc(), newReadType, op.getBase(), op.getIndices(),
386 AffineMapAttr::get(newMap), op.getPadding(), op.getMask(),
387 newInBoundsAttr);
388 return vector::BroadcastOp::create(rewriter, op.getLoc(), originalVecType,
389 newRead)
390 .getVector();
391 }
392};
393
394} // namespace
395
397 RewritePatternSet &patterns, PatternBenefit benefit) {
398 patterns
399 .add<TransferReadPermutationLowering, TransferWritePermutationLowering,
400 TransferOpReduceRank, TransferWriteNonPermutationLowering>(
401 patterns.getContext(), benefit);
402}
403
404//===----------------------------------------------------------------------===//
405// populateVectorTransferLoweringPatterns
406//===----------------------------------------------------------------------===//
407
408namespace {
409/// Progressive lowering of transfer_read. This pattern supports lowering of
410/// `vector.transfer_read` to a combination of `vector.load` and
411/// `vector.broadcast` if all of the following hold:
412/// - Stride of most minor memref dimension must be 1.
413/// - Out-of-bounds masking is not required.
414/// - If the memref's element type is a vector type then it coincides with the
415/// result type.
416/// - The permutation map doesn't perform permutation (broadcasting is allowed).
417struct TransferReadToVectorLoadLowering
418 : public MaskableOpRewritePattern<vector::TransferReadOp> {
419 TransferReadToVectorLoadLowering(MLIRContext *context,
420 std::optional<unsigned> maxRank,
421 PatternBenefit benefit = 1)
422 : MaskableOpRewritePattern<vector::TransferReadOp>(context, benefit),
423 maxTransferRank(maxRank) {}
424
425 FailureOr<mlir::Value>
426 matchAndRewriteMaskableOp(vector::TransferReadOp read,
427 MaskingOpInterface maskOp,
428 PatternRewriter &rewriter) const override {
429 if (maxTransferRank && read.getVectorType().getRank() > *maxTransferRank) {
430 return rewriter.notifyMatchFailure(
431 read, "vector type is greater than max transfer rank");
432 }
433
434 if (maskOp)
435 return rewriter.notifyMatchFailure(read, "Masked case not supported");
436 SmallVector<unsigned> broadcastedDims;
437 // Permutations are handled by VectorToSCF or
438 // populateVectorTransferPermutationMapLoweringPatterns.
439 // We let the 0-d corner case pass-through as it is supported.
440 if (!read.getPermutationMap().isMinorIdentityWithBroadcasting(
441 &broadcastedDims))
442 return rewriter.notifyMatchFailure(read, "not minor identity + bcast");
443
444 auto memRefType = dyn_cast<MemRefType>(read.getShapedType());
445 if (!memRefType)
446 return rewriter.notifyMatchFailure(read, "not a memref source");
447
448 // Non-unit strides are handled by VectorToSCF.
449 if (!memRefType.isLastDimUnitStride())
450 return rewriter.notifyMatchFailure(read, "!= 1 stride needs VectorToSCF");
451
452 // If there is broadcasting involved then we first load the unbroadcasted
453 // vector, and then broadcast it with `vector.broadcast`.
454 ArrayRef<int64_t> vectorShape = read.getVectorType().getShape();
455 SmallVector<int64_t> unbroadcastedVectorShape(vectorShape);
456 for (unsigned i : broadcastedDims)
457 unbroadcastedVectorShape[i] = 1;
458 VectorType unbroadcastedVectorType = read.getVectorType().cloneWith(
459 unbroadcastedVectorShape, read.getVectorType().getElementType());
460
461 // `vector.load` supports vector types as memref's elements only when the
462 // resulting vector type is the same as the element type.
463 auto memrefElTy = memRefType.getElementType();
464 if (isa<VectorType>(memrefElTy) && memrefElTy != unbroadcastedVectorType)
465 return rewriter.notifyMatchFailure(read, "incompatible element type");
466
467 // Otherwise, element types of the memref and the vector must match.
468 if (!isa<VectorType>(memrefElTy) &&
469 memrefElTy != read.getVectorType().getElementType())
470 return rewriter.notifyMatchFailure(read, "non-matching element type");
471
472 // Out-of-bounds dims are handled by MaterializeTransferMask.
473 if (read.hasOutOfBoundsDim())
474 return rewriter.notifyMatchFailure(read, "out-of-bounds needs mask");
475
476 // Create vector load op.
477 Operation *res;
478 if (read.getMask()) {
479 if (read.getVectorType().getRank() != 1)
480 // vector.maskedload operates on 1-D vectors.
481 return rewriter.notifyMatchFailure(
482 read, "vector type is not rank 1, can't create masked load, needs "
483 "VectorToSCF");
484
485 Value fill = vector::BroadcastOp::create(
486 rewriter, read.getLoc(), unbroadcastedVectorType, read.getPadding());
487 res = vector::MaskedLoadOp::create(
488 rewriter, read.getLoc(), unbroadcastedVectorType, read.getBase(),
489 read.getIndices(), read.getMask(), fill);
490 } else {
491 res = vector::LoadOp::create(rewriter, read.getLoc(),
492 unbroadcastedVectorType, read.getBase(),
493 read.getIndices());
494 }
495
496 // Insert a broadcasting op if required.
497 if (!broadcastedDims.empty())
498 res = vector::BroadcastOp::create(
499 rewriter, read.getLoc(), read.getVectorType(), res->getResult(0));
500 return res->getResult(0);
501 }
502
503 std::optional<unsigned> maxTransferRank;
504};
505
506/// Progressive lowering of transfer_write. This pattern supports lowering of
507/// `vector.transfer_write` to `vector.store` if all of the following hold:
508/// - Stride of most minor memref dimension must be 1.
509/// - Out-of-bounds masking is not required.
510/// - If the memref's element type is a vector type then it coincides with the
511/// type of the written value.
512/// - The permutation map is the minor identity map (neither permutation nor
513/// broadcasting is allowed).
514struct TransferWriteToVectorStoreLowering
515 : public MaskableOpRewritePattern<vector::TransferWriteOp> {
516 TransferWriteToVectorStoreLowering(MLIRContext *context,
517 std::optional<unsigned> maxRank,
518 PatternBenefit benefit = 1)
519 : MaskableOpRewritePattern<vector::TransferWriteOp>(context, benefit),
520 maxTransferRank(maxRank) {}
521
522 FailureOr<mlir::Value>
523 matchAndRewriteMaskableOp(vector::TransferWriteOp write,
524 MaskingOpInterface maskOp,
525 PatternRewriter &rewriter) const override {
526 if (maxTransferRank && write.getVectorType().getRank() > *maxTransferRank) {
527 return rewriter.notifyMatchFailure(
528 write, "vector type is greater than max transfer rank");
529 }
530 if (maskOp)
531 return rewriter.notifyMatchFailure(write, "Masked case not supported");
532
533 // Permutations are handled by VectorToSCF or
534 // populateVectorTransferPermutationMapLoweringPatterns.
535 if ( // pass-through for the 0-d corner case.
536 !write.getPermutationMap().isMinorIdentity())
537 return rewriter.notifyMatchFailure(write.getLoc(), [=](Diagnostic &diag) {
538 diag << "permutation map is not minor identity: " << write;
539 });
540
541 auto memRefType = dyn_cast<MemRefType>(write.getShapedType());
542 if (!memRefType)
543 return rewriter.notifyMatchFailure(write.getLoc(), [=](Diagnostic &diag) {
544 diag << "not a memref type: " << write;
545 });
546
547 // Non-unit strides are handled by VectorToSCF.
548 if (!memRefType.isLastDimUnitStride())
549 return rewriter.notifyMatchFailure(write.getLoc(), [=](Diagnostic &diag) {
550 diag << "most minor stride is not 1: " << write;
551 });
552
553 // `vector.store` supports vector types as memref's elements only when the
554 // type of the vector value being written is the same as the element type.
555 auto memrefElTy = memRefType.getElementType();
556 if (isa<VectorType>(memrefElTy) && memrefElTy != write.getVectorType())
557 return rewriter.notifyMatchFailure(write.getLoc(), [=](Diagnostic &diag) {
558 diag << "elemental type mismatch: " << write;
559 });
560
561 // Otherwise, element types of the memref and the vector must match.
562 if (!isa<VectorType>(memrefElTy) &&
563 memrefElTy != write.getVectorType().getElementType())
564 return rewriter.notifyMatchFailure(write.getLoc(), [=](Diagnostic &diag) {
565 diag << "elemental type mismatch: " << write;
566 });
567
568 // Out-of-bounds dims are handled by MaterializeTransferMask.
569 if (write.hasOutOfBoundsDim())
570 return rewriter.notifyMatchFailure(write.getLoc(), [=](Diagnostic &diag) {
571 diag << "out of bounds dim: " << write;
572 });
573 if (write.getMask()) {
574 if (write.getVectorType().getRank() != 1)
575 // vector.maskedstore operates on 1-D vectors.
576 return rewriter.notifyMatchFailure(
577 write.getLoc(), [=](Diagnostic &diag) {
578 diag << "vector type is not rank 1, can't create masked store, "
579 "needs VectorToSCF: "
580 << write;
581 });
582
583 vector::MaskedStoreOp::create(rewriter, write.getLoc(), write.getBase(),
584 write.getIndices(), write.getMask(),
585 write.getVector());
586 } else {
587 vector::StoreOp::create(rewriter, write.getLoc(), write.getVector(),
588 write.getBase(), write.getIndices());
589 }
590 // There's no return value for StoreOps. Use Value() to signal success to
591 // matchAndRewrite.
592 return Value();
593 }
594
595 std::optional<unsigned> maxTransferRank;
596};
597} // namespace
598
600 RewritePatternSet &patterns, std::optional<unsigned> maxTransferRank,
601 PatternBenefit benefit) {
602 patterns.add<TransferReadToVectorLoadLowering,
603 TransferWriteToVectorStoreLowering>(patterns.getContext(),
604 maxTransferRank, benefit);
605}
ArrayAttr()
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 std::optional< VectorShape > vectorShape(Type type)
static AffineMap getMinorIdentityMap(unsigned dims, unsigned results, MLIRContext *context)
Returns an identity affine map (d0, ..., dn) -> (dp, ..., dn) on the most minor dimensions.
bool isMinorIdentity() const
Returns true if this affine map is a minor identity, i.e.
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 ...
unsigned getNumDims() const
ArrayRef< AffineExpr > getResults() const
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...
unsigned getNumResults() const
static AffineMap getPermutationMap(ArrayRef< unsigned > permutation, MLIRContext *context)
Returns an AffineMap representing a permutation.
AffineMap compose(AffineMap map) const
Returns the AffineMap resulting from composing this with map.
bool isIdentity() const
Returns true if this affine map is an identity affine map.
AffineExpr getAffineDimExpr(unsigned position)
Definition Builders.cpp:373
ArrayAttr getArrayAttr(ArrayRef< Attribute > value)
Definition Builders.cpp:275
MLIRContext * getContext() const
Definition Builders.h:56
ArrayAttr getBoolArrayAttr(ArrayRef< bool > values)
Definition Builders.cpp:279
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
Definition Location.h:76
MLIRContext is the top-level object for a collection of MLIR operations.
Definition MLIRContext.h:63
This class helps build Operations.
Definition Builders.h:210
OpResult getResult(unsigned idx)
Get the 'idx'th result of this operation.
Definition Operation.h:432
This class represents the benefit of a pattern match in a unitless scheme that ranges from 0 (very li...
A special type of RewriterBase that coordinates the application of a rewrite pattern on the current I...
MLIRContext * getContext() const
RewritePatternSet & add(ConstructorArg &&arg, ConstructorArgs &&...args)
Add an instance of each of the pattern types 'Ts' to the pattern list with the given arguments.
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,...
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:105
void populateVectorTransferPermutationMapLoweringPatterns(RewritePatternSet &patterns, PatternBenefit benefit=1)
Collect a set of transfer read/write lowering patterns that simplify the permutation map (e....
Operation * maskOperation(OpBuilder &builder, Operation *maskableOp, Value mask, Value passthru=Value())
Creates a vector.mask operation around a maskable operation.
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.
AffineMap inversePermutation(AffineMap map)
Returns a map of codomain to domain dimensions such that the first codomain dimension for a particula...
AffineMap compressUnusedDims(AffineMap map)
Drop the dims that are not used.
SmallVector< int64_t > invertPermutationVector(ArrayRef< int64_t > permutation)
Helper method to apply to inverse a permutation.
A pattern for ops that implement MaskableOpInterface and that might be masked (i.e.