[Mlir-commits] [mlir] [MLIR][Linalg] Add support for folding reduce_or and reduce_and of a broadcast (PR #227013)
Chibuoyim Ogbonna
llvmlistbot at llvm.org
Tue Sep 29 01:24:37 PDT 2026
https://github.com/bruteforceboy updated https://github.com/llvm/llvm-project/pull/227013
>From 91adb377708c6d82de0cb855e2a6b8d9eefc79f3 Mon Sep 17 00:00:00 2001
From: bruteforceboy <chibuoyim.faith.ogbonna at huawei.com>
Date: Mon, 28 Sep 2026 15:33:27 +0100
Subject: [PATCH 1/2] [MLIR][Linalg] Add support for folding reduce_or and
reduce_and of a broadcast
---
mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp | 14 ++++-
mlir/test/Dialect/Linalg/canonicalize.mlir | 67 ++++++++++++++++++++++
2 files changed, 79 insertions(+), 2 deletions(-)
diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
index 87b14439597b4..a40f99b9fdbfb 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -2044,9 +2044,11 @@ enum class BroadcastReduceKind {
MaxUI,
MinSI,
MinUI,
+ AndI,
+ OrI,
};
-/// Match a supported max/min reduction body and return its reduction kind.
+/// Match a supported reduction body and return its reduction kind.
static std::optional<BroadcastReduceKind>
matchBroadcastReduceBody(ReduceOp reduceOp) {
if (reduceOp.getNumDpsInputs() != 1 || reduceOp.getNumDpsInits() != 1 ||
@@ -2087,6 +2089,9 @@ matchBroadcastReduceBody(ReduceOp reduceOp) {
[](arith::MinSIOp) { return BroadcastReduceKind::MinSI; })
.Case<arith::MinUIOp>(
[](arith::MinUIOp) { return BroadcastReduceKind::MinUI; })
+ .Case<arith::AndIOp>(
+ [](arith::AndIOp) { return BroadcastReduceKind::AndI; })
+ .Case<arith::OrIOp>([](arith::OrIOp) { return BroadcastReduceKind::OrI; })
.Default([](Operation *) -> std::optional<BroadcastReduceKind> {
return std::nullopt;
});
@@ -2111,7 +2116,10 @@ static bool hasBroadcastReduceIdentity(Value init, BroadcastReduceKind kind) {
case BroadcastReduceKind::MinSI:
return value.isMaxSignedValue();
case BroadcastReduceKind::MinUI:
+ case BroadcastReduceKind::AndI:
return value.isAllOnes();
+ case BroadcastReduceKind::OrI:
+ return value.isZero();
}
llvm_unreachable("unknown broadcast reduction kind");
}
@@ -2120,8 +2128,10 @@ static bool hasBroadcastReduceIdentity(Value init, BroadcastReduceKind kind) {
///
/// maxsi(broadcast(x)) -> x
/// minsi(broadcast(y)) -> y
+/// andi(broadcast(z)) -> z
+/// ori(broadcast(w)) -> w
///
-// TODO: We can add other op. e.g., add, mul, and, or, xor.
+/// provided the reduction init is the identity value of the combiner.
struct FoldReduceBroadcast : public OpRewritePattern<linalg::ReduceOp> {
using OpRewritePattern<linalg::ReduceOp>::OpRewritePattern;
diff --git a/mlir/test/Dialect/Linalg/canonicalize.mlir b/mlir/test/Dialect/Linalg/canonicalize.mlir
index 3fd6cb54b2a35..e5e0be01aece1 100644
--- a/mlir/test/Dialect/Linalg/canonicalize.mlir
+++ b/mlir/test/Dialect/Linalg/canonicalize.mlir
@@ -109,6 +109,73 @@ func.func @negative_reduce_broadcast_empty_dimension(%input: tensor<3xi32>,
// -----
+// CHECK-LABEL: func.func @reduce_broadcast_andi(
+// CHECK-SAME: %[[INPUT:[a-zA-Z0-9]+]]: tensor<3xi32>
+// CHECK-NEXT: return %[[INPUT]] : tensor<3xi32>
+func.func @reduce_broadcast_andi(%input: tensor<3xi32>,
+ %broadcast_init: tensor<3x4xi32>)
+ -> tensor<3xi32> {
+ %reduce_init = arith.constant dense<-1> : tensor<3xi32>
+ %broadcasted = linalg.broadcast
+ ins(%input : tensor<3xi32>)
+ outs(%broadcast_init : tensor<3x4xi32>) dimensions = [1]
+ %result = linalg.reduce
+ ins(%broadcasted : tensor<3x4xi32>)
+ outs(%reduce_init : tensor<3xi32>) dimensions = [1]
+ (%in: i32, %out: i32) {
+ %and = arith.andi %in, %out : i32
+ linalg.yield %and : i32
+ }
+ return %result : tensor<3xi32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @reduce_broadcast_ori(
+// CHECK-SAME: %[[INPUT:[a-zA-Z0-9]+]]: tensor<3xi32>
+// CHECK-NEXT: return %[[INPUT]] : tensor<3xi32>
+func.func @reduce_broadcast_ori(%input: tensor<3xi32>,
+ %broadcast_init: tensor<2x3xi32>)
+ -> tensor<3xi32> {
+ %reduce_init = arith.constant dense<0> : tensor<3xi32>
+ %broadcasted = linalg.broadcast
+ ins(%input : tensor<3xi32>)
+ outs(%broadcast_init : tensor<2x3xi32>) dimensions = [0]
+ %result = linalg.reduce
+ ins(%broadcasted : tensor<2x3xi32>)
+ outs(%reduce_init : tensor<3xi32>) dimensions = [0]
+ (%in: i32, %out: i32) {
+ %or = arith.ori %out, %in : i32
+ linalg.yield %or : i32
+ }
+ return %result : tensor<3xi32>
+}
+
+// -----
+
+// Init 5 is not the ori identity, so no fold.
+// CHECK-LABEL: func.func @negative_reduce_broadcast_ori_non_identity_init
+// CHECK: linalg.broadcast
+// CHECK: linalg.reduce
+func.func @negative_reduce_broadcast_ori_non_identity_init(
+ %input: tensor<3xi32>, %broadcast_init: tensor<3x4xi32>)
+ -> tensor<3xi32> {
+ %reduce_init = arith.constant dense<5> : tensor<3xi32>
+ %broadcasted = linalg.broadcast
+ ins(%input : tensor<3xi32>)
+ outs(%broadcast_init : tensor<3x4xi32>) dimensions = [1]
+ %result = linalg.reduce
+ ins(%broadcasted : tensor<3x4xi32>)
+ outs(%reduce_init : tensor<3xi32>) dimensions = [1]
+ (%in: i32, %out: i32) {
+ %or = arith.ori %in, %out : i32
+ linalg.yield %or : i32
+ }
+ return %result : tensor<3xi32>
+}
+
+// -----
+
// CHECK-LABEL: func @memref_cast(
func.func @memref_cast(%a: index, %b: index) -> memref<?x?xf32> {
%c0 = arith.constant 0 : index
>From 30d2340b5c1b93a6a03bd68cd9002d31cad742a5 Mon Sep 17 00:00:00 2001
From: bruteforceboy <chibuoyim.faith.ogbonna at huawei.com>
Date: Tue, 29 Sep 2026 09:24:01 +0100
Subject: [PATCH 2/2] Apply review comment: add negative test for andi
---
mlir/test/Dialect/Linalg/canonicalize.mlir | 23 ++++++++++++++++++++++
1 file changed, 23 insertions(+)
diff --git a/mlir/test/Dialect/Linalg/canonicalize.mlir b/mlir/test/Dialect/Linalg/canonicalize.mlir
index e5e0be01aece1..26f4d92fa0e04 100644
--- a/mlir/test/Dialect/Linalg/canonicalize.mlir
+++ b/mlir/test/Dialect/Linalg/canonicalize.mlir
@@ -176,6 +176,29 @@ func.func @negative_reduce_broadcast_ori_non_identity_init(
// -----
+// Init 1 is not the andi identity, so no fold.
+// CHECK-LABEL: func.func @negative_reduce_broadcast_andi_non_identity_init
+// CHECK: linalg.broadcast
+// CHECK: linalg.reduce
+func.func @negative_reduce_broadcast_andi_non_identity_init(
+ %input: tensor<3xi32>, %broadcast_init: tensor<3x4xi32>)
+ -> tensor<3xi32> {
+ %reduce_init = arith.constant dense<1> : tensor<3xi32>
+ %broadcasted = linalg.broadcast
+ ins(%input : tensor<3xi32>)
+ outs(%broadcast_init : tensor<3x4xi32>) dimensions = [1]
+ %result = linalg.reduce
+ ins(%broadcasted : tensor<3x4xi32>)
+ outs(%reduce_init : tensor<3xi32>) dimensions = [1]
+ (%in: i32, %out: i32) {
+ %and = arith.andi %in, %out : i32
+ linalg.yield %and : i32
+ }
+ return %result : tensor<3xi32>
+}
+
+// -----
+
// CHECK-LABEL: func @memref_cast(
func.func @memref_cast(%a: index, %b: index) -> memref<?x?xf32> {
%c0 = arith.constant 0 : index
More information about the Mlir-commits
mailing list