[Mlir-commits] [mlir] Revert "[mlir][tensor] Refine the semantics of `createPadHighOp`" (PR #110153)

[Han-Chung Wang]

https://github.com/hanhanW created https://github.com/llvm/llvm-project/pull/110153

Reverts llvm/llvm-project#109667

It breaks the behavior of [linalg::padAndHoistLinalgOp](https://github.com/llvm/llvm-project/blob/fbec1c2a08ce2ae9750ddf3cecc86c5dd2bbc9d8/mlir/lib/Dialect/Linalg/Transforms/Padding.cpp#L265C7-L265C34), which calls [makeComposedPadHighOp here](https://github.com/llvm/llvm-project/blob/fbec1c2a08ce2ae9750ddf3cecc86c5dd2bbc9d8/mlir/lib/Dialect/Linalg/Transforms/Padding.cpp#L148-L149). The implementation of [makeComposedPadHighOp](https://github.com/llvm/llvm-project/blob/83368191a21340a6c3a8f88b01ecae6433640957/mlir/lib/Dialect/Linalg/Utils/Utils.cpp#L192) relies on the behavior to create valid pad ops.

To repro: `mlir-opt --transform-interpreter ~/repro.mlir`

```mlir
module {
  func.func @batch_matmul_f16(%arg0: tensor<1x?x1281xf16>, %arg1: tensor<1x1281x?xf16>, %arg2: tensor<1x?x?xf16>) -> tensor<1x?x?xf16> {
    %cst = arith.constant 0.000000e+00 : f16
    %c0 = arith.constant 0 : index
    %0 = linalg.fill ins(%cst : f16) outs(%arg2 : tensor<1x?x?xf16>) -> tensor<1x?x?xf16>
    %1 = linalg.batch_matmul ins(%arg0, %arg1 : tensor<1x?x1281xf16>, tensor<1x1281x?xf16>) outs(%0 : tensor<1x?x?xf16>) -> tensor<1x?x?xf16>
    return %1 : tensor<1x?x?xf16>
  }
  module attributes {transform.with_named_sequence} {
    transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
      %0 = transform.structured.match ops{["linalg.batch_matmul"]} in %arg0 : (!transform.any_op) -> !transform.any_op
      %padded, %pad, %copy = transform.structured.pad %0 {pack_paddings = [1, 1], padding_dimensions = [3], padding_values = [0.000000e+00 : f16, 0.000000e+00 : f16, 0.000000e+00 : f16]} : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.op<"bufferization.materialize_in_destination">)
      %1 = transform.num_associations %copy : (!transform.op<"bufferization.materialize_in_destination">) -> !transform.param<i64>
      transform.debug.emit_param_as_remark %1 : !transform.param<i64>
      transform.yield
    }
  }
}
```

>From 61e03b68511862ace5bc361565b2535e27fde833 Mon Sep 17 00:00:00 2001
From: Han-Chung Wang <hanhan0912 at gmail.com>
Date: Thu, 26 Sep 2024 11:17:47 -0700
Subject: [PATCH] Revert "[mlir][tensor] Refine the semantics of
 `createPadHighOp` (#109667)"

This reverts commit 9c48a04328f1dfa739985f64b33f20b67e085277.
---
 mlir/include/mlir/Dialect/Tensor/Utils/Utils.h |  8 ++++----
 mlir/lib/Dialect/Tensor/Utils/Utils.cpp        | 11 +++--------
 2 files changed, 7 insertions(+), 12 deletions(-)

diff --git a/mlir/include/mlir/Dialect/Tensor/Utils/Utils.h b/mlir/include/mlir/Dialect/Tensor/Utils/Utils.h
index e63749eb384316..84d06d456bb689 100644
--- a/mlir/include/mlir/Dialect/Tensor/Utils/Utils.h
+++ b/mlir/include/mlir/Dialect/Tensor/Utils/Utils.h
@@ -14,10 +14,10 @@
 namespace mlir {
 namespace tensor {
 
-// Return a PadOp that pads `source` to `type` size. Output sizes (from `type`)
-// are assumed to be static and greater than the potentially dynamic input sizes
-// (from `source). The op performs "high" padding (i.e. it adds trailing padding
-// values until the desired size is met).
+// Return a PadOp that pads `source` to `type` size where the static
+// sizes are assumed to be greater than the dynamic sizes. If `type` has dynamic
+// dimensions the padding width is set to zero. The op performs "high" padding
+// (i.e. it adds trailing padding values until the desired size is met).
 PadOp createPadHighOp(RankedTensorType type, Value source, Value pad,
                       bool nofold, Location loc, OpBuilder &builder);
 
diff --git a/mlir/lib/Dialect/Tensor/Utils/Utils.cpp b/mlir/lib/Dialect/Tensor/Utils/Utils.cpp
index 0cb16c28b829c2..a0d8a08fc6ba47 100644
--- a/mlir/lib/Dialect/Tensor/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/Tensor/Utils/Utils.cpp
@@ -24,17 +24,12 @@ using namespace mlir::tensor;
 PadOp mlir::tensor::createPadHighOp(RankedTensorType type, Value source,
                                     Value pad, bool nofold, Location loc,
                                     OpBuilder &b) {
-
-  // TODO: Either relax or turn this into a failure
-  assert(!ShapedType::isDynamicShape(type.getShape()) &&
-         "The output type is dynamic - that's not supported ATM.");
-
-  // Init "low" and "high" padding values ("low" is kept as is, "high" is
-  // computed below).
   SmallVector<OpFoldResult> low(type.getRank(), b.getIndexAttr(0));
   SmallVector<OpFoldResult> high(type.getRank(), b.getIndexAttr(0));
-
   for (const auto &en : enumerate(type.getShape())) {
+    // Pad only the static dimensions of the result tensor type.
+    if (ShapedType::isDynamic(en.value()))
+      continue;
     // Compute the padding width.
     AffineExpr d0;
     bindDims(b.getContext(), d0);