[Mlir-commits] [mlir] [mlir][vector] Fold transpose(broadcast(shape_cast)) to broadcast (PR #215940)
Jianhui Li
llvmlistbot at llvm.org
Thu Aug 20 23:03:00 PDT 2026
https://github.com/Jianhui-Li updated https://github.com/llvm/llvm-project/pull/215940
>From a56d8bfbfcea7bf6ff61809d7f8ee42fae453584 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Thu, 13 Aug 2026 03:52:44 +0000
Subject: [PATCH 1/2] [mlir][vector] Fold transpose(broadcast(shape_cast)) to
broadcast
FoldTransposeBroadcast only folds transpose(broadcast(y)) when the
transpose permutes within y's own broadcast groups. When a unit-dim-only
shape_cast sits between the broadcast and its source x, the equivalent
broadcast is of x rather than of y, so that check fails even though the
chain is a plain broadcast of x.
Add FoldTransposeShapeCastBroadcast, which looks through the shape_cast
and folds the chain to a single broadcast of x when the non-unit dims are
preserved and land where a direct broadcast would put them. Scalable dims
(including a scalable [1]) are treated as non-unit.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply at anthropic.com>
---
mlir/lib/Dialect/Vector/IR/VectorOps.cpp | 90 ++++++++++++++++-
.../Vector/canonicalize/vector-transpose.mlir | 98 +++++++++++++++++++
2 files changed, 187 insertions(+), 1 deletion(-)
diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index 35e93ef81516d..82697e4bef5f8 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -7618,13 +7618,101 @@ class FoldTransposeBroadcast : public OpRewritePattern<vector::TransposeOp> {
}
};
+/// Folds transpose(broadcast(shape_cast(x))) to broadcast(x) when the chain is
+/// equivalent to a single broadcast of x to the transpose result type.
+///
+/// FoldTransposeBroadcast only folds transpose(broadcast(y)) when the transpose
+/// permutes within y's own broadcast groups. Here the equivalent broadcast is
+/// of x, not of the shape_cast result y, so that check fails even though the
+/// chain is a plain broadcast of x. Looking through the shape_cast recovers it.
+///
+/// Example 1, broadcast prepends a dim that the transpose moves to the back:
+/// ```
+/// %0 = vector.shape_cast %x : vector<1x32x1xf32> to vector<1x32xf32>
+/// %1 = vector.broadcast %0 : vector<1x32xf32> to vector<64x1x32xf32>
+/// %2 = vector.transpose %1, [1, 2, 0] : vector<64x1x32xf32>
+/// to vector<1x32x64xf32>
+/// ```
+/// rewrites to broadcast %x : vector<1x32x1xf32> to vector<1x32x64xf32>.
+///
+/// Example 2, broadcast stretches an existing size-1 dim:
+/// ```
+/// %0 = vector.shape_cast %x : vector<1x4xf32> to vector<4x1xf32>
+/// %1 = vector.broadcast %0 : vector<4x1xf32> to vector<4x3xf32>
+/// %2 = vector.transpose %1, [1, 0] : vector<4x3xf32> to vector<3x4xf32>
+/// ```
+/// rewrites to broadcast %x : vector<1x4xf32> to vector<3x4xf32>.
+///
+/// The fold is valid when two things hold. First, the only difference between x
+/// and the shape_cast result is in unit (size-1) dims; the rest of the dims are
+/// the same size in the same order. Second, the transpose puts each non-unit
+/// dim where a plain broadcast of x would put it. The rest are size-1 or
+/// broadcast dims, and a broadcast fills those the same way wherever they land.
+class FoldTransposeShapeCastBroadcast
+ : public OpRewritePattern<vector::TransposeOp> {
+public:
+ using Base::Base;
+
+ LogicalResult matchAndRewrite(vector::TransposeOp transpose,
+ PatternRewriter &rewriter) const override {
+ auto broadcast = transpose.getVector().getDefiningOp<vector::BroadcastOp>();
+ if (!broadcast)
+ return rewriter.notifyMatchFailure(transpose, "not a broadcast source");
+ auto shapeCast = broadcast.getSource().getDefiningOp<vector::ShapeCastOp>();
+ if (!shapeCast)
+ return rewriter.notifyMatchFailure(transpose, "not a shape_cast source");
+
+ VectorType srcType = shapeCast.getSourceVectorType();
+ VectorType midType = shapeCast.getResultVectorType();
+ VectorType bcastType = broadcast.getResultVectorType();
+ VectorType outType = transpose.getResultVectorType();
+
+ // Non-unit axis positions, in order. A scalable [1] is not a unit dim.
+ auto nonUnitAxes = [](VectorType ty) {
+ SmallVector<int64_t> axes;
+ for (auto [i, d] : llvm::enumerate(ty.getShape()))
+ if (d != 1 || ty.getScalableDims()[i])
+ axes.push_back(i);
+ return axes;
+ };
+ SmallVector<int64_t> srcAxes = nonUnitAxes(srcType);
+ SmallVector<int64_t> midAxes = nonUnitAxes(midType);
+
+ if (srcAxes.size() != midAxes.size() ||
+ vector::isBroadcastableTo(srcType, outType) !=
+ vector::BroadcastableToResult::Success)
+ return rewriter.notifyMatchFailure(transpose, "not a plain broadcast");
+
+ // Check that non-unit dims are preserved (same size and scalability) in
+ // order and land where a direct broadcast of x would.
+ SmallVector<int64_t> invPerm =
+ invertPermutationVector(transpose.getPermutation());
+ for (auto [srcAxis, midAxis] : llvm::zip_equal(srcAxes, midAxes)) {
+ if (srcType.getDimSize(srcAxis) != midType.getDimSize(midAxis) ||
+ srcType.getScalableDims()[srcAxis] !=
+ midType.getScalableDims()[midAxis])
+ return rewriter.notifyMatchFailure(transpose,
+ "reshapes a non-unit dim");
+ int64_t bcastAxis = midAxis + bcastType.getRank() - midType.getRank();
+ int64_t directAxis = srcAxis + outType.getRank() - srcType.getRank();
+ if (invPerm[bcastAxis] != directAxis)
+ return rewriter.notifyMatchFailure(transpose,
+ "reorders a broadcast axis");
+ }
+
+ rewriter.replaceOpWithNewOp<vector::BroadcastOp>(transpose, outType,
+ shapeCast.getSource());
+ return success();
+ }
+};
+
} // namespace
void vector::TransposeOp::getCanonicalizationPatterns(
RewritePatternSet &results, MLIRContext *context) {
results.add<FoldTransposeCreateMask, FoldTransposeShapeCast, TransposeFolder,
FoldTransposeSplat, FoldTransposeFromElements,
- FoldTransposeBroadcast>(context);
+ FoldTransposeBroadcast, FoldTransposeShapeCastBroadcast>(context);
}
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/Vector/canonicalize/vector-transpose.mlir b/mlir/test/Dialect/Vector/canonicalize/vector-transpose.mlir
index f1e1c5e896c66..237c66218ffd9 100644
--- a/mlir/test/Dialect/Vector/canonicalize/vector-transpose.mlir
+++ b/mlir/test/Dialect/Vector/canonicalize/vector-transpose.mlir
@@ -304,3 +304,101 @@ func.func @negative_transpose_fold(%arg : vector<2x2xi8>) -> vector<2x2xi8> {
%0 = vector.transpose %arg, [1, 0] : vector<2x2xi8> to vector<2x2xi8>
return %0 : vector<2x2xi8>
}
+
+// -----
+
+// +----------------------------------------------------------------------------
+// Tests of FoldTransposeShapeCastBroadcast:
+// transpose(broadcast(shape_cast)) -> broadcast
+// +----------------------------------------------------------------------------
+
+// The shape_cast drops a trailing unit dim, so the broadcast must prepend the
+// new dim and the transpose moves it back to the trailing position. Peeking
+// through the shape_cast recovers a single direct broadcast.
+// CHECK-LABEL: func @transpose_shape_cast_broadcast
+// CHECK-SAME: (%[[ARG:.+]]: vector<1x32x1xf32>)
+// CHECK: %[[V:.+]] = vector.broadcast %[[ARG]] : vector<1x32x1xf32> to vector<1x32x64xf32>
+// CHECK: return %[[V]] : vector<1x32x64xf32>
+func.func @transpose_shape_cast_broadcast(%arg: vector<1x32x1xf32>) -> vector<1x32x64xf32> {
+ %sc = vector.shape_cast %arg : vector<1x32x1xf32> to vector<1x32xf32>
+ %bc = vector.broadcast %sc : vector<1x32xf32> to vector<64x1x32xf32>
+ %t = vector.transpose %bc, [1, 2, 0] : vector<64x1x32xf32> to vector<1x32x64xf32>
+ return %t : vector<1x32x64xf32>
+}
+
+// -----
+
+// The broadcast stretches an existing size-1 dim rather than prepending one;
+// still equivalent to a single broadcast (no leading/trailing dim rule).
+// CHECK-LABEL: func @transpose_shape_cast_broadcast_stretch
+// CHECK-SAME: (%[[ARG:.+]]: vector<1x4xf32>)
+// CHECK: %[[V:.+]] = vector.broadcast %[[ARG]] : vector<1x4xf32> to vector<3x4xf32>
+// CHECK: return %[[V]] : vector<3x4xf32>
+func.func @transpose_shape_cast_broadcast_stretch(%arg: vector<1x4xf32>) -> vector<3x4xf32> {
+ %sc = vector.shape_cast %arg : vector<1x4xf32> to vector<4x1xf32>
+ %bc = vector.broadcast %sc : vector<4x1xf32> to vector<4x3xf32>
+ %t = vector.transpose %bc, [1, 0] : vector<4x3xf32> to vector<3x4xf32>
+ return %t : vector<3x4xf32>
+}
+
+// -----
+
+// The transpose reorders the two non-unit dims (2 and 4), so the chain is not a
+// plain broadcast and must not be folded.
+// CHECK-LABEL: func @negative_transpose_shape_cast_broadcast_reorder
+// CHECK: vector.shape_cast
+// CHECK: vector.broadcast
+// CHECK: %[[T:.+]] = vector.transpose
+// CHECK: return %[[T]]
+func.func @negative_transpose_shape_cast_broadcast_reorder(%arg: vector<2x1x4xf32>) -> vector<4x8x2xf32> {
+ %sc = vector.shape_cast %arg : vector<2x1x4xf32> to vector<2x4xf32>
+ %bc = vector.broadcast %sc : vector<2x4xf32> to vector<8x2x4xf32>
+ %t = vector.transpose %bc, [2, 0, 1] : vector<8x2x4xf32> to vector<4x8x2xf32>
+ return %t : vector<4x8x2xf32>
+}
+
+// -----
+
+// The shape_cast merges two non-unit dims (not a unit-dim-only reshape), so the
+// look-through does not apply and nothing is folded.
+// CHECK-LABEL: func @negative_transpose_shape_cast_broadcast_nonunit_reshape
+// CHECK: vector.shape_cast
+// CHECK: vector.broadcast
+// CHECK: %[[T:.+]] = vector.transpose
+// CHECK: return %[[T]]
+func.func @negative_transpose_shape_cast_broadcast_nonunit_reshape(%arg: vector<2x4xf32>) -> vector<8x3xf32> {
+ %sc = vector.shape_cast %arg : vector<2x4xf32> to vector<8xf32>
+ %bc = vector.broadcast %sc : vector<8xf32> to vector<3x8xf32>
+ %t = vector.transpose %bc, [1, 0] : vector<3x8xf32> to vector<8x3xf32>
+ return %t : vector<8x3xf32>
+}
+
+// -----
+
+// Scalable non-unit dims ([4]) are handled like any other non-unit dim.
+// CHECK-LABEL: func @transpose_shape_cast_broadcast_scalable
+// CHECK-SAME: (%[[ARG:.+]]: vector<[4]x1xf32>)
+// CHECK: %[[V:.+]] = vector.broadcast %[[ARG]] : vector<[4]x1xf32> to vector<1x[4]x8xf32>
+// CHECK: return %[[V]] : vector<1x[4]x8xf32>
+func.func @transpose_shape_cast_broadcast_scalable(%arg: vector<[4]x1xf32>) -> vector<1x[4]x8xf32> {
+ %sc = vector.shape_cast %arg : vector<[4]x1xf32> to vector<[4]xf32>
+ %bc = vector.broadcast %sc : vector<[4]xf32> to vector<8x1x[4]xf32>
+ %t = vector.transpose %bc, [1, 2, 0] : vector<8x1x[4]xf32> to vector<1x[4]x8xf32>
+ return %t : vector<1x[4]x8xf32>
+}
+
+// -----
+
+// A scalable [1] is not a fixed unit dim, so the shape_cast that folds it into
+// the [4] reshapes a scalable dim and must not be looked through.
+// CHECK-LABEL: func @negative_transpose_shape_cast_broadcast_scalable_unit
+// CHECK: vector.shape_cast
+// CHECK: vector.broadcast
+// CHECK: %[[T:.+]] = vector.transpose
+// CHECK: return %[[T]]
+func.func @negative_transpose_shape_cast_broadcast_scalable_unit(%arg: vector<[1]x4xf32>) -> vector<[4]x8xf32> {
+ %sc = vector.shape_cast %arg : vector<[1]x4xf32> to vector<[4]xf32>
+ %bc = vector.broadcast %sc : vector<[4]xf32> to vector<8x[4]xf32>
+ %t = vector.transpose %bc, [1, 0] : vector<8x[4]xf32> to vector<[4]x8xf32>
+ return %t : vector<[4]x8xf32>
+}
>From 6c22bf08f58152ea7073518925c7f62a8115735b Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Fri, 21 Aug 2026 06:00:24 +0000
Subject: [PATCH 2/2] [mlir][vector] Share core folding logic between
transpose(broadcast) patterns
Address review feedback on #215940: FoldTransposeBroadcast and
FoldTransposeShapeCastBroadcast duplicated the validity/rewrite logic.
Extract two small shared utilities:
- nonBroadcastAxes(type): positions of the data-carrying (size != 1, or
scalable) dims.
- transposeMapsAxes(from, to, permutation): whether the transpose moves the
axis at each input position from[i] to output position to[i].
Both patterns now compute their non-broadcast axis positions and delegate the
order-preservation check to transposeMapsAxes:
- FoldTransposeBroadcast checks each non-broadcast dim maps to itself
(from == to).
- FoldTransposeShapeCastBroadcast adjusts for the size-1-dim-only shape_cast,
mapping each non-broadcast axis from its position in y (the broadcast input)
to its direct-broadcast position in x.
The shared check is the exact order-preservation condition, which folds a few
extra valid cases the previous group-based algorithm conservatively rejected
(e.g. store_to_load_tensor_perm_broadcast now folds fully); all existing folds
are unchanged. Adds positive/negative tests, including a size-1-dim-only reshape
that reorders non-broadcast dims (5x4 -> 4x5x1).
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply at anthropic.com>
---
mlir/lib/Dialect/Vector/IR/VectorOps.cpp | 188 ++++++++----------
mlir/test/Dialect/Vector/canonicalize.mlir | 5 +-
.../Vector/canonicalize/vector-transpose.mlir | 62 +++---
3 files changed, 122 insertions(+), 133 deletions(-)
diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index c64a4130cbabc..f962810a4afea 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -7523,9 +7523,32 @@ class FoldTransposeFromElements final : public OpRewritePattern<TransposeOp> {
}
};
-/// 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.
+/// Positions (in order) of `type`'s non-broadcast dims: the data-carrying dims
+/// (size != 1, or a scalable [1]) a broadcast copies verbatim. Size-1 dims are
+/// replicated, so their position is irrelevant.
+static SmallVector<int64_t> nonBroadcastAxes(VectorType type) {
+ SmallVector<int64_t> axes;
+ for (auto [i, size] : llvm::enumerate(type.getShape()))
+ if (size != 1 || type.getScalableDims()[i])
+ axes.push_back(i);
+ return axes;
+}
+
+/// Returns true if `permutation` moves the axis at each input position `from[i]`
+/// to output position `to[i]`.
+static bool transposeMapsAxes(ArrayRef<int64_t> from, ArrayRef<int64_t> to,
+ ArrayRef<int64_t> permutation) {
+ SmallVector<int64_t> invPerm = invertPermutationVector(permutation);
+ for (auto [f, t] : llvm::zip_equal(from, to))
+ if (invPerm[f] != t)
+ return false;
+ return true;
+}
+
+/// Folds transpose(broadcast(x)) to broadcast(x) when the transpose is order
+/// preserving, i.e. it only reorders broadcast/size-1 dims and leaves every
+/// non-broadcast dim of x where a direct broadcast to the transpose result would
+/// place it.
///
/// Example:
/// ```
@@ -7537,17 +7560,6 @@ class FoldTransposeFromElements final : public OpRewritePattern<TransposeOp> {
/// ```
/// %0 = vector.broadcast %input : vector<1x1xi32> to vector<8x1xi32>.
/// ```
-/// The algorithm works by partitioning dimensions into groups that can be
-/// locally permuted while preserving order, and checks that the transpose
-/// only permutes within these groups.
-///
-/// Groups are either contiguous sequences of 1s, or non-1s (1-element groups).
-/// Consider broadcasting 4x1x1x7 to 2x3x4x5x6x7. This is equivalent to
-/// broadcasting from 1x1x4x1x1x7.
-/// ^^^ ^ ^^^ ^
-/// groups: 0 1 2 3
-/// Order preserving permutations for this example are ones that only permute
-/// within the groups [0,1] and [3,4], like (1 0 2 4 3 5 6).
class FoldTransposeBroadcast : public OpRewritePattern<vector::TransposeOp> {
public:
using Base::Base;
@@ -7556,75 +7568,46 @@ class FoldTransposeBroadcast : public OpRewritePattern<vector::TransposeOp> {
LogicalResult matchAndRewrite(vector::TransposeOp transpose,
PatternRewriter &rewriter) const override {
-
vector::BroadcastOp broadcast =
transpose.getVector().getDefiningOp<vector::BroadcastOp>();
- if (!broadcast) {
+ if (!broadcast)
return rewriter.notifyMatchFailure(transpose,
"not preceded by a broadcast");
- }
- auto inputType = dyn_cast<VectorType>(broadcast.getSourceType());
- VectorType outputType = transpose.getResultVectorType();
+ VectorType outType = transpose.getResultVectorType();
+ auto srcType = dyn_cast<VectorType>(broadcast.getSourceType());
- // transpose(broadcast(scalar)) -> broadcast(scalar) is always valid
- bool inputIsScalar = !inputType;
- if (inputIsScalar) {
- rewriter.replaceOpWithNewOp<vector::BroadcastOp>(transpose, outputType,
- broadcast.getSource());
- return success();
- }
+ // transpose(broadcast(scalar)) always folds. Otherwise the source must be
+ // broadcastable to the result and the transpose must leave its non-broadcast
+ // dims in place, i.e. map each to itself (from == to).
+ if (srcType) {
+ if (vector::isBroadcastableTo(srcType, outType) !=
+ vector::BroadcastableToResult::Success)
+ return rewriter.notifyMatchFailure(transpose, "not broadcastable");
- ArrayRef<int64_t> permutation = transpose.getPermutation();
- ArrayRef<int64_t> inputShape = inputType.getShape();
- int64_t inputRank = inputType.getRank();
- int64_t outputRank = transpose.getType().getRank();
- int64_t deltaRank = outputRank - inputRank;
-
- int low = 0;
- for (int inputIndex = 0; inputIndex < inputRank; ++inputIndex) {
- bool notOne = inputShape[inputIndex] != 1;
- bool prevNotOne = (inputIndex != 0 && inputShape[inputIndex - 1] != 1);
- bool groupEndFound = notOne || prevNotOne;
- if (groupEndFound) {
- int high = inputIndex + deltaRank;
- // Return failure if not all permutation destinations for indices in
- // [low, high) are in [low, high), i.e. the permutation is not local to
- // the group.
- for (int i = low; i < high; ++i) {
- if (permutation[i] < low || permutation[i] >= high) {
- return rewriter.notifyMatchFailure(
- transpose, "permutation not local to group");
- }
- }
- low = high;
- }
+ int64_t rankDelta = outType.getRank() - srcType.getRank();
+ SmallVector<int64_t> axes;
+ for (int64_t axis : nonBroadcastAxes(srcType))
+ axes.push_back(axis + rankDelta);
+ if (!transposeMapsAxes(axes, axes, transpose.getPermutation()))
+ return rewriter.notifyMatchFailure(transpose,
+ "not an order-preserving broadcast");
}
- // We don't need to check the final group [low, outputRank) because if it is
- // not locally bound, there must be a preceding group that already failed
- // the check (impossible to have just 1 non-locally bound group).
-
- // The preceding logic also ensures that at this point, the output of the
- // transpose is definitely broadcastable from the input shape, assert so:
- assert(vector::isBroadcastableTo(inputType, outputType) ==
- vector::BroadcastableToResult::Success &&
- "not broadcastable directly to transpose output");
-
- rewriter.replaceOpWithNewOp<vector::BroadcastOp>(transpose, outputType,
+ rewriter.replaceOpWithNewOp<vector::BroadcastOp>(transpose, outType,
broadcast.getSource());
-
return success();
}
};
-/// Folds transpose(broadcast(shape_cast(x))) to broadcast(x) when the chain is
-/// equivalent to a single broadcast of x to the transpose result type.
+/// Folds transpose(broadcast(shape_cast(x))) to broadcast(x) when a
+/// size-1-dim-only shape_cast `x -> y` sits between the broadcast and its
+/// source, i.e. y (the broadcast's input) is x with only size-1 dims rearranged.
///
-/// FoldTransposeBroadcast only folds transpose(broadcast(y)) when the transpose
-/// permutes within y's own broadcast groups. Here the equivalent broadcast is
-/// of x, not of the shape_cast result y, so that check fails even though the
-/// chain is a plain broadcast of x. Looking through the shape_cast recovers it.
+/// Such a shape_cast leaves x and y with the same non-broadcast dims (in order),
+/// so the chain is still a plain broadcast of x even when the transpose moves
+/// things relative to y. Reusing transposeMapsAxes, each non-broadcast axis is
+/// mapped from its position in y to its direct-broadcast position in x.
///
/// Example 1, broadcast prepends a dim that the transpose moves to the back:
/// ```
@@ -7642,12 +7625,6 @@ class FoldTransposeBroadcast : public OpRewritePattern<vector::TransposeOp> {
/// %2 = vector.transpose %1, [1, 0] : vector<4x3xf32> to vector<3x4xf32>
/// ```
/// rewrites to broadcast %x : vector<1x4xf32> to vector<3x4xf32>.
-///
-/// The fold is valid when two things hold. First, the only difference between x
-/// and the shape_cast result is in unit (size-1) dims; the rest of the dims are
-/// the same size in the same order. Second, the transpose puts each non-unit
-/// dim where a plain broadcast of x would put it. The rest are size-1 or
-/// broadcast dims, and a broadcast fills those the same way wherever they land.
class FoldTransposeShapeCastBroadcast
: public OpRewritePattern<vector::TransposeOp> {
public:
@@ -7662,46 +7639,45 @@ class FoldTransposeShapeCastBroadcast
if (!shapeCast)
return rewriter.notifyMatchFailure(transpose, "not a shape_cast source");
- VectorType srcType = shapeCast.getSourceVectorType();
- VectorType midType = shapeCast.getResultVectorType();
- VectorType bcastType = broadcast.getResultVectorType();
+ VectorType srcType = shapeCast.getSourceVectorType(); // x
+ VectorType midType = shapeCast.getResultVectorType(); // y, the broadcast input
VectorType outType = transpose.getResultVectorType();
- // Non-unit axis positions, in order. A scalable [1] is not a unit dim.
- auto nonUnitAxes = [](VectorType ty) {
- SmallVector<int64_t> axes;
- for (auto [i, d] : llvm::enumerate(ty.getShape()))
- if (d != 1 || ty.getScalableDims()[i])
- axes.push_back(i);
- return axes;
- };
- SmallVector<int64_t> srcAxes = nonUnitAxes(srcType);
- SmallVector<int64_t> midAxes = nonUnitAxes(midType);
-
- if (srcAxes.size() != midAxes.size() ||
- vector::isBroadcastableTo(srcType, outType) !=
- vector::BroadcastableToResult::Success)
- return rewriter.notifyMatchFailure(transpose, "not a plain broadcast");
-
- // Check that non-unit dims are preserved (same size and scalability) in
- // order and land where a direct broadcast of x would.
- SmallVector<int64_t> invPerm =
- invertPermutationVector(transpose.getPermutation());
- for (auto [srcAxis, midAxis] : llvm::zip_equal(srcAxes, midAxes)) {
+ if (vector::isBroadcastableTo(srcType, outType) !=
+ vector::BroadcastableToResult::Success)
+ return rewriter.notifyMatchFailure(transpose, "not broadcastable");
+
+ // Size-1-dim-only shape_cast: x and y must have identical non-broadcast dims
+ // in the same order. E.g. rejects 5x4 -> 4x5x1, which reorders 5 and 4.
+ SmallVector<int64_t> srcAxes = nonBroadcastAxes(srcType);
+ SmallVector<int64_t> midAxes = nonBroadcastAxes(midType);
+ if (srcAxes.size() != midAxes.size())
+ return rewriter.notifyMatchFailure(transpose,
+ "reshapes a non-broadcast dim");
+ for (auto [srcAxis, midAxis] : llvm::zip_equal(srcAxes, midAxes))
if (srcType.getDimSize(srcAxis) != midType.getDimSize(midAxis) ||
srcType.getScalableDims()[srcAxis] !=
midType.getScalableDims()[midAxis])
return rewriter.notifyMatchFailure(transpose,
- "reshapes a non-unit dim");
- int64_t bcastAxis = midAxis + bcastType.getRank() - midType.getRank();
- int64_t directAxis = srcAxis + outType.getRank() - srcType.getRank();
- if (invPerm[bcastAxis] != directAxis)
- return rewriter.notifyMatchFailure(transpose,
- "reorders a broadcast axis");
+ "reshapes a non-broadcast dim");
+
+ // For each non-broadcast axis, find its index in the broadcast result
+ // (`from`) and in a direct broadcast of x (`to`), then check the transpose
+ // maps `from` -> `to`. E.g. dim 32 in Example 1:
+ // bcast = 64x1x32 -> 32 at index 2 (from)
+ // bcast(x) = 1x32x64 -> 32 at index 1 (to)
+ int64_t bcastRank = broadcast.getResultVectorType().getRank();
+ SmallVector<int64_t> from, to;
+ for (auto [srcAxis, midAxis] : llvm::zip_equal(srcAxes, midAxes)) {
+ from.push_back(midAxis + bcastRank - midType.getRank());
+ to.push_back(srcAxis + outType.getRank() - srcType.getRank());
}
+ if (!transposeMapsAxes(from, to, transpose.getPermutation()))
+ return rewriter.notifyMatchFailure(transpose,
+ "not a plain broadcast of the source");
- rewriter.replaceOpWithNewOp<vector::BroadcastOp>(transpose, outType,
- shapeCast.getSource());
+ rewriter.replaceOpWithNewOp<vector::BroadcastOp>(
+ transpose, outType, shapeCast.getSource());
return success();
}
};
diff --git a/mlir/test/Dialect/Vector/canonicalize.mlir b/mlir/test/Dialect/Vector/canonicalize.mlir
index 20ce397f10b22..bbdb3749ca3ce 100644
--- a/mlir/test/Dialect/Vector/canonicalize.mlir
+++ b/mlir/test/Dialect/Vector/canonicalize.mlir
@@ -2100,9 +2100,8 @@ func.func @store_to_load_tensor_broadcast_scalable(%arg0 : tensor<?xf32>,
// CHECK-LABEL: func @store_to_load_tensor_perm_broadcast
// CHECK-SAME: (%[[ARG:.*]]: tensor<4x4x4xf32>, %[[V0:.*]]: vector<4x1xf32>)
-// CHECK: %[[B:.*]] = vector.broadcast %[[V0]] : vector<4x1xf32> to vector<100x5x4x1xf32>
-// CHECK: %[[T:.*]] = vector.transpose %[[B]], [3, 0, 2, 1] : vector<100x5x4x1xf32> to vector<1x100x4x5xf32>
-// CHECK: return %[[T]] : vector<1x100x4x5xf32>
+// CHECK: %[[B:.*]] = vector.broadcast %[[V0]] : vector<4x1xf32> to vector<1x100x4x5xf32>
+// CHECK: return %[[B]] : vector<1x100x4x5xf32>
func.func @store_to_load_tensor_perm_broadcast(%arg0 : tensor<4x4x4xf32>,
%v0 : vector<4x1xf32>) -> vector<1x100x4x5xf32> {
%c0 = arith.constant 0 : index
diff --git a/mlir/test/Dialect/Vector/canonicalize/vector-transpose.mlir b/mlir/test/Dialect/Vector/canonicalize/vector-transpose.mlir
index 237c66218ffd9..12823fe6fa21a 100644
--- a/mlir/test/Dialect/Vector/canonicalize/vector-transpose.mlir
+++ b/mlir/test/Dialect/Vector/canonicalize/vector-transpose.mlir
@@ -91,6 +91,18 @@ func.func @broadcast_transpose_final_group(%arg0 : vector<4x7x1x1xi8>) -> vector
// -----
+// CHECK-LABEL: broadcast_transpose_reorders_broadcast_dims
+// CHECK-SAME: %[[ARG:.*]]: vector<4x1xf32>) -> vector<1x100x4x5xf32> {
+// CHECK: %[[RES:.*]] = vector.broadcast %[[ARG]] : vector<4x1xf32> to vector<1x100x4x5xf32>
+// CHECK: return %[[RES]] : vector<1x100x4x5xf32>
+func.func @broadcast_transpose_reorders_broadcast_dims(%arg0 : vector<4x1xf32>) -> vector<1x100x4x5xf32> {
+ %0 = vector.broadcast %arg0 : vector<4x1xf32> to vector<100x5x4x1xf32>
+ %1 = vector.transpose %0, [3, 0, 2, 1] : vector<100x5x4x1xf32> to vector<1x100x4x5xf32>
+ return %1 : vector<1x100x4x5xf32>
+}
+
+// -----
+
// CHECK-LABEL: negative_broadcast_transpose_square
// CHECK-SAME: %[[ARG:.*]]:
// CHECK: %[[BCT:.*]] = vector.broadcast %[[ARG]]
@@ -312,14 +324,11 @@ func.func @negative_transpose_fold(%arg : vector<2x2xi8>) -> vector<2x2xi8> {
// transpose(broadcast(shape_cast)) -> broadcast
// +----------------------------------------------------------------------------
-// The shape_cast drops a trailing unit dim, so the broadcast must prepend the
-// new dim and the transpose moves it back to the trailing position. Peeking
-// through the shape_cast recovers a single direct broadcast.
-// CHECK-LABEL: func @transpose_shape_cast_broadcast
+// CHECK-LABEL: func @transpose_shape_cast_broadcast_drop_unit_dim
// CHECK-SAME: (%[[ARG:.+]]: vector<1x32x1xf32>)
// CHECK: %[[V:.+]] = vector.broadcast %[[ARG]] : vector<1x32x1xf32> to vector<1x32x64xf32>
// CHECK: return %[[V]] : vector<1x32x64xf32>
-func.func @transpose_shape_cast_broadcast(%arg: vector<1x32x1xf32>) -> vector<1x32x64xf32> {
+func.func @transpose_shape_cast_broadcast_drop_unit_dim(%arg: vector<1x32x1xf32>) -> vector<1x32x64xf32> {
%sc = vector.shape_cast %arg : vector<1x32x1xf32> to vector<1x32xf32>
%bc = vector.broadcast %sc : vector<1x32xf32> to vector<64x1x32xf32>
%t = vector.transpose %bc, [1, 2, 0] : vector<64x1x32xf32> to vector<1x32x64xf32>
@@ -328,13 +337,11 @@ func.func @transpose_shape_cast_broadcast(%arg: vector<1x32x1xf32>) -> vector<1x
// -----
-// The broadcast stretches an existing size-1 dim rather than prepending one;
-// still equivalent to a single broadcast (no leading/trailing dim rule).
-// CHECK-LABEL: func @transpose_shape_cast_broadcast_stretch
+// CHECK-LABEL: func @transpose_shape_cast_broadcast_stretch_unit_dim
// CHECK-SAME: (%[[ARG:.+]]: vector<1x4xf32>)
// CHECK: %[[V:.+]] = vector.broadcast %[[ARG]] : vector<1x4xf32> to vector<3x4xf32>
// CHECK: return %[[V]] : vector<3x4xf32>
-func.func @transpose_shape_cast_broadcast_stretch(%arg: vector<1x4xf32>) -> vector<3x4xf32> {
+func.func @transpose_shape_cast_broadcast_stretch_unit_dim(%arg: vector<1x4xf32>) -> vector<3x4xf32> {
%sc = vector.shape_cast %arg : vector<1x4xf32> to vector<4x1xf32>
%bc = vector.broadcast %sc : vector<4x1xf32> to vector<4x3xf32>
%t = vector.transpose %bc, [1, 0] : vector<4x3xf32> to vector<3x4xf32>
@@ -343,14 +350,12 @@ func.func @transpose_shape_cast_broadcast_stretch(%arg: vector<1x4xf32>) -> vect
// -----
-// The transpose reorders the two non-unit dims (2 and 4), so the chain is not a
-// plain broadcast and must not be folded.
-// CHECK-LABEL: func @negative_transpose_shape_cast_broadcast_reorder
+// CHECK-LABEL: func @negative_transpose_reorders_nonbroadcast_dims
// CHECK: vector.shape_cast
// CHECK: vector.broadcast
// CHECK: %[[T:.+]] = vector.transpose
// CHECK: return %[[T]]
-func.func @negative_transpose_shape_cast_broadcast_reorder(%arg: vector<2x1x4xf32>) -> vector<4x8x2xf32> {
+func.func @negative_transpose_reorders_nonbroadcast_dims(%arg: vector<2x1x4xf32>) -> vector<4x8x2xf32> {
%sc = vector.shape_cast %arg : vector<2x1x4xf32> to vector<2x4xf32>
%bc = vector.broadcast %sc : vector<2x4xf32> to vector<8x2x4xf32>
%t = vector.transpose %bc, [2, 0, 1] : vector<8x2x4xf32> to vector<4x8x2xf32>
@@ -359,14 +364,12 @@ func.func @negative_transpose_shape_cast_broadcast_reorder(%arg: vector<2x1x4xf3
// -----
-// The shape_cast merges two non-unit dims (not a unit-dim-only reshape), so the
-// look-through does not apply and nothing is folded.
-// CHECK-LABEL: func @negative_transpose_shape_cast_broadcast_nonunit_reshape
+// CHECK-LABEL: func @negative_shape_cast_merges_nonbroadcast_dims
// CHECK: vector.shape_cast
// CHECK: vector.broadcast
// CHECK: %[[T:.+]] = vector.transpose
// CHECK: return %[[T]]
-func.func @negative_transpose_shape_cast_broadcast_nonunit_reshape(%arg: vector<2x4xf32>) -> vector<8x3xf32> {
+func.func @negative_shape_cast_merges_nonbroadcast_dims(%arg: vector<2x4xf32>) -> vector<8x3xf32> {
%sc = vector.shape_cast %arg : vector<2x4xf32> to vector<8xf32>
%bc = vector.broadcast %sc : vector<8xf32> to vector<3x8xf32>
%t = vector.transpose %bc, [1, 0] : vector<3x8xf32> to vector<8x3xf32>
@@ -375,12 +378,25 @@ func.func @negative_transpose_shape_cast_broadcast_nonunit_reshape(%arg: vector<
// -----
-// Scalable non-unit dims ([4]) are handled like any other non-unit dim.
-// CHECK-LABEL: func @transpose_shape_cast_broadcast_scalable
+// CHECK-LABEL: func @negative_shape_cast_reorders_nonbroadcast_dims
+// CHECK: vector.shape_cast
+// CHECK: vector.broadcast
+// CHECK: %[[T:.+]] = vector.transpose
+// CHECK: return %[[T]]
+func.func @negative_shape_cast_reorders_nonbroadcast_dims(%arg: vector<5x4xf32>) -> vector<5x4x3xf32> {
+ %sc = vector.shape_cast %arg : vector<5x4xf32> to vector<4x5x1xf32>
+ %bc = vector.broadcast %sc : vector<4x5x1xf32> to vector<4x5x3xf32>
+ %t = vector.transpose %bc, [1, 0, 2] : vector<4x5x3xf32> to vector<5x4x3xf32>
+ return %t : vector<5x4x3xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @transpose_shape_cast_broadcast_scalable_dim
// CHECK-SAME: (%[[ARG:.+]]: vector<[4]x1xf32>)
// CHECK: %[[V:.+]] = vector.broadcast %[[ARG]] : vector<[4]x1xf32> to vector<1x[4]x8xf32>
// CHECK: return %[[V]] : vector<1x[4]x8xf32>
-func.func @transpose_shape_cast_broadcast_scalable(%arg: vector<[4]x1xf32>) -> vector<1x[4]x8xf32> {
+func.func @transpose_shape_cast_broadcast_scalable_dim(%arg: vector<[4]x1xf32>) -> vector<1x[4]x8xf32> {
%sc = vector.shape_cast %arg : vector<[4]x1xf32> to vector<[4]xf32>
%bc = vector.broadcast %sc : vector<[4]xf32> to vector<8x1x[4]xf32>
%t = vector.transpose %bc, [1, 2, 0] : vector<8x1x[4]xf32> to vector<1x[4]x8xf32>
@@ -389,14 +405,12 @@ func.func @transpose_shape_cast_broadcast_scalable(%arg: vector<[4]x1xf32>) -> v
// -----
-// A scalable [1] is not a fixed unit dim, so the shape_cast that folds it into
-// the [4] reshapes a scalable dim and must not be looked through.
-// CHECK-LABEL: func @negative_transpose_shape_cast_broadcast_scalable_unit
+// CHECK-LABEL: func @negative_shape_cast_reshapes_scalable_unit_dim
// CHECK: vector.shape_cast
// CHECK: vector.broadcast
// CHECK: %[[T:.+]] = vector.transpose
// CHECK: return %[[T]]
-func.func @negative_transpose_shape_cast_broadcast_scalable_unit(%arg: vector<[1]x4xf32>) -> vector<[4]x8xf32> {
+func.func @negative_shape_cast_reshapes_scalable_unit_dim(%arg: vector<[1]x4xf32>) -> vector<[4]x8xf32> {
%sc = vector.shape_cast %arg : vector<[1]x4xf32> to vector<[4]xf32>
%bc = vector.broadcast %sc : vector<[4]xf32> to vector<8x[4]xf32>
%t = vector.transpose %bc, [1, 0] : vector<8x[4]xf32> to vector<[4]x8xf32>
More information about the Mlir-commits
mailing list