[Mlir-commits] [mlir] [mlir][tosa] Preserve all-NaN max-pool windows (PR #225744)
Jack Frankland
llvmlistbot at llvm.org
Wed Sep 30 08:51:55 PDT 2026
https://github.com/FranklandJack updated https://github.com/llvm/llvm-project/pull/225744
>From 0e333f3a201687465c4fd23129207c22e29be2ac 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 | 34 +++++---
.../TosaToLinalg/tosa-to-linalg-named.mlir | 82 +++++++++++++++++++
2 files changed, 105 insertions(+), 11 deletions(-)
diff --git a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp
index 4a11a927090bf6..ae5390b36ddc7b 100644
--- a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp
+++ b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp
@@ -742,16 +742,22 @@ 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())
- initialAttr = rewriter.getFloatAttr(
- resultETy, APFloat::getLargest(
- cast<FloatType>(resultETy).getFloatSemantics(), true));
+ if (resultETy.isF32() || resultETy.isBF16() || resultETy.isF16()) {
+ const llvm::fltSemantics &semantics =
+ cast<FloatType>(resultETy).getFloatSemantics();
+ if (llvm::APFloat::semanticsHasNaN(semantics))
+ initialAttr = rewriter.getFloatAttr(
+ resultETy, nanMode == NanPropagationMode::IGNORE
+ ? APFloat::getNaN(semantics)
+ : APFloat::getLargest(semantics, true));
+ }
else if (isUnsigned)
initialAttr = rewriter.getIntegerAttr(
@@ -805,7 +811,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 +822,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 +837,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 a4abc687bfe387..2034db67e9e3d3 100644
--- a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-named.mlir
+++ b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-named.mlir
@@ -106,6 +106,88 @@ 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: %[[NAN:.*]] = arith.constant 0x7FC00000 : f32
+ // CHECK: %[[FILL:.*]] = linalg.fill ins(%[[NAN]] : f32)
+ // CHECK: %[[POOL:.*]] = linalg.generic
+ // CHECK-SAME: outs(%[[FILL]] : tensor<1x3x3x1xf32>)
+ // CHECK: ^bb0(%[[INPUT:.*]]: f32, %[[WINDOW:.*]]: f32, %[[ACC:.*]]: f32):
+ // CHECK-DAG: %[[MAX:.*]] = arith.maximumf %[[ACC]], %[[INPUT]] : f32
+ // CHECK-DAG: %[[INPUT_NAN:.*]] = arith.cmpf uno, %[[INPUT]], %[[INPUT]] : f32
+ // CHECK-DAG: %[[ACC_NAN:.*]] = arith.cmpf uno, %[[ACC]], %[[ACC]] : f32
+ // CHECK: %[[CANDIDATE:.*]] = arith.select %[[ACC_NAN]], %[[INPUT]], %[[MAX]] : f32
+ // CHECK: %[[RESULT:.*]] = arith.select %[[INPUT_NAN]], %[[ACC]], %[[CANDIDATE]] : f32
+ // CHECK: linalg.yield %[[RESULT]] : f32
+ // CHECK: return %[[POOL]] : tensor<1x3x3x1xf32>
+ %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 that IGNORE mode also uses NaN to pad the input.
+// CHECK-LABEL: @max_pool_ignore_padded
+func.func @max_pool_ignore_padded(%arg0: tensor<1x4x4x1xf32>) -> tensor<1x3x4x1xf32> {
+ // CHECK: %[[PAD_NAN:.*]] = arith.constant 0x7FC00000 : f32
+ // CHECK: %[[PADDED:.*]] = tensor.pad %arg0 low[0, 0, 0, 0] high[0, 0, 1, 0]
+ // CHECK: tensor.yield %[[PAD_NAN]] : f32
+ // CHECK: %[[INIT_NAN:.*]] = arith.constant 0x7FC00000 : f32
+ // CHECK: %[[FILL:.*]] = linalg.fill ins(%[[INIT_NAN]] : f32)
+ // CHECK: %[[POOL:.*]] = linalg.generic
+ // CHECK-SAME: ins(%[[PADDED]]
+ // CHECK: return %[[POOL]] : tensor<1x3x4x1xf32>
+ %0 = tosa.max_pool2d %arg0 kernel([2, 2]) stride([1, 1]) pad([0, 0, 0, 1]) nan_mode<IGNORE> :
+ (tensor<1x4x4x1xf32>) -> tensor<1x3x4x1xf32>
+ return %0 : tensor<1x3x4x1xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @max_pool_ignore_bf16
+func.func @max_pool_ignore_bf16(%arg0: tensor<1x4x4x1xbf16>) -> tensor<1x3x3x1xbf16> {
+ // CHECK: %[[NAN:.*]] = arith.constant 0x7FC0 : bf16
+ // CHECK: %[[FILL:.*]] = linalg.fill ins(%[[NAN]] : bf16)
+ // CHECK: %[[POOL:.*]] = linalg.generic
+ // CHECK-SAME: outs(%[[FILL]] : tensor<1x3x3x1xbf16>)
+ // CHECK: ^bb0(%[[INPUT:.*]]: bf16, %[[WINDOW:.*]]: bf16, %[[ACC:.*]]: bf16):
+ // CHECK-DAG: %[[MAX:.*]] = arith.maximumf %[[ACC]], %[[INPUT]] : bf16
+ // CHECK-DAG: %[[INPUT_NAN:.*]] = arith.cmpf uno, %[[INPUT]], %[[INPUT]] : bf16
+ // CHECK-DAG: %[[ACC_NAN:.*]] = arith.cmpf uno, %[[ACC]], %[[ACC]] : bf16
+ // CHECK: %[[CANDIDATE:.*]] = arith.select %[[ACC_NAN]], %[[INPUT]], %[[MAX]] : bf16
+ // CHECK: %[[RESULT:.*]] = arith.select %[[INPUT_NAN]], %[[ACC]], %[[CANDIDATE]] : bf16
+ // CHECK: linalg.yield %[[RESULT]] : bf16
+ // CHECK: return %[[POOL]] : tensor<1x3x3x1xbf16>
+ %0 = tosa.max_pool2d %arg0 kernel([2, 2]) stride([1, 1]) pad([0, 0, 0, 0]) nan_mode<IGNORE> :
+ (tensor<1x4x4x1xbf16>) -> tensor<1x3x3x1xbf16>
+ return %0 : tensor<1x3x3x1xbf16>
+}
+
+// -----
+
+// CHECK-LABEL: @max_pool_ignore_f16
+func.func @max_pool_ignore_f16(%arg0: tensor<1x4x4x1xf16>) -> tensor<1x3x3x1xf16> {
+ // CHECK: %[[NAN:.*]] = arith.constant 0x7E00 : f16
+ // CHECK: %[[FILL:.*]] = linalg.fill ins(%[[NAN]] : f16)
+ // CHECK: %[[POOL:.*]] = linalg.generic
+ // CHECK-SAME: outs(%[[FILL]] : tensor<1x3x3x1xf16>)
+ // CHECK: ^bb0(%[[INPUT:.*]]: f16, %[[WINDOW:.*]]: f16, %[[ACC:.*]]: f16):
+ // CHECK-DAG: %[[MAX:.*]] = arith.maximumf %[[ACC]], %[[INPUT]] : f16
+ // CHECK-DAG: %[[INPUT_NAN:.*]] = arith.cmpf uno, %[[INPUT]], %[[INPUT]] : f16
+ // CHECK-DAG: %[[ACC_NAN:.*]] = arith.cmpf uno, %[[ACC]], %[[ACC]] : f16
+ // CHECK: %[[CANDIDATE:.*]] = arith.select %[[ACC_NAN]], %[[INPUT]], %[[MAX]] : f16
+ // CHECK: %[[RESULT:.*]] = arith.select %[[INPUT_NAN]], %[[ACC]], %[[CANDIDATE]] : f16
+ // CHECK: linalg.yield %[[RESULT]] : f16
+ // CHECK: return %[[POOL]] : tensor<1x3x3x1xf16>
+ %0 = tosa.max_pool2d %arg0 kernel([2, 2]) stride([1, 1]) pad([0, 0, 0, 0]) nan_mode<IGNORE> :
+ (tensor<1x4x4x1xf16>) -> tensor<1x3x3x1xf16>
+ return %0 : tensor<1x3x3x1xf16>
+}
+
// CHECK-LABEL: @max_pool_padded
func.func @max_pool_padded(%arg0: tensor<1x6x34x62xf32>) -> () {
// CHECK-DAG: [[CONST:%.+]] = arith.constant -3.40282347E+38 : f32
More information about the Mlir-commits
mailing list