[Mlir-commits] [mlir] [mlir][Func] Preserve attribute when converting CallOp/ReturnOp signature (PR #127772)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-func

Author: Hongren Zheng (ZenithalHourlyRate)

<details>
<summary>Changes</summary>

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
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"()

``````````

</details>


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