[mlir][vector] Canonicalize/fold 'order preserving' transposes #135841
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -5575,6 +5575,34 @@ LogicalResult ShapeCastOp::verify() { | |
| return success(); | ||
| } | ||
|
|
||
| namespace { | ||
|
|
||
| /// Return true if `transpose` does not permute a pair of non-unit dims. | ||
| /// By `order preserving` we mean that the flattened versions of the input and | ||
| /// output vectors are (numerically) identical. In other words `transpose` is | ||
| /// effectively a shape cast. | ||
| bool isOrderPreserving(TransposeOp transpose) { | ||
| ArrayRef<int64_t> permutation = transpose.getPermutation(); | ||
| VectorType sourceType = transpose.getSourceVectorType(); | ||
| ArrayRef<int64_t> inShape = sourceType.getShape(); | ||
| ArrayRef<bool> inDimIsScalable = sourceType.getScalableDims(); | ||
| auto isNonScalableUnitDim = [&](int64_t dim) { | ||
| return inShape[dim] == 1 && !inDimIsScalable[dim]; | ||
| }; | ||
| int64_t current = 0; | ||
| for (auto p : permutation) { | ||
| if (!isNonScalableUnitDim(p)) { | ||
| if (p < current) { | ||
| return false; | ||
| } | ||
| current = p; | ||
| } | ||
| } | ||
| return true; | ||
| } | ||
|
|
||
| } // namespace | ||
|
|
||
| OpFoldResult ShapeCastOp::fold(FoldAdaptor adaptor) { | ||
|
|
||
| VectorType resultType = getType(); | ||
|
|
@@ -5583,17 +5611,32 @@ OpFoldResult ShapeCastOp::fold(FoldAdaptor adaptor) { | |
| if (getSource().getType() == resultType) | ||
| return getSource(); | ||
|
|
||
| // shape_cast(shape_cast(x)) -> shape_cast(x) | ||
| if (auto precedingShapeCast = getSource().getDefiningOp<ShapeCastOp>()) { | ||
| setOperand(precedingShapeCast.getSource()); | ||
| return getResult(); | ||
| } | ||
|
|
||
| // shape_cast(transpose(x)) -> shape_cast(x) | ||
| if (auto transpose = getSource().getDefiningOp<TransposeOp>()) { | ||
| // This folder does | ||
| // shape_cast(transpose) -> shape_cast | ||
| // But another pattern, ConvertIllegalShapeCastOpsToTransposes, does | ||
| // shape_cast -> shape_cast(transpose) | ||
| // i.e. the complete opposite. When paired, these 2 patterns can cause | ||
| // infinite cycles in pattern rewriting. | ||
| // ConvertIllegalShapeCastOpsToTransposes only matches on scalable | ||
| // vectors, so by disabling this folder for scalable vectors the | ||
| // cycle is avoided. | ||
| // TODO: Check if ConvertIllegalShapeCastOpsToTransposes is | ||
|
There was a problem hiding this comment. Do you know why we generate illegal shape cast ops in first place? It sounds like something that shouldn't happen... There was a problem hiding this comment. They're illegal according to an arm specific lowering target, which I'm not familiar with I think @banach-space suspects that actually they're not illegal and will investigate the removal of this constraint here (and the pattern ConvertIllegalShapeCastOpsToTransposes). There was a problem hiding this comment. "Illegal" in this context means "Illegal as an SME target - transform before invoking Tidying this up is high on my TODO list (will look into it this week). There was a problem hiding this comment. Looks like There was a problem hiding this comment. 🎉 thanks @banach-space |
||
| // still needed. If it's not, then we can fold here. | ||
| if (!transpose.getType().isScalable() && isOrderPreserving(transpose)) { | ||
| setOperand(transpose.getVector()); | ||
| return getResult(); | ||
| } | ||
| return {}; | ||
| } | ||
|
|
||
| // Y = shape_cast(broadcast(X)) | ||
| // -> X, if X and Y have same type | ||
| if (auto bcastOp = getSource().getDefiningOp<BroadcastOp>()) { | ||
|
|
@@ -5619,7 +5662,7 @@ namespace { | |
| /// Helper function that computes a new vector type based on the input vector | ||
| /// type by removing the trailing one dims: | ||
| /// | ||
| /// vector<4x1x1xi1> --> vector<4x1xi1> | ||
| /// | ||
| static VectorType trimTrailingOneDims(VectorType oldType) { | ||
| ArrayRef<int64_t> oldShape = oldType.getShape(); | ||
|
|
@@ -6086,6 +6129,32 @@ class FoldTransposeCreateMask final : public OpRewritePattern<TransposeOp> { | |
| } | ||
| }; | ||
|
|
||
| /// Folds transpose(shape_cast) into a new shape_cast. | ||
| class FoldTransposeShapeCast final : public OpRewritePattern<TransposeOp> { | ||
| public: | ||
| using OpRewritePattern::OpRewritePattern; | ||
|
|
||
| LogicalResult matchAndRewrite(TransposeOp transposeOp, | ||
| PatternRewriter &rewriter) const override { | ||
| auto shapeCastOp = | ||
| transposeOp.getVector().getDefiningOp<vector::ShapeCastOp>(); | ||
| if (!shapeCastOp) | ||
| return failure(); | ||
| if (!isOrderPreserving(transposeOp)) | ||
| return failure(); | ||
|
|
||
| VectorType resultType = transposeOp.getType(); | ||
|
|
||
| // We don't need to check isValidShapeCast at this point, because it is | ||
| // guaranteed that merging the transpose into the the shape_cast is a valid | ||
| // shape_cast, because the transpose just inserts/removes ones. | ||
|
|
||
| rewriter.replaceOpWithNewOp<vector::ShapeCastOp>(transposeOp, resultType, | ||
| shapeCastOp.getSource()); | ||
| return success(); | ||
| } | ||
| }; | ||
|
|
||
| /// Folds transpose(broadcast(x)) to broadcast(x) if the transpose is | ||
| /// 'order preserving', where 'order preserving' means the flattened | ||
| /// inputs and outputs of the transpose have identical (numerical) values. | ||
|
|
@@ -6184,8 +6253,8 @@ class FoldTransposeBroadcast : public OpRewritePattern<vector::TransposeOp> { | |
|
|
||
| void vector::TransposeOp::getCanonicalizationPatterns( | ||
| RewritePatternSet &results, MLIRContext *context) { | ||
| results.add<FoldTransposeCreateMask, FoldTransposeShapeCast, TransposeFolder, | ||
| FoldTransposeSplat, FoldTransposeBroadcast>(context); | ||
| } | ||
|
|
||
| //===----------------------------------------------------------------------===// | ||
|
|
||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Probably a case for
static: https://llvm.org/docs/CodingStandards.html#restrict-visibilityThere was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks. I intend to change this in a follow-up PR.