[Mlir-commits] [mlir] [[mlir][linalg] Infer reduction-neutral padding values in rewriteAsPaddedOp (PR #216517)

[Andrzej Warzyński]

================
@@ -0,0 +1,367 @@
+// RUN: mlir-opt --transform-interpreter -canonicalize -split-input-file --verify-diagnostics %s | FileCheck %s
+
+// When padding values are not specified, `pad_tiling_interface` infers them.
+// For an operand that is reduced, a zero pad value would corrupt the result,
+// so the neutral element of the reduction combiner must be inferred instead.
+
+///--------------------------------------------------------------------------------------
+/// Tests for linalg.reduce
+///--------------------------------------------------------------------------------------
+
+// CHECK-LABEL: @pad_reduce_maximumf
+func.func @pad_reduce_maximumf(%input: tensor<8x30xf32>, %init: tensor<8xf32>)
+    -> tensor<8xf32> {
+  // maximumf neutral element is -inf.
+  // 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>)
+  %0 = linalg.reduce ins(%input : tensor<8x30xf32>) outs(%init : tensor<8xf32>)
+      dimensions = [1]
+    (%in: f32, %out: f32) {
+      %m = arith.maximumf %in, %out : f32
+      linalg.yield %m : f32
+    }
+  return %0 : tensor<8xf32>
+}
+
+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
+    // No padding_values: they are inferred.
+    %padded, %pad = transform.structured.pad_tiling_interface %red to padding_sizes [8, 32]
+      : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+    transform.yield
+  }
+}
+
+// -----
+
+// CHECK-LABEL: @pad_reduce_minimumf
+func.func @pad_reduce_minimumf(%input: tensor<8x30xf32>, %init: tensor<8xf32>)
+    -> tensor<8xf32> {
+  // minimumf neutral element is +inf.
+  // CHECK-DAG:   %[[NEUTRAL:.*]] = arith.constant 0x7F800000 : 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>)
+  %0 = linalg.reduce ins(%input : tensor<8x30xf32>) outs(%init : tensor<8xf32>)
+      dimensions = [1]
+    (%in: f32, %out: f32) {
+      %m = arith.minimumf %in, %out : f32
+      linalg.yield %m : f32
+    }
+  return %0 : tensor<8xf32>
+}
+
+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
+    %padded, %pad = transform.structured.pad_tiling_interface %red to padding_sizes [8, 32]
+      : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+    transform.yield
+  }
+}
+
+// -----
+
+// CHECK-LABEL: @pad_reduce_mulf
+func.func @pad_reduce_mulf(%input: tensor<8x30xf32>, %init: tensor<8xf32>)
+    -> tensor<8xf32> {
+  // mulf neutral element is 1.0 (a zero pad would zero out the product).
+  // CHECK-DAG:   %[[NEUTRAL:.*]] = arith.constant 1.000000e+00 : 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>)
+  %0 = linalg.reduce ins(%input : tensor<8x30xf32>) outs(%init : tensor<8xf32>)
+      dimensions = [1]
+    (%in: f32, %out: f32) {
+      %m = arith.mulf %in, %out : f32
+      linalg.yield %m : f32
+    }
+  return %0 : tensor<8xf32>
+}
+
+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
+    %padded, %pad = transform.structured.pad_tiling_interface %red to padding_sizes [8, 32]
+      : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+    transform.yield
+  }
+}
+
+// -----
+
+// maxnumf is the NaN-ignoring variant, so its neutral element is NaN.
+
+// CHECK-LABEL: @pad_reduce_maxnumf
+func.func @pad_reduce_maxnumf(%input: tensor<8x30xf32>, %init: tensor<8xf32>)
+    -> tensor<8xf32> {
+  // CHECK-DAG:   %[[NAN:.*]] = arith.constant 0xFFC00000 : f32
+  // CHECK:       tensor.pad %{{.*}} low[0, 0] high[0, 2]
+  // CHECK:         tensor.yield %[[NAN]] : f32
+  // CHECK:       } : tensor<8x30xf32> to tensor<8x32xf32>
+  %0 = linalg.reduce ins(%input : tensor<8x30xf32>) outs(%init : tensor<8xf32>) dimensions = [1]
+    (%in: f32, %out: f32) {
+      %m = arith.maxnumf %in, %out : f32
+      linalg.yield %m : f32
+    }
+  return %0 : tensor<8xf32>
+}
+
+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
+    %padded, %pad = transform.structured.pad_tiling_interface %red to padding_sizes [8, 32]
+      : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+    transform.yield
+  }
+}
+
+// -----
+
+// Only the reduced input takes the neutral (-inf); the non-reduced init operand
----------------
banach-space wrote:

[nit] 
```suggestion
// Only the reduced input takes `-inf`, the reduction Op neutral value (`maximumf`), as the pad value; the non-reduced init operand
```

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