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

[ZenithalHourlyRate]

For func.call and func.return, the attribute should be preserved during function type conversion.

Other func ops are converted using rewriter.modifyOpInplace while these two ops are converted using rewriter.replaceOpWithNewOp so the attribute is left out.

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-func

Author: Hongren Zheng (ZenithalHourlyRate)

Changes

For func.call and func.return, the attribute should be preserved during function type conversion.

Other func ops are converted using rewriter.modifyOpInplace while these two ops are converted using rewriter.replaceOpWithNewOp so the attribute is left out.

---

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

2 Files Affected:

• (modified) mlir/lib/Dialect/Func/Transforms/FuncConversions.cpp (+4-2)
• (modified) mlir/test/Transforms/test-legalizer.mlir (+16)

diff --git a/mlir/lib/Dialect/Func/Transforms/FuncConversions.cpp b/mlir/lib/Dialect/Func/Transforms/FuncConversions.cpp
index a3638c8766a5c..41a083fb78ff4 100644
--- a/mlir/lib/Dialect/Func/Transforms/FuncConversions.cpp
+++ b/mlir/lib/Dialect/Func/Transforms/FuncConversions.cpp
@@ -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) {
@@ -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());
     return success();
   }
 };
diff --git a/mlir/test/Transforms/test-legalizer.mlir b/mlir/test/Transforms/test-legalizer.mlir
index ae7d344b7167f..4ed01b7392073 100644
--- a/mlir/test/Transforms/test-legalizer.mlir
+++ b/mlir/test/Transforms/test-legalizer.mlir
@@ -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"()

[joker-eph]

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.

[ZenithalHourlyRate]

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.

[joker-eph]

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

[ZenithalHourlyRate]

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.

[joker-eph]

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.

[joker-eph]

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

[matthias-springer]

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.

[ZenithalHourlyRate]

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

[matthias-springer]

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.

[joker-eph]

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

[ZenithalHourlyRate]

ping for update