[AArch64][WinCFI] Handle cases where no SEH opcodes in the prologue

but there are some in the epilogue.

Make a decision whether or not to have a startepilogue/endepilogue
based on whether we actually insert SEH opcodes in the epilogue,
rather than whether we had SEH opcodes in the prologue or not.

This fixes an assert failure when there are no SEH opcodes in the
prologue but there are SEH opcodes in the epilogue (for example, when
there is no stack frame but there are stack arguments) which was not
covered in https://reviews.llvm.org/D88641.

Assertion failed: HasWinCFI == MF.hasWinCFI(), file C:\Users\hiroshi\llvm-project\llvm\lib\Target\AArch64\AArch64FrameLowering.cpp, line 1988

Differential Revision: https://reviews.llvm.org/D159238

Author: hjyamauchi

llvm/lib/Target/AArch64/AArch64FrameLowering.cpp

@@ -1948,23 +1948,32 @@ void AArch64FrameLowering::emitEpilogue(MachineFunction &MF,
   bool EmitCFI = MF.getInfo<AArch64FunctionInfo>()->needsAsyncDwarfUnwindInfo();
   bool HasWinCFI = false;
   bool IsFunclet = false;
-  auto WinCFI = make_scope_exit([&]() { assert(HasWinCFI == MF.hasWinCFI()); });
 
   if (MBB.end() != MBBI) {
     DL = MBBI->getDebugLoc();
     IsFunclet = isFuncletReturnInstr(*MBBI);
   }
 
+  MachineBasicBlock::iterator EpilogStartI = MBB.end();
+
   auto FinishingTouches = make_scope_exit([&]() {
     InsertReturnAddressAuth(MF, MBB, NeedsWinCFI, &HasWinCFI);
     if (needsShadowCallStackPrologueEpilogue(MF))
       emitShadowCallStackEpilogue(*TII, MF, MBB, MBB.getFirstTerminator(), DL);
     if (EmitCFI)
       emitCalleeSavedGPRRestores(MBB, MBB.getFirstTerminator());
-    if (HasWinCFI)
+    if (HasWinCFI) {
       BuildMI(MBB, MBB.getFirstTerminator(), DL,
               TII->get(AArch64::SEH_EpilogEnd))
           .setMIFlag(MachineInstr::FrameDestroy);
+      if (!MF.hasWinCFI())
+        MF.setHasWinCFI(true);
+    }
+    if (NeedsWinCFI) {
+      assert(EpilogStartI != MBB.end());
+      if (!HasWinCFI)
+        MBB.erase(EpilogStartI);
+    }
   });
 
   int64_t NumBytes = IsFunclet ? getWinEHFuncletFrameSize(MF)
@@ -2029,7 +2038,7 @@ void AArch64FrameLowering::emitEpilogue(MachineFunction &MF,
     // Adjust local stack
     emitFrameOffset(MBB, LastPopI, DL, AArch64::SP, AArch64::SP,
                     StackOffset::getFixed(AFI->getLocalStackSize()), TII,
-                    MachineInstr::FrameDestroy, false, NeedsWinCFI);
+                    MachineInstr::FrameDestroy, false, NeedsWinCFI, &HasWinCFI);
 
     // SP has been already adjusted while restoring callee save regs.
     // We've bailed-out the case with adjusting SP for arguments.
@@ -2081,16 +2090,17 @@ void AArch64FrameLowering::emitEpilogue(MachineFunction &MF,
                                         NeedsWinCFI, &HasWinCFI);
   }
 
-  if (MF.hasWinCFI()) {
-    // If the prologue didn't contain any SEH opcodes and didn't set the
-    // MF.hasWinCFI() flag, assume the epilogue won't either, and skip the
-    // EpilogStart - to avoid generating CFI for functions that don't need it.
-    // (And as we didn't generate any prologue at all, it would be asymmetrical
-    // to the epilogue.) By the end of the function, we assert that
-    // HasWinCFI is equal to MF.hasWinCFI(), to verify this assumption.
-    HasWinCFI = true;
+  if (NeedsWinCFI) {
+    // Note that there are cases where we insert SEH opcodes in the
+    // epilogue when we had no SEH opcodes in the prologue. For
+    // example, when there is no stack frame but there are stack
+    // arguments. Insert the SEH_EpilogStart and remove it later if it
+    // we didn't emit any SEH opcodes to avoid generating WinCFI for
+    // functions that don't need it.
     BuildMI(MBB, LastPopI, DL, TII->get(AArch64::SEH_EpilogStart))
         .setMIFlag(MachineInstr::FrameDestroy);
+    EpilogStartI = LastPopI;
+    --EpilogStartI;
   }
 
   if (hasFP(MF) && AFI->hasSwiftAsyncContext()) {
@@ -2244,11 +2254,11 @@ void AArch64FrameLowering::emitEpilogue(MachineFunction &MF,
     emitFrameOffset(
         MBB, LastPopI, DL, AArch64::SP, AArch64::FP,
         StackOffset::getFixed(-AFI->getCalleeSaveBaseToFrameRecordOffset()),
-        TII, MachineInstr::FrameDestroy, false, NeedsWinCFI);
+        TII, MachineInstr::FrameDestroy, false, NeedsWinCFI, &HasWinCFI);
   } else if (NumBytes)
     emitFrameOffset(MBB, LastPopI, DL, AArch64::SP, AArch64::SP,
                     StackOffset::getFixed(NumBytes), TII,
-                    MachineInstr::FrameDestroy, false, NeedsWinCFI);
+                    MachineInstr::FrameDestroy, false, NeedsWinCFI, &HasWinCFI);
 
   // When we are about to restore the CSRs, the CFA register is SP again.
   if (EmitCFI && hasFP(MF)) {

llvm/test/CodeGen/AArch64/wincfi-seh-only-in-epilogue.ll

@@ -0,0 +1,13 @@
+; RUN: llc -mtriple=aarch64-windows %s -o - | FileCheck %s
+
+define hidden swifttailcc void @test(ptr noalias nocapture %0, ptr swiftasync %1, ptr %2, ptr noalias nocapture %3, ptr nocapture dereferenceable(8) %4, ptr %5, ptr %6, ptr %Act, ptr %Err, ptr %Res, ptr %Act.DistributedActor, ptr %Err.Error, ptr %Res.Decodable, ptr %Res.Encodable, ptr swiftself %7) #0 {
+entry:
+  ret void
+}
+
+; Check that there is no .seh_endprologue but there is seh_startepilogue/seh_endepilogue.
+; CHECK-NOT: .seh_endprologue
+; CHECK:     .seh_startepilogue
+; CHECK:     add sp, sp, #48
+; CHECK:     .seh_stackalloc 48
+; CHECK:     .seh_endepilogue