[AArch64][Windows] Fix the slot offset of the swift async context register.

[hjyamauchi]

Summary

This fixes a code gen issue where savings the swift async context register (x22) accidentally overwrites the saved value of another callee-saved register, corrupts its value and causes a crash.

[compnerd]

Ah right, the ABI mandates the FP, LR spill. Thanks for fixing this!

llvm/lib/Target/AArch64/AArch64FrameLowering.cpp

for (unsigned i = FirstReg; i < Count; i += RegInc) {
     if (RPI.isScalable())
       ScalableByteOffset += StackFillDir * Scale;
     else
       ByteOffset += StackFillDir * (RPI.isPaired() ? 2 * Scale : Scale);
 
     // Swift's async context is directly before FP, so allocate an extra
     // 8 bytes for it.
     if (NeedsFrameRecord && AFI->hasSwiftAsyncContext() &&
-        RPI.Reg2 == AArch64::FP)
+        ((!IsWindows && RPI.Reg2 == AArch64::FP) ||
+         (IsWindows && RPI.Reg2 == AArch64::LR)))
       ByteOffset += StackFillDir * 8;
 
     assert(!(RPI.isScalable() && RPI.isPaired()) &&
            "Paired spill/fill instructions don't exist for SVE vectors");
 
     // Round up size of non-pair to pair size if we need to pad the
     // callee-save area to ensure 16-byte alignment.
     if (NeedGapToAlignStack && !NeedsWinCFI &&
for (unsigned i = FirstReg; i < Count; i += RegInc) {
     assert(OffsetPost % Scale == 0);
     // If filling top down (default), we want the offset after incrementing it.
     // If fillibg bootom up (WinCFI) we need the original offset.
     int Offset = NeedsWinCFI ? OffsetPre : OffsetPost;
 
     // The FP, LR pair goes 8 bytes into our expanded 24-byte slot so that the
     // Swift context can directly precede FP.
     if (NeedsFrameRecord && AFI->hasSwiftAsyncContext() &&
-        RPI.Reg2 == AArch64::FP)
+        ((!IsWindows && RPI.Reg2 == AArch64::FP) ||
+         (IsWindows && RPI.Reg2 == AArch64::LR)))
       Offset += 8;
     RPI.Offset = Offset / Scale;
 
     assert(((!RPI.isScalable() && RPI.Offset >= -64 && RPI.Offset <= 63) ||
             (RPI.isScalable() && RPI.Offset >= -256 && RPI.Offset <= 255)) &&
            "Offset out of bounds for LDP/STP immediate");
 
     // Save the offset to frame record so that the FP register can point to the

llvm/test/CodeGen/AArch64/swift-async-context-slot-offset-win.ll

+; RUN: llc -mtriple aarch64-unknown-windows-msvc -filetype asm %s -o - | FileCheck %s
+
+; Check that the slot offset of the async context (x22) doesn't
+; conflict with that of another callee-saved register (x21 here) and
+; saving it won't overwrite the saved value of the callee-saved
+; register.
+;
+; CHECK:        sub     sp, sp, #64
+; CHECK:        str     x19, [sp, #16]
+; CHECK:        str     x21, [sp, #24]
+; CHECK-NOT:    stp     x29, x30, [sp, #32]
+; CHECK:        stp     x29, x30, [sp, #40]
+; CHECK-NOT:    str     x22, [sp, #24]
+; CHECK:        str     x22, [sp, #32]
+
+declare ptr @llvm.swift.async.context.addr()
+declare swiftcc i64 @foo(i64 %0, i64 %1)
+declare swifttailcc void @tail(ptr swiftasync %0, ptr swiftself dereferenceable(8) %1, i64 %2)
+define internal swifttailcc void @test(ptr swiftasync %0, ptr swiftself %1, i64 %2) {
+entry:
+  %3 = load ptr, ptr %0, align 8
+  %4 = call ptr @llvm.swift.async.context.addr()
+  store ptr %3, ptr %4, align 8
+  %5 = call swiftcc i64 @foo(i64 %2, i64 %2)
+  %6 = call swiftcc i64 @foo(i64 %2, i64 %5)
+  %7 = call swiftcc i64 @foo(i64 %5, i64 %2)
+  %8 = call swiftcc i64 @foo(i64 %7, i64 %6)
+  %9 = call swiftcc i64 @foo(i64 %2, i64 %8)
+  %10 = call ptr @llvm.swift.async.context.addr()
+  musttail call swifttailcc void @tail(ptr swiftasync %10, ptr swiftself %1, i64 %2)
+  ret void
+}