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