Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Sep 26, 2024
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
7a5a875
[mlir][tensor] Refine the semantics of `createPadHighOp`
banach-space Sep 23, 2024
8751e42
fixup! [mlir][tensor] Refine the semantics of `createPadHighOp`
banach-space Sep 26, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 4 additions & 4 deletions mlir/include/mlir/Dialect/Tensor/Utils/Utils.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);

Expand Down
11 changes: 8 additions & 3 deletions mlir/lib/Dialect/Tensor/Utils/Utils.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,12 +24,17 @@ 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()) &&
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 :)

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe dont assert and instead return an error?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

"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);
Expand Down
Loading