[flang] Set attribute at the correct position

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.

Reviewed By: PeteSteinfeld

Differential Revision: https://reviews.llvm.org/D132113

diff --git a/flang/lib/Optimizer/CodeGen/TargetRewrite.cpp b/flang/lib/Optimizer/CodeGen/TargetRewrite.cpp
index 30df227..d87701d 100644 (file)
--- a/flang/lib/Optimizer/CodeGen/TargetRewrite.cpp
+++ b/flang/lib/Optimizer/CodeGen/TargetRewrite.cpp
@@ -475,6 +475,7 @@ public:
       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)
@@ -552,9 +553,12 @@ public:
             }
           })
           .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())});
       }
     }
 
@@ -716,6 +720,10 @@ public:
     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);

diff --git a/flang/test/Fir/target-rewrite-arg-position.fir b/flang/test/Fir/target-rewrite-arg-position.fir
new file mode 100644 (file)
index 0000000..1b01e73
--- /dev/null
+++ b/flang/test/Fir/target-rewrite-arg-position.fir
@@ -0,0 +1,17 @@
+// 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}) {