[Mlir-commits] [mlir] [mlir][vector] Fold transpose(broadcast(shape_cast)) to broadcast (PR #215940)
Jianhui Li
llvmlistbot at llvm.org
Wed Aug 12 20:59:32 PDT 2026
https://github.com/Jianhui-Li created https://github.com/llvm/llvm-project/pull/215940
This PR folds shape_cast, broadcast, and transpose to one broadcast op.
Example:
```mlir
%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>
```
folds to
```mlir
%2 = vector.broadcast %x : vector<1x32x1xf32> to vector<1x32x64xf32>
```
>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] [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>
+}
More information about the Mlir-commits
mailing list