[Mlir-commits] [mlir] [mlir][arith] Honor finite min/max reduction identities (PR #232257)

[Anutosh Bhat]

https://github.com/anutosh491 created https://github.com/llvm/llvm-project/pull/232257

I hit this while compiling a JAX transformer block through IREE/WebGPU in JupyterLite. `linalg.softmax` requests a finite reduction initializer, but `getIdentityValueAttr` returns NaN for `maxnumf`, which Tint rejects.

Honor `useOnlyFiniteValue` for `maxnumf` and `minnumf`, preserving the existing NaN identities when the flag is false. Update the existing softmax check and add finite min/max reduction cases.

>From d02e28cbcd168eb05b79a35c17bc4a962e3dc994 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/lib/Dialect/Arith/IR/ArithOps.cpp        |  8 +++-
 .../Linalg/transform-op-decompose.mlir        |  2 +-
 .../Linalg/transform-op-split-reduction.mlir  | 44 +++++++++++++++++++
 3 files changed, 51 insertions(+), 3 deletions(-)

diff --git a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
index 40fcfa5c458819..09aea5c8eba837 100644
--- a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
+++ b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
@@ -3316,7 +3316,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:
@@ -3345,7 +3347,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:
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..de23fc1af6dc07 100644
--- a/mlir/test/Dialect/Linalg/transform-op-split-reduction.mlir
+++ b/mlir/test/Dialect/Linalg/transform-op-split-reduction.mlir
@@ -499,3 +499,47 @@ module attributes {transform.with_named_sequence} {
       transform.yield
   }
 }
+
+// -----
+
+// Finite identities must also be used for NaN-ignoring min/max reductions.
+// CHECK-LABEL: func.func @maxnumf_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_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_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_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
+  }
+}