27 for (
unsigned pos : permutation)
28 newInBoundsValues[pos] =
29 cast<BoolAttr>(attr.getValue()[
index++]).getValue();
37 auto vecType = cast<VectorType>(vec.
getType());
38 int64_t numUnitDims = unitDimPositions.size();
41 llvm::append_range(newDims, vecType.getShape());
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);
49 for (
auto [i, pos] : llvm::enumerate(unitDimPositions))
52 for (
int64_t &dim : permutation)
57 return vector::TransposeOp::create(builder, loc,
broadcast, permutation);
84struct TransferReadPermutationLowering
86 using MaskableOpRewritePattern::MaskableOpRewritePattern;
88 FailureOr<mlir::Value>
89 matchAndRewriteMaskableOp(vector::TransferReadOp op,
90 MaskingOpInterface maskOp,
91 PatternRewriter &rewriter)
const override {
93 if (op.getTransferRank() == 0)
96 SmallVector<unsigned> permutation;
97 AffineMap map = op.getPermutationMap();
102 op,
"map is not permutable to minor identity, apply another pattern");
104 AffineMap permutationMap =
112 AffineMap newMap = permutationMap.
compose(map);
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()];
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(),
135 SmallVector<int64_t> transposePerm(permutation.begin(), permutation.end());
140 Value passthru = maskOp.getPassthru();
143 vector::TransposeOp::create(rewriter, op.getLoc(), passthru,
150 return vector::TransposeOp::create(rewriter, op.getLoc(),
172struct TransferWritePermutationLowering
174 using MaskableOpRewritePattern::MaskableOpRewritePattern;
176 FailureOr<mlir::Value>
177 matchAndRewriteMaskableOp(vector::TransferWriteOp op,
178 MaskingOpInterface maskOp,
179 PatternRewriter &rewriter)
const override {
181 if (op.getTransferRank() == 0)
184 SmallVector<unsigned> permutation;
191 op,
"map is not permutable to minor identity, apply another pattern");
202 [](AffineExpr expr) {
203 return dyn_cast<AffineDimExpr>(expr).getPosition();
211 Value newVec = vector::TransposeOp::create(rewriter, op.getLoc(),
215 auto newWrite = vector::TransferWriteOp::create(
216 rewriter, op.getLoc(), newVec, op.getBase(), op.getIndices(),
217 AffineMapAttr::get(newMap), op.getMask(), newInBoundsAttr);
221 Operation *rewritten = newWrite;
225 if (newWrite.hasPureTensorSemantics())
248struct TransferWriteNonPermutationLowering
250 using MaskableOpRewritePattern::MaskableOpRewritePattern;
252 FailureOr<mlir::Value>
253 matchAndRewriteMaskableOp(vector::TransferWriteOp op,
254 MaskingOpInterface maskOp,
255 PatternRewriter &rewriter)
const override {
257 if (op.getTransferRank() == 0)
263 SmallVector<unsigned> permutation;
268 "map is already permutable to minor identity, apply another pattern");
273 SmallVector<bool> foundDim(map.
getNumDims(),
false);
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++) {
289 missingInnerDim.push_back(i - *firstDim);
294 rewriter, op.getLoc(), op.getVector(),
295 llvm::to_vector(llvm::seq<int64_t>(missingInnerDim.size())));
301 insertUnitDims(rewriter, op.getLoc(), op.getMask(), missingInnerDim);
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));
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();
330struct TransferOpReduceRank
332 using MaskableOpRewritePattern::MaskableOpRewritePattern;
334 FailureOr<mlir::Value>
335 matchAndRewriteMaskableOp(vector::TransferReadOp op,
336 MaskingOpInterface maskOp,
337 PatternRewriter &rewriter)
const override {
339 if (op.getTransferRank() == 0)
346 unsigned numLeadingBroadcast = 0;
348 auto dimExpr = dyn_cast<AffineConstantExpr>(expr);
349 if (!dimExpr || dimExpr.getValue() != 0)
351 numLeadingBroadcast++;
354 if (numLeadingBroadcast == 0)
357 VectorType originalVecType = op.getVectorType();
358 unsigned reducedShapeRank = originalVecType.getRank() - numLeadingBroadcast;
367 op,
"map is not a minor identity with broadcasting");
370 SmallVector<int64_t> newShape(
371 originalVecType.getShape().take_back(reducedShapeRank));
372 SmallVector<bool> newScalableDims(
373 originalVecType.getScalableDims().take_back(reducedShapeRank));
375 VectorType newReadType = VectorType::get(
376 newShape, originalVecType.getElementType(), newScalableDims);
380 op.getInBoundsAttr().getValue().take_back(reducedShapeRank))
382 Value newRead = vector::TransferReadOp::create(
383 rewriter, op.getLoc(), newReadType, op.getBase(), op.getIndices(),
384 AffineMapAttr::get(newMap), op.getPadding(), op.getMask(),
386 return vector::BroadcastOp::create(rewriter, op.getLoc(), originalVecType,
397 .
add<TransferReadPermutationLowering, TransferWritePermutationLowering,
398 TransferOpReduceRank, TransferWriteNonPermutationLowering>(
415struct TransferReadToVectorLoadLowering
417 TransferReadToVectorLoadLowering(
MLIRContext *context,
418 std::optional<unsigned> maxRank,
421 maxTransferRank(maxRank) {}
423 FailureOr<mlir::Value>
424 matchAndRewriteMaskableOp(vector::TransferReadOp read,
425 MaskingOpInterface maskOp,
427 if (maxTransferRank && read.getVectorType().getRank() > *maxTransferRank) {
429 read,
"vector type is greater than max transfer rank");
434 SmallVector<unsigned> broadcastedDims;
438 if (!read.getPermutationMap().isMinorIdentityWithBroadcasting(
442 auto memRefType = dyn_cast<MemRefType>(read.getShapedType());
447 if (!memRefType.isLastDimUnitStride())
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());
461 auto memrefElTy = memRefType.getElementType();
462 if (isa<VectorType>(memrefElTy) && memrefElTy != unbroadcastedVectorType)
466 if (!isa<VectorType>(memrefElTy) &&
467 memrefElTy != read.getVectorType().getElementType())
471 if (read.hasOutOfBoundsDim())
476 if (read.getMask()) {
477 if (read.getVectorType().getRank() != 1)
480 read,
"vector type is not rank 1, can't create masked load, needs "
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);
489 res = vector::LoadOp::create(rewriter, read.getLoc(),
490 unbroadcastedVectorType, read.getBase(),
495 if (!broadcastedDims.empty())
496 res = vector::BroadcastOp::create(
497 rewriter, read.getLoc(), read.getVectorType(), res->getResult(0));
498 return res->getResult(0);
501 std::optional<unsigned> maxTransferRank;
512struct TransferWriteToVectorStoreLowering
514 TransferWriteToVectorStoreLowering(MLIRContext *context,
515 std::optional<unsigned> maxRank,
516 PatternBenefit benefit = 1)
517 : MaskableOpRewritePattern<vector::TransferWriteOp>(context, benefit),
518 maxTransferRank(maxRank) {}
520 FailureOr<mlir::Value>
521 matchAndRewriteMaskableOp(vector::TransferWriteOp write,
522 MaskingOpInterface maskOp,
523 PatternRewriter &rewriter)
const override {
524 if (maxTransferRank && write.getVectorType().getRank() > *maxTransferRank) {
526 write,
"vector type is greater than max transfer rank");
534 !write.getPermutationMap().isMinorIdentity())
536 diag <<
"permutation map is not minor identity: " << write;
539 auto memRefType = dyn_cast<MemRefType>(write.getShapedType());
542 diag <<
"not a memref type: " << write;
546 if (!memRefType.isLastDimUnitStride())
548 diag <<
"most minor stride is not 1: " << write;
553 auto memrefElTy = memRefType.getElementType();
554 if (isa<VectorType>(memrefElTy) && memrefElTy != write.getVectorType())
556 diag <<
"elemental type mismatch: " << write;
560 if (!isa<VectorType>(memrefElTy) &&
561 memrefElTy != write.getVectorType().getElementType())
563 diag <<
"elemental type mismatch: " << write;
567 if (write.hasOutOfBoundsDim())
569 diag <<
"out of bounds dim: " << write;
571 if (write.getMask()) {
572 if (write.getVectorType().getRank() != 1)
575 write.getLoc(), [=](Diagnostic &diag) {
576 diag <<
"vector type is not rank 1, can't create masked store, "
577 "needs VectorToSCF: "
581 vector::MaskedStoreOp::create(rewriter, write.getLoc(), write.getBase(),
582 write.getIndices(), write.getMask(),
585 vector::StoreOp::create(rewriter, write.getLoc(), write.getVector(),
586 write.getBase(), write.getIndices());
593 std::optional<unsigned> maxTransferRank;
600 patterns.
add<TransferReadToVectorLoadLowering,
601 TransferWriteToVectorStoreLowering>(patterns.
getContext(),
602 maxTransferRank, benefit);
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)
ArrayAttr getArrayAttr(ArrayRef< Attribute > value)
MLIRContext * getContext() const
ArrayAttr getBoolArrayAttr(ArrayRef< bool > values)
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
MLIRContext is the top-level object for a collection of MLIR operations.
This class helps build Operations.
OpResult getResult(unsigned idx)
Get the 'idx'th result of this operation.
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...
Type getType() const
Return the type of this value.
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.