[mlir][vector] Canonicalize broadcast of shape_cast #150523

mshockwave · 2025-07-24T21:01:16Z

Fold broadcast(shape_cast(x)) into broadcast(x) if the type of x is compatible with broadcast's result type.

Fold `broadcast(shape_cast(x))` into `broadcast(x)` if the type of x is compatible with broadcast's result type.

llvmbot · 2025-07-24T21:01:47Z

@llvm/pr-subscribers-mlir-vector

@llvm/pr-subscribers-mlir

Author: Min-Yih Hsu (mshockwave)

Changes

Fold broadcast(shape_cast(x)) into broadcast(x) if the type of x is compatible with broadcast's result type.

Full diff: https://github.com/llvm/llvm-project/pull/150523.diff

2 Files Affected:

(modified) mlir/lib/Dialect/Vector/IR/VectorOps.cpp (+23-1)
(modified) mlir/test/Dialect/Vector/canonicalize.mlir (+22)

diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index 8c97aed6e7742..ad908319d8584 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -2938,13 +2938,35 @@ struct BroadcastFolder : public OpRewritePattern<BroadcastOp> {
     return success();
   }
 };
+
+// Fold broadcast(shape_cast(x)) into broadcast(x) if x's type is compatible
+// with broadcast's result type.
+struct FoldBroadcastOfShapeCast : public OpRewritePattern<BroadcastOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(BroadcastOp broadcastOp,
+                                PatternRewriter &rewriter) const override {
+    if (auto srcShapeCast =
+            broadcastOp.getSource().getDefiningOp<ShapeCastOp>()) {
+      VectorType srcType = srcShapeCast.getSourceVectorType();
+      VectorType destType = broadcastOp.getResultVectorType();
+      if (vector::isBroadcastableTo(srcType, destType) ==
+          BroadcastableToResult::Success) {
+        rewriter.replaceOpWithNewOp<BroadcastOp>(broadcastOp, destType,
+                                                 srcShapeCast.getSource());
+        return success();
+      }
+    }
+    return failure();
+  }
+};
 } // 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, FoldBroadcastOfShapeCast>(context);
 }
 
 //===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/Vector/canonicalize.mlir b/mlir/test/Dialect/Vector/canonicalize.mlir
index 1461c30162c5f..0fd2acd06c8ec 100644
--- a/mlir/test/Dialect/Vector/canonicalize.mlir
+++ b/mlir/test/Dialect/Vector/canonicalize.mlir
@@ -1168,6 +1168,28 @@ func.func @canonicalize_broadcast_shapecast_both_possible(%arg0: vector<1xf32>)
 
 // -----
 
+// CHECK-LABEL: func @canonicalize_shapecast_broadcast_to_broadcast
+//   CHECK-NOT:   vector.shape_cast
+//       CHECK:   vector.broadcast {{.+}} : vector<2xf32> to vector<32x2xf32>
+func.func @canonicalize_shapecast_broadcast_to_broadcast(%arg0 : vector<2xf32>) -> vector<32x2xf32> {
+  %0 = vector.shape_cast %arg0 : vector<2xf32> to vector<1x2xf32>
+  %1 = vector.broadcast %0 : vector<1x2xf32> to vector<32x2xf32>
+  return %1 : vector<32x2xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @canonicalize_shapecast_broadcast_invalid_shape
+//       CHECK:   vector.shape_cast {{.+}} : vector<64xf32> to vector<4x16xf32
+//       CHECK:   vector.broadcast {{.+}} : vector<4x16xf32> to vector<2x4x16xf32>
+func.func @canonicalize_shapecast_broadcast_invalid_shape(%arg0 : vector<64xf32>) -> vector<2x4x16xf32> {
+  %0 = vector.shape_cast %arg0 : vector<64xf32> to vector<4x16xf32>
+  %1 = vector.broadcast %0 : vector<4x16xf32> to vector<2x4x16xf32>
+  return %1 : vector<2x4x16xf32>
+}
+
+// -----
+
 // 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

banach-space

Thanks, makes sense % minor suggestions.

banach-space · 2025-07-25T15:20:00Z

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

+    if (auto srcShapeCast =
+            broadcastOp.getSource().getDefiningOp<ShapeCastOp>()) {
+      VectorType srcType = srcShapeCast.getSourceVectorType();
+      VectorType destType = broadcastOp.getResultVectorType();
+      if (vector::isBroadcastableTo(srcType, destType) ==
+          BroadcastableToResult::Success) {
+        rewriter.replaceOpWithNewOp<BroadcastOp>(broadcastOp, destType,
+                                                 srcShapeCast.getSource());
+        return success();
+      }
+    }
+    return failure();


[nit] Prefer early exits - helps reduce indentation.

Suggested change

if (auto srcShapeCast =

broadcastOp.getSource().getDefiningOp<ShapeCastOp>()) {

VectorType srcType = srcShapeCast.getSourceVectorType();

VectorType destType = broadcastOp.getResultVectorType();

if (vector::isBroadcastableTo(srcType, destType) ==

BroadcastableToResult::Success) {

rewriter.replaceOpWithNewOp<BroadcastOp>(broadcastOp, destType,

srcShapeCast.getSource());

return success();

}

}

return failure();

auto srcShapeCast =

broadcastOp.getSource().getDefiningOp<ShapeCastOp>();

if (!srcShapeCast)

return failure();

VectorType srcType = srcShapeCast.getSourceVectorType();

VectorType destType = broadcastOp.getResultVectorType();

if (vector::isBroadcastableTo(srcType, destType) !=

BroadcastableToResult::Success)

return failure();

rewriter.replaceOpWithNewOp<BroadcastOp>(broadcastOp, destType,

srcShapeCast.getSource());

return success();

mlir/test/Dialect/Vector/canonicalize.mlir

Co-authored-by: Andrzej Warzyński <[email protected]>

newling · 2025-07-25T16:55:09Z

Can you reuse isBroadcastLike

llvm-project/mlir/lib/Dialect/Vector/IR/VectorOps.cpp

Line 1712 in 2571924

static bool isBroadcastLike(Operation *op) {

?

General rule is that is something can be a folder, it should be (i.e. on BroadcastOp::fold) https://mlir.llvm.org/docs/Canonicalization/#when-to-use-the-fold-method-vs-rewriterpatterns-for-canonicalizations

newling

I don't think this is always valid?

(2,1) -> shape_cast -> (1,2) -> broadcast (2,2)

and

(2,1) -> broadcast (2,2)

are different.

Example. If input is [[5], [6]]. then first one's output is [[5, 6], [5, 6]] but second one's is [[5, 5], [6, 6]].

mshockwave · 2025-07-25T21:38:42Z

I don't think this is always valid?

(2,1) -> shape_cast -> (1,2) -> broadcast (2,2)

and

(2,1) -> broadcast (2,2)

are different.

Example. If input is [[5], [6]]. then first one's output is [[5, 6], [5, 6]] but second one's is [[5, 5], [6, 6]].

Yeah you're right. Let me turn this PR draft and think about this.

Groverkss · 2025-08-01T09:17:34Z

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

+// Fold broadcast(shape_cast(x)) into broadcast(x) if x's type is compatible
+// with broadcast's result type.
+struct FoldBroadcastOfShapeCast : public OpRewritePattern<BroadcastOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(BroadcastOp broadcastOp,
+                                PatternRewriter &rewriter) const override {
+    auto srcShapeCast = broadcastOp.getSource().getDefiningOp<ShapeCastOp>();
+    if (!srcShapeCast)
+      return failure();
+
+    VectorType srcType = srcShapeCast.getSourceVectorType();
+    VectorType destType = broadcastOp.getResultVectorType();
+    if (vector::isBroadcastableTo(srcType, destType) !=
+        BroadcastableToResult::Success)
+      return failure();
+
+    rewriter.replaceOpWithNewOp<BroadcastOp>(broadcastOp, destType,
+                                             srcShapeCast.getSource());
+    return success();
+  }
+};


This should be a folder, not a rewrite pattern.

I just rewrote it into a folder

mshockwave · 2025-08-06T00:21:59Z

I don't think this is always valid?

(2,1) -> shape_cast -> (1,2) -> broadcast (2,2)

and

(2,1) -> broadcast (2,2)

are different.

Example. If input is [[5], [6]]. then first one's output is [[5, 6], [5, 6]] but second one's is [[5, 5], [6, 6]].

I updated the algorithm to add a condition that the replicating dimensions have to be the same before and after the transformations.

newling

Thanks! I think this is correct now. But added a suggestion which might simplify it.

newling · 2025-08-06T16:30:37Z

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

+// Fold broadcast(shape_cast(x)) into broadcast(x) if x's type is compatible
+// with broadcast's result type and the broadcasted dimensions are the same.
+static LogicalResult foldBroadcastOfShapeCast(BroadcastOp broadcastOp) {
+  auto srcShapeCast = broadcastOp.getSource().getDefiningOp<ShapeCastOp>();


I think this is the same as saying (where srcShape -> shapecastShape -> destShape)

rank(srcShape) <= rank(destShape)

srcShape and shapeCastShape are the same, except that one has some 1's prepended. i.e. where R = min(srcShape.rank, shapeCastShape.rank), last R dimensions of srcShape and shapeCastCast are the same.

If so, would be more intuitive I think. If not, can you please provided a counterexample?

[mlir][vector] Canonicalize broadcast of shape_cast

9ca07a1

Fold `broadcast(shape_cast(x))` into `broadcast(x)` if the type of x is compatible with broadcast's result type.

mshockwave requested review from antiagainst and newling July 24, 2025 21:01

mshockwave requested review from banach-space, dcaballe, nicolasvasilache and Groverkss as code owners July 24, 2025 21:01

llvmbot added mlir:vectorops mlir mlir:vector labels Jul 24, 2025

banach-space approved these changes Jul 25, 2025

View reviewed changes

mshockwave and others added 2 commits July 25, 2025 09:06

fixup! Address review comments

10a914e

fixup! Update mlir/test/Dialect/Vector/canonicalize.mlir

067f115

Co-authored-by: Andrzej Warzyński <[email protected]>

newling reviewed Jul 25, 2025

View reviewed changes

fixup! fixup! Update mlir/test/Dialect/Vector/canonicalize.mlir

32c870b

mshockwave marked this pull request as draft July 25, 2025 21:38

Groverkss requested changes Aug 1, 2025

View reviewed changes

mshockwave added 3 commits August 5, 2025 15:18

Merge branch 'main' into patch/mlir/shapecast-broadcast

5517462

fixup! Fix invalid folding on mismatching broadcast dimensions

0cf5cc1

fixup! Rewrite as a folding pattern

236c545

mshockwave marked this pull request as ready for review August 6, 2025 00:19

mshockwave requested review from Groverkss and newling August 6, 2025 00:19

newling reviewed Aug 6, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[mlir][vector] Canonicalize broadcast of shape_cast #150523

[mlir][vector] Canonicalize broadcast of shape_cast #150523

Uh oh!

mshockwave commented Jul 24, 2025

Uh oh!

llvmbot commented Jul 24, 2025 •

edited

Loading

Uh oh!

banach-space left a comment

Uh oh!

banach-space Jul 25, 2025

Uh oh!

mshockwave Jul 25, 2025

Uh oh!

Uh oh!

newling commented Jul 25, 2025

Uh oh!

newling left a comment

Uh oh!

mshockwave commented Jul 25, 2025

Uh oh!

Groverkss Aug 1, 2025

Uh oh!

mshockwave Aug 6, 2025

Uh oh!

mshockwave commented Aug 6, 2025

Uh oh!

newling left a comment

Uh oh!

newling Aug 6, 2025

Uh oh!

Uh oh!

[mlir][vector] Canonicalize broadcast of shape_cast #150523

Are you sure you want to change the base?

[mlir][vector] Canonicalize broadcast of shape_cast #150523

Uh oh!

Conversation

mshockwave commented Jul 24, 2025

Uh oh!

llvmbot commented Jul 24, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

banach-space left a comment

Choose a reason for hiding this comment

Uh oh!

banach-space Jul 25, 2025

Choose a reason for hiding this comment

Uh oh!

mshockwave Jul 25, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

newling commented Jul 25, 2025

Uh oh!

newling left a comment

Choose a reason for hiding this comment

Uh oh!

mshockwave commented Jul 25, 2025

Uh oh!

Groverkss Aug 1, 2025

Choose a reason for hiding this comment

Uh oh!

mshockwave Aug 6, 2025

Choose a reason for hiding this comment

Uh oh!

mshockwave commented Aug 6, 2025

Uh oh!

newling left a comment

Choose a reason for hiding this comment

Uh oh!

newling Aug 6, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

llvmbot commented Jul 24, 2025 •

edited

Loading