[llvm] [AArch64] Skip over shadow space for ARM64EC entry thunk variadic calls (PR #80994)

[Billy Laws via llvm-commits]

https://github.com/bylaws created https://github.com/llvm/llvm-project/pull/80994

When in an entry thunk the x64 SP is passed in x4 but this cannot be directly passed through since x64 varargs calls have a 32 byte shadow store at SP followed by the in-stack parameters. ARM64EC varargs calls on the other hand expect x4 to point to the first in-stack parameter.

CC: @efriedma-quic @cjacek @mstorsjo 

>From dff9bb50c75bcc781a6b8414fa33c4c8a0dfd2da Mon Sep 17 00:00:00 2001
From: Billy Laws <blaws05 at gmail.com>
Date: Sun, 28 Jan 2024 22:51:08 +0000
Subject: [PATCH] [AArch64] Skip over shadow space for ARM64EC entry thunk
 variadic calls

When in an entry thunk the x64 SP is passed in x4 but this cannot be directly
passed through since x64 varargs calls have a 32 byte shadow store at SP
followed by the in-stack parameters. ARM64EC varargs calls on the other hand
expect x4 to point to the first in-stack parameter.
---
 llvm/lib/Target/AArch64/AArch64ISelLowering.cpp   | 12 +++++++++++-
 llvm/test/CodeGen/AArch64/arm64ec-entry-thunks.ll |  2 +-
 2 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
index 8573939b04389f..b7e1c96c237281 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -8010,7 +8010,17 @@ AArch64TargetLowering::LowerCall(CallLoweringInfo &CLI,
 
   if (IsVarArg && Subtarget->isWindowsArm64EC()) {
     SDValue ParamPtr = StackPtr;
-    if (IsTailCall) {
+    if (MF.getFunction().getCallingConv() == CallingConv::ARM64EC_Thunk_X64) {
+      // When in an entry thunk the x64 SP is passed via x4. This cannot
+      // be directly passed through since x64 varargs calls have a 32 byte
+      // shadow store at SP followed by the in-stack parameters,
+      // Arm64EC varargs calls on the other hand expect x4
+      // to point to the first in-stack parameter.
+      Register VReg = MF.addLiveIn(AArch64::X4, &AArch64::GPR64RegClass);
+      SDValue Val = DAG.getCopyFromReg(Chain, DL, VReg, MVT::i64);
+      SDValue PtrOff = DAG.getIntPtrConstant(32, DL);
+      ParamPtr = DAG.getNode(ISD::ADD, DL, PtrVT, Val, PtrOff);
+    } else if (IsTailCall) {
       // Create a dummy object at the top of the stack that can be used to get
       // the SP after the epilogue
       int FI = MF.getFrameInfo().CreateFixedObject(1, FPDiff, true);
diff --git a/llvm/test/CodeGen/AArch64/arm64ec-entry-thunks.ll b/llvm/test/CodeGen/AArch64/arm64ec-entry-thunks.ll
index 0083818def1514..bb9ba05f7a2724 100644
--- a/llvm/test/CodeGen/AArch64/arm64ec-entry-thunks.ll
+++ b/llvm/test/CodeGen/AArch64/arm64ec-entry-thunks.ll
@@ -147,7 +147,7 @@ define void @has_varargs(...) nounwind {
 ; CHECK-NEXT:     add     x29, sp, #160
 ; CHECK-NEXT:     .seh_add_fp     160
 ; CHECK-NEXT:     .seh_endprologue
-; CHECK-NEXT:     mov     x4, sp
+; CHECK-NEXT:     add     x4, x4, #32
 ; CHECK-NEXT:     mov     x5, xzr
 ; CHECK-NEXT:     blr     x9
 ; CHECK-NEXT:     adrp    x8, __os_arm64x_dispatch_ret