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

[Jianhui Li]

================
@@ -0,0 +1,320 @@
+// 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.
+
+// 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
+  }
+}
+
+// -----
+
+// An operand reduced along several dimensions is padded with the neutral on all
+// of them (here d1 and d2); one neutral suffices since the combiner is the same.
+
+#in3  = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+#out3 = affine_map<(d0, d1, d2) -> (d0)>
+// CHECK-LABEL: @pad_reduce_two_reduction_dims
+func.func @pad_reduce_two_reduction_dims(%in: tensor<8x6x10xf32>, %init: tensor<8xf32>)
----------------
Jianhui-Li wrote:

changed. 

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