[X86] Handle EAX being live when calling chkstk for x86_64

[mstorsjo]

Summary

EAX can turn out to be alive here, when shrink wrapping is done (which is allowed when using dwarf exceptions, contrary to the normal case with WinCFI).

This fixes PR36487.

[mstorsjo]

This is just a first draft - do you have any suggestions on a testcase for it? Just including the reduced IR from the bug report, or something more synthetic?

[RKSimon]

In D43968#1024404, @mstorsjo wrote:

This is just a first draft - do you have any suggestions on a testcase for it? Just including the reduced IR from the bug report, or something more synthetic?

Yes you should at least include that, if you can try to reduce it further - what is the coverage like for the 32-bit path?

[rnk]

I see, shrink-wrapping makes it so that EAX may be live into the prologue.

I don't like that the CFI becomes imprecise, though. The user won't be able to break on __chkstk and take a backtrace. Does that work today when shrink wrapping isn't in use? If not, then let's just do this. If it does work, we should consider emitting an 8 byte .seh_stackalloc for the push or RAX and a smaller update after __chkstk.

This needs a test case. I think this might be a good time to use an MIR test case.

[mstorsjo]

In D43968#1024556, @rnk wrote:

I see, shrink-wrapping makes it so that EAX may be live into the prologue.

I don't like that the CFI becomes imprecise, though. The user won't be able to break on __chkstk and take a backtrace. Does that work today when shrink wrapping isn't in use?

It doesn't seem so; both with dwarf and SEH, and GCC with SEH, I get an incorrect backtrace in gdb when breaking on ___chkstk_ms in this case. (Breakpoints in other places give proper backtraces though.)

If not, then let's just do this. If it does work, we should consider emitting an 8 byte .seh_stackalloc for the push or RAX and a smaller update after __chkstk.

Well, this case isn't too relevant for SEH at all, since shrink wrapping is disabled altogether when SEH is enabled - in lib/CodeGen/ShrinkWrap.cpp, ShrinkWrap::isShrinkWrapEnabled:

bool ShrinkWrap::isShrinkWrapEnabled(const MachineFunction &MF) {
  const TargetFrameLowering *TFI = MF.getSubtarget().getFrameLowering();

  switch (EnableShrinkWrapOpt) {
  case cl::BOU_UNSET:
    return TFI->enableShrinkWrapping(MF) &&
           // Windows with CFI has some limitations that make it impossible
           // to use shrink-wrapping.
           !MF.getTarget().getMCAsmInfo()->usesWindowsCFI() &&

[mstorsjo]

Added a testcase based on the reduced IR from the bug.

[rnk]

In D43968#1024973, @mstorsjo wrote:

Well, this case isn't too relevant for SEH at all, since shrink wrapping is disabled altogether when SEH is enabled - in lib/CodeGen/ShrinkWrap.cpp, ShrinkWrap::isShrinkWrapEnabled:

Fair enough.

llvm/trunk/lib/Target/X86/X86FrameLowering.cpp

void X86FrameLowering::emitPrologue(MachineFunction &MF,
   if (AlignedNumBytes >= StackProbeSize && UseStackProbe) {
     assert(!X86FI->getUsesRedZone() &&
            "The Red Zone is not accounted for in stack probes");
 
     // Check whether EAX is livein for this block.
     bool isEAXAlive = isEAXLiveIn(MBB);
 
     if (isEAXAlive) {
-      // Sanity check that EAX is not livein for this function.
-      // It should not be, so throw an assert.
-      assert(!Is64Bit && "EAX is livein in x64 case!");
-
+      if (Is64Bit) {
+        // Save RAX
+        BuildMI(MBB, MBBI, DL, TII.get(X86::PUSH64r))
+          .addReg(X86::RAX, RegState::Kill)
+          .setMIFlag(MachineInstr::FrameSetup);
+      } else {
         // Save EAX
         BuildMI(MBB, MBBI, DL, TII.get(X86::PUSH32r))
           .addReg(X86::EAX, RegState::Kill)
           .setMIFlag(MachineInstr::FrameSetup);
       }
+    }
 
     if (Is64Bit) {
       // Handle the 64-bit Windows ABI case where we need to call __chkstk.
       // Function prologue is responsible for adjusting the stack pointer.
-      if (isUInt<32>(NumBytes)) {
+      int Alloc = isEAXAlive ? NumBytes - 8 : NumBytes;
+      if (isUInt<32>(Alloc)) {
         BuildMI(MBB, MBBI, DL, TII.get(X86::MOV32ri), X86::EAX)
-            .addImm(NumBytes)
+            .addImm(Alloc)
             .setMIFlag(MachineInstr::FrameSetup);
-      } else if (isInt<32>(NumBytes)) {
+      } else if (isInt<32>(Alloc)) {
         BuildMI(MBB, MBBI, DL, TII.get(X86::MOV64ri32), X86::RAX)
-            .addImm(NumBytes)
+            .addImm(Alloc)
             .setMIFlag(MachineInstr::FrameSetup);
       } else {
         BuildMI(MBB, MBBI, DL, TII.get(X86::MOV64ri), X86::RAX)
-            .addImm(NumBytes)
+            .addImm(Alloc)
             .setMIFlag(MachineInstr::FrameSetup);
       }
     } else {
       // Allocate NumBytes-4 bytes on stack in case of isEAXAlive.
       // We'll also use 4 already allocated bytes for EAX.
       BuildMI(MBB, MBBI, DL, TII.get(X86::MOV32ri), X86::EAX)
           .addImm(isEAXAlive ? NumBytes - 4 : NumBytes)
           .setMIFlag(MachineInstr::FrameSetup);
     }
 
     // Call __chkstk, __chkstk_ms, or __alloca.
     emitStackProbe(MF, MBB, MBBI, DL, true);
 
     if (isEAXAlive) {
-      // Restore EAX
-      MachineInstr *MI =
-          addRegOffset(BuildMI(MF, DL, TII.get(X86::MOV32rm), X86::EAX),
+      // Restore RAX/EAX
+      MachineInstr *MI;
+      if (Is64Bit)
+        MI = addRegOffset(BuildMI(MF, DL, TII.get(X86::MOV64rm), X86::RAX),
+                          StackPtr, false, NumBytes - 8);
+      else
+        MI = addRegOffset(BuildMI(MF, DL, TII.get(X86::MOV32rm), X86::EAX),
                           StackPtr, false, NumBytes - 4);
       MI->setFlag(MachineInstr::FrameSetup);
       MBB.insert(MBBI, MI);
     }
   } else if (NumBytes) {
     emitSPUpdate(MBB, MBBI, DL, -(int64_t)NumBytes, /*InEpilogue=*/false);
   }
 
   if (NeedsWinCFI && NumBytes) {

llvm/trunk/test/CodeGen/X86/shrink-wrap-chkstk-x86_64.ll

+; RUN: llc -mtriple=x86_64-windows-gnu -exception-model=dwarf < %s | FileCheck %s
+
+%struct.A = type { [4096 x i8] }
+
+@a = common global i32 0, align 4
+@b = common global i32 0, align 4
+
+define void @fn1() nounwind uwtable {
+entry:
+  %ctx = alloca %struct.A, align 1
+  %0 = load i32, i32* @a, align 4
+  %tobool = icmp eq i32 %0, 0
+  %div = sdiv i32 %0, 6
+  %cond = select i1 %tobool, i32 %div, i32 %0
+  store i32 %cond, i32* @b, align 4
+  %1 = getelementptr inbounds %struct.A, %struct.A* %ctx, i64 0, i32 0, i64 0
+  call void @llvm.lifetime.start.p0i8(i64 4096, i8* nonnull %1)
+  %2 = ptrtoint %struct.A* %ctx to i64
+  %3 = trunc i64 %2 to i32
+  call void @fn2(i32 %3)
+  call void @llvm.lifetime.end.p0i8(i64 4096, i8* nonnull %1)
+  ret void
+}
+
+declare void @fn2(i32)
+declare void @llvm.lifetime.start.p0i8(i64, i8* nocapture)
+declare void @llvm.lifetime.end.p0i8(i64, i8* nocapture)
+
+; CHECK-LABEL: fn1:
+; CHECK: pushq %rax
+; CHECK: movl $4128, %eax
+; CHECK: callq ___chkstk_ms
+; CHECK: subq %rax, %rsp
+; CHECK: movq 4128(%rsp), %rax
+
+; CHECK: addq $4136, %rsp