Skip to content

[mlir][tensor] Fix tensor::PackOp fold() handling of padding value #87296

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 3 commits into from
Apr 2, 2024
Merged

[mlir][tensor] Fix tensor::PackOp fold() handling of padding value #87296

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
dc6004e
[mlir][tensor] Fix a bug in constant packing folding.
hanhanW Apr 2, 2024
ac9d0f7
fix compile error
hanhanW Apr 2, 2024
e1d25b0
address comments
hanhanW Apr 2, 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
16 changes: 11 additions & 5 deletions mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1068,10 +1068,13 @@ void EmptyOp::getCanonicalizationPatterns(RewritePatternSet &results,

/// Try to remove a tensor operation if it would only reshape a constant.
/// Removes the op and replaces the constant with a new constant of the result
/// shape.
static OpFoldResult reshapeConstantSource(DenseElementsAttr source,
TensorType result) {
if (source && source.isSplat() && result.hasStaticShape())
/// shape. When an optional cst attribute is passed, it is reshaped only if the
/// splat value matches the value in the attribute.
static OpFoldResult
reshapeConstantSource(DenseElementsAttr source, TensorType result,
std::optional<Attribute> cst = std::nullopt) {
Copy link
Collaborator

Choose a reason for hiding this comment

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

Please update the documentation for this.

if (source && source.isSplat() && result.hasStaticShape() &&
(!cst.has_value() || source.getSplatValue<Attribute>() == cst.value()))
return source.resizeSplat(result);

return {};
Expand Down Expand Up @@ -4143,9 +4146,12 @@ bool PackOp::isLikePad() {
}

OpFoldResult PackOp::fold(FoldAdaptor adaptor) {
std::optional<Attribute> paddingValue;
if (auto pad = adaptor.getPaddingValue())
paddingValue = pad;
if (OpFoldResult reshapedSource = reshapeConstantSource(
llvm::dyn_cast_if_present<DenseElementsAttr>(adaptor.getSource()),
getResult().getType()))
getDestType(), paddingValue))
return reshapedSource;
return {};
}
Expand Down
33 changes: 33 additions & 0 deletions mlir/test/Dialect/Tensor/canonicalize.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -830,6 +830,39 @@ func.func @fold_pack_constant_splat(%dest : tensor<8x16x8x32xf32>) -> tensor<8x1

// -----

// CHECK-LABEL: func @fold_padding_value_pack_constant_splat
// CHECK-NOT: tensor.pack
// CHECK: arith.constant dense<1.000000e-01> : tensor<8x16x8x32xf32>
func.func @fold_padding_value_pack_constant_splat(%dest : tensor<8x16x8x32xf32>) -> tensor<8x16x8x32xf32> {
%pad = arith.constant 1.000000e-01 : f32
%cst = arith.constant dense<1.000000e-01> : tensor<63x127xf32>
%0 = tensor.pack %cst
padding_value(%pad : f32)
outer_dims_perm = [1, 0] inner_dims_pos = [0, 1]
inner_tiles = [8, 32] into %dest : tensor<63x127xf32> -> tensor<8x16x8x32xf32>
return %0 : tensor<8x16x8x32xf32>
}


// -----

// CHECK-LABEL: func @nofold_padding_value_pack_constant_splat
// CHECK: arith.constant dense<1.000000e-01> : tensor<63x127xf32>
// CHECK: tensor.pack
func.func @nofold_padding_value_pack_constant_splat(%dest : tensor<8x16x8x32xf32>) -> tensor<8x16x8x32xf32> {
%pad = arith.constant 0.0 : f32
%cst = arith.constant dense<1.000000e-01> : tensor<63x127xf32>
%0 = tensor.pack %cst
padding_value(%pad : f32)
outer_dims_perm = [1, 0]
inner_dims_pos = [0, 1]
inner_tiles = [8, 32]
into %dest : tensor<63x127xf32> -> tensor<8x16x8x32xf32>
return %0 : tensor<8x16x8x32xf32>
}

// -----

func.func @fold_padding_value_pack(%arg0: tensor<1200x500000xf32>) -> tensor<31250x1200x16x1xf32> {
%cst = arith.constant 0.000000e+00 : f32
%0 = tensor.empty() : tensor<31250x1200x16x1xf32>
Expand Down