[Mlir-commits] [mlir] [mlir][tosa] Preserve all-NaN max-pool windows (PR #225744)

[Jack Frankland]

https://github.com/FranklandJack updated https://github.com/llvm/llvm-project/pull/225744

>From d03592f45eb5f0f899592398736d7f6dfcf00630 Mon Sep 17 00:00:00 2001
From: Jack Frankland <jack.frankland at arm.com>
Date: Wed, 23 Sep 2026 09:39:09 +0100
Subject: [PATCH] [mlir][tosa] Preserve all-NaN max-pool windows

Initialize IGNORE-mode floating-point max pooling with NaN and select
the first finite input. This keeps an all-NaN window NaN instead of
returning the lowest finite value.

Assited-by: Codex
---
 .../TosaToLinalg/TosaToLinalgNamed.cpp        | 31 +++++++++++++------
 .../TosaToLinalg/tosa-to-linalg-named.mlir    | 18 +++++++++++
 2 files changed, 39 insertions(+), 10 deletions(-)

diff --git a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp
index 4a11a927090bf..c6a931ac4f528 100644
--- a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp
+++ b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp
@@ -742,16 +742,21 @@ class MaxPool2dConverter : public OpConversionPattern<tosa::MaxPool2dOp> {
     if (!resultTy)
       return rewriter.notifyMatchFailure(op, "failed to convert type");
     Type resultETy = inputTy.getElementType();
+    NanPropagationMode nanMode = op.getNanMode();
 
     SmallVector<Value> dynamicDims =
         computeDynamicOutputSizes(op, adaptor, rewriter);
 
     // Determine what the initial value needs to be for the max pool op.
     TypedAttr initialAttr;
-    if (resultETy.isF32() || resultETy.isBF16() || resultETy.isF16())
+    if (resultETy.isF32() || resultETy.isBF16() || resultETy.isF16()) {
+      const llvm::fltSemantics &semantics =
+          cast<FloatType>(resultETy).getFloatSemantics();
       initialAttr = rewriter.getFloatAttr(
-          resultETy, APFloat::getLargest(
-                         cast<FloatType>(resultETy).getFloatSemantics(), true));
+          resultETy, nanMode == NanPropagationMode::IGNORE
+                         ? APFloat::getNaN(semantics)
+                         : APFloat::getLargest(semantics, true));
+    }
 
     else if (isUnsigned)
       initialAttr = rewriter.getIntegerAttr(
@@ -805,7 +810,6 @@ class MaxPool2dConverter : public OpConversionPattern<tosa::MaxPool2dOp> {
         ValueRange{paddedInput, fakeWindowDims}, filledEmptyTensor, strideAttr,
         dilationAttr);
 
-    NanPropagationMode nanMode = op.getNanMode();
     rewriter.replaceOp(op, resultOp);
 
     // NaN propagation has no meaning for non floating point types.
@@ -817,8 +821,9 @@ class MaxPool2dConverter : public OpConversionPattern<tosa::MaxPool2dOp> {
     //
     // In the case of "IGNORE" we need to insert a compare and select. Since
     // we've already produced a named op we will just take its body and modify
-    // it to include the appropriate checks. If the current value is NaN the
-    // old value of pool will be taken otherwise we use the result.
+    // it to include the appropriate checks. A NaN accumulator represents a
+    // window with no finite values yet. The first finite input replaces it;
+    // input NaNs are ignored. An all-NaN window therefore remains NaN.
     if (nanMode == NanPropagationMode::IGNORE) {
       auto genericOp = linalg::GenericOp::create(
           rewriter, loc, resultOp.getType(0), resultOp.getInputs(),
@@ -831,12 +836,18 @@ class MaxPool2dConverter : public OpConversionPattern<tosa::MaxPool2dOp> {
             auto &oldMaxOp = *resultOp.getBlock()->begin();
             map.map(oldArgs, blockArgs);
             auto *newOp = opBuilder.clone(oldMaxOp, map);
-            Value isNaN =
+            Value inputIsNaN =
                 arith::CmpFOp::create(opBuilder, loc, arith::CmpFPredicate::UNO,
                                       blockArgs.front(), blockArgs.front());
-            auto selectOp = arith::SelectOp::create(
-                opBuilder, loc, isNaN, blockArgs.back(), newOp->getResult(0));
-            linalg::YieldOp::create(opBuilder, loc, selectOp.getResult());
+            Value accumulatorIsNaN =
+                arith::CmpFOp::create(opBuilder, loc, arith::CmpFPredicate::UNO,
+                                      blockArgs.back(), blockArgs.back());
+            Value finiteMaximum =
+                arith::SelectOp::create(opBuilder, loc, accumulatorIsNaN,
+                                        blockArgs.front(), newOp->getResult(0));
+            Value result = arith::SelectOp::create(
+                opBuilder, loc, inputIsNaN, blockArgs.back(), finiteMaximum);
+            linalg::YieldOp::create(opBuilder, loc, result);
           });
       rewriter.replaceOp(resultOp, genericOp);
     }
diff --git a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-named.mlir b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-named.mlir
index a4abc687bfe38..effb2dd38c132 100644
--- a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-named.mlir
+++ b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-named.mlir
@@ -106,6 +106,24 @@ func.func @max_pool(%arg0: tensor<1x6x34x62xf32>) -> () {
   return
 }
 
+// -----
+
+// Lower IGNORE mode with a NaN accumulator so an all-NaN window stays NaN.
+// CHECK-LABEL: @max_pool_ignore
+func.func @max_pool_ignore(%arg0: tensor<1x4x4x1xf32>) -> tensor<1x3x3x1xf32> {
+  // CHECK: arith.constant {{.*}} : f32
+  // CHECK: linalg.fill
+  // CHECK: linalg.generic
+  // CHECK: arith.cmpf uno
+  // CHECK: arith.cmpf uno
+  // CHECK: arith.select
+  // CHECK: arith.select
+  // CHECK-NOT: tosa.max_pool2d
+  %0 = tosa.max_pool2d %arg0 kernel([2, 2]) stride([1, 1]) pad([0, 0, 0, 0]) nan_mode<IGNORE> :
+    (tensor<1x4x4x1xf32>) -> tensor<1x3x3x1xf32>
+  return %0 : tensor<1x3x3x1xf32>
+}
+
 // CHECK-LABEL: @max_pool_padded
 func.func @max_pool_padded(%arg0: tensor<1x6x34x62xf32>) -> () {
   // CHECK-DAG: [[CONST:%.+]] = arith.constant -3.40282347E+38 : f32