[AArch64] [Windows] Don't skip constructing UnwindHelp.

[efriedma]

Summary

In certain cases, the first non-frame-setup instruction in a function is a branch. For example, it could be a cbz on an argument. Make sure we correctly allocate the UnwindHelp, and find an appropriate register to use to initialize it.

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

[rnk]

lgtm

llvm/trunk/lib/Target/AArch64/AArch64FrameLowering.cpp

void AArch64FrameLowering::processFunctionBeforeFrameFinalized(
   MachineFrameInfo &MFI = MF.getFrameInfo();
   WinEHFuncInfo &EHInfo = *MF.getWinEHFuncInfo();
 
   MachineBasicBlock &MBB = MF.front();
   auto MBBI = MBB.begin();
   while (MBBI != MBB.end() && MBBI->getFlag(MachineInstr::FrameSetup))
     ++MBBI;
 
-  if (MBBI->isTerminator())
-    return;
-
   // Create an UnwindHelp object.
   int UnwindHelpFI =
       MFI.CreateStackObject(/*size*/8, /*alignment*/16, false);
   EHInfo.UnwindHelpFrameIdx = UnwindHelpFI;
   // We need to store -2 into the UnwindHelp object at the start of the
   // function.
   DebugLoc DL;
-  RS->enterBasicBlock(MBB);
-  unsigned DstReg = RS->scavengeRegister(&AArch64::GPR64RegClass, MBBI, 0);
+  RS->enterBasicBlockEnd(MBB);
+  RS->backward(std::prev(MBBI));
+  unsigned DstReg = RS->FindUnusedReg(&AArch64::GPR64commonRegClass);
+  assert(DstReg && "There must be a free register after frame setup");
   BuildMI(MBB, MBBI, DL, TII.get(AArch64::MOVi64imm), DstReg).addImm(-2);
   BuildMI(MBB, MBBI, DL, TII.get(AArch64::STURXi))
       .addReg(DstReg, getKillRegState(true))
       .addFrameIndex(UnwindHelpFI)
       .addImm(0);
 }
 
 /// For Win64 AArch64 EH, the offset to the Unwind object is from the SP before

llvm/trunk/test/CodeGen/AArch64/seh-finally.ll

 
 %struct.S = type { i32 }
 
 ; Test simple SEH (__try/__finally).
 define void @simple_seh() #0 personality i8* bitcast (i32 (...)* @__C_specific_handler to i8*) {
 entry:
 ; CHECK-LABEL: simple_seh
 ; CHECK: add     x29, sp, #16
-; CHECK: orr     x1, xzr, #0xfffffffffffffffe
-; CHECK: stur    x1, [x29, #-16]
+; CHECK: orr     x0, xzr, #0xfffffffffffffffe
+; CHECK: stur    x0, [x29, #-16]
 ; CHECK: .set .Lsimple_seh$frame_escape_0, -8
 ; CHECK: ldur    w0, [x29, #-8]
 ; CHECK: bl      foo
 
   %o = alloca %struct.S, align 4
   call void (...) @llvm.localescape(%struct.S* %o)
   %x = getelementptr inbounds %struct.S, %struct.S* %o, i32 0, i32 0
   %0 = load i32, i32* %x, align 4

 ; Test SEH when stack realignment is needed in case highly aligned stack objects are present.
 define void @stack_realign() #0 personality i8* bitcast (i32 (...)* @__C_specific_handler to i8*) {
 entry:
 ; CHECK-LABEL: stack_realign
 ; CHECK: add     x29, sp, #16
 ; CHECK: sub     x9, sp, #64
 ; CHECK: and     sp, x9, #0xffffffffffffffe0
 ; CHECK: mov     x19, sp
-; CHECK: orr     x1, xzr, #0xfffffffffffffffe
-; CHECK: stur    x1, [x19, #16]
+; CHECK: orr     x0, xzr, #0xfffffffffffffffe
+; CHECK: stur    x0, [x19, #16]
 ; CHECK: .set .Lstack_realign$frame_escape_0, 32
 ; CHECK: ldr     w0, [x19, #32]
 ; CHECK: bl      foo
 
   %o = alloca %struct.S, align 32
   call void (...) @llvm.localescape(%struct.S* %o)
   %x = getelementptr inbounds %struct.S, %struct.S* %o, i32 0, i32 0
   %0 = load i32, i32* %x, align 32

llvm/trunk/test/CodeGen/AArch64/wineh-try-catch-cbz.ll

+; RUN: llc < %s | FileCheck %s
+
+; Make sure the prologue is sane.  (Doesn't need to exactly match this,
+; but the original issue only reproduced if the cbz was immediately
+; after the frame setup.)
+
+; CHECK:      sub     sp, sp, #32
+; CHECK-NEXT: stp     x29, x30, [sp, #16]
+; CHECK-NEXT: add     x29, sp, #16
+; CHECK-NEXT: orr     x1, xzr, #0xfffffffffffffffe
+; CHECK-NEXT: stur    x1, [x29, #-16]
+; CHECK-NEXT: cbz     w0, .LBB0_2
+
+target datalayout = "e-m:w-p:64:64-i32:32-i64:64-i128:128-n32:64-S128"
+target triple = "aarch64-unknown-windows-msvc19.11.0"
+
+; Function Attrs: uwtable
+define dso_local void @"?f@@YAXH@Z"(i32 %x) local_unnamed_addr #0 personality i8* bitcast (i32 (...)* @__CxxFrameHandler3 to i8*) {
+entry:
+  %cmp = icmp eq i32 %x, 0
+  br i1 %cmp, label %try.cont, label %if.then
+
+if.then:                                          ; preds = %entry
+  invoke void @"?g@@YAXXZ"()
+          to label %try.cont unwind label %catch.dispatch
+
+catch.dispatch:                                   ; preds = %if.then
+  %0 = catchswitch within none [label %catch] unwind to caller
+
+catch:                                            ; preds = %catch.dispatch
+  %1 = catchpad within %0 [i8* null, i32 64, i8* null]
+  catchret from %1 to label %try.cont
+
+try.cont:                                         ; preds = %entry, %if.then, %catch
+  ret void
+}
+
+declare dso_local void @"?g@@YAXXZ"() local_unnamed_addr #1
+
+declare dso_local i32 @__CxxFrameHandler3(...)

llvm/trunk/test/CodeGen/AArch64/wineh-try-catch.ll

 ; CHECK-LABEL: "?func@@YAHXZ":
 ; CHECK:       str     x28, [sp, #-48]!
 ; CHECK:       str     x21, [sp, #8]
 ; CHECK:       stp     x19, x20, [sp, #16]
 ; CHECK:       stp     x29, x30, [sp, #32]
 ; CHECK:       add     x29, sp, #32
 ; CHECK:       sub     sp, sp, #624
 ; CHECK:       mov     x19, sp
-; CHECK:       orr     x1, xzr, #0xfffffffffffffffe
-; CHECK:       stur    x1, [x19]
+; CHECK:       orr     x0, xzr, #0xfffffffffffffffe
+; CHECK:       stur    x0, [x19]
 
 ; Now check that x is stored at fp - 20.  We check that this is the same
 ; location accessed from the funclet to retrieve x.
 ; CHECK:       orr     w8, wzr, #0x1
 ; CHECK:       stur    w8, [x29, [[X_OFFSET:#-[1-9][0-9]+]]
 
 ; Check the offset off the frame pointer at which B is located.
 ; Check the same offset is used to pass the address of B to init2 in the