[Mlir-commits] [mlir] [mlir][vector] Add DialectReductionPatternInterface and broadcast reduction (PR #183181)

[Omer Farkash]

https://github.com/OmerFarkash updated https://github.com/llvm/llvm-project/pull/183181

>From 211560fe72509744ca1ece2e5398452eb598fea7 Mon Sep 17 00:00:00 2001
From: omer farkash <omerfarkash at gmail.com>
Date: Tue, 24 Feb 2026 15:09:25 +0200
Subject: [PATCH] [mlir][vector] Add reduction pattern for vector.broadcast

This commit introduces a DialectReductionPatternInterface for the Vector dialect and implements a reduction pattern for vector.broadcast.
The pattern replaces the broadcast operation with a ub.poison value of the same type. This breaks the dependency on the input scalar/vector, allowing the reducer to eliminate upstream operations while maintaining valid IR types.

Helps with generating smaller reproducers in mlir-reduce.
---
 mlir/lib/Dialect/Vector/IR/VectorOps.cpp | 28 ++++++++++++++++++++++++
 mlir/test/mlir-reduce/vector-reduce.mlir | 13 +++++++++++
 2 files changed, 41 insertions(+)
 create mode 100644 mlir/test/mlir-reduce/vector-reduce.mlir

diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index 613adeb5eeaaf..4376ddd18a3f4 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -54,6 +54,7 @@
 #include "mlir/Dialect/Vector/IR/VectorDialect.cpp.inc"
 // Pull in all enum type and utility function definitions.
 #include "mlir/Dialect/Vector/IR/VectorEnums.cpp.inc"
+#include "mlir/Reducer/ReductionPatternInterface.h"
 
 using namespace mlir;
 using namespace mlir::vector;
@@ -467,6 +468,31 @@ struct VectorInlinerInterface : public DialectInlinerInterface {
     return true;
   }
 };
+// Reduction pattern for vector.broadcast.
+struct ReduceBroadcastPattern : public OpRewritePattern<vector::BroadcastOp> {
+  using OpRewritePattern<vector::BroadcastOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(vector::BroadcastOp op,
+                                PatternRewriter &rewriter) const override {
+    // Replace the broadcast operation with a poison value.
+    // This breaks the dependency on the input and removes the operation,
+    // which helps the reducer generate smaller reproducers.
+    rewriter.replaceOpWithNewOp<ub::PoisonOp>(op, op.getType());
+    return success();
+  }
+};
+
+// Dialect reduction pattern interface for the Vector dialect.
+struct VectorReductionPatternInterface
+    : public DialectReductionPatternInterface {
+  explicit VectorReductionPatternInterface(Dialect *dialect)
+      : DialectReductionPatternInterface(dialect) {}
+
+  void populateReductionPatterns(RewritePatternSet &patterns) const override {
+    // Register the reduction patterns.
+    patterns.add<ReduceBroadcastPattern>(patterns.getContext());
+  }
+};
 } // namespace
 
 void VectorDialect::initialize() {
@@ -482,6 +508,8 @@ void VectorDialect::initialize() {
 
   addInterfaces<VectorInlinerInterface>();
 
+  addInterfaces<VectorReductionPatternInterface>();
+
   declarePromisedInterfaces<bufferization::BufferizableOpInterface,
                             TransferReadOp, TransferWriteOp, GatherOp, MaskOp,
                             YieldOp>();
diff --git a/mlir/test/mlir-reduce/vector-reduce.mlir b/mlir/test/mlir-reduce/vector-reduce.mlir
new file mode 100644
index 0000000000000..2513820cd435f
--- /dev/null
+++ b/mlir/test/mlir-reduce/vector-reduce.mlir
@@ -0,0 +1,13 @@
+// RUN: rm -f %t.sh
+// RUN: echo "#!/bin/bash" > %t.sh
+// RUN: echo "exit 1" >> %t.sh
+// RUN: chmod +x %t.sh
+// RUN: mlir-reduce %s --reduction-tree="test=%t.sh" | FileCheck %s
+
+// CHECK-LABEL: func.func @test_broadcast
+func.func @test_broadcast(%arg0 : f32)->vector<4xf32> {
+  // CHECK-NOT: vector.broadcast
+  // CHECK: %[[POISON:.*]] = ub.poison : vector<4xf32>
+  // CHECK: return %[[POISON]] : vector<4xf32>
+  %0 = vector.broadcast %arg0 : f32 to vector<4xf32> return %0 : vector<4xf32>
+}