[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 updated https://github.com/llvm/llvm-project/pull/66967

>From 73213adae4b12335436056f29e7c35ba25ce84a1 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   | 31 ++++++++++++++++++-
 .../AArch64/swift-async-context-seh.ll        | 24 ++++++++++++++
 2 files changed, 54 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..d65f0ea71846f83 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,20 @@ 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_Nop))
+            .setMIFlags(MachineInstr::FrameSetup);
+        HasWinCFI = true;
+      }
     }
 
     if (HomPrologEpilog) {
@@ -2132,6 +2156,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..852c97743a10cfd
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/swift-async-context-seh.ll
@@ -0,0 +1,24 @@
+; 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
+
+; Check that the prologue/epilogue instructions for the swift async
+; context have an associated SEH instruction and that it doesn't error
+; when the output is an object file.
+
+; CHECK: orr     x29, x29, #0x1000000000000000
+; CHECK-NEXT: .seh_nop
+; CHECK:  str     x22, [sp, #16]
+; CHECK-NEXT: .seh_nop
+; CHECK: and     x29, x29, #0xefffffffffffffff
+; CHECK-NEXT: .seh_nop
+
+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
+}