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

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.

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

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@326764 91177308-0d34-0410-b5e6-96231b3b80d8

diff --git a/lib/Target/X86/X86FrameLowering.cpp b/lib/Target/X86/X86FrameLowering.cpp
index 34febe6..d052998 100644 (file)
--- a/lib/Target/X86/X86FrameLowering.cpp
+++ b/lib/Target/X86/X86FrameLowering.cpp
@@ -1213,30 +1213,34 @@ void X86FrameLowering::emitPrologue(MachineFunction &MF,
     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!");
-
-      // Save EAX
-      BuildMI(MBB, MBBI, DL, TII.get(X86::PUSH32r))
-        .addReg(X86::EAX, RegState::Kill)
-        .setMIFlag(MachineInstr::FrameSetup);
+      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 {
@@ -1251,10 +1255,14 @@ void X86FrameLowering::emitPrologue(MachineFunction &MF,
     emitStackProbe(MF, MBB, MBBI, DL, true);
 
     if (isEAXAlive) {
-      // Restore EAX
-      MachineInstr *MI =
-          addRegOffset(BuildMI(MF, DL, TII.get(X86::MOV32rm), X86::EAX),
-                       StackPtr, false, NumBytes - 4);
+      // 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);
     }

diff --git a/test/CodeGen/X86/shrink-wrap-chkstk-x86_64.ll b/test/CodeGen/X86/shrink-wrap-chkstk-x86_64.ll
new file mode 100644 (file)
index 0000000..b908200
--- /dev/null
+++ b/test/CodeGen/X86/shrink-wrap-chkstk-x86_64.ll
@@ -0,0 +1,36 @@
+; 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