Skip to content

[mlir][Func] Preserve attribute when converting CallOp/ReturnOp signature #127772

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

[mlir][Func] Preserve attribute when converting CallOp/ReturnOp signature #127772

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
6 changes: 4 additions & 2 deletions mlir/lib/Dialect/Func/Transforms/FuncConversions.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,7 @@ struct CallOpSignatureConversion : public OpConversionPattern<CallOp> {
auto newCallOp = rewriter.create<CallOp>(
callOp.getLoc(), callOp.getCallee(), convertedResults,
flattenValues(adaptor.getOperands()));
newCallOp->setAttrs(callOp->getAttrs());
SmallVector<ValueRange> replacements;
size_t offset = 0;
for (int i = 0, e = callOp->getNumResults(); i < e; ++i) {
Expand Down Expand Up @@ -126,8 +127,9 @@ class ReturnOpTypeConversion : public OpConversionPattern<ReturnOp> {
LogicalResult
matchAndRewrite(ReturnOp op, OneToNOpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const final {
rewriter.replaceOpWithNewOp<ReturnOp>(op,
flattenValues(adaptor.getOperands()));
rewriter
.replaceOpWithNewOp<ReturnOp>(op, flattenValues(adaptor.getOperands()))
->setAttrs(op->getAttrs());
Copy link
Collaborator

Choose a reason for hiding this comment

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

It's not clear to me that this is always safe: can't the attribute be irrelevant or incorrect post-transformation? For example the number of operands may change but the attribute refers to them positionally.

Copy link
Member Author

Choose a reason for hiding this comment

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

I think it is the responsibility of downstream dialect conversion user to maintain what they want: it is them that provide the type converter and call this type conversion pass. If the attribute accompany the call op does change its meaning after type conversion, the user could add another rewrite pattern to correct them.

The pass itself should not silently discard the attribute, instead it should let the user decide. It is more convenient for the user to see the original attribute after the conversion, than to save these attributes in case they might be lost, and recover it after the standard type conversion.

Copy link
Collaborator

@joker-eph joker-eph Feb 19, 2025 •
edited
Loading

Choose a reason for hiding this comment

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

I think it is the responsibility of downstream dialect conversion user to maintain what they want: it is them that provide the type converter and call this type conversion pass. If the attribute accompany the call op does change its meaning after type conversion, the user could add another rewrite pattern to correct them.

From this point of view, you could say that it is downstream responsibility to gather their attributes before the pass and reapply them afterward.

The pass itself should not silently discard the attribute,

I'm not sure why you believe that: these are called "discardable attributes" because the intent is that transformations should discard them when they don't "understand" them (through interfaces or other mechanism).

Copy link
Member Author

Choose a reason for hiding this comment

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

Or more conversatively we can preserve some dialect attribute?

newOp->setDialectAttrs(op->getDialectAttrs())

transformations should discard them when they don't "understand" them (through interfaces or other mechanism).

I have not seen some mechanism for the transformation to query that. For dialect attribute the transformation could query the dialect namedAttr.getName().getDialect() but I have not found dialect interface / attribute interface that describe such information.

Copy link
Collaborator

Choose a reason for hiding this comment

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

Or more conversatively we can preserve some dialect attribute?

"dialect attributes" are just a different way to say "discardable attributes" I believe. The API you mention don't exist, it is called setDiscardableAttrs() in the codebase.

Copy link
Collaborator

Choose a reason for hiding this comment

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

@matthias-springer : that's an interesting idea, but note that this wouldn't work with most passes which don't use the greedy driver.

Copy link
Member

@matthias-springer matthias-springer Mar 5, 2025 •
edited
Loading

Choose a reason for hiding this comment

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

It wouldn't work with any upstream MLIR passes. This can only be done in a downstream pass. E.g., there's no way to inject attribute conversion rules into the upstream -canonicalize pass. (Unless we expose it as a pass option that can only be set from C++ but not from the command line.)

Instead, users would write their own canonicalizer pass that has a custom greedy config with these conversion rules.

Interestingly, so far, this issue has only come up with patterns.

Copy link
Member Author

Choose a reason for hiding this comment

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

FYI, downstream I am working on now unsafely store and load back the attribute in func.call to work around this issue, with the implicit assumption that walk order is unchanged and no call op is inserted/deleted.

https://github.com/google/heir/blob/b2ab9dcaac644ab8131ed06c4767251f46cf02d8/lib/Dialect/LWE/Conversions/LWEToLattigo/LWEToLattigo.cpp#L511-L540

Copy link
Member

Choose a reason for hiding this comment

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

I've seen too many workarounds related to discardable attribute handling. We also have some in our code base. I will discuss this with a few more people and send a prototype+RFC for the attribute converter idea if it works out.

Copy link
Collaborator

Choose a reason for hiding this comment

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

If you do so, please account for the previous RFC and discussions on the topic on Discourse, in particular we had an ODM on the topic: https://discourse.llvm.org/t/odm-tomorrow-2-4-discuss-the-role-of-dialect-attributes-merging-and-propagation/2725

return success();
}
};
Expand Down
16 changes: 16 additions & 0 deletions mlir/test/Transforms/test-legalizer.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -398,6 +398,22 @@ func.func @caller() {

// -----

module {
// CHECK-LABEL: func.func private @callee()
func.func private @callee() -> (i24)

// CHECK: func.func @call_op_attr_preserved()
func.func @call_op_attr_preserved() {
// i24 is converted to ().
// CHECK: call @callee() {dialect.attr = 1 : i64} : () -> ()
%0 = func.call @callee() {dialect.attr = 1 : i64} : () -> (i24)

"test.return"() : () -> ()
}
}

// -----

// CHECK-LABEL: func @test_move_op_before_rollback()
func.func @test_move_op_before_rollback() {
// CHECK: "test.one_region_op"()
Expand Down