[Mlir-commits] [mlir] [mlir][linalg] Add fill to broadcast canonicalization pattern (PR #195619)

[Hocky Yudhiono]

https://github.com/hockyy created https://github.com/llvm/llvm-project/pull/195619

Add fill to broadcast linalg canonicalization pattern.

>From fded6a75ca25c0082883d38102e29733cb99d8ed Mon Sep 17 00:00:00 2001
From: Hocky Yudhiono <hocky.yudhiono at gmail.com>
Date: Mon, 4 May 2026 17:13:27 +0800
Subject: [PATCH] [mlir][linalg] Add fill to broadcast canonicalization pattern

---
 mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp   | 28 ++++++++++++++++++----
 mlir/test/Dialect/Linalg/canonicalize.mlir | 27 +++++++++++++++++++++
 2 files changed, 50 insertions(+), 5 deletions(-)

diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
index 27988a451173c..7048b74232955 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -1012,6 +1012,22 @@ struct FoldFillWithTranspose : OpRewritePattern<linalg::TransposeOp> {
   }
 };
 
+/// Fold fill with broadcast.
+struct FoldFillWithBroadcast : OpRewritePattern<linalg::BroadcastOp> {
+  using OpRewritePattern<linalg::BroadcastOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(linalg::BroadcastOp broadcastOp,
+                                PatternRewriter &rewriter) const override {
+    if (auto fillOp = broadcastOp.getInput().getDefiningOp<FillOp>()) {
+      rewriter.replaceOpWithNewOp<FillOp>(
+          broadcastOp, broadcastOp.getResultTypes(), fillOp.getInputs(),
+          broadcastOp.getDpsInitOperand(0)->get());
+      return success();
+    }
+    return failure();
+  }
+};
+
 /// Fold a concat with all elements being fills of the same value
 /// into a fill of the concat result shape.
 struct FoldConcatsOfFill : public OpRewritePattern<tensor::ConcatOp> {
@@ -1067,11 +1083,13 @@ struct FoldConcatsOfFill : public OpRewritePattern<tensor::ConcatOp> {
 
 void FillOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                          MLIRContext *context) {
-  results.add<FoldConcatsOfFill, FoldFillWithCopy, FoldFillWithTensorExtract,
-              FoldFillWithPack, FoldFillWithPad,
-              FoldFillWithTensorReshape<tensor::CollapseShapeOp>,
-              FoldFillWithTensorReshape<tensor::ExpandShapeOp>,
-              FoldInsertPadIntoFill, FoldFillWithTranspose>(context);
+  results
+      .add<FoldConcatsOfFill, FoldFillWithCopy, FoldFillWithTensorExtract,
+           FoldFillWithPack, FoldFillWithPad,
+           FoldFillWithTensorReshape<tensor::CollapseShapeOp>,
+           FoldFillWithTensorReshape<tensor::ExpandShapeOp>,
+           FoldInsertPadIntoFill, FoldFillWithTranspose, FoldFillWithBroadcast>(
+          context);
 }
 
 //===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/Linalg/canonicalize.mlir b/mlir/test/Dialect/Linalg/canonicalize.mlir
index 019b7433b2777..ed9be08e24e9e 100644
--- a/mlir/test/Dialect/Linalg/canonicalize.mlir
+++ b/mlir/test/Dialect/Linalg/canonicalize.mlir
@@ -1164,6 +1164,33 @@ func.func @canonicalize_fill_to_transpose_input(%arg0 : tensor<?x?xf32>, %arg1 :
 
 // -----
 
+// CHECK-LABEL: func @canonicalize_fill_to_broadcast_dyn(
+//  CHECK-SAME:     %[[ARG0:[a-zA-Z0-9]+]]: tensor<?xf32>
+//  CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]+]]: tensor<?x?xf32>)
+//       CHECK:   %[[ZERO:.+]] = arith.constant 0.0
+//       CHECK:   linalg.fill ins(%[[ZERO]] : f32) outs(%[[ARG1]] : tensor<?x?xf32>)
+func.func @canonicalize_fill_to_broadcast_dyn(%arg0 : tensor<?xf32>, %arg1 : tensor<?x?xf32>) -> tensor<?x?xf32> {
+  %c0 = arith.constant 0.0 : f32
+  %fill = linalg.fill ins(%c0 : f32) outs(%arg0 : tensor<?xf32>) -> tensor<?xf32>
+  %broadcast = linalg.broadcast ins(%fill : tensor<?xf32>) outs(%arg1 : tensor<?x?xf32>) dimensions = [1]
+  return %broadcast : tensor<?x?xf32>
+}
+// -----
+
+// CHECK-LABEL: func @canonicalize_fill_to_broadcast_input(
+//  CHECK-SAME:     %[[ARG0:[a-zA-Z0-9]+]]: tensor<6x7xf32>
+//  CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]+]]: tensor<6x1x7xf32>)
+//       CHECK:   %[[ZERO:.+]] = arith.constant 0.0
+//       CHECK:   linalg.fill ins(%[[ZERO]] : f32) outs(%[[ARG1]] : tensor<6x1x7xf32>)
+func.func @canonicalize_fill_to_broadcast_input(%arg0 : tensor<6x7xf32>, %arg1 : tensor<6x1x7xf32>) -> tensor<6x1x7xf32> {
+  %c0 = arith.constant 0.0 : f32
+  %fill = linalg.fill ins(%c0 : f32) outs(%arg0 : tensor<6x7xf32>) -> tensor<6x7xf32>
+  %broadcast = linalg.broadcast ins(%fill : tensor<6x7xf32>) outs(%arg1 : tensor<6x1x7xf32>) dimensions = [1]
+  return %broadcast : tensor<6x1x7xf32>
+}
+
+// -----
+
 // CHECK-LABEL: func @broadcast_same_shape(
 //  CHECK-SAME:     %[[ARG0:[a-zA-Z0-9]+]]: tensor<2x3xf32>
 //  CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]+]]: tensor<2x3xf32>)