[llvm] 6e6d2b7 - [AArch64][Windows] Fix the callee-saved registers for swiftcc on Windows/AArch64

[Hiroshi Yamauchi via llvm-commits]

Author: Hiroshi Yamauchi
Date: 2023-07-12T13:13:21-07:00
New Revision: 6e6d2b7859a36cefacb1322540cfc9f90d4f19f1

URL: https://github.com/llvm/llvm-project/commit/6e6d2b7859a36cefacb1322540cfc9f90d4f19f1
DIFF: https://github.com/llvm/llvm-project/commit/6e6d2b7859a36cefacb1322540cfc9f90d4f19f1.diff

LOG: [AArch64][Windows] Fix the callee-saved registers for swiftcc on Windows/AArch64

Fix a miscompilation crash where swiftself on x20 gets corrupted due
to incorrect save/restore at prologue/epilogue.

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

Added: 
    llvm/test/CodeGen/AArch64/swift-csr-win.ll

Modified: 
    llvm/lib/Target/AArch64/AArch64CallingConvention.td
    llvm/lib/Target/AArch64/AArch64RegisterInfo.cpp
    llvm/test/CodeGen/AArch64/swift-async-win.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Target/AArch64/AArch64CallingConvention.td b/llvm/lib/Target/AArch64/AArch64CallingConvention.td
index ce087be79202e3..37976a222783b2 100644
--- a/llvm/lib/Target/AArch64/AArch64CallingConvention.td
+++ b/llvm/lib/Target/AArch64/AArch64CallingConvention.td
@@ -416,6 +416,12 @@ def CSR_Win_AArch64_AAPCS : CalleeSavedRegs<(add X19, X20, X21, X22, X23, X24,
                                                D8, D9, D10, D11,
                                                D12, D13, D14, D15)>;
 
+def CSR_Win_AArch64_AAPCS_SwiftError
+    : CalleeSavedRegs<(sub CSR_Win_AArch64_AAPCS, X21)>;
+
+def CSR_Win_AArch64_AAPCS_SwiftTail
+    : CalleeSavedRegs<(sub CSR_Win_AArch64_AAPCS, X20, X22)>;
+
 // The Control Flow Guard check call uses a custom calling convention that also
 // preserves X0-X8 and Q0-Q7.
 def CSR_Win_AArch64_CFGuard_Check : CalleeSavedRegs<(add CSR_Win_AArch64_AAPCS,

diff  --git a/llvm/lib/Target/AArch64/AArch64RegisterInfo.cpp b/llvm/lib/Target/AArch64/AArch64RegisterInfo.cpp
index f0bdce6938536e..d1ddf6d7697547 100644
--- a/llvm/lib/Target/AArch64/AArch64RegisterInfo.cpp
+++ b/llvm/lib/Target/AArch64/AArch64RegisterInfo.cpp
@@ -85,8 +85,16 @@ AArch64RegisterInfo::getCalleeSavedRegs(const MachineFunction *MF) const {
 
   if (MF->getFunction().getCallingConv() == CallingConv::CFGuard_Check)
     return CSR_Win_AArch64_CFGuard_Check_SaveList;
-  if (MF->getSubtarget<AArch64Subtarget>().isTargetWindows())
+  if (MF->getSubtarget<AArch64Subtarget>().isTargetWindows()) {
+    if (MF->getSubtarget<AArch64Subtarget>().getTargetLowering()
+            ->supportSwiftError() &&
+        MF->getFunction().getAttributes().hasAttrSomewhere(
+            Attribute::SwiftError))
+      return CSR_Win_AArch64_AAPCS_SwiftError_SaveList;
+    if (MF->getFunction().getCallingConv() == CallingConv::SwiftTail)
+      return CSR_Win_AArch64_AAPCS_SwiftTail_SaveList;
     return CSR_Win_AArch64_AAPCS_SaveList;
+  }
   if (MF->getFunction().getCallingConv() == CallingConv::AArch64_VectorCall)
     return CSR_AArch64_AAVPCS_SaveList;
   if (MF->getFunction().getCallingConv() == CallingConv::AArch64_SVE_VectorCall)

diff  --git a/llvm/test/CodeGen/AArch64/swift-async-win.ll b/llvm/test/CodeGen/AArch64/swift-async-win.ll
index 995953df29c15a..75a27445675436 100644
--- a/llvm/test/CodeGen/AArch64/swift-async-win.ll
+++ b/llvm/test/CodeGen/AArch64/swift-async-win.ll
@@ -34,12 +34,11 @@ entryresume.0:
 ; NOTE: we do not see the canonical windows frame setup due to the `nounwind`
 ; attribtue on the function.
 
-; CHECK: sub sp, sp, #64
-; CHECK: stp x30, x29, [sp, #16]
-; CHECK: add x29, sp, #16
-; CHECK: stp x22, x21, [sp, #32]
+; CHECK: sub sp, sp, #48
+; CHECK: stp x30, x29, [sp, #24]
+; CHECK: add x29, sp, #24
+; CHECK: str x19, [sp, #40]
 ; CHECK: sub x8, x29, #8
-; CHECK: stp x20, x19, [sp, #48]
 ; CHECK: ldr x9, [x0]
 ; CHECK: str x9, [x8]
 

diff  --git a/llvm/test/CodeGen/AArch64/swift-csr-win.ll b/llvm/test/CodeGen/AArch64/swift-csr-win.ll
new file mode 100644
index 00000000000000..06f9986f510903
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/swift-csr-win.ll
@@ -0,0 +1,22 @@
+; RUN: llc -mtriple aarch64-unknown-windows-msvc %s -o - | FileCheck %s
+
+define internal swifttailcc void @"?future_adapter@@YWXPEAVAsyncContext at swift@@@Z"(ptr noundef swiftasync %_context) #0 {
+entry:
+  %add.ptr = getelementptr inbounds i8, ptr %_context, i64 -32
+  %asyncEntryPoint = getelementptr inbounds i8, ptr %_context, i64 -24
+  %0 = load ptr, ptr %asyncEntryPoint, align 8
+  %closureContext = getelementptr inbounds i8, ptr %_context, i64 -16
+  %1 = load ptr, ptr %closureContext, align 8
+  %2 = load ptr, ptr %add.ptr, align 8
+  musttail call swifttailcc void %0(ptr noundef %2, ptr noundef swiftasync %_context, ptr noundef swiftself %1) #17
+  ret void
+}
+
+; Check that x20 isn't saved/restored at the prologue/epilogue which
+; would interfere with the outgoing self parameter on x20 at the tail
+; call.
+
+; CHECK-NOT:  st{{.*}}x20
+; CHECK:      ldp x1, x20, [x22, #-24]
+; CHECK-NEXT: ldur x0, [x22, #-32]
+; CHECK-NEXT: br x1