[vector][mlir] Canonicalize to shape_cast where possible #140583
Conversation
newling
changed the title
newling
changed the title
newling
force-pushed
the
canonicalize_to_shape_cast
branch
from
d546ab3 to
29d41d8
Compare
newling
force-pushed
the
canonicalize_to_shape_cast
branch
from
29d41d8 to
f2e5417
Compare
|
✅ With the latest revision this PR passed the C/C++ code formatter. |
There was a problem hiding this comment.
@banach-space I'm getting back to this PR. Peephole question: is this operation ok? i.e. is
vector.shape_cast %a vector<[4]x1xf32> to vector<1x[4]xf32>
an acceptable operation to have after running mlir-opt -arm-sme-vector-legalization -cse -canonicalize ?
There was a problem hiding this comment.
In general, yes. But I can't guarantee there's no logic that expects vector<[4]x1xf32> instead of vector<1x[4]xf32> ;-) If that's the case, we will fix it and I will be grateful for uncovering this :)
newling
force-pushed
the
canonicalize_to_shape_cast
branch
from
7bc5da0 to
e673522
Compare
There was a problem hiding this comment.
Author note: I've removed this, as now it happens in 2 steps during canonicalization. The first converts the Broadcast to a ShapeCast. The second combines the 2 ShapeCasts.
There was a problem hiding this comment.
Author note: I've removed this, as it now happens in 2 steps during canonicalization. The first (new) step is to rewrite the transpose as a shape_cast. The second step is to fold shape_cast(shape_cast) to shape_cast.
There was a problem hiding this comment.
Author note: I've removed this pattern, as it is a special case of TransposeToShapeCast
There was a problem hiding this comment.
Author note: removed these tests, as the pattern they are testing is removed
There was a problem hiding this comment.
Shouldn't we keep them? shouldn't they still be canonicalized?
There was a problem hiding this comment.
I'll add them back, yes they're still canonicalized
There was a problem hiding this comment.
Author note: as the vector.transpose is canonicalized to a vector.shape_cast, the lowering test is now moved to shape_cast lowering
newling
requested review from
Groverkss,
dcaballe,
hanhanW and
nicolasvasilache
as code owners
|
@llvm/pr-subscribers-mlir @llvm/pr-subscribers-mlir-sme Author: James Newling (newling) ChangesDiscussions suggest that we should use shape_cast as a canonical form of broadcast/transpose/extract where possible (see #138777) For example these can all be expressed as shape casts: %0 = vector.broadcast %arg0 : vector<4xi8> to vector<1x1x4xi8>
%1 = vector.transpose %arg1, [1, 0] : vector<2x1xi8> to vector<1x2xi8>
%2 = vector.extract %arg2[0] : vector<4xi8> from vector<1x4xi8>This PR adds canonicalizes to convert the above 3 examples to shape_casts. I've added some more comments as review comments. I'm happy to split this PR up and add the new patterns separately. Patch is 41.81 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/140583.diff 10 Files Affected:
diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index 862ed7bae1fbb..08cc4af158e10 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -2351,11 +2351,41 @@ LogicalResult foldExtractFromFromElements(ExtractOp extractOp,
return success();
}
+/// BEFORE:
+/// %0 = vector.extract %arg0[0] : vector<4xf32> from vector<1x4xf32>
+/// AFTER:
+/// %0 = vector.shape_cast %arg0 : vector<1x4xf32> to vector<4xf32>
+struct ExtractToShapeCast final : public OpRewritePattern<vector::ExtractOp> {
+ using OpRewritePattern::OpRewritePattern;
+ LogicalResult matchAndRewrite(vector::ExtractOp extractOp,
+ PatternRewriter &rewriter) const override {
+ VectorType sourceType = extractOp.getSourceVectorType();
+ VectorType outType = dyn_cast<VectorType>(extractOp.getType());
+ if (!outType)
+ return failure();
+
+ // Negative values in `position` indicates poison, which cannot be
+ // represented with a shape_cast
+ if (llvm::any_of(extractOp.getMixedPosition(),
+ [](OpFoldResult v) { return !isConstantIntValue(v, 0); }))
+ return failure();
+
+ if (sourceType.getNumElements() != outType.getNumElements())
+ return failure();
+
+ rewriter.replaceOpWithNewOp<vector::ShapeCastOp>(extractOp, outType,
+ extractOp.getVector());
+ return success();
+ }
+};
+
} // namespace
void ExtractOp::getCanonicalizationPatterns(RewritePatternSet &results,
MLIRContext *context) {
- results.add<ExtractOpFromBroadcast, ExtractOpFromCreateMask>(context);
+ results
+ .add<ExtractOpFromBroadcast, ExtractOpFromCreateMask, ExtractToShapeCast>(
+ context);
results.add(foldExtractFromShapeCastToShapeCast);
results.add(foldExtractFromFromElements);
}
@@ -2867,13 +2897,36 @@ struct BroadcastFolder : public OpRewritePattern<BroadcastOp> {
return success();
}
};
+
+/// BEFORE:
+/// %0 = vector.broadcast %arg0 : vector<4xi8> to vector<1x1x4xi8>
+/// AFTER:
+/// %0 = vector.shape_cast %arg0 : vector<4xi8> to vector<1x1x4xi8>
+struct BroadcastToShapeCast final
+ : public OpRewritePattern<vector::BroadcastOp> {
+ using OpRewritePattern::OpRewritePattern;
+ LogicalResult matchAndRewrite(vector::BroadcastOp broadcast,
+ PatternRewriter &rewriter) const override {
+ auto sourceType = dyn_cast<VectorType>(broadcast.getSourceType());
+ if (!sourceType) {
+ return rewriter.notifyMatchFailure(
+ broadcast, "source is a scalar, shape_cast doesn't support scalar");
+ }
+
+ VectorType outType = broadcast.getType();
+ if (sourceType.getNumElements() != outType.getNumElements())
+ return failure();
+
+ rewriter.replaceOpWithNewOp<vector::ShapeCastOp>(broadcast, outType,
+ broadcast.getSource());
+ return success();
+ }
+};
} // namespace
void BroadcastOp::getCanonicalizationPatterns(RewritePatternSet &results,
MLIRContext *context) {
- // BroadcastToShapeCast is not a default canonicalization, it is opt-in by
- // calling `populateCastAwayVectorLeadingOneDimPatterns`
- results.add<BroadcastFolder>(context);
+ results.add<BroadcastFolder, BroadcastToShapeCast>(context);
}
//===----------------------------------------------------------------------===//
@@ -5991,10 +6044,7 @@ class ShapeCastCreateMaskFolderTrailingOneDim final
}
};
-/// Pattern to rewrite Y = ShapeCast(Broadcast(X)) as either
-/// i) Y = ShapeCast(X), or
-/// ii) Y = Broadcast(X)
-/// If both (i) and (ii) are possible, (i) is chosen.
+/// Pattern to rewrite Y = ShapeCast(Broadcast(X)) as Y = Broadcast(X)
class ShapeCastBroadcastFolder final : public OpRewritePattern<ShapeCastOp> {
public:
using OpRewritePattern::OpRewritePattern;
@@ -6009,22 +6059,6 @@ class ShapeCastBroadcastFolder final : public OpRewritePattern<ShapeCastOp> {
auto srcVectorType = dyn_cast<VectorType>(broadcastOp.getSourceType());
bool srcIsScalar = !srcVectorType;
- // Replace Y = ShapeCast(Broadcast(X)) with Y = ShapeCast(X).
- // Example:
- // %0 = vector.broadcast %in : vector<3x4xf32> to vector<1x3x4xf32>
- // %1 = vector.shape_cast %0 : vector<1x3x4xf32> to vector<12xf32>
- // to
- // %1 = vector.shape_cast %in : vector<3x4xf32> to vector<12xf32>
- if (srcVectorType) {
- if (srcVectorType.getNumElements() ==
- shapeCastOp.getResultVectorType().getNumElements()) {
- rewriter.replaceOpWithNewOp<vector::ShapeCastOp>(
- shapeCastOp, shapeCastOp.getResultVectorType(),
- broadcastOp.getSource());
- return success();
- }
- }
-
// Replace Y = ShapeCast(Broadcast(X)) with Y = Broadcast(X)
// Example
// %0 = vector.broadcast %in : vector<3xf32> to vector<2x4x3xf32>
@@ -6233,7 +6267,7 @@ OpFoldResult vector::TransposeOp::fold(FoldAdaptor adaptor) {
// %0 = vector.transpose %arg, [0, 1] : vector<2x2xi8> to vector<2x2xi8>
// %0 = vector.transpose %arg, [1, 0] : vector<1x1xi8> to vector<1x1xi8>
//
- // Example of what NOT to fold:
+ // Example of what not to fold:
// %0 = vector.transpose %arg, [1, 0] : vector<2x2xi8> to vector<2x2xi8>
//
if (getSourceVectorType() == getResultVectorType() &&
@@ -6359,32 +6393,6 @@ 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.
@@ -6480,12 +6488,35 @@ class FoldTransposeBroadcast : public OpRewritePattern<vector::TransposeOp> {
}
};
+/// BEFORE:
+/// %0 = vector.transpose %arg0, [0, 2, 1] :
+/// vector<2x1x2xf32> to vector<2x2x1xf32>
+/// AFTER:
+/// %0 = vector.shape_cast %arg0 :
+/// vector<2x1x2xf32> to vector<2x2x1xf32>
+struct TransposeToShapeCast final
+ : public OpRewritePattern<vector::TransposeOp> {
+ using OpRewritePattern::OpRewritePattern;
+ LogicalResult matchAndRewrite(vector::TransposeOp transpose,
+ PatternRewriter &rewriter) const override {
+
+ if (!isOrderPreserving(transpose)) {
+ return rewriter.notifyMatchFailure(
+ transpose, "not order preserving, so not semantically a 'copy'");
+ }
+ rewriter.replaceOpWithNewOp<vector::ShapeCastOp>(
+ transpose, transpose.getType(), transpose.getVector());
+ return success();
+ }
+};
+
} // namespace
void vector::TransposeOp::getCanonicalizationPatterns(
RewritePatternSet &results, MLIRContext *context) {
- results.add<FoldTransposeCreateMask, FoldTransposeShapeCast, TransposeFolder,
- FoldTransposeSplat, FoldTransposeBroadcast>(context);
+ results.add<FoldTransposeBroadcast, FoldTransposeCreateMask,
+ FoldTransposeSplat, TransposeFolder, TransposeToShapeCast>(
+ context);
}
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/Vector/Transforms/LowerVectorTranspose.cpp b/mlir/lib/Dialect/Vector/Transforms/LowerVectorTranspose.cpp
index 732e316c93381..71410eda28297 100644
--- a/mlir/lib/Dialect/Vector/Transforms/LowerVectorTranspose.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/LowerVectorTranspose.cpp
@@ -11,7 +11,6 @@
//
//===----------------------------------------------------------------------===//
-#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/UB/IR/UBOps.h"
#include "mlir/Dialect/Utils/IndexingUtils.h"
@@ -382,64 +381,6 @@ class TransposeOpLowering : public OpRewritePattern<vector::TransposeOp> {
vector::VectorTransposeLowering vectorTransposeLowering;
};
-/// Rewrites vector.transpose as vector.shape_cast. This pattern is only applied
-/// to 2D vectors with at least one unit dim. For example:
-///
-/// Replace:
-/// vector.transpose %0, [1, 0] : vector<4x1xi32>> to
-/// vector<1x4xi32>
-/// with:
-/// vector.shape_cast %0 : vector<4x1xi32> to vector<1x4xi32>
-///
-/// Source with leading unit dim (inverse) is also replaced. Unit dim must
-/// be fixed. Non-unit dim can be scalable.
-///
-/// TODO: This pattern was introduced specifically to help lower scalable
-/// vectors. In hindsight, a more specialised canonicalization (for shape_cast's
-/// to cancel out) would be preferable:
-///
-/// BEFORE:
-/// %0 = some_op
-/// %1 = vector.shape_cast %0 : vector<[4]xf32> to vector<[4]x1xf32>
-/// %2 = vector.transpose %1 [1, 0] : vector<[4]x1xf32> to vector<1x[4]xf32>
-/// AFTER:
-/// %0 = some_op
-/// %1 = vector.shape_cast %0 : vector<[4]xf32> to vector<1x[4]xf32>
-///
-/// Given the context above, we may want to consider (re-)moving this pattern
-/// at some later time. I am leaving it for now in case there are other users
-/// that I am not aware of.
-class Transpose2DWithUnitDimToShapeCast
- : public OpRewritePattern<vector::TransposeOp> {
-public:
- using OpRewritePattern::OpRewritePattern;
-
- Transpose2DWithUnitDimToShapeCast(MLIRContext *context,
- PatternBenefit benefit = 1)
- : OpRewritePattern<vector::TransposeOp>(context, benefit) {}
-
- LogicalResult matchAndRewrite(vector::TransposeOp op,
- PatternRewriter &rewriter) const override {
- Value input = op.getVector();
- VectorType resType = op.getResultVectorType();
-
- // Set up convenience transposition table.
- ArrayRef<int64_t> transp = op.getPermutation();
-
- if (resType.getRank() == 2 &&
- ((resType.getShape().front() == 1 &&
- !resType.getScalableDims().front()) ||
- (resType.getShape().back() == 1 &&
- !resType.getScalableDims().back())) &&
- transp == ArrayRef<int64_t>({1, 0})) {
- rewriter.replaceOpWithNewOp<vector::ShapeCastOp>(op, resType, input);
- return success();
- }
-
- return failure();
- }
-};
-
/// Rewrite a 2-D vector.transpose as a sequence of shuffle ops.
/// If the strategy is Shuffle1D, it will be lowered to:
/// vector.shape_cast 2D -> 1D
@@ -511,8 +452,6 @@ class TransposeOp2DToShuffleLowering
void mlir::vector::populateVectorTransposeLoweringPatterns(
RewritePatternSet &patterns,
VectorTransposeLowering vectorTransposeLowering, PatternBenefit benefit) {
- patterns.add<Transpose2DWithUnitDimToShapeCast>(patterns.getContext(),
- benefit);
patterns.add<TransposeOpLowering, TransposeOp2DToShuffleLowering>(
vectorTransposeLowering, patterns.getContext(), benefit);
}
diff --git a/mlir/test/Dialect/ArmSME/vector-legalization.mlir b/mlir/test/Dialect/ArmSME/vector-legalization.mlir
index 6cdf576272ebc..a9a2fdccdd82f 100644
--- a/mlir/test/Dialect/ArmSME/vector-legalization.mlir
+++ b/mlir/test/Dialect/ArmSME/vector-legalization.mlir
@@ -480,11 +480,11 @@ func.func @lift_illegal_transpose_to_memory_with_in_bounds_attr(%a: index, %b: i
// -----
-// The pass should do nothing (and not crash).
-// CHECK-LABEL: @illegal_transpose_no_defining_source_op
-func.func @illegal_transpose_no_defining_source_op(%vec: vector<[4]x1xf32>) -> vector<1x[4]xf32>
+// CHECK-LABEL: @transpose_no_defining_source_op
+func.func @transpose_no_defining_source_op(%vec: vector<[4]x1xf32>) -> vector<1x[4]xf32>
{
- // CHECK: vector.transpose
+ // CHECK: vector.shape_cast
+ // CHECK-SAME: vector<[4]x1xf32> to vector<1x[4]xf32>
%0 = vector.transpose %vec, [1, 0] : vector<[4]x1xf32> to vector<1x[4]xf32>
return %0 : vector<1x[4]xf32>
}
diff --git a/mlir/test/Dialect/Vector/canonicalize.mlir b/mlir/test/Dialect/Vector/canonicalize.mlir
index 65b73375831da..374c71c814e89 100644
--- a/mlir/test/Dialect/Vector/canonicalize.mlir
+++ b/mlir/test/Dialect/Vector/canonicalize.mlir
@@ -451,16 +451,25 @@ func.func @extract_strided_fold_insert(%a: vector<2x8xf32>, %b: vector<1x4xf32>,
// -----
// CHECK-LABEL: transpose_3D_identity
-// CHECK-SAME: ([[ARG:%.*]]: vector<4x3x2xf32>)
+// CHECK-SAME: ([[ARG:%.*]]: vector<4x3x2xf32>)
+// CHECK-NEXT: return [[ARG]]
func.func @transpose_3D_identity(%arg : vector<4x3x2xf32>) -> vector<4x3x2xf32> {
- // CHECK-NOT: transpose
%0 = vector.transpose %arg, [0, 1, 2] : vector<4x3x2xf32> to vector<4x3x2xf32>
- // CHECK-NEXT: return [[ARG]]
return %0 : vector<4x3x2xf32>
}
// -----
+// CHECK-LABEL: transpose_0D_identity
+// CHECK-SAME: ([[ARG:%.*]]: vector<i8>)
+// CHECK-NEXT: return [[ARG]]
+func.func @transpose_0D_identity(%arg : vector<i8>) -> vector<i8> {
+ %0 = vector.transpose %arg, [] : vector<i8> to vector<i8>
+ return %0 : vector<i8>
+}
+
+// -----
+
// CHECK-LABEL: transpose_2D_sequence
// CHECK-SAME: ([[ARG:%.*]]: vector<4x3xf32>)
func.func @transpose_2D_sequence(%arg : vector<4x3xf32>) -> vector<4x3xf32> {
@@ -753,12 +762,13 @@ func.func @fold_extract_broadcast_0dvec_input_scalar_output(%a : vector<f32>,
// -----
+
// CHECK-LABEL: negative_fold_extract_broadcast
-// CHECK: vector.broadcast %{{.*}} : vector<1x1xf32> to vector<1x1x4xf32>
-// CHECK: vector.extract %{{.*}}[0, 0] : vector<4xf32> from vector<1x1x4xf32>
+// CHECK: vector.broadcast %{{.*}} : vector<1x1xf32> to vector<1x2x4xf32>
+// CHECK: vector.extract %{{.*}}[0, 0] : vector<4xf32> from vector<1x2x4xf32>
func.func @negative_fold_extract_broadcast(%a : vector<1x1xf32>) -> vector<4xf32> {
- %b = vector.broadcast %a : vector<1x1xf32> to vector<1x1x4xf32>
- %r = vector.extract %b[0, 0] : vector<4xf32> from vector<1x1x4xf32>
+ %b = vector.broadcast %a : vector<1x1xf32> to vector<1x2x4xf32>
+ %r = vector.extract %b[0, 0] : vector<4xf32> from vector<1x2x4xf32>
return %r : vector<4xf32>
}
@@ -797,8 +807,8 @@ func.func @fold_extract_broadcast_dim1_broadcasting(%a : vector<2x1xf32>,
// rank(extract_output) < rank(broadcast_input)
func.func @fold_extract_broadcast_to_lower_rank(%a : vector<2x4xf32>,
%idx0 : index, %idx1 : index) -> vector<4xf32> {
- %b = vector.broadcast %a : vector<2x4xf32> to vector<1x2x4xf32>
- %r = vector.extract %b[%idx0, %idx1] : vector<4xf32> from vector<1x2x4xf32>
+ %b = vector.broadcast %a : vector<2x4xf32> to vector<2x2x4xf32>
+ %r = vector.extract %b[%idx0, %idx1] : vector<4xf32> from vector<2x2x4xf32>
return %r : vector<4xf32>
}
@@ -1033,30 +1043,6 @@ func.func @canonicalize_broadcast_shapecast_to_broadcast_scalar(%arg0: f32) -> v
// -----
-// In this test, broadcast (2)->(1,2,1) is not legal, but shape_cast (2)->(1,2,1) is.
-// CHECK-LABEL: func @canonicalize_broadcast_shapecast_to_shapcast
-// CHECK-NOT: vector.broadcast
-// CHECK: vector.shape_cast {{.+}} : vector<2xf32> to vector<1x2x1xf32>
-func.func @canonicalize_broadcast_shapecast_to_shapcast(%arg0 : vector<2xf32>) -> vector<1x2x1xf32> {
- %0 = vector.broadcast %arg0 : vector<2xf32> to vector<1x2xf32>
- %1 = vector.shape_cast %0 : vector<1x2xf32> to vector<1x2x1xf32>
- return %1 : vector<1x2x1xf32>
-}
-
-// -----
-
-// In this test, broadcast (1)->(1,1) and shape_cast (1)->(1,1) are both legal. shape_cast is chosen.
-// CHECK-LABEL: func @canonicalize_broadcast_shapecast_both_possible
-// CHECK-NOT: vector.broadcast
-// CHECK: vector.shape_cast {{.+}} : vector<1xf32> to vector<1x1xf32>
-func.func @canonicalize_broadcast_shapecast_both_possible(%arg0: vector<1xf32>) -> vector<1x1xf32> {
- %0 = vector.broadcast %arg0 : vector<1xf32> to vector<1x1x1xf32>
- %1 = vector.shape_cast %0 : vector<1x1x1xf32> to vector<1x1xf32>
- return %1 : vector<1x1xf32>
-}
-
-// -----
-
// CHECK-LABEL: fold_vector_transfer_masks
func.func @fold_vector_transfer_masks(%A: memref<?x?xf32>) -> (vector<4x8xf32>, vector<4x[4]xf32>) {
// CHECK: %[[C0:.+]] = arith.constant 0 : index
@@ -1920,12 +1906,12 @@ func.func @extract_strided_splat(%arg0: f16) -> vector<2x4xf16> {
// -----
-// CHECK-LABEL: func @insert_extract_to_broadcast
+// CHECK-LABEL: func @insert_extract_to_shape_cast
// CHECK-SAME: (%[[ARG0:.*]]: vector<1x1x4xf32>, %[[ARG1:.*]]: vector<4xf32>)
-// CHECK: %[[V0:.*]] = vector.extract %[[ARG0]][0, 0] : vector<4xf32> from vector<1x1x4xf32>
-// CHECK: %[[V1:.*]] = vector.broadcast %[[ARG1]] : vector<4xf32> to vector<1x1x4xf32>
+// CHECK: %[[V0:.*]] = vector.shape_cast %[[ARG0]] : vector<1x1x4xf32> to vector<4xf32>
+// CHECK: %[[V1:.*]] = vector.shape_cast %[[ARG1]] : vector<4xf32> to vector<1x1x4xf32>
// CHECK: return %[[V0]], %[[V1]] : vector<4xf32>, vector<1x1x4xf32>
-func.func @insert_extract_to_broadcast(%arg0 : vector<1x1x4xf32>,
+func.func @insert_extract_to_shape_cast(%arg0 : vector<1x1x4xf32>,
%arg1 : vector<4xf32>) -> (vector<4xf32>, vector<1x1x4xf32>) {
%0 = vector.extract %arg0[0, 0] : vector<4xf32> from vector<1x1x4xf32>
%1 = vector.insert %arg1, %arg0 [0, 0] : vector<4xf32> into vector<1x1x4xf32>
@@ -2277,7 +2263,7 @@ func.func @shuffle_1d_rhs_poison() -> vector<4xi32> {
// CHECK-LABEL: func @shuffle_canonicalize_0d
func.func @shuffle_canonicalize_0d(%v0 : vector<i32>, %v1 : vector<i32>) -> vector<1xi32> {
- // CHECK: vector.broadcast %{{.*}} : vector<i32> to vector<1xi32>
+ // CHECK: vector.shape_cast %{{.*}} : vector<i32> to vector<1xi32>
%shuffle = vector.shuffle %v0, %v1 [0] : vector<i32>, vector<i32>
return %shuffle : vector<1xi32>
}
@@ -2764,9 +2750,8 @@ func.func @transfer_read_from_rank_reducing_extract_slice(%src: tensor<1x8x8x8xf
// CHECK-LABEL: func.func @extract_from_broadcast
func.func @extract_from_broadcast(%src: vector<1x1x1xf32>) -> vector<1xf32> {
%0 = vector.broadcast %src : vector<1x1x1xf32> to vector<1x1x32x1xf32>
-
- // CHECK-NEXT: %0 = vector.extract {{.*}}[0, 0] : vector<1xf32> from vector<1x1x1xf32>
- // CHECK-NEXT: return %0 : vector<1xf32>
+ // CHECK-NEXT: %[[RES:.*]] = vector.shape_cast{{.*}} vector<1x1x1xf32> to vector<1xf32>
+ // CHECK-NEXT: return %[[RES]] : vector<1xf32>
%1 = vector.extract %0[0, 0, 31] : vector<1xf32> from vector<1x1x32x1xf32>
return %1: vector<1xf32>
}
diff --git a/mlir/test/Dialect/Vector/canonicalize/vector-from-elements.mlir b/mlir/test/Dialect/Vector/canonicalize/vector-from-elements.mlir
index fdab2a8918a2e..d5f96a8928770 100644
--- a/mlir/test/Dialect/Vector/canonicalize/vector-from-elements.mlir
+++ b/mlir/test/Dialect/Vector/canonicalize/vector-from-elements.mlir
@@ -81,8 +81,8 @@ func.func @from_elements_to_splat(%a: f32, %b: f32) -> (vector<2x3xf32>, vector<
// CHECK-LABEL: func @to_shape_cast_rank2_to_rank1(
// CHECK-SAME: %[[A:.*]]: vector<1x2xi8>)
-// CHECK: %[[EXTRACT:.*]] = vector.extract %[[A]][0] : vector<2xi8> from vector<1x2xi8>
-// CHECK: return %[[EXTRACT]] : vector<2xi8>
+// CHECK: %[[SC:.*]] = vector.shape_cast %[[A]] : vector<1x2xi8> to vector<2xi8>
+// CHECK: return %[[SC]] : vector<2xi8>
func.func @to_shape_cast_rank2_to_rank1(%arg0: vector<1x2xi8>) -> vector<2xi8> {
%0 = vector.extract %arg0[0, 0] : i8 from vector<1x2xi8>
%1 = vector.extract %arg0[0, 1] : i8 from vector<1x2xi8>
diff --git...
[truncated]
|
|
Hi @banach-space and @dcaballe, I've pulled this PR out of draft mode, so please feel free to comment on it whenever! |
There was a problem hiding this comment.
Keep both tests, one with the original shape and one with the new ones?
Unrelated: it looks like we are missing a canonicalization patter here? This should be turned into a single vector.broadcast to vector<4xf32>?
There was a problem hiding this comment.
Keep both tests, one with the original shape and one with the new ones?
Makes sense, will do.
Unrelated: it looks like we are missing a canonicalization patter here? This should be turned into a single vector.broadcast to vector<4xf32>?
No because you can't broadcast <1x1xf32> to <4xf32> -- broadcasts can never reduce rank in Vector. FWIW slightly related to my comment here where this would be simpler if ops didn't do implicit shape casting. In this case if it was something like
%s = vector.shape_cast %a : vector<1x1xf32> to vector<1x1x1xf32>
%b = vector.broadcast %s : vector<1x1x1xf32> to vector<1x2x4xf32>
%r = vector.extract %b[0, 0] : vector<1x1x4xf32> from vector<1x2x4xf32>
%s = vector.shape_cast %r : vector<1x1x4> to vector<4>
ie if we constrained broadcasts and extracts to be rank retaining, then this would be canonicalized to
%s = vector.shape_cast %a : vector<1x1xf32> to vector<1x1x1xf32>
%b = vector.broadcast %s : vector<1x1x1xf32> to vector<1x1x4xf32>
%s = vector.shape_cast %b : vector<1x1x4> to vector<4>
which, if you have faith that the shape_casts will vanish at a later point, is simpler!
p.s. I plan to reply in #145740 later today
There was a problem hiding this comment.
Shouldn't we keep them? shouldn't they still be canonicalized?
|
Thanks! I run the SME e2e tests and all pass. I wasn't able to cherry-pick this in IREE though, getting weird compilation errors. Though upstream tests should be sufficient to surface all potential issues. @newling , why not name all "folding" patterns as |
I'll give this a spin with IREE
Yes, I think so. Actually fe3933d made me wonder if we should split canonicalize.mlir into 2 files (the new one with name fold.mlir containing everything in canonicalize.mlir that only depends on 1-time folds). @banach-space and @dcaballe thanks for your feedback! Unfortunately I'm going to put this on hold again temporarily, as I've uncovered some other things which should be done before this. Moving back into draft mode, will ping when I think it's ready again. |
+1 |
newling
force-pushed
the
canonicalize_to_shape_cast
branch
from
1ff3399 to
92e809e
Compare
|
This PR is back, and ready for review! Let me summarize the previous concerns as this is quite old now: @dcaballe raised concerns about removing tests. I have reinstated all canonicalization tests. |
There was a problem hiding this comment.
Sorry, I didn't follow the discussion in the previous PR, but to me something seems wrong, that we are moving towards shape_cast being the canonical form for removing unit dimensions.
shape_cast by itself, is a more general operation and requires inferring what the shape_cast actually did. We are throwing away information in each of these examples:
1. %0 = vector.broadcast %arg0 : vector<4xi8> to vector<1x1x4xi8>
2. %1 = vector.transpose %arg1, [1, 0] : vector<2x1xi8> to vector<1x2xi8>
3. %2 = vector.extract %arg2[0] : vector<4xi8> from vector<1x4xi8>
For each of these cases, we now need to analyze the shape_cast to see what it was doing:
- It is clear that broadcast always does a leading dim broadcast, and we already know its a unit dim, there is no possibility of collapsing
- transpose makes it clear that we are doing a permutation, there is no expansion/collapse of rank happening
- extract makes it clear we are rank reducing the vector, there is no possibility of expansion
We are throwing away information, this is in no way a canonicalization. In fact, this should be preprocessing pattern to prepare for conversion / unrolling / flattening so that we have less ops to handle and can cancel out more operations, if we actually want to run something like this.
This can clearly cause problems where we cannot fold things properly because of a shape_cast in between:
%0 = vector.extract %arg0[0] : vector<1x4> from vector<8x1x4>
%1 = vector.extract %0[0] : vector<4> from vector<1x4>
Now, depending on how the patterns run, we could end up with:
%1 = vector.extract %arg0[0, 0] : vector<4> from vector<8x1x4>
or
%0 = vector.extract %arg0[0] : vector<1x4> from vector<8x1x4>
%1 = vector.shape_cast %0 : vector<4> from vector<1x4>
To actually write the canonicalizer of shape_cast(extract()), we will end up actually having to infer if the shape_cast was a vector.extract, which defeats the purpose of this canonicalization.
Also note that each of the cases mentioned are disjoint. They do not overlap in terms of what they can do and the ops carry restrictions on what they can do.
|
@Groverkss that's fair, I probably should have done this as an RFC on discourse as there's been some debate and drift since the original #138777 Maybe a less controversial approach is to use functions like |
I would really argue that the opposite pattern should be a canonicalization, since we are caching the analyzed information, by using a more restricted op. |
Possibly, but I think "more restricted" depends on how you measure volumes in op-space (I discussed this here). But yeah, I think this canonicalization (in either direction) is not obvious enough to include until we have a formal definition (energy function for fixed point convergence). |
|
The point of having a canonical representation is to choose a unique representation for a computation that may have multiple alternative representations. That form doesn't have to be the "best" representation (whatever that means: it doesn’t have to be the most efficient, the most compact, the most aesthetically pleasing, the most restrictive…) It has to be unique and hold all the information that the original form held. We currently have 3+ ways to drop unit dimensions and perform vector reshapes that don’t imply data movement in the virtual vector world and we are paying the cost of maintaining all those redundant patterns. I think the big misalignment here is that we expect the canonical form to be the ideal input for all our use cases without doing any pre-processing and that is something we need to change for the shake of the overall project health. For me this is similar to other well-established canonical forms in LLVM. For example, look at
This sounds contradictory to me. If we can infer the information from the shape_cast itself, what is the information that is being thrown away? There shouldn't be a problem with analyzing the shape_cast op to get the information that we need, or convert that shape cast into something else before a specific pass/transformation. We are talking about comparing the rank of the input and the result types to determine if the reshape is collapsing/expanding/preserving dimensions. I fail to see the problem here.
Retrieving this information from a shape_cast is just about
For unit dimension transposes, there is no data permutation. Again, retrieving the expansion collapse/information is just about
Same.
This is just about having folder for |
Introducing new ops in the canonical form to cache |
What is this PR achieving?
To me, this change isn't really about removing unit dims. Its about improving the state of In the distant past (~4yrs ago) we did use Without a crisp contract,
Context + previous PRsNote, James has already proposed/implemented a number of steps to improve
This PR simply takes things further (and was discussed before: #138777). I just wanted to highlight - there has been a number of PRs through which we concluded that this would be beneficial. Why take this approach?
Indeed! With multiple ways of expressing identical computation, lets just agree on one. I don't see any information being lost, so to me this a matter of simply agreeing on what the canonical form should be and then working with that. From #138777:
Yes! Clarifying this takes us to a cleaner, better defined state - that's beneficial for As a specific argument for this change, replacing: %0 = vector.broadcast %arg0 : vector<4xi8> to vector<1x1x4xi8>with a Is this really a canonicalization?Yes. The original Ops are mechanically recoverable in every case, so I don't see any loss of information here.
I believe so, yes. Do we need an RFC?
We are discussing fine implementation details within LGTM?Yes |
|
I don't mind which way this PR goes. The original goal of the PR was to generate better code in a simpler way, not to define a canonical form. Converting to shape_cast achieves the original goal, but we don't have to say that shape_cast is a canonical form. As already mentioned in previous comments, technically it is a canonicalization if we want it to be, but until we formalize a canonical form (ideally with an energy function with guarantees of converging to a fixed point, and clear documentation and motivation in a single obvious place) maybe we're better off leaving things that don't directly reduce op count or provide some other obvious generic improvement, out. As someone who's tried to unravel complex decisions make 4+ years ago in this dialect (and often not followed through on) I'd rather not leave a whole new paper trail for future developers. What I've actually been considering for a while is making a 2-step pass that
I think that would be quite powerful. The nice thing about |
I'm sorry, but this response is a bit confusing to me. I'm very aware of the previous prs by James, because James and I are aligned on what we are working towards. I didn't chime in on those prs because I didn't want two people with the same ideas , from the same team and working on similar things to just review things between them and land things, given that we have multiple maintainers for Vector dialect. I do skim by those prs, I just didn't read until the title of this PR that we came to a conclusion that these patterns should be canonicalizations. In fact, i'm more confused that we talk about context here, given that the same change was reverted in the past and I distinctly remember we never came to conclusion what is the right form here: #72918 . I do not see that we resolved any of the concerns in that discussion, so I don't think concluding this is beneficial is correct, without previous context. There was even an RFC where @banach-space and @dcaballe both were present: https://discourse.llvm.org/t/improving-handling-of-unit-dimensions-in-the-vector-dialect/75216/29 .I don't see anybody followup on any of those discussions, so we cannot throw those conclusions away and proceed with something a lot of reviewers disagreed with.
Given that there was already an RFC about this, with no followup on things, yes we need an RFC on this. A lot of vector dialect users were present in that discussion whose concerns need to be addressed. I'm not okay with this proceeding as a PR given that there was an RFC about this where this PR was not the conclusion.
I don't know when we concluded that this "experiment" failed. I see that in that RFC @nicolasvasilache had the same point about vector.shape_cast https://discourse.llvm.org/t/improving-handling-of-unit-dimensions-in-the-vector-dialect/75216/10 . This is why we need to address the points in the RFC before we start a discussion on these patterns being canonicalizations. |
This is exactly the situation we are trying to improve! There have been numerous complaints along the years about the arbitrary set of disconnected patterns that we have in the Vector dialect and how difficult is for someone new (and I would say for someone with experience starting a new compiler) to put all of them together and in the right order to produce an optimized output. I’m ok with having different passes that are able to convert one form to another but we need a common starting point, which should be the canonical one.
That RFC is from 2023. Quite a few things have changed since then. and The RFC referenced above is also a good example of what can happen when we pull too many voices into a discussion. The discussion derailed and there were people talking over each other with counter-proposals and no conclusion was reached. I know not everyone agrees with this but the Tensor WG was built with the intent to improve on this situation and help build certain level of consensus before having a broader discussion. IMO, we should start there. Hopefully @rengolin or @rolfmorel can help organize. I think some people are on vacation right now so we may have to wait for a few days/weeks. |
|
Returning after some OOO + sick leave. I wanted to clarify my earlier feedback.
On context: @Groverkss ,I wasn’t suggesting you weren’t aware of previous PRs. My intent was to provide context for other reviewers (including, but not limited to you) so they could see that this work is part of a longer effort, not an isolated change. And where my support was coming from.
On the 2023 concerns: The PR you linked was ~2 years ago. It’s not clear to me whether the issues raised there still apply. If they do, could someone summarize them for us here? The earlier discussions:
involved many people and spanned months, but there’s been little follow-up since. That suggests (to me) that the concerns may no longer be relevant, especially given the changes in the ecosystem that Diego pointed out.
On process: I share the worry that RFCs can stall progress. We’ve seen this before (e.g. the 0-D vectors RFC that took 6 months and required an ODM). That said, I’m not against RFCs, just cautious that we don’t end up in the same situation again. The Tensor WG seems like a good venue to first align and then escalate if needed. Next step: If anyone believes the 2023 concerns still block this direction, could you restate them here? That would help us decide whether an RFC is warranted, or whether we can move forward with this canonicalization. In the meantime, I added this to the Tensor WG agenda. Thank you, |
While this was on the agenda for our call today, we didn't get a chance to discuss it (notes from today: https://discourse.llvm.org/t/tcdg-tensor-compiler-design-group-meeting-notes-2025-09-17) - we prioritised the unrolling vs flattening topic that was raised earlier. We'll get back to this in the next Tensor WG meeting! |
While we lacked The overall sentiment was that we should proceed with this unless specific and actionable concerns are raised. Regarding the earlier discussion on Discourse, much has changed since then, and no concrete follow-ups occurred. Given the ~2-year gap, it’s reasonable to assume that the previous concerns no longer hold. Previously, Given LLVM’s liberal revert policy, I propose that we land this and, if issues arise, revert. We’ve been converging toward this change for some time. -Andrzej |
|
|
||
| if (!isOrderPreserving(transpose)) { | ||
| return rewriter.notifyMatchFailure( | ||
| transpose, "not order preserving, so not semantically a 'copy'"); |
There was a problem hiding this comment.
Did you mean "semantically" without "not"?
| transpose, "not order preserving, so not semantically a 'copy'"); | |
| transpose, "not order preserving, so semantically a 'copy'"); |
| /// to fold shape_cast into create_mask. | ||
| /// | ||
| /// BEFORE: | ||
| /// to fold shape_cast into creatto a greater number of BEFORE: |
There was a problem hiding this comment.
Something has gone wrong here :)
Hi, I'll try to write up concerns by Monday next week and ask for an RFC on how those concerns are addressed. I have raised concerns about this, other people have raised concerns about this, and if you want to proceed please send a RFC. I am happy to summarise all problems with this change by Monday. Landing this without a RFC even when requested would be against the LLVM policy. |
There was a problem hiding this comment.
This change has had problems in the past with concerns raised from @MaheshRavishankar @kuhar @qedawkins and me. I have talked with them privately recently, and we still have the same concerns. I will summarise the previous concerns by Monday.
Can you please send a RFC summarising what has changed and how you plan to handle the previously raised concerns? Happy to send a summary of the concerns by Monday which you can use.
Please do not land this without sending a RFC and getting explicit approval on the RFC.
|
I did write a summary of previous concerns today, but I'm talking to others who were on previous threads to make sure if the concerns are still valid (some of them don't look valid to me). I'll post an update tommorow, due to timezone differences I couldn't sync with some people from US timezone today. |
|
The main issue with this canonicalization comes down to the fact that it drops The crux of the issue raised in the previous prs, reverts and discourse This is visible when describing each of these operations based on how their Note that transpose, broadcast and extract represent a projected permutation When coming from tensor land via vectorization, the vector code produced This PR tries to canonicalize different ops which can do different things but have The reasonings I'm seeing in this PR are:
I don't see canonicalizing towards a less restricted op to be a goal of Infact, that is only just the initial bar for something to be considered as a
It's not the bar for something to be a canonicalization. The goal that is relevant here, that I do see is:
There is no justification in this PR that canonicalizing towards a single operation Let's say I'm writing a transformation (a commonly written one), which folds It is easy to write, with just matching operations that have transpose like Not only does my transformation has to look for transpose-like ops, it has to This is making subsequent transformations harder to apply, this is actively going In practice, we have number of patterns in SPIRV backend in our downstream Some examples (from @kuhar): https://github.com/iree-org/iree/blob/main/compiler/src/iree/compiler/Codegen/SPIRV/SPIRVInitialVectorLowering.cpp#L148 I'm sure I could find more examples from upstream, but these are just from Keeping this in mind, I don't agree this transformation is a canonicalization If you think otherwise, please send a RFC and we can discuss further there. This The reason I ask for an RFC is that there are justifications given in this PR
A simple example of why this isn't a good arguement is we also have linalg.contraction, linalg.generic and linalg.matmul in the same dialect, but we don't choose to canonicalize them to each other. We have 3+ ways of writing a matmul in linalg. We should be consistent with this argument and apply it everywhere if so. What I would like to see in the RFC is:
|
"In general" having to match a single form for the same program is what makes subsequent analysis more effective: that is we can match a single "thing" instead of needing to match many. So this sentence as-is is a bit odd to me (but you're elaborating pretty well in the rest of your post, so don't take this as a rebuttal of your post as a whole but just a minor nitpick of this particular sentence). |
6 participants
Discussions suggest that we should use shape_cast as a canonical form of broadcast/transpose/extract where possible (see #138777)
For example these can all be expressed as shape casts:
This PR adds canonicalizes to convert the above 3 examples to shape_casts.
I've added some more comments as review comments.
I'm happy to split this PR up and add the new patterns separately.