[flang-commits] [flang] [flang][TargetRewrite] Keep argument attributes consistent after ABI arg shift (PR #208124)

[Slava Zakharin via flang-commits]

https://github.com/vzakhari created https://github.com/llvm/llvm-project/pull/208124

When target-rewrite expands an argument into several arguments
(for example splitting a complex value into two scalars),
the arguments after it shift to the right.
Argument attributes were only moved for a shift introduced
by result lowering, so an attribute on a later argument
(e.g. fir.host_assoc) could stay on its old index and end up
on an unrelated argument. Remap every saved argument
attribute to the new index of its argument.

This is a consistency fix and may not cause a miscompile today.

>From d37a208cacbd54f0a189d088a948d4d97e649d5f Mon Sep 17 00:00:00 2001
From: Slava Zakharin <szakharin at nvidia.com>
Date: Tue, 7 Jul 2026 16:38:53 -0700
Subject: [PATCH] [flang][TargetRewrite] Keep argument attributes consistent
 after ABI arg shift

When target-rewrite expands an argument into several arguments (for example
splitting a complex value into two scalars), the arguments after it shift to
the right. Argument attributes were only moved for a shift introduced by
result lowering, so an attribute on a later argument (e.g. fir.host_assoc)
could stay on its old index and end up on an unrelated argument. Remap every
saved argument attribute to the new index of its argument. This is a
consistency fix and may not cause a miscompile today.
---
 flang/lib/Optimizer/CodeGen/TargetRewrite.cpp | 30 +++++++++++--------
 .../test/Fir/target-rewrite-arg-position.fir  | 14 +++++++++
 2 files changed, 31 insertions(+), 13 deletions(-)

diff --git a/flang/lib/Optimizer/CodeGen/TargetRewrite.cpp b/flang/lib/Optimizer/CodeGen/TargetRewrite.cpp
index 3701e5abff6df..0ce27abce870b 100644
--- a/flang/lib/Optimizer/CodeGen/TargetRewrite.cpp
+++ b/flang/lib/Optimizer/CodeGen/TargetRewrite.cpp
@@ -916,15 +916,16 @@ class TargetRewrite : public fir::impl::TargetRewritePassBase<TargetRewrite> {
           })
           .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 = newInTyAndAttrs.size();
+    // New index of each original argument, so that saved attributes can follow
+    // their argument when target expansion shifts arguments to the right.
+    llvm::SmallVector<unsigned> newArgIndex;
 
     // Convert arguments
     llvm::SmallVector<mlir::Type> trailingTys;
     for (auto e : llvm::enumerate(funcTy.getInputs())) {
       auto ty = e.value();
       unsigned index = e.index();
+      newArgIndex.push_back(newInTyAndAttrs.size());
       llvm::TypeSwitch<mlir::Type>(ty)
           .Case([&](fir::BoxCharType boxTy) {
             if (noCharacterConversion) {
@@ -1187,16 +1188,19 @@ class TargetRewrite : public fir::impl::TargetRewritePassBase<TargetRewrite> {
       for (mlir::NamedAttribute resAttr : resAttrList)
         func.setResultAttr(resId, resAttr.getName(), resAttr.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());
-      }
-    }
+    // Move each saved argument attribute to the index its argument now
+    // occupies. Target expansion (e.g. splitting a complex value into two
+    // scalars) can insert arguments mid-list and shift the following ones; an
+    // attribute such as fir.host_assoc must follow its argument rather than
+    // stay on the old index and land on an unrelated one. This keeps the
+    // attributes consistent (it may not cause a miscompile today). Remove all
+    // before setting so a moved attribute does not clobber another argument's
+    // attribute.
+    for (std::pair<unsigned, mlir::NamedAttribute> savedAttr : savedAttrs)
+      func.removeArgAttr(savedAttr.first, savedAttr.second.getName());
+    for (std::pair<unsigned, mlir::NamedAttribute> savedAttr : savedAttrs)
+      func.setArgAttr(newArgIndex[savedAttr.first], savedAttr.second.getName(),
+                      savedAttr.second.getValue());
 
     for (auto &fixup : fixups) {
       if constexpr (std::is_same_v<FuncOpTy, mlir::func::FuncOp>)
diff --git a/flang/test/Fir/target-rewrite-arg-position.fir b/flang/test/Fir/target-rewrite-arg-position.fir
index 45e783662fb01..34770e43d5b12 100644
--- a/flang/test/Fir/target-rewrite-arg-position.fir
+++ b/flang/test/Fir/target-rewrite-arg-position.fir
@@ -27,3 +27,17 @@ func.func @_QFPs(%arg0: !fir.ref<i32> {fir.host_assoc}) {
 }
 
 // CHECK: func.func @_QFPs(%arg0: !fir.ref<i32> {fir.host_assoc, llvm.nest})
+
+// -----
+
+// Test with a mid-list argument shift: the target ABI splits the complex value
+// into two scalars, inserting an argument. fir.host_assoc on the following
+// argument must move with the shift and stay together with the llvm.nest that
+// is derived from it, not land on one of the split scalars.
+
+func.func @_QFPg(%arg0: complex<f64>, %arg1: !fir.ref<tuple<!fir.ref<i32>>> {fir.host_assoc}) {
+  return
+}
+
+// CHECK-LABEL: func.func @_QFPg
+// CHECK-SAME:    %{{.*}}: f64, %{{.*}}: f64, %{{.*}}: !fir.ref<tuple<!fir.ref<i32>>> {fir.host_assoc, llvm.nest})