Fix mismatch in prologue and epilogue for funclets on AArch64 Windows

[danielframpton]

Summary

The generated code for a funclet can have an add to sp in the epilogue for which there is no corresponding sub in the prologue.

This patch removes the early return from emitPrologue that was preventing the sub to sp, and instead conditionalizes the appropriate parts of the rest of the function.

Test case included.

Fixes: https://bugs.llvm.org/show_bug.cgi?id=45345

[efriedma]

If I'm understanding correctly, the space that wasn't getting allocated in the prologue is the space for arguments to calls inside the funclet?

llvm/lib/Target/AArch64/AArch64FrameLowering.cpp

1038

I think this COPY still needs to happen after the SEH_PrologEnd, or else we need to emit an SEH opcode for it. Otherwise, the SEH opcodes won't match the instructions. (The mismatch probably doesn't have many visible side-effects here, but still.)

[danielframpton]

In D77015#1948788, @efriedma wrote:

If I'm understanding correctly, the space that wasn't getting allocated in the prologue is the space for arguments to calls inside the funclet?

Correct, although the relevant value (MaxCallFrameSize) is computed for the function, so the call does not actually need to be in the funclet.

[danielframpton]

Move the COPY to after the end of the prologue (SEH_PrologEnd) as suggested in the review.

[danielframpton]

llvm/lib/Target/AArch64/AArch64FrameLowering.cpp

1038

Done

[danielframpton]

Updated to include context.

[efriedma]

LGTM

I assume you don't have commit access, so I'll commit it for you.

[danielframpton]

In D77015#1950924, @efriedma wrote:

LGTM

I assume you don't have commit access, so I'll commit it for you.

That would be great. I am doing some final validation of the version you accepted (using it to self-host rustc). I will let you know that is complete.

[danielframpton]

In D77015#1951716, @danielframpton wrote:

In D77015#1950924, @efriedma wrote:

LGTM

I assume you don't have commit access, so I'll commit it for you.

That would be great. I am doing some final validation of the version you accepted (using it to self-host rustc). I will let you know that is complete.

My additional validation showed no issues. Please let me know if there is anything you need me to do.

[efriedma]

Both patches are merged now. Thanks for the patches!

llvm/lib/Target/AArch64/AArch64FrameLowering.cpp

void AArch64FrameLowering::emitPrologue(MachineFunction &MF,
   while (MBBI != End && MBBI->getFlag(MachineInstr::FrameSetup) &&
          !IsSVECalleeSave(MBBI)) {
     if (CombineSPBump)
       fixupCalleeSaveRestoreStackOffset(*MBBI, AFI->getLocalStackSize(),
                                         NeedsWinCFI, &HasWinCFI);
     ++MBBI;
   }
 
-  // The code below is not applicable to funclets. We have emitted all the SEH
-  // opcodes that we needed to emit.  The FP and BP belong to the containing
-  // function.
-  if (IsFunclet) {
-    if (NeedsWinCFI) {
-      HasWinCFI = true;
-      BuildMI(MBB, MBBI, DL, TII->get(AArch64::SEH_PrologEnd))
-          .setMIFlag(MachineInstr::FrameSetup);
-    }
-
-    // SEH funclets are passed the frame pointer in X1.  If the parent
-    // function uses the base register, then the base register is used
-    // directly, and is not retrieved from X1.
-    if (F.hasPersonalityFn()) {
-      EHPersonality Per = classifyEHPersonality(F.getPersonalityFn());
-      if (isAsynchronousEHPersonality(Per)) {
-        BuildMI(MBB, MBBI, DL, TII->get(TargetOpcode::COPY), AArch64::FP)
-            .addReg(AArch64::X1).setMIFlag(MachineInstr::FrameSetup);
-        MBB.addLiveIn(AArch64::X1);
-      }
-    }
-
-    return;
-  }
-
-  if (HasFP) {
+  // For funclets the FP belongs to the containing function.
+  if (!IsFunclet && HasFP) {
     // Only set up FP if we actually need to.
     int64_t FPOffset = isTargetDarwin(MF) ? (AFI->getCalleeSavedStackSize() - 16) : 0;
 
     if (CombineSPBump)
       FPOffset += AFI->getLocalStackSize();
 
     // Issue    sub fp, sp, FPOffset or
     //          mov fp,sp          when FPOffset is zero.
     // Note: All stores of callee-saved registers are marked as "FrameSetup".
     // This code marks the instruction(s) that set the FP also.
     emitFrameOffset(MBB, MBBI, DL, AArch64::FP, AArch64::SP,
                     {FPOffset, MVT::i8}, TII, MachineInstr::FrameSetup, false,
                     NeedsWinCFI, &HasWinCFI);
void AArch64FrameLowering::emitPrologue(MachineFunction &MF,
 
   // Finally allocate remaining SVE stack space.
   emitFrameOffset(MBB, CalleeSavesEnd, DL, AArch64::SP, AArch64::SP,
                   -AllocateAfter, TII,
                   MachineInstr::FrameSetup);
 
   // Allocate space for the rest of the frame.
   if (NumBytes) {
-    const bool NeedsRealignment = RegInfo->needsStackRealignment(MF);
+    // Alignment is required for the parent frame, not the funclet
+    const bool NeedsRealignment =
+        !IsFunclet && RegInfo->needsStackRealignment(MF);
     unsigned scratchSPReg = AArch64::SP;
 
     if (NeedsRealignment) {
       scratchSPReg = findScratchNonCalleeSaveRegister(&MBB);
       assert(scratchSPReg != AArch64::NoRegister);
     }
 
     // If we're a leaf function, try using the red zone.
void AArch64FrameLowering::emitPrologue(MachineFunction &MF,
 
   // If we need a base pointer, set it up here. It's whatever the value of the
   // stack pointer is at this point. Any variable size objects will be allocated
   // after this, so we can still use the base pointer to reference locals.
   //
   // FIXME: Clarify FrameSetup flags here.
   // Note: Use emitFrameOffset() like above for FP if the FrameSetup flag is
   // needed.
-  if (RegInfo->hasBasePointer(MF)) {
+  // For funclets the BP belongs to the containing function.
+  if (!IsFunclet && RegInfo->hasBasePointer(MF)) {
     TII->copyPhysReg(MBB, MBBI, DL, RegInfo->getBaseRegister(), AArch64::SP,
                      false);
     if (NeedsWinCFI) {
       HasWinCFI = true;
       BuildMI(MBB, MBBI, DL, TII->get(AArch64::SEH_Nop))
           .setMIFlag(MachineInstr::FrameSetup);
     }
   }
 
   // The very last FrameSetup instruction indicates the end of prologue. Emit a
   // SEH opcode indicating the prologue end.
   if (NeedsWinCFI && HasWinCFI) {
     BuildMI(MBB, MBBI, DL, TII->get(AArch64::SEH_PrologEnd))
         .setMIFlag(MachineInstr::FrameSetup);
   }
 
+  // SEH funclets are passed the frame pointer in X1.  If the parent
+  // function uses the base register, then the base register is used
+  // directly, and is not retrieved from X1.
+  if (IsFunclet && F.hasPersonalityFn()) {
+    EHPersonality Per = classifyEHPersonality(F.getPersonalityFn());
+    if (isAsynchronousEHPersonality(Per)) {
+      BuildMI(MBB, MBBI, DL, TII->get(TargetOpcode::COPY), AArch64::FP)
+          .addReg(AArch64::X1)
+          .setMIFlag(MachineInstr::FrameSetup);
+      MBB.addLiveIn(AArch64::X1);
+    }
+  }
+
   if (needsFrameMoves) {
     const DataLayout &TD = MF.getDataLayout();
     const int StackGrowth = isTargetDarwin(MF)
                                 ? (2 * -TD.getPointerSize(0))
                                 : -AFI->getCalleeSavedStackSize();
     Register FramePtr = RegInfo->getFrameRegister(MF);
     // An example of the prologue:
     //

llvm/test/CodeGen/AArch64/funclet-match-add-sub-stack.ll

+; RUN: llc -o - %s -mtriple=aarch64-windows | FileCheck %s
+; Check that the stack bump around a funclet is computed correctly in both the
+; prologue and epilogue in the case we have a MaxCallFrameSize > 0 and are doing alloca
+target datalayout = "e-m:w-p:64:64-i32:32-i64:64-i128:128-n32:64-S128"
+target triple = "aarch64-pc-windows-msvc19.25.28611"
+
+; // requires passing arguments on the stack
+; void test2(void*, int, int, int, int, int, int, int, int);
+;
+; // function with the funclet being checked
+; void test1(size_t bytes)
+; {
+;   // alloca forces a separate callee save bump and stack bump
+;   void *data = _alloca(bytes);
+;   try {
+;     test2(data, 0, 1, 2, 3, 4, 5, 6, 7);
+;   } catch (...) {
+;     // the funclet being checked
+;   }
+; }
+
+; CHECK-LABEL: ?catch$2@?0??test1@@YAX_K@Z@4HA
+; CHECK: sub sp, sp, #16
+; CHECK: add sp, sp, #16
+; Function Attrs: uwtable
+define dso_local void @"?test1@@YAX_K@Z"(i64 %0) #0 personality i8* bitcast (i32 (...)* @__CxxFrameHandler3 to i8*) {
+  %2 = alloca i64, align 8
+  %3 = alloca i8*, align 8
+  store i64 %0, i64* %2, align 8
+  %4 = load i64, i64* %2, align 8
+  %5 = alloca i8, i64 %4, align 16
+  store i8* %5, i8** %3, align 8
+  %6 = load i8*, i8** %3, align 8
+  invoke void @"?test2@@YAXPEAXHHHHHHHH@Z"(i8* %6, i32 0, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7)
+          to label %13 unwind label %7
+
+7:                                                ; preds = %1
+  %8 = catchswitch within none [label %9] unwind to caller
+
+9:                                                ; preds = %7
+  %10 = catchpad within %8 [i8* null, i32 64, i8* null]
+  catchret from %10 to label %11
+
+11:                                               ; preds = %9
+  br label %12
+
+12:                                               ; preds = %11, %13
+  ret void
+
+13:                                               ; preds = %1
+  br label %12
+}
+
+declare dso_local void @"?test2@@YAXPEAXHHHHHHHH@Z"(i8*, i32, i32, i32, i32, i32, i32, i32, i32) #1
+
+declare dso_local i32 @__CxxFrameHandler3(...)
+
+attributes #0 = { uwtable }
+
+!llvm.module.flags = !{!0}
+
+!0 = !{i32 1, !"wchar_size", i32 2}