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

[Jack Frankland]

================
@@ -106,6 +106,28 @@ 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> :
----------------
FranklandJack wrote:

Yep good idea, added a test for this.

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