[llvm] [X86] Use StackArgTokenFactor for all stores when setting up tail calls. (PR #126244)

[Florian Hahn via llvm-commits]

https://github.com/fhahn created https://github.com/llvm/llvm-project/pull/126244

Before this patch, the stack argument token factor was used if any outgoing stack slots needed to be written, but not when writing the return address to a stack slot. Writing the return address stack slot can also alias a stack slot for an input argument.

Always use StackArgumentTokenFactor, to ensure the store of the return address will always use it, i.e. happen after all input arguments are loaded.

>From 347f82e8306e982f582e961eb2b4d476037cb4c1 Mon Sep 17 00:00:00 2001
From: Florian Hahn <flo at fhahn.com>
Date: Fri, 7 Feb 2025 13:39:44 +0000
Subject: [PATCH] [X86] Use StackArgTokenFactor for all stores when setting up
 tail calls.

Before this patch, the stack argument token factor was used if any
outgoing stack slots needed to be written, but not when writing the
return address to a stack slot. Writing the return address stack slot
can also alias a stack slot for an input argument.

Always use StackArgumentTokenFactor, to ensure the store of the return
address will always use it, i.e. happen after all input arguments are
loaded.
---
 llvm/lib/Target/X86/X86ISelLoweringCall.cpp   | 20 +++++++++----------
 ...c-store-ret-address-aliasing-stack-slot.ll |  4 ++--
 2 files changed, 12 insertions(+), 12 deletions(-)

diff --git a/llvm/lib/Target/X86/X86ISelLoweringCall.cpp b/llvm/lib/Target/X86/X86ISelLoweringCall.cpp
index 6835c7e336a5cb1..ee4bb758102f42d 100644
--- a/llvm/lib/Target/X86/X86ISelLoweringCall.cpp
+++ b/llvm/lib/Target/X86/X86ISelLoweringCall.cpp
@@ -2326,12 +2326,13 @@ X86TargetLowering::LowerCall(TargetLowering::CallLoweringInfo &CLI,
   // shuffling arguments passed in memory.
   if (!IsSibcall && isTailCall) {
     // Force all the incoming stack arguments to be loaded from the stack
-    // before any new outgoing arguments are stored to the stack, because the
-    // outgoing stack slots may alias the incoming argument stack slots, and
-    // the alias isn't otherwise explicit. This is slightly more conservative
-    // than necessary, because it means that each store effectively depends
-    // on every argument instead of just those arguments it would clobber.
-    SDValue ArgChain = DAG.getStackArgumentTokenFactor(Chain);
+    // before any new outgoing arguments or the return address are stored to the
+    // stack, because the outgoing stack slots may alias the incoming argument
+    // stack slots, and the alias isn't otherwise explicit. This is slightly
+    // more conservative than necessary, because it means that each store
+    // effectively depends on every argument instead of just those arguments it
+    // would clobber.
+    Chain = DAG.getStackArgumentTokenFactor(Chain);
 
     SmallVector<SDValue, 8> MemOpChains2;
     SDValue FIN;
@@ -2373,13 +2374,12 @@ X86TargetLowering::LowerCall(TargetLowering::CallLoweringInfo &CLI,
         Source = DAG.getNode(ISD::ADD, dl, getPointerTy(DAG.getDataLayout()),
                              StackPtr, Source);
 
-        MemOpChains2.push_back(CreateCopyOfByValArgument(Source, FIN,
-                                                         ArgChain,
-                                                         Flags, DAG, dl));
+        MemOpChains2.push_back(
+            CreateCopyOfByValArgument(Source, FIN, Chain, Flags, DAG, dl));
       } else {
         // Store relative to framepointer.
         MemOpChains2.push_back(DAG.getStore(
-            ArgChain, dl, Arg, FIN,
+            Chain, dl, Arg, FIN,
             MachinePointerInfo::getFixedStack(DAG.getMachineFunction(), FI)));
       }
     }
diff --git a/llvm/test/CodeGen/X86/swifttailcc-store-ret-address-aliasing-stack-slot.ll b/llvm/test/CodeGen/X86/swifttailcc-store-ret-address-aliasing-stack-slot.ll
index 78e810bb67f45c9..cd669768705e53b 100644
--- a/llvm/test/CodeGen/X86/swifttailcc-store-ret-address-aliasing-stack-slot.ll
+++ b/llvm/test/CodeGen/X86/swifttailcc-store-ret-address-aliasing-stack-slot.ll
@@ -24,9 +24,9 @@ define swifttailcc void @test(ptr %0, ptr swiftasync %1, i64 %2, i64 %3, ptr %4,
 ; CHECK-NEXT:    movq {{[0-9]+}}(%rsp), %r15
 ; CHECK-NEXT:    callq _foo
 ; CHECK-NEXT:    movq %r14, (%rax)
+; CHECK-NEXT:    movl [[OFF:[0-9]+]](%rsp), %edx
 ; CHECK-NEXT:    movq {{[0-9]+}}(%rsp), %rcx
-; CHECK-NEXT:    movq %rcx, [[OFF:[0-9]+]](%rsp)
-; CHECK-NEXT:    movl [[OFF]](%rsp), %edx
+; CHECK-NEXT:    movq %rcx, [[OFF]](%rsp)
 ; CHECK-NEXT:    movq %rax, %r14
 ; CHECK-NEXT:    movq %r13, %rdi
 ; CHECK-NEXT:    movq %r15, %rsi