[Mlir-commits] [mlir] [mlir][arith] Honor finite min/max reduction identities (PR #232257)
Anutosh Bhat
llvmlistbot at llvm.org
Wed Oct 21 23:07:04 PDT 2026
https://github.com/anutosh491 updated https://github.com/llvm/llvm-project/pull/232257
>From d33b4d0417b62d92fc2c548cffcb2b1ffb915f1c Mon Sep 17 00:00:00 2001
From: anutosh491 <andersonbhat491 at gmail.com>
Date: Thu, 15 Oct 2026 15:44:54 +0530
Subject: [PATCH] [mlir][arith] Honor finite min/max reduction identities
---
mlir/include/mlir/Dialect/Arith/IR/Arith.h | 8 +-
mlir/lib/Dialect/Arith/IR/ArithOps.cpp | 18 ++-
mlir/test/Dialect/Linalg/pad-neutral.mlir | 14 ++-
.../Linalg/transform-op-decompose.mlir | 2 +-
.../Linalg/transform-op-split-reduction.mlir | 107 ++++++++++++++++++
5 files changed, 137 insertions(+), 12 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Arith/IR/Arith.h b/mlir/include/mlir/Dialect/Arith/IR/Arith.h
index 319ea68c2f2cdd..c940a36b5aa9a2 100644
--- a/mlir/include/mlir/Dialect/Arith/IR/Arith.h
+++ b/mlir/include/mlir/Dialect/Arith/IR/Arith.h
@@ -134,15 +134,17 @@ bool applyCmpPredicate(arith::CmpFPredicate predicate, const APFloat &lhs,
/// E.g., The identity value for maxf is in theory `-Inf`, but if we want to
/// stay in the finite range, it would be `BiggestRepresentableNegativeFloat`.
/// The purpose of this boolean is to offer constants that will play nice
-/// with fast math related optimizations.
+/// with fast math related optimizations. The caller must ensure that using a
+/// finite value preserves the reduction semantics.
TypedAttr getIdentityValueAttr(AtomicRMWKind kind, Type resultType,
OpBuilder &builder, Location loc,
bool useOnlyFiniteValue = false);
/// Return the identity numeric value associated to the give op. Return
/// std::nullopt if there is no known neutral element.
-/// If `op` has `FastMathFlags::ninf`, only finite values will be used
-/// as neutral element.
+/// Finite min/max identities require `FastMathFlags::ninf`. NaN-ignoring
+/// min/max operations additionally require `FastMathFlags::nnan`; otherwise,
+/// their neutral element is NaN.
std::optional<TypedAttr> getNeutralElement(Operation *op);
/// Returns the identity value associated with an AtomicRMWKind op.
diff --git a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
index f3a03a6a05c7fd..19e9caca6567a6 100644
--- a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
+++ b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
@@ -3385,7 +3385,9 @@ TypedAttr mlir::arith::getIdentityValueAttr(AtomicRMWKind kind, Type resultType,
case AtomicRMWKind::maxnumf: {
const llvm::fltSemantics &semantic =
llvm::cast<FloatType>(resultType).getFloatSemantics();
- APFloat identity = APFloat::getNaN(semantic, /*Negative=*/true);
+ APFloat identity = useOnlyFiniteValue
+ ? APFloat::getLargest(semantic, /*Negative=*/true)
+ : APFloat::getNaN(semantic, /*Negative=*/true);
return builder.getFloatAttr(resultType, identity);
}
case AtomicRMWKind::addf:
@@ -3414,7 +3416,9 @@ TypedAttr mlir::arith::getIdentityValueAttr(AtomicRMWKind kind, Type resultType,
case AtomicRMWKind::minnumf: {
const llvm::fltSemantics &semantic =
llvm::cast<FloatType>(resultType).getFloatSemantics();
- APFloat identity = APFloat::getNaN(semantic, /*Negative=*/false);
+ APFloat identity = useOnlyFiniteValue
+ ? APFloat::getLargest(semantic, /*Negative=*/false)
+ : APFloat::getNaN(semantic, /*Negative=*/false);
return builder.getFloatAttr(resultType, identity);
}
case AtomicRMWKind::mins:
@@ -3469,6 +3473,16 @@ std::optional<TypedAttr> mlir::arith::getNeutralElement(Operation *op) {
arith::FastMathFlagsAttr fmfAttr = fmfOpInterface.getFastMathFlagsAttr();
useOnlyFiniteValue =
bitEnumContainsAny(fmfAttr.getValue(), arith::FastMathFlags::ninf);
+ if (*maybeKind == AtomicRMWKind::maxnumf ||
+ *maybeKind == AtomicRMWKind::minnumf) {
+ if (bitEnumContainsAny(fmfAttr.getValue(), arith::FastMathFlags::nnan)) {
+ *maybeKind = *maybeKind == AtomicRMWKind::maxnumf
+ ? AtomicRMWKind::maximumf
+ : AtomicRMWKind::minimumf;
+ } else {
+ useOnlyFiniteValue = false;
+ }
+ }
}
// Builder only used as helper for attribute creation.
diff --git a/mlir/test/Dialect/Linalg/pad-neutral.mlir b/mlir/test/Dialect/Linalg/pad-neutral.mlir
index c85ad4d73c7059..2309d44a1f7bfd 100644
--- a/mlir/test/Dialect/Linalg/pad-neutral.mlir
+++ b/mlir/test/Dialect/Linalg/pad-neutral.mlir
@@ -166,12 +166,15 @@ module attributes {transform.with_named_sequence} {
// -----
-// Negative: `nnan` declares that NaN operands do not occur, so the NaN neutral
-// element of maxnumf is not a usable pad value and inference fails instead.
-
-func.func @negative_pad_reduce_maxnumf_nnan(%input: tensor<8x30xf32>,
+// CHECK-LABEL: @pad_reduce_maxnumf_nnan
+func.func @pad_reduce_maxnumf_nnan(%input: tensor<8x30xf32>,
%init: tensor<8xf32>) -> tensor<8xf32> {
- // expected-note @below {{target op}}
+ // CHECK-DAG: %[[NEUTRAL:.*]] = arith.constant 0xFF800000 : f32
+ // CHECK: %[[PAD:.*]] = tensor.pad %{{.*}} low[0, 0] high[0, 2]
+ // CHECK: tensor.yield %[[NEUTRAL]] : f32
+ // CHECK: } : tensor<8x30xf32> to tensor<8x32xf32>
+ // CHECK: linalg.reduce ins(%[[PAD]] : tensor<8x32xf32>)
+ // CHECK: arith.maxnumf {{.*}} fastmath<nnan> : f32
%0 = linalg.reduce ins(%input : tensor<8x30xf32>) outs(%init : tensor<8xf32>)
dimensions = [1]
(%in: f32, %out: f32) {
@@ -185,7 +188,6 @@ module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
%red = transform.structured.match ops{["linalg.reduce"]} in %arg1
: (!transform.any_op) -> !transform.any_op
- // expected-error @below {{failed to pad op}}
%padded, %pad = transform.structured.pad_tiling_interface %red to padding_sizes [8, 32]
: (!transform.any_op) -> (!transform.any_op, !transform.any_op)
transform.yield
diff --git a/mlir/test/Dialect/Linalg/transform-op-decompose.mlir b/mlir/test/Dialect/Linalg/transform-op-decompose.mlir
index 3897f8502bb04e..05a01b0a20eb36 100644
--- a/mlir/test/Dialect/Linalg/transform-op-decompose.mlir
+++ b/mlir/test/Dialect/Linalg/transform-op-decompose.mlir
@@ -369,7 +369,7 @@ func.func @softmax(%arg0: tensor<2x16x32xf32>, %dst: tensor<2x16x32xf32>) -> ten
// CHECK-LABEL: func.func @softmax(
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<2x16x32xf32>, %[[DST:[a-zA-Z0-9_]+]]: tensor<2x16x32xf32>) -> tensor<2x16x32xf32> {
// CHECK-DAG: %[[D1:.+]] = tensor.empty() : tensor<2x16xf32>
-// CHECK-DAG: %[[CST:.+]] = arith.constant 0xFFC00000 : f32
+// CHECK-DAG: %[[CST:.+]] = arith.constant -3.40282347E+38 : f32
// CHECK: %[[D2:.+]] = linalg.fill ins(%[[CST]] : f32) outs(%[[D1]] : tensor<2x16xf32>) -> tensor<2x16xf32>
// CHECK: %[[D3:.+]] = linalg.generic {indexing_maps = [#[[$MAP]], #[[$MAP1]]], iterator_types = ["parallel",
// CHECK-SAME: "parallel", "reduction"]} ins(%[[ARG0]] : tensor<2x16x32xf32>) outs(%[[D2]] : tensor<2x16xf32>) {
diff --git a/mlir/test/Dialect/Linalg/transform-op-split-reduction.mlir b/mlir/test/Dialect/Linalg/transform-op-split-reduction.mlir
index 82111f911aa244..dde3a425b9e175 100644
--- a/mlir/test/Dialect/Linalg/transform-op-split-reduction.mlir
+++ b/mlir/test/Dialect/Linalg/transform-op-split-reduction.mlir
@@ -499,3 +499,110 @@ module attributes {transform.with_named_sequence} {
transform.yield
}
}
+
+// -----
+
+// CHECK-LABEL: func.func @maxnumf_ninf(
+// CHECK: %[[INIT:.*]] = arith.constant 0xFFC00000 : f32
+// CHECK: linalg.fill ins(%[[INIT]] : f32)
+// CHECK: arith.maxnumf {{.*}} fastmath<ninf> : f32
+func.func @maxnumf_ninf(%in: tensor<32xf32>, %out: tensor<f32>) -> tensor<f32> {
+ %r = linalg.generic {
+ indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> ()>],
+ iterator_types = ["reduction"]
+ } ins(%in : tensor<32xf32>) outs(%out : tensor<f32>) {
+ ^bb0(%value: f32, %acc: f32):
+ %result = arith.maxnumf %value, %acc fastmath<ninf> : f32
+ linalg.yield %result : f32
+ } -> tensor<f32>
+ return %r : tensor<f32>
+}
+
+// CHECK-LABEL: func.func @maxnumf_nnan(
+// CHECK: %[[INIT:.*]] = arith.constant 0xFF800000 : f32
+// CHECK: linalg.fill ins(%[[INIT]] : f32)
+// CHECK: arith.maxnumf {{.*}} fastmath<nnan> : f32
+func.func @maxnumf_nnan(%in: tensor<32xf32>, %out: tensor<f32>) -> tensor<f32> {
+ %r = linalg.generic {
+ indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> ()>],
+ iterator_types = ["reduction"]
+ } ins(%in : tensor<32xf32>) outs(%out : tensor<f32>) {
+ ^bb0(%value: f32, %acc: f32):
+ %result = arith.maxnumf %value, %acc fastmath<nnan> : f32
+ linalg.yield %result : f32
+ } -> tensor<f32>
+ return %r : tensor<f32>
+}
+
+// CHECK-LABEL: func.func @minnumf_ninf(
+// CHECK: %[[INIT:.*]] = arith.constant 0x7FC00000 : f32
+// CHECK: linalg.fill ins(%[[INIT]] : f32)
+// CHECK: arith.minnumf {{.*}} fastmath<ninf> : f32
+func.func @minnumf_ninf(%in: tensor<32xf32>, %out: tensor<f32>) -> tensor<f32> {
+ %r = linalg.generic {
+ indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> ()>],
+ iterator_types = ["reduction"]
+ } ins(%in : tensor<32xf32>) outs(%out : tensor<f32>) {
+ ^bb0(%value: f32, %acc: f32):
+ %result = arith.minnumf %value, %acc fastmath<ninf> : f32
+ linalg.yield %result : f32
+ } -> tensor<f32>
+ return %r : tensor<f32>
+}
+
+// CHECK-LABEL: func.func @minnumf_nnan(
+// CHECK: %[[INIT:.*]] = arith.constant 0x7F800000 : f32
+// CHECK: linalg.fill ins(%[[INIT]] : f32)
+// CHECK: arith.minnumf {{.*}} fastmath<nnan> : f32
+func.func @minnumf_nnan(%in: tensor<32xf32>, %out: tensor<f32>) -> tensor<f32> {
+ %r = linalg.generic {
+ indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> ()>],
+ iterator_types = ["reduction"]
+ } ins(%in : tensor<32xf32>) outs(%out : tensor<f32>) {
+ ^bb0(%value: f32, %acc: f32):
+ %result = arith.minnumf %value, %acc fastmath<nnan> : f32
+ linalg.yield %result : f32
+ } -> tensor<f32>
+ return %r : tensor<f32>
+}
+
+// CHECK-LABEL: func.func @maxnumf_nnan_ninf
+// CHECK: %[[MAX_INIT:.*]] = arith.constant -3.40282347E+38 : f32
+// CHECK: linalg.fill ins(%[[MAX_INIT]] : f32)
+// CHECK: arith.maxnumf {{.*}} fastmath<nnan,ninf> : f32
+func.func @maxnumf_nnan_ninf(%in: tensor<32xf32>, %out: tensor<f32>) -> tensor<f32> {
+ %r = linalg.generic {
+ indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> ()>],
+ iterator_types = ["reduction"]
+ } ins(%in : tensor<32xf32>) outs(%out : tensor<f32>) {
+ ^bb0(%value: f32, %acc: f32):
+ %max = arith.maxnumf %value, %acc fastmath<nnan,ninf> : f32
+ linalg.yield %max : f32
+ } -> tensor<f32>
+ return %r : tensor<f32>
+}
+
+// CHECK-LABEL: func.func @minnumf_nnan_ninf
+// CHECK: %[[MIN_INIT:.*]] = arith.constant 3.40282347E+38 : f32
+// CHECK: linalg.fill ins(%[[MIN_INIT]] : f32)
+// CHECK: arith.minnumf {{.*}} fastmath<nnan,ninf> : f32
+func.func @minnumf_nnan_ninf(%in: tensor<32xf32>, %out: tensor<f32>) -> tensor<f32> {
+ %r = linalg.generic {
+ indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> ()>],
+ iterator_types = ["reduction"]
+ } ins(%in : tensor<32xf32>) outs(%out : tensor<f32>) {
+ ^bb0(%value: f32, %acc: f32):
+ %min = arith.minnumf %value, %acc fastmath<nnan,ninf> : f32
+ linalg.yield %min : f32
+ } -> tensor<f32>
+ return %r : tensor<f32>
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%root: !transform.any_op {transform.readonly}) {
+ %ops = transform.structured.match ops{["linalg.generic"]} in %root : (!transform.any_op) -> !transform.any_op
+ %split:4 = transform.structured.split_reduction %ops <split_factor = 4, insert_split_dimension = 0, inner_parallel>
+ : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)
+ transform.yield
+ }
+}
More information about the Mlir-commits
mailing list