diff --git a/flang/lib/Optimizer/CodeGen/TargetRewrite.cpp b/flang/lib/Optimizer/CodeGen/TargetRewrite.cpp --- a/flang/lib/Optimizer/CodeGen/TargetRewrite.cpp +++ b/flang/lib/Optimizer/CodeGen/TargetRewrite.cpp @@ -475,6 +475,7 @@ return; llvm::SmallVector newResTys; llvm::SmallVector newInTys; + llvm::SmallVector> extraAttrs; llvm::SmallVector fixups; // Convert return value(s) @@ -552,9 +553,12 @@ } }) .Default([&](mlir::Type ty) { newInTys.push_back(ty); }); + if (func.getArgAttrOfType(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 @@ LLVM_DEBUG(llvm::dbgs() << "new func: " << newFuncTy << '\n'); func.setType(newFuncTy); + for (std::pair 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 --- /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>> {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> + %5 = fir.load %0 : !fir.ref> + return %5 : !fir.complex<16> +} + +// CHECK-LABEL: func.func @_QFPf +// CHECK-SAME: %{{.*}}: !fir.ref, !fir.real<16>>> {fir.host_assoc, llvm.align = 16 : i32, llvm.sret}, %arg1: !fir.ref>> {llvm.nest}) {