[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.

Author: hjyamauchi

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,

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)

llvm/test/CodeGen/AArch64/swift-async-win.ll

@@ -39,12 +39,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]
 

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@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