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

[Martin Storsjo via llvm-commits]

Author: mstorsjo
Date: Mon Mar  5 22:00:13 2018
New Revision: 326764

URL: http://llvm.org/viewvc/llvm-project?rev=326764&view=rev
Log:
[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

Added:
    llvm/trunk/test/CodeGen/X86/shrink-wrap-chkstk-x86_64.ll
Modified:
    llvm/trunk/lib/Target/X86/X86FrameLowering.cpp

Modified: llvm/trunk/lib/Target/X86/X86FrameLowering.cpp
URL: http://llvm.org/viewvc/llvm-project/llvm/trunk/lib/Target/X86/X86FrameLowering.cpp?rev=326764&r1=326763&r2=326764&view=diff
==============================================================================
--- llvm/trunk/lib/Target/X86/X86FrameLowering.cpp (original)
+++ llvm/trunk/lib/Target/X86/X86FrameLowering.cpp Mon Mar  5 22:00:13 2018
@@ -1213,30 +1213,34 @@ void X86FrameLowering::emitPrologue(Mach
     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(Mach
     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);
     }

Added: llvm/trunk/test/CodeGen/X86/shrink-wrap-chkstk-x86_64.ll
URL: http://llvm.org/viewvc/llvm-project/llvm/trunk/test/CodeGen/X86/shrink-wrap-chkstk-x86_64.ll?rev=326764&view=auto
==============================================================================
--- llvm/trunk/test/CodeGen/X86/shrink-wrap-chkstk-x86_64.ll (added)
+++ llvm/trunk/test/CodeGen/X86/shrink-wrap-chkstk-x86_64.ll Mon Mar  5 22:00:13 2018
@@ -0,0 +1,36 @@
+; RUN: llc -mtriple=x86_64-windows-gnu -exception-model=dwarf < %s | FileCheck %s
+
+%struct.A = type { [4096 x i8] }
+
+ at a = common global i32 0, align 4
+ at 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