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

[Hongren Zheng]

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

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.

>From a5375e7e37961ee2b30e0a426efe8a8bf7fc872b Mon Sep 17 00:00:00 2001
From: Zenithal <i at zenithal.me>
Date: Wed, 19 Feb 2025 09:17:41 +0000
Subject: [PATCH] [mlir][Func] Preserve attribute when converting
 CallOp/ReturnOp signature

---
 .../Dialect/Func/Transforms/FuncConversions.cpp  |  6 ++++--
 mlir/test/Transforms/test-legalizer.mlir         | 16 ++++++++++++++++
 2 files changed, 20 insertions(+), 2 deletions(-)

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