[Mlir-commits] [mlir] aaa2c08 - [mlir][linalg] Extend FoldReduceBroadcast to float maximumf/minimumf (#233417)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sun Oct 25 05:10:17 PDT 2026
Author: Wijkqwe
Date: 2026-10-25T21:10:12+09:00
New Revision: aaa2c083d57b597673e95eccba536f74c588ae86
URL: https://github.com/llvm/llvm-project/commit/aaa2c083d57b597673e95eccba536f74c588ae86
DIFF: https://github.com/llvm/llvm-project/commit/aaa2c083d57b597673e95eccba536f74c588ae86.diff
LOG: [mlir][linalg] Extend FoldReduceBroadcast to float maximumf/minimumf (#233417)
Extend `FoldReduceBroadcast` to fold `reduce(broadcast(x))` into `x`
for `arith.maximumf`/`arith.minimumf` when the init is the identity
(`-inf`/`+inf`). The identity check gains a `FloatAttr` branch.
Safety: `maximumf(-inf, x) == x` for all values — finite, ±inf,
-0.0, and NaN (NaN-propagating semantics). Only an exact ±inf init
passes the check; finite/NaN inits are rejected by negative tests.
`maxnumf`/`minnumf` are deliberately excluded: they ignore NaN, so
`reduce(broadcast(NaN))` with init `-inf` would yield `-inf`, not NaN
(covered by a negative test).
Assisted-by: AI.
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 a7580196e10c0c..3c9797cc6acb5c 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -2046,6 +2046,8 @@ enum class BroadcastReduceKind {
MinUI,
AndI,
OrI,
+ MaximumF,
+ MinimumF,
};
/// Match a supported reduction body and return its reduction kind.
@@ -2092,6 +2094,10 @@ matchBroadcastReduceBody(ReduceOp reduceOp) {
.Case<arith::AndIOp>(
[](arith::AndIOp) { return BroadcastReduceKind::AndI; })
.Case<arith::OrIOp>([](arith::OrIOp) { return BroadcastReduceKind::OrI; })
+ .Case<arith::MaximumFOp>(
+ [](arith::MaximumFOp) { return BroadcastReduceKind::MaximumF; })
+ .Case<arith::MinimumFOp>(
+ [](arith::MinimumFOp) { return BroadcastReduceKind::MinimumF; })
.Default([](Operation *) -> std::optional<BroadcastReduceKind> {
return std::nullopt;
});
@@ -2103,6 +2109,17 @@ static bool hasBroadcastReduceIdentity(Value init, BroadcastReduceKind kind) {
if (!initAttr)
return false;
+ // Float kinds: identity is the corresponding infinity.
+ if (kind == BroadcastReduceKind::MaximumF ||
+ kind == BroadcastReduceKind::MinimumF) {
+ auto floatAttr = dyn_cast<FloatAttr>(*initAttr);
+ if (!floatAttr)
+ return false;
+ const APFloat &value = floatAttr.getValue();
+ return kind == BroadcastReduceKind::MaximumF ? value.isNegInfinity()
+ : value.isPosInfinity();
+ }
+
auto integerAttr = dyn_cast<IntegerAttr>(*initAttr);
if (!integerAttr)
return false;
@@ -2120,6 +2137,9 @@ static bool hasBroadcastReduceIdentity(Value init, BroadcastReduceKind kind) {
return value.isAllOnes();
case BroadcastReduceKind::OrI:
return value.isZero();
+ case BroadcastReduceKind::MaximumF:
+ case BroadcastReduceKind::MinimumF:
+ llvm_unreachable("float kinds handled by the branch above");
}
llvm_unreachable("unknown broadcast reduction kind");
}
@@ -2130,6 +2150,8 @@ static bool hasBroadcastReduceIdentity(Value init, BroadcastReduceKind kind) {
/// minsi(broadcast(y)) -> y
/// andi(broadcast(z)) -> z
/// ori(broadcast(w)) -> w
+/// maximumf(broadcast(x)) -> x
+/// minimumf(broadcast(y)) -> y
///
/// provided the reduction init is the identity value of the combiner.
struct FoldReduceBroadcast : public OpRewritePattern<linalg::ReduceOp> {
diff --git a/mlir/test/Dialect/Linalg/canonicalize.mlir b/mlir/test/Dialect/Linalg/canonicalize.mlir
index 26f4d92fa0e04e..6d8ce938168329 100644
--- a/mlir/test/Dialect/Linalg/canonicalize.mlir
+++ b/mlir/test/Dialect/Linalg/canonicalize.mlir
@@ -199,6 +199,120 @@ func.func @negative_reduce_broadcast_andi_non_identity_init(
// -----
+// CHECK-LABEL: func.func @reduce_broadcast_maximumf(
+// CHECK-SAME: %[[INPUT:[a-zA-Z0-9]+]]: tensor<3xf32>
+// CHECK-NEXT: return %[[INPUT]] : tensor<3xf32>
+func.func @reduce_broadcast_maximumf(%input: tensor<3xf32>,
+ %broadcast_init: tensor<3x4xf32>)
+ -> tensor<3xf32> {
+ %reduce_init = arith.constant dense<0xFF800000> : tensor<3xf32>
+ %broadcasted = linalg.broadcast
+ ins(%input : tensor<3xf32>)
+ outs(%broadcast_init : tensor<3x4xf32>) dimensions = [1]
+ %result = linalg.reduce
+ ins(%broadcasted : tensor<3x4xf32>)
+ outs(%reduce_init : tensor<3xf32>) dimensions = [1]
+ (%in: f32, %out: f32) {
+ %max = arith.maximumf %in, %out : f32
+ linalg.yield %max : f32
+ }
+ return %result : tensor<3xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @reduce_broadcast_minimumf(
+// CHECK-SAME: %[[INPUT:[a-zA-Z0-9]+]]: tensor<3xf32>
+// CHECK-NEXT: return %[[INPUT]] : tensor<3xf32>
+func.func @reduce_broadcast_minimumf(%input: tensor<3xf32>,
+ %broadcast_init: tensor<2x3xf32>)
+ -> tensor<3xf32> {
+ %reduce_init = arith.constant dense<0x7F800000> : tensor<3xf32>
+ %broadcasted = linalg.broadcast
+ ins(%input : tensor<3xf32>)
+ outs(%broadcast_init : tensor<2x3xf32>) dimensions = [0]
+ %result = linalg.reduce
+ ins(%broadcasted : tensor<2x3xf32>)
+ outs(%reduce_init : tensor<3xf32>) dimensions = [0]
+ (%in: f32, %out: f32) {
+ %min = arith.minimumf %out, %in : f32
+ linalg.yield %min : f32
+ }
+ return %result : tensor<3xf32>
+}
+
+// -----
+
+// Init 0.0 is not the maximumf identity (-inf), so no fold.
+// CHECK-LABEL: func.func @negative_reduce_broadcast_maximumf_finite_init
+// CHECK: linalg.broadcast
+// CHECK: linalg.reduce
+func.func @negative_reduce_broadcast_maximumf_finite_init(
+ %input: tensor<3xf32>, %broadcast_init: tensor<3x4xf32>)
+ -> tensor<3xf32> {
+ %reduce_init = arith.constant dense<0.0> : tensor<3xf32>
+ %broadcasted = linalg.broadcast
+ ins(%input : tensor<3xf32>)
+ outs(%broadcast_init : tensor<3x4xf32>) dimensions = [1]
+ %result = linalg.reduce
+ ins(%broadcasted : tensor<3x4xf32>)
+ outs(%reduce_init : tensor<3xf32>) dimensions = [1]
+ (%in: f32, %out: f32) {
+ %max = arith.maximumf %in, %out : f32
+ linalg.yield %max : f32
+ }
+ return %result : tensor<3xf32>
+}
+
+// -----
+
+// maxnumf ignores NaN, so reduce(broadcast(NaN)) with init=-inf yields -inf,
+// not NaN: the fold would be incorrect. Not folded even with an -inf init.
+// CHECK-LABEL: func.func @negative_reduce_broadcast_maxnumf_identity_init
+// CHECK: linalg.broadcast
+// CHECK: linalg.reduce
+func.func @negative_reduce_broadcast_maxnumf_identity_init(
+ %input: tensor<3xf32>, %broadcast_init: tensor<3x4xf32>)
+ -> tensor<3xf32> {
+ %reduce_init = arith.constant dense<0xFF800000> : tensor<3xf32>
+ %broadcasted = linalg.broadcast
+ ins(%input : tensor<3xf32>)
+ outs(%broadcast_init : tensor<3x4xf32>) dimensions = [1]
+ %result = linalg.reduce
+ ins(%broadcasted : tensor<3x4xf32>)
+ outs(%reduce_init : tensor<3xf32>) dimensions = [1]
+ (%in: f32, %out: f32) {
+ %max = arith.maxnumf %in, %out : f32
+ linalg.yield %max : f32
+ }
+ return %result : tensor<3xf32>
+}
+
+// -----
+
+// NaN is not the maximumf identity (maximumf(NaN, x) = NaN, not x), so no fold.
+// CHECK-LABEL: func.func @negative_reduce_broadcast_maximumf_nan_init
+// CHECK: linalg.broadcast
+// CHECK: linalg.reduce
+func.func @negative_reduce_broadcast_maximumf_nan_init(
+ %input: tensor<3xf32>, %broadcast_init: tensor<3x4xf32>)
+ -> tensor<3xf32> {
+ %reduce_init = arith.constant dense<0x7FC00000> : tensor<3xf32>
+ %broadcasted = linalg.broadcast
+ ins(%input : tensor<3xf32>)
+ outs(%broadcast_init : tensor<3x4xf32>) dimensions = [1]
+ %result = linalg.reduce
+ ins(%broadcasted : tensor<3x4xf32>)
+ outs(%reduce_init : tensor<3xf32>) dimensions = [1]
+ (%in: f32, %out: f32) {
+ %max = arith.maximumf %in, %out : f32
+ linalg.yield %max : f32
+ }
+ return %result : tensor<3xf32>
+}
+
+// -----
+
// CHECK-LABEL: func @memref_cast(
func.func @memref_cast(%a: index, %b: index) -> memref<?x?xf32> {
%c0 = arith.constant 0 : index
@@ -711,7 +825,7 @@ func.func @fold_self_copy(%0 : memref<4x16xf32>) {
// -----
// CHECK-LABEL: func @negative_fold_fill_like_memref
-// CHECK-NEXT: linalg.generic
+// CHECK-NEXT: linalg.generic
func.func @negative_fold_fill_like_memref(%in_out : memref<4x16xf32>, %fill_val : f32) {
linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
affine_map<(d0, d1) -> (d0, d1)>],
@@ -727,7 +841,7 @@ func.func @negative_fold_fill_like_memref(%in_out : memref<4x16xf32>, %fill_val
// -----
// CHECK-LABEL: func @negative_fold_fill_like_tensor
-// CHECK-NEXT: linalg.generic
+// CHECK-NEXT: linalg.generic
func.func @negative_fold_fill_like_tensor(%in_out : tensor<4x16xf32>, %fill_val : f32) -> tensor<4x16xf32> {
%result = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
affine_map<(d0, d1) -> (d0, d1)>],
More information about the Mlir-commits
mailing list