[Mlir-commits] [mlir] 81a8d38 - [MLIR][Linalg] Add support for folding reduce_or and reduce_and of a broadcast (#227013)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Oct 15 02:28:05 PDT 2026
Author: Chibuoyim (Wilson) Ogbonna
Date: 2026-10-15T17:27:58+08:00
New Revision: 81a8d38739fb3e705ff046648bb15dee5095676e
URL: https://github.com/llvm/llvm-project/commit/81a8d38739fb3e705ff046648bb15dee5095676e
DIFF: https://github.com/llvm/llvm-project/commit/81a8d38739fb3e705ff046648bb15dee5095676e.diff
LOG: [MLIR][Linalg] Add support for folding reduce_or and reduce_and of a broadcast (#227013)
This PR adds support for folding the following patterns:
```
reduce_or(broadcast(x), init=0) -> x
reduce_and(broadcast(x), init=allOnes) -> x
```
These are not currently folded.
Repros:
```mlir
func.func @reduce_ori_broadcast(%x: tensor<8xi32>) -> tensor<8xi32> {
%e = tensor.empty() : tensor<8x4xi32>
%b = linalg.broadcast ins(%x : tensor<8xi32>) outs(%e : tensor<8x4xi32>) dimensions = [1]
%init = arith.constant dense<0> : tensor<8xi32>
%r = linalg.reduce ins(%b : tensor<8x4xi32>) outs(%init : tensor<8xi32>) dimensions = [1]
(%in: i32, %acc: i32) {
%o = arith.ori %in, %acc : i32
linalg.yield %o : i32
}
return %r : tensor<8xi32>
}
```
and
```mlir
func.func @reduce_andi_broadcast(%x: tensor<8xi32>) -> tensor<8xi32> {
%e = tensor.empty() : tensor<8x4xi32>
%b = linalg.broadcast ins(%x : tensor<8xi32>) outs(%e : tensor<8x4xi32>) dimensions = [1]
%init = arith.constant dense<-1> : tensor<8xi32> // all ones in binary
%r = linalg.reduce ins(%b : tensor<8x4xi32>) outs(%init : tensor<8xi32>) dimensions = [1]
(%in: i32, %acc: i32) {
%a = arith.andi %in, %acc : i32
linalg.yield %a : i32
}
return %r : tensor<8xi32>
}
```
Both of these can be folded to a single `return %x : tensor<8xi32>`
using the trivially provable patterns above.
Added lit tests.
Follow up to
[#213190](https://github.com/llvm/llvm-project/pull/213190).
Added:
Modified:
mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
mlir/test/Dialect/Linalg/canonicalize.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
index 9e63ccf47872eb..a7580196e10c0c 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 3fd6cb54b2a357..26f4d92fa0e04e 100644
--- a/mlir/test/Dialect/Linalg/canonicalize.mlir
+++ b/mlir/test/Dialect/Linalg/canonicalize.mlir
@@ -109,6 +109,96 @@ 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>
+}
+
+// -----
+
+// 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