[flang] Set attribute at the correct position

[clementval]

Summary

The TargetRewrite pass can change the number of argument of a function.
An extra llvm.nest attribute is added and was not set at the correct position
if an extra argument was inserted before.

[PeteSteinfeld]

All builds and tests correctly and looks good.

flang/lib/Optimizer/CodeGen/TargetRewrite.cpp

public:
   /// Rewrite the signatures and body of the `FuncOp`s in the module for
   /// the immediately subsequent target code gen.
   void convertSignature(mlir::func::FuncOp func) {
     auto funcTy = func.getFunctionType().cast<mlir::FunctionType>();
     if (hasPortableSignature(funcTy) && !hasHostAssociations(func))
       return;
     llvm::SmallVector<mlir::Type> newResTys;
     llvm::SmallVector<mlir::Type> newInTys;
+    llvm::SmallVector<std::pair<unsigned, mlir::NamedAttribute>> extraAttrs;
     llvm::SmallVector<FixupTy> fixups;
 
     // Convert return value(s)
     for (auto ty : funcTy.getResults())
       llvm::TypeSwitch<mlir::Type>(ty)
           .Case<fir::ComplexType>([&](fir::ComplexType cmplx) {
             if (noComplexConversion)
               newResTys.push_back(cmplx);
for (auto e : llvm::enumerate(funcTy.getInputs())) {
                                   newInTys.size(), trailingTys.size());
               newInTys.push_back(tuple.getType(0));
               trailingTys.push_back(tuple.getType(1));
             } else {
               newInTys.push_back(ty);
             }
           })
           .Default([&](mlir::Type ty) { newInTys.push_back(ty); });
+
       if (func.getArgAttrOfType<mlir::UnitAttr>(index,
                                                 fir::getHostAssocAttrName())) {
-        func.setArgAttr(index, "llvm.nest", rewriter->getUnitAttr());
+        extraAttrs.push_back(
+            {newInTys.size() - 1,
+             rewriter->getNamedAttr("llvm.nest", rewriter->getUnitAttr())});
       }
     }
 
     if (!func.empty()) {
       // If the function has a body, then apply the fixups to the arguments and
       // return ops as required. These fixups are done in place.
       auto loc = func.getLoc();
       const auto fixupSize = fixups.size();
void convertSignature(mlir::func::FuncOp func) {
 
     // Set the new type and finalize the arguments, etc.
     newInTys.insert(newInTys.end(), trailingTys.begin(), trailingTys.end());
     auto newFuncTy =
         mlir::FunctionType::get(func.getContext(), newInTys, newResTys);
     LLVM_DEBUG(llvm::dbgs() << "new func: " << newFuncTy << '\n');
     func.setType(newFuncTy);
 
+    for (std::pair<unsigned, mlir::NamedAttribute> extraAttr : extraAttrs)
+      func.setArgAttr(extraAttr.first, extraAttr.second.getName(),
+                      extraAttr.second.getValue());
+
     for (auto &fixup : fixups)
       if (fixup.finalizer)
         (*fixup.finalizer)(func);
   }
 
   inline bool functionArgIsSRet(unsigned index, mlir::func::FuncOp func) {
     if (auto attr = func.getArgAttrOfType<mlir::UnitAttr>(index, "llvm.sret"))
       return true;

flang/test/Fir/target-rewrite-arg-position.fir

+// RUN: fir-opt --target-rewrite="target=x86_64-unknown-linux-gnu" %s | FileCheck %s
+
+func.func @_QFPf(%arg0: !fir.ref<tuple<!fir.ref<i32>>> {fir.host_assoc}) -> !fir.complex<16> {
+  %0 = fir.alloca !fir.complex<16> {bindc_name = "f", uniq_name = "_QFfEf"}
+  %c2_i32 = arith.constant 2 : i32
+  %1 = fir.convert %c2_i32 : (i32) -> f128
+  %cst = arith.constant 0.000000e+00 : f128
+  %2 = fir.undefined !fir.complex<16>
+  %3 = fir.insert_value %2, %1, [0 : index] : (!fir.complex<16>, f128) -> !fir.complex<16>
+  %4 = fir.insert_value %3, %cst, [1 : index] : (!fir.complex<16>, f128) -> !fir.complex<16>
+  fir.store %4 to %0 : !fir.ref<!fir.complex<16>>
+  %5 = fir.load %0 : !fir.ref<!fir.complex<16>>
+  return %5 : !fir.complex<16>
+}
+
+// CHECK-LABEL: func.func @_QFPf
+// CHECK-SAME:    %{{.*}}: !fir.ref<tuple<!fir.real<16>, !fir.real<16>>> {fir.host_assoc, llvm.align = 16 : i32, llvm.sret}, %arg1: !fir.ref<tuple<!fir.ref<i32>>> {llvm.nest}) {