[llvm] [AArch64][Win] Emit SEH instructions for the swift async context-related instructions in the prologue and the epilogue. (PR #66967)

[Hiroshi Yamauchi via llvm-commits]

https://github.com/hjyamauchi created https://github.com/llvm/llvm-project/pull/66967

This fixes an error from checkARM64Instructions() in MCWin64EH.cpp.

>From 23e6e1f30ab2109bdcf06e4fc5a2df7181c13b44 Mon Sep 17 00:00:00 2001
From: Hiroshi Yamauchi <hjyamauchi at gmail.com>
Date: Wed, 20 Sep 2023 15:36:12 -0700
Subject: [PATCH] [AArch64][Win] Emit SEH instructions for the swift async
 context-related instructions in the prologue and the epilogue.

This fixes an error from checkARM64Instructions() in MCWin64EH.cpp.
---
 .../Target/AArch64/AArch64FrameLowering.cpp   | 33 ++++++++++++++++++-
 .../AArch64/swift-async-context-seh.ll        | 26 +++++++++++++++
 2 files changed, 58 insertions(+), 1 deletion(-)
 create mode 100644 llvm/test/CodeGen/AArch64/swift-async-context-seh.ll

diff --git a/llvm/lib/Target/AArch64/AArch64FrameLowering.cpp b/llvm/lib/Target/AArch64/AArch64FrameLowering.cpp
index 68e68449d4073b2..a05bf976d9ebb84 100644
--- a/llvm/lib/Target/AArch64/AArch64FrameLowering.cpp
+++ b/llvm/lib/Target/AArch64/AArch64FrameLowering.cpp
@@ -1476,10 +1476,20 @@ void AArch64FrameLowering::emitPrologue(MachineFunction &MF,
         BuildMI(MBB, MBBI, DL, TII->get(AArch64::LOADgot), AArch64::X16)
             .addExternalSymbol("swift_async_extendedFramePointerFlags",
                                AArch64II::MO_GOT);
+        if (NeedsWinCFI) {
+          BuildMI(MBB, MBBI, DL, TII->get(AArch64::SEH_Nop))
+              .setMIFlags(MachineInstr::FrameSetup);
+          HasWinCFI = true;
+        }
         BuildMI(MBB, MBBI, DL, TII->get(AArch64::ORRXrs), AArch64::FP)
             .addUse(AArch64::FP)
             .addUse(AArch64::X16)
             .addImm(Subtarget.isTargetILP32() ? 32 : 0);
+        if (NeedsWinCFI) {
+          BuildMI(MBB, MBBI, DL, TII->get(AArch64::SEH_Nop))
+              .setMIFlags(MachineInstr::FrameSetup);
+          HasWinCFI = true;
+        }
         break;
       }
       [[fallthrough]];
@@ -1490,6 +1500,11 @@ void AArch64FrameLowering::emitPrologue(MachineFunction &MF,
           .addUse(AArch64::FP)
           .addImm(0x1100)
           .setMIFlag(MachineInstr::FrameSetup);
+      if (NeedsWinCFI) {
+        BuildMI(MBB, MBBI, DL, TII->get(AArch64::SEH_Nop))
+            .setMIFlags(MachineInstr::FrameSetup);
+        HasWinCFI = true;
+      }
       break;
 
     case SwiftAsyncFramePointerMode::Never:
@@ -1613,11 +1628,22 @@ void AArch64FrameLowering::emitPrologue(MachineFunction &MF,
       bool HaveInitialContext = Attrs.hasAttrSomewhere(Attribute::SwiftAsync);
       if (HaveInitialContext)
         MBB.addLiveIn(AArch64::X22);
+      Register Reg = HaveInitialContext ? AArch64::X22 : AArch64::XZR;
       BuildMI(MBB, MBBI, DL, TII->get(AArch64::StoreSwiftAsyncContext))
-          .addUse(HaveInitialContext ? AArch64::X22 : AArch64::XZR)
+          .addUse(Reg)
           .addUse(AArch64::SP)
           .addImm(FPOffset - 8)
           .setMIFlags(MachineInstr::FrameSetup);
+      if (NeedsWinCFI) {
+        // WinCFI and arm64e, where StoreSwiftAsyncContext is expanded
+        // to multiple instructions, should be mutually-exclusive.
+        assert(Subtarget.getTargetTriple().getArchName() != "arm64e");
+        BuildMI(MBB, MBBI, DL, TII->get(AArch64::SEH_SaveReg))
+            .addImm(RegInfo->getSEHRegNum(Reg))
+            .addImm(FPOffset - 8)
+            .setMIFlags(MachineInstr::FrameSetup);
+        HasWinCFI = true;
+      }
     }
 
     if (HomPrologEpilog) {
@@ -2132,6 +2158,11 @@ void AArch64FrameLowering::emitEpilogue(MachineFunction &MF,
           .addUse(AArch64::FP)
           .addImm(0x10fe)
           .setMIFlag(MachineInstr::FrameDestroy);
+      if (NeedsWinCFI) {
+        BuildMI(MBB, MBBI, DL, TII->get(AArch64::SEH_Nop))
+            .setMIFlags(MachineInstr::FrameDestroy);
+        HasWinCFI = true;
+      }
       break;
 
     case SwiftAsyncFramePointerMode::Never:
diff --git a/llvm/test/CodeGen/AArch64/swift-async-context-seh.ll b/llvm/test/CodeGen/AArch64/swift-async-context-seh.ll
new file mode 100644
index 000000000000000..101f6f7623d7a50
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/swift-async-context-seh.ll
@@ -0,0 +1,26 @@
+; RUN: rm -rf %t && mkdir -p %t
+; RUN: llc -mtriple aarch64-unknown-windows-msvc %s -o - | FileCheck %s
+; RUN: llc -mtriple aarch64-unknown-windows-msvc -filetype obj %s -o %t/a.o 2>&1 | FileCheck %s -check-prefix=MC --allow-empty
+
+; Check that the prologue/epilogue instructions for the swift async context
+; have an associated SEH instruction.
+
+; CHECK: orr     x29, x29, #0x1000000000000000
+; CHECK-NEXT: .seh_nop
+; CHECK:  str     x22, [sp, #16]
+; CHECK-NEXT: .seh_save_reg   x22, 16
+; CHECK: and     x29, x29, #0xefffffffffffffff
+; CHECK-NEXT: .seh_nop
+
+; MC-NOT: error: Incorrect size for test prologue: {{[0-9]+}} bytes of instructions in range, but .seh directives corresponding to {{[0-9]+}} bytes
+; MC-NOT: error: Incorrect size for test epilogue: {{[0-9]+}} bytes of instructions in range, but .seh directives corresponding to {{[0-9]+}} bytes
+
+declare ptr @llvm.swift.async.context.addr()
+
+define internal swifttailcc void @test(ptr nocapture readonly swiftasync %0) {
+entryresume.0:
+  %1 = load ptr, ptr %0, align 8
+  %2 = tail call ptr @llvm.swift.async.context.addr()
+  store ptr %1, ptr %2, align 8
+  ret void
+}
\ No newline at end of file