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

[hjyamauchi]

Summary

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

[compnerd]

Thanks for this fix! I think that as long as we have a slightly stronger check to ensure that a store isn't occurring in the case, this should be good.

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

18

Should this have something like CHECK-NOT for a stp?

[hjyamauchi]

Added a CHECK-NOT.

llvm/lib/Target/AArch64/AArch64CallingConvention.td

 // Win64 has unwinding codes for an (FP,LR) pair, save_fplr and save_fplr_x.
 // We put FP before LR, so that frame lowering logic generates (FP,LR) pairs,
 // and not (LR,FP) pairs.
 def CSR_Win_AArch64_AAPCS : CalleeSavedRegs<(add X19, X20, X21, X22, X23, X24,
                                                X25, X26, X27, X28, FP, LR,
                                                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,
                                                (sequence "X%u", 0, 8),
                                                (sequence "Q%u", 0, 7))>;
 
 // AArch64 PCS for vector functions (VPCS)
 // must (additionally) preserve full Q8-Q23 registers

llvm/lib/Target/AArch64/AArch64RegisterInfo.cpp

AArch64RegisterInfo::getCalleeSavedRegs(const MachineFunction *MF) const {
 
   // Darwin has its own CSR_AArch64_AAPCS_SaveList, which means most CSR save
   // lists depending on that will need to have their Darwin variant as well.
   if (MF->getSubtarget<AArch64Subtarget>().isTargetDarwin())
     return getDarwinCalleeSavedRegs(MF);
 
   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)
     return CSR_AArch64_SVE_AAPCS_SaveList;
   if (MF->getFunction().getCallingConv() ==
           CallingConv::AArch64_SME_ABI_Support_Routines_PreserveMost_From_X0)
     report_fatal_error(
         "Calling convention AArch64_SME_ABI_Support_Routines_PreserveMost_From_X0 is "

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

entryresume.0:
   %5 = load ptr, ptr %4, align 8
   musttail call swifttailcc void %5(ptr %3) #4
   ret void
 }
 
 ; 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]
 
 ; Function Attrs: nounwind readnone
 declare ptr @llvm.swift.async.context.addr() #2
 
 ; Function Attrs: argmemonly nounwind
 declare dllimport swiftcc void @swift_task_dealloc(ptr) local_unnamed_addr #3
 
 attributes #0 = { argmemonly nofree nosync nounwind willreturn }
 attributes #1 = { nounwind "frame-pointer"="none" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-cpu"="generic" "target-features"="+neon" }
 attributes #2 = { nounwind readnone }
 attributes #3 = { argmemonly nounwind }
 attributes #4 = { nounwind }
 

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

+; 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