[flang] Shift argument attributes when result operand is inserted

[clementval]

Summary

The TargetRewrite pass convert the signature of the function.
In some cases it adds operands to the function to hanlde the result of it.
This patch makes sure the argument attributes present before the conversion
are replaced with the correct arguments after the conversion is performed.

Depends D132113

[clementval]

Rebase

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>> savedAttrs;
     llvm::SmallVector<std::pair<unsigned, mlir::NamedAttribute>> extraAttrs;
     llvm::SmallVector<FixupTy> fixups;
 
+    // Save argument attributes in case there is a shift so we can replace them
+    // correctly.
+    for (auto e : llvm::enumerate(funcTy.getInputs())) {
+      unsigned index = e.index();
+      llvm::ArrayRef<mlir::NamedAttribute> attrs = func.getArgAttrs(index);
+      for (mlir::NamedAttribute attr : attrs) {
+        savedAttrs.push_back({index, attr});
+      }
+    }
+
     // 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);
             else
               doComplexReturn(func, cmplx, newResTys, newInTys, fixups);
           })
           .Case<mlir::ComplexType>([&](mlir::ComplexType cmplx) {
             if (noComplexConversion)
               newResTys.push_back(cmplx);
             else
               doComplexReturn(func, cmplx, newResTys, newInTys, fixups);
           })
           .Default([&](mlir::Type ty) { newResTys.push_back(ty); });
 
+    // Saved potential shift in argument. Handling of result can add arguments
+    // at the beginning of the function signature.
+    unsigned argumentShift = newInTys.size();
+
     // Convert arguments
     llvm::SmallVector<mlir::Type> trailingTys;
     for (auto e : llvm::enumerate(funcTy.getInputs())) {
       auto ty = e.value();
       unsigned index = e.index();
       llvm::TypeSwitch<mlir::Type>(ty)
           .Case<fir::BoxCharType>([&](fir::BoxCharType boxTy) {
             if (noCharacterConversion) {
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());
 
+    // Replace attributes to the correct argument if there was an argument shift
+    // to the right.
+    if (argumentShift > 0) {
+      for (std::pair<unsigned, mlir::NamedAttribute> savedAttr : savedAttrs) {
+        func.removeArgAttr(savedAttr.first, savedAttr.second.getName());
+        func.setArgAttr(savedAttr.first + argumentShift,
+                        savedAttr.second.getName(),
+                        savedAttr.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
 
+// Test with an argument shift.
+
 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}) {
+// CHECK-SAME:    %{{.*}}: !fir.ref<tuple<!fir.real<16>, !fir.real<16>>> {llvm.align = 16 : i32, llvm.sret}, %arg1: !fir.ref<tuple<!fir.ref<i32>>> {fir.host_assoc, llvm.nest}) {
+
+// -----
+
+// Test with no shift.
+
+func.func @_QFPs(%arg0: !fir.ref<i32> {fir.host_assoc}) {
+  return
+}
+
+// CHECK: func.func @_QFPs(%arg0: !fir.ref<i32> {fir.host_assoc, llvm.nest})