Skip to content

Fix support for complex types in transform.structured.pad #139841

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

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Fix support for complex types in transform.structured.pad #139841

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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: 12 additions & 4 deletions mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1901,9 +1901,10 @@ transform::PadOp::apply(transform::TransformRewriter &rewriter,
SmallVector<Attribute> paddingValues;
for (auto const &it :
llvm::zip(getPaddingValues(), linalgTarget->getOperandTypes())) {
auto attr = dyn_cast<TypedAttr>(std::get<0>(it));
if (!attr) {
emitOpError("expects padding values to be typed attributes");
Attribute attr = std::get<0>(it);
if (!llvm::isa<TypedAttr, ArrayAttr>(attr)) {
emitOpError("expects padding values to be typed attributes or array "
"attributes (for complex numbers)");
return DiagnosedSilenceableFailure::definiteFailure();
}
Type elementType = getElementTypeOrSelf(std::get<1>(it));
Expand All @@ -1922,7 +1923,14 @@ transform::PadOp::apply(transform::TransformRewriter &rewriter,
continue;
}
// Otherwise, add the attribute directly.
if (attr.getType() != elementType) {
if (isa<TypedAttr>(attr) &&
cast<TypedAttr>(attr).getType() != elementType) {
Comment on lines +1926 to +1927
Copy link
Member

Choose a reason for hiding this comment

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

isa followed by cast is an explicit anti-pattern in LLVM, please use dyn_cast combined with c++17 feature of declaring a value inside if.

auto diag = this->emitOpError("expects a padding value of type ")
<< elementType << ", got " << attr;
diag.attachNote(linalgTarget.getLoc()) << "when applied to this op";
return DiagnosedSilenceableFailure::definiteFailure();
}
if (isa<ArrayAttr>(attr) && !isa<ComplexType>(elementType)) {
auto diag = this->emitOpError("expects a padding value of type ")
<< elementType << ", got " << attr;
diag.attachNote(linalgTarget.getLoc()) << "when applied to this op";
Expand Down
37 changes: 37 additions & 0 deletions mlir/test/Dialect/Linalg/transform-op-pad.mlir
View file
Open in desktop
Copy link
Contributor

Choose a reason for hiding this comment

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

Rather than adding a completely new test case, could you re-use or modify one of the existing ones? Otherwise, it’s harder to see what makes the new test "special" or "unique" - modifying an existing test often makes it easier to highlight the impact of a change.

That said, I don’t see any truly “basic” cases for padding linalg.matmul, so it’s reasonable if you need to create your own. Alternatively, you might be able to re-use one of the pad_to_multiple_of tests. I'm not very familiar with transform.structured.pad, but do you actually rely on pad_to_multiple_of to verify this change?

Separately, please make sure your test is self-documenting - for example, consider renaming @pad_matmul to something more descriptive. You can find relevant testing guidelines here:

Thanks!

Original file line number Diff line number Diff line change
Expand Up @@ -419,3 +419,40 @@ module attributes {transform.with_named_sequence} {
transform.yield
}
}

// -----

!type = tensor<10x10xcomplex<f32>>
// CHECK-LABEL: @pad_matmul
func.func @pad_matmul(%arg0: !type,
%arg1: !type,
%arg2: !type
) -> !type {
// CHECK: complex.constant [{{.*}}] : complex<f32>
// CHECK: tensor.pad
// CHECK: tensor.yield
// CHECK: complex.constant [{{.*}}] : complex<f32>
// CHECK: tensor.pad
// CHECK: tensor.yield
// CHECK: complex.constant [{{.*}}] : complex<f32>
// CHECK: tensor.pad
// CHECK: tensor.yield
// CHECK: linalg.matmul
%0 = linalg.matmul ins(%arg0, %arg1 : !type, !type) outs(%arg2 : !type) -> !type
func.return %0 : !type
}

module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
%0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op
%padded, %pad, %copy_back = transform.structured.pad %0 pad_to_multiple_of [3, 3, 3] {
padding_values=[
[0.1 : f32, 0.2 : f32],
[0.3 : f32, 0.4 : f32],
[0.5 : f32, 0.6 : f32]
],
padding_dimensions = [0, 1, 2]
} : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
transform.yield
}
}