[Mlir-commits] [mlir] [mlir][Transforms] Handle attributes in 1-to-many function conversions (PR #162579)

[llvmlistbot at llvm.org]

https://github.com/peterbell10 updated https://github.com/llvm/llvm-project/pull/162579

>From 732ebeda0f1328f0aea1a9d3cc78c728e659cef8 Mon Sep 17 00:00:00 2001
From: Peter Bell <peterbell10 at openai.com>
Date: Thu, 9 Oct 2025 02:13:11 +0100
Subject: [PATCH 1/2] [mlir][Transforms] Handle attributes in 1-to-many
 function conversions

Function attributes are associated to the argument index, which may
change during a 1-to-many type conversion but is currently not being
updated by the conversion pattern.

The results is that attributes get applied to the wrong arguments after
the conversion.

Happy to add a lit test, if you can suggest a good place for it.
---
 .../Transforms/Utils/DialectConversion.cpp    | 32 ++++++++++++++++++-
 1 file changed, 31 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Transforms/Utils/DialectConversion.cpp b/mlir/lib/Transforms/Utils/DialectConversion.cpp
index 3a23bbfd70eac..ec8d0971f2bbe 100644
--- a/mlir/lib/Transforms/Utils/DialectConversion.cpp
+++ b/mlir/lib/Transforms/Utils/DialectConversion.cpp
@@ -3789,6 +3789,27 @@ TypeConverter::convertTypeAttribute(Type type, Attribute attr) const {
 // FunctionOpInterfaceSignatureConversion
 //===----------------------------------------------------------------------===//
 
+static SmallVector<Attribute>
+convertFuncOpAttrs(FunctionOpInterface funcOp,
+                   TypeConverter::SignatureConversion &sigConv,
+                   FunctionType newType) {
+  if (newType.getNumInputs() == funcOp.getNumArguments()) {
+    return {};
+  }
+  ArrayAttr allArgAttrs = funcOp.getAllArgAttrs();
+  if (!allArgAttrs)
+    return {};
+
+  SmallVector<Attribute> newAttrs(newType.getNumInputs());
+  for (auto i : llvm::seq(allArgAttrs.size())) {
+    auto mapping = sigConv.getInputMapping(i);
+    assert(mapping.has_value());
+    auto outIdx = mapping->inputNo;
+    newAttrs[outIdx] = allArgAttrs[i];
+  }
+  return newAttrs;
+}
+
 static LogicalResult convertFuncOpTypes(FunctionOpInterface funcOp,
                                         const TypeConverter &typeConverter,
                                         ConversionPatternRewriter &rewriter) {
@@ -3809,7 +3830,16 @@ static LogicalResult convertFuncOpTypes(FunctionOpInterface funcOp,
   auto newType = FunctionType::get(rewriter.getContext(),
                                    result.getConvertedTypes(), newResults);
 
-  rewriter.modifyOpInPlace(funcOp, [&] { funcOp.setType(newType); });
+  // If using 1-to-n type conversion, we must re-map argument attributes 
+  // to the corresponding new argument index.
+  auto newArgAttrs = convertFuncOpAttrs(funcOp, result, newType);
+
+  rewriter.modifyOpInPlace(funcOp, [&] {
+    funcOp.setType(newType);
+    if (!newArgAttrs.empty()) {
+      funcOp.setAllArgAttrs(newArgAttrs);
+    }
+  });
 
   return success();
 }

>From ed9821a62eb26b16977874c0ac3ed42a825aebdb Mon Sep 17 00:00:00 2001
From: Peter Bell <peterbell10 at openai.com>
Date: Thu, 9 Oct 2025 09:24:28 +0100
Subject: [PATCH 2/2] Lint

---
 mlir/lib/Transforms/Utils/DialectConversion.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlir/lib/Transforms/Utils/DialectConversion.cpp b/mlir/lib/Transforms/Utils/DialectConversion.cpp
index ec8d0971f2bbe..c7fdef74d33bf 100644
--- a/mlir/lib/Transforms/Utils/DialectConversion.cpp
+++ b/mlir/lib/Transforms/Utils/DialectConversion.cpp
@@ -3830,7 +3830,7 @@ static LogicalResult convertFuncOpTypes(FunctionOpInterface funcOp,
   auto newType = FunctionType::get(rewriter.getContext(),
                                    result.getConvertedTypes(), newResults);
 
-  // If using 1-to-n type conversion, we must re-map argument attributes 
+  // If using 1-to-n type conversion, we must re-map argument attributes
   // to the corresponding new argument index.
   auto newArgAttrs = convertFuncOpAttrs(funcOp, result, newType);