From 732ebeda0f1328f0aea1a9d3cc78c728e659cef8 Mon Sep 17 00:00:00 2001 From: Peter Bell 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 +convertFuncOpAttrs(FunctionOpInterface funcOp, + TypeConverter::SignatureConversion &sigConv, + FunctionType newType) { + if (newType.getNumInputs() == funcOp.getNumArguments()) { + return {}; + } + ArrayAttr allArgAttrs = funcOp.getAllArgAttrs(); + if (!allArgAttrs) + return {}; + + SmallVector 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 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);