[mlir][tensor] Refine the semantics of `createPadHighOp` #109667

banach-space · 2024-09-23T14:22:29Z

Refine createPadHighOp so that the output tensor is required to be
statically shaped. This is to prevent the current behaviour, which is
incorrect:

// If type has dynamic dimensions the padding width is set to zero.

The actual padding width should be set to: %new_dim - %old_dim, where
%new_dimand%old_dimare defined via e.g.tensor.dim` Op applied to
output and input tensors, respectively.

This PR is an attempt to clarify the semantics surrounding dynamic
shapes in preparation for adding support for scalable vectors to the
pack/unpack logic in Tensor/Linalg (dynamic shapes is what we use to
model scalable (*) sizes at the Tensor/MemRef level).

(*) Scalable as in Arm's Scalable Vector Extension (SVE)

llvmbot · 2024-09-23T14:23:04Z

@llvm/pr-subscribers-mlir

Author: Andrzej Warzyński (banach-space)

Changes

Refine createPadHighOp so that the output tensor is required to be
statically shaped. This is to prevent the current behaviour, which is
incorrect:

> // If type has dynamic dimensions the padding width is set to zero.

The actual padding width should be set to: %new_dim - %old_dim, where
%new_dimand%old_dimare defined via e.g.tensor.dim` Op applied to
output and input tensors, respectively.

This PR is an attempt to clarify the semantics surrounding dynamic
shapes in preparation for adding support for scalable vectors to the
pack/unpack logic in Tensor/Linalg (dynamic shapes is what we use to
model scalable (*) sizes at the Tensor/MemRef level).

(*) Scalable as in Arm's Scalable Vector Extension (SVE)

Full diff: https://github.com/llvm/llvm-project/pull/109667.diff

2 Files Affected:

(modified) mlir/include/mlir/Dialect/Tensor/Utils/Utils.h (+4-4)
(modified) mlir/lib/Dialect/Tensor/Utils/Utils.cpp (+7-3)

diff --git a/mlir/include/mlir/Dialect/Tensor/Utils/Utils.h b/mlir/include/mlir/Dialect/Tensor/Utils/Utils.h
index 84d06d456bb689..e63749eb384316 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 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).
+// 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).
 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 a0d8a08fc6ba47..c8e0c05bfb2b87 100644
--- a/mlir/lib/Dialect/Tensor/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/Tensor/Utils/Utils.cpp
@@ -24,12 +24,16 @@ using namespace mlir::tensor;
 PadOp mlir::tensor::createPadHighOp(RankedTensorType type, Value source,
                                     Value pad, bool nofold, Location loc,
                                     OpBuilder &b) {
+
+  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);

llvmbot · 2024-09-23T14:23:05Z

@llvm/pr-subscribers-mlir-tensor

Author: Andrzej Warzyński (banach-space)

Changes

Refine createPadHighOp so that the output tensor is required to be
statically shaped. This is to prevent the current behaviour, which is
incorrect:

> // If type has dynamic dimensions the padding width is set to zero.

The actual padding width should be set to: %new_dim - %old_dim, where
%new_dimand%old_dimare defined via e.g.tensor.dim` Op applied to
output and input tensors, respectively.

This PR is an attempt to clarify the semantics surrounding dynamic
shapes in preparation for adding support for scalable vectors to the
pack/unpack logic in Tensor/Linalg (dynamic shapes is what we use to
model scalable (*) sizes at the Tensor/MemRef level).

(*) Scalable as in Arm's Scalable Vector Extension (SVE)

Full diff: https://github.com/llvm/llvm-project/pull/109667.diff

2 Files Affected:

(modified) mlir/include/mlir/Dialect/Tensor/Utils/Utils.h (+4-4)
(modified) mlir/lib/Dialect/Tensor/Utils/Utils.cpp (+7-3)

diff --git a/mlir/include/mlir/Dialect/Tensor/Utils/Utils.h b/mlir/include/mlir/Dialect/Tensor/Utils/Utils.h
index 84d06d456bb689..e63749eb384316 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 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).
+// 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).
 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 a0d8a08fc6ba47..c8e0c05bfb2b87 100644
--- a/mlir/lib/Dialect/Tensor/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/Tensor/Utils/Utils.cpp
@@ -24,12 +24,16 @@ using namespace mlir::tensor;
 PadOp mlir::tensor::createPadHighOp(RankedTensorType type, Value source,
                                     Value pad, bool nofold, Location loc,
                                     OpBuilder &b) {
+
+  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);

MaheshRavishankar · 2024-09-24T13:24:18Z

mlir/lib/Dialect/Tensor/Utils/Utils.cpp

@@ -24,12 +24,16 @@ using namespace mlir::tensor;
 PadOp mlir::tensor::createPadHighOp(RankedTensorType type, Value source,
                                    Value pad, bool nofold, Location loc,
                                    OpBuilder &b) {
+
+  assert(!ShapedType::isDynamicShape(type.getShape()) &&


Instead cant we make the semantics here that the dynamic shapes of the output are same as the dynamic shapes of the input?

Instead cant we make the semantics here that the dynamic shapes of the output are same as the dynamic shapes of the input?

How would we verify that these are indeed identical?

Also, what should happen when the input is static, and the output is dynamic? This is one specific example that I have in mind right now:

func.func @simple_pad_and_pack_scalable(%input: tensor<5x1xf32>, %output: tensor<1x1x?x2xf32>, %pad: f32) -> tensor<1x1x?x2xf32> { %c8 = arith.constant 8 : index %vscale = vector.vscale %c8_vscale = arith.muli %vscale, %c8 : index %0 = tensor.pack %input padding_value(%pad : f32) inner_dims_pos = [0, 1] inner_tiles = [%c8_vscale, 2] into %output : tensor<5x1xf32> -> tensor<1x1x?x2xf32> return %0 : tensor<1x1x?x2xf32> }

That's based on the example from #109815. Overall, there's quite a few cases to consider (and to test):

input static, output static (easy)

input dynamic, output static (tricky)

input static, output dynamic (tricky)

input dynamic, output dynamic (tricky)

I'd rather restrict the semantics to begin with, and then gradually relax them when there's need (e.g. for #109815). No strong opinion though, whatever works and is correct :)

Instead cant we make the semantics here that the dynamic shapes of the output are same as the dynamic shapes of the input?

How would we verify that these are indeed identical?

I think its a matter of defining it. You can say that the createPadHighOp will effectively do low and high for dynamic shape to be 0. So for dynamic dims the input dim extent and output dim extent would be the same.

Also, what should happen when the input is static, and the output is dynamic? This is one specific example that I have in mind right now:

func.func @simple_pad_and_pack_scalable(%input: tensor<5x1xf32>, %output: tensor<1x1x?x2xf32>, %pad: f32) -> tensor<1x1x?x2xf32> { %c8 = arith.constant 8 : index %vscale = vector.vscale %c8_vscale = arith.muli %vscale, %c8 : index %0 = tensor.pack %input padding_value(%pad : f32) inner_dims_pos = [0, 1] inner_tiles = [%c8_vscale, 2] into %output : tensor<5x1xf32> -> tensor<1x1x?x2xf32> return %0 : tensor<1x1x?x2xf32> }

Thats a pack operation? I am not following your point.

That's based on the example from #109815. Overall, there's quite a few cases to consider (and to test):

input static, output static (easy)
+1

input dynamic, output static (tricky)
You can compute the high padding needed as output dim - input dim, and the high padding can be dynamic.

input static, output dynamic (tricky)

input dynamic, output dynamic (tricky)

I'd rather restrict the semantics to begin with, and then gradually relax them when there's need (e.g. for #109815). No strong opinion though, whatever works and is correct :)

Ok, that makes sense. Let me re-review.

Maybe dont assert and instead return an error?

Adding an error complicates things for the clients of this API, and I will be relaxing this shortly. Let me leave a TODO for myself and I'll follow-up shortly.

Refine `createPadHighOp` so that the output tensor is required to be statically shaped. This is to prevent the current behaviour, which is incorrect: > // If `type` has dynamic dimensions the padding width is set to zero. The actual padding width should be set to: `%new_dim - %old_dim`, where %new_dim` and `%old_dim` are defined via e.g. `tensor.dim` Op applied to output and input tensors, respectively. This PR is an attempt to clarify the semantics surrounding dynamic shapes in preparation for adding support for scalable vectors to the pack/unpack logic in Tensor/Linalg (dynamic shapes is what we use to model scalable (*) sizes at the Tensor/MemRef level). (*) Scalable as in Arm's Scalable Vector Extension (SVE)

MaheshRavishankar · 2024-09-25T16:53:20Z

mlir/lib/Dialect/Tensor/Utils/Utils.cpp

@@ -24,12 +24,16 @@ using namespace mlir::tensor;
 PadOp mlir::tensor::createPadHighOp(RankedTensorType type, Value source,
                                    Value pad, bool nofold, Location loc,
                                    OpBuilder &b) {
+
+  assert(!ShapedType::isDynamicShape(type.getShape()) &&


Maybe dont assert and instead return an error?

Add a TODO

hanhanW · 2024-09-26T18:15:37Z

This breaks the behavior of linalg::padAndHoistLinalgOp, which calls makeComposedPadHighOp here. The implementation of makeComposedPadHighOp relies on the behavior to create valid pad ops.

In the transform script, it asks for padding dim 3, so it does not care about the padding values of other dimension and use the createPadHighOp method to generate the pad op.

Can we revert it until it is fixed?

To repro: mlir-opt --transform-interpreter ~/repro.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
    }
  }
}

…m#109667)" This reverts commit 9c48a04.

…9667)" This reverts commit 9c48a04.

Refine `createPadHighOp` so that the output tensor is required to be statically shaped. This is to prevent the current behaviour, which is incorrect: > // If `type` has dynamic dimensions the padding width is set to zero. The actual padding width should be set to: `%new_dim - %old_dim`, where %new_dim` and `%old_dim` are defined via e.g. `tensor.dim` Op applied to output and input tensors, respectively. This PR is an attempt to clarify the semantics surrounding dynamic shapes in preparation for adding support for scalable vectors to the pack/unpack logic in Tensor/Linalg (dynamic shapes is what we use to model scalable (*) sizes at the Tensor/MemRef level). (*) Scalable as in Arm's Scalable Vector Extension (SVE)

Allow multiple dynamic dims and remove the dependency on llvm#109667

llvmbot added mlir mlir:tensor labels Sep 23, 2024

MaheshRavishankar reviewed Sep 24, 2024

View reviewed changes

banach-space force-pushed the andrzej/restrict_create_pad_high branch from 84fcd1a to 7a5a875 Compare September 24, 2024 14:23

banach-space mentioned this pull request Sep 24, 2024

[mlir][tensor] Extend the logic to generalise tensor.pack #109815

Merged

banach-space requested a review from qedawkins September 24, 2024 15:55

MaheshRavishankar approved these changes Sep 25, 2024

View reviewed changes

fixup! [mlir][tensor] Refine the semantics of createPadHighOp

8751e42

Add a TODO

banach-space merged commit 9c48a04 into llvm:main Sep 26, 2024
6 of 7 checks passed

hanhanW added a commit to iree-org/llvm-project that referenced this pull request Sep 26, 2024

Revert "[mlir][tensor] Refine the semantics of createPadHighOp (llv…

6580227

…m#109667)" This reverts commit 9c48a04.

hanhanW added a commit that referenced this pull request Sep 26, 2024

Revert "[mlir][tensor] Refine the semantics of createPadHighOp (#10…

61e03b6

…9667)" This reverts commit 9c48a04.

hanhanW mentioned this pull request Sep 26, 2024

Revert "[mlir][tensor] Refine the semantics of createPadHighOp" #110153

Merged

banach-space added a commit to banach-space/llvm-project that referenced this pull request Sep 30, 2024

fixup! fixup! [mlir][tensor] Extend the logic to generalise tensor.pack

b3da00e

Allow multiple dynamic dims and remove the dependency on llvm#109667

banach-space deleted the andrzej/restrict_create_pad_high branch October 1, 2024 18:54

banach-space added a commit to banach-space/llvm-project that referenced this pull request Oct 2, 2024

fixup! fixup! [mlir][tensor] Extend the logic to generalise tensor.pack

bfa6a2d

Allow multiple dynamic dims and remove the dependency on llvm#109667

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[mlir][tensor] Refine the semantics of `createPadHighOp` #109667

[mlir][tensor] Refine the semantics of `createPadHighOp` #109667

Uh oh!

banach-space commented Sep 23, 2024

Uh oh!

llvmbot commented Sep 23, 2024

Uh oh!

llvmbot commented Sep 23, 2024

Uh oh!

MaheshRavishankar Sep 24, 2024

Uh oh!

banach-space Sep 24, 2024

Uh oh!

MaheshRavishankar Sep 25, 2024

Uh oh!

MaheshRavishankar Sep 25, 2024

Uh oh!

banach-space Sep 26, 2024

Uh oh!

MaheshRavishankar Sep 25, 2024

Uh oh!

Uh oh!

hanhanW commented Sep 26, 2024

Uh oh!

Uh oh!

[mlir][tensor] Refine the semantics of createPadHighOp #109667

[mlir][tensor] Refine the semantics of createPadHighOp #109667

Uh oh!

Conversation

banach-space commented Sep 23, 2024

Uh oh!

llvmbot commented Sep 23, 2024

Uh oh!

llvmbot commented Sep 23, 2024

Uh oh!

MaheshRavishankar Sep 24, 2024

Choose a reason for hiding this comment

Uh oh!

banach-space Sep 24, 2024

Choose a reason for hiding this comment

Uh oh!

MaheshRavishankar Sep 25, 2024

Choose a reason for hiding this comment

Uh oh!

MaheshRavishankar Sep 25, 2024

Choose a reason for hiding this comment

Uh oh!

banach-space Sep 26, 2024

Choose a reason for hiding this comment

Uh oh!

MaheshRavishankar Sep 25, 2024

Choose a reason for hiding this comment

Uh oh!

Uh oh!

hanhanW commented Sep 26, 2024

Uh oh!

Uh oh!

[mlir][tensor] Refine the semantics of `createPadHighOp` #109667

[mlir][tensor] Refine the semantics of `createPadHighOp` #109667