[PATCH 2/2] [x86] Add support for "probe-stack"

[Zoxc]

Summary

This is an alternative to http://reviews.llvm.org/D9654

Stack probes are done in an loop in the prologue instead of calling external functions.

Note that when modifying basic blocks in the emitPrologue the epilog will be inserted into the wrong basic block. This is because
the basic blocks to insert epilogs into is calculated before
emitPrologue is called.
I fixed this by adding the following code after the emitPrologue call
in PEI::insertPrologEpilogCode:

RestoreBlocks.clear();
calculateSets(Fn);

This doesn't seem like a very nice solution.

Another thing is that I'm not sure that I modified all the test for stack probing correctly (there were quite a few cases of it).

[Zoxc]

Code generation on Windows is preserved in this version.

[Zoxc]

Rebased tests and added a test for variable length stack allocation.

[AndyAyers]

We have some changes coming in our fork (see https://github.com/Microsoft/llvm/pull/79) where for the inline probes we emit a call to a stub during prolog generation, and then we come back later after epilogs are created to expand the probe sequence inline. We should be able to upstream this soon.

This seems cleaner to us then trying to repair state after prolog generation...

[Zoxc]

Is said state actually used before prolog generation? Just moving the generation to after prolog generation seems like the nicest solution if it's possible.

[Zoxc]

This now uses the same method as the CoreCLR inline stack probes.

[Zoxc]

Rebased patch.

include/llvm/CodeGen/MachineFunction.h

MachineBasicBlock *getBlockNumbered(unsigned N) const {
     assert(N < MBBNumbering.size() && "Illegal block number");
     assert(MBBNumbering[N] && "Block was removed from the machine function!");
     return MBBNumbering[N];
   }
 
   /// Should we be emitting segmented stack stuff for the function
   bool shouldSplitStack() const;
 
+  /// \brief Should we be probing the stack for the function.
+  ///
+  /// Probing the stack means that we must read or write to the stack on every
+  /// page. This is to ensure that a guard page will be hit and stack overflow
+  /// can be detected. We insert instructions to do this when allocating from
+  /// the stack.
+  bool shouldProbeStack() const;
+
   /// getNumBlockIDs - Return the number of MBB ID's allocated.
   ///
   unsigned getNumBlockIDs() const { return (unsigned)MBBNumbering.size(); }
 
   /// RenumberBlocks - This discards all of the MachineBasicBlock numbers and
   /// recomputes them.  This guarantees that the MBB numbers are sequential,
   /// dense, and match the ordering of the blocks within the function.  If a
   /// specific MachineBasicBlock is specified, only that block and those after

lib/CodeGen/MachineFunction.cpp

getOrCreateJumpTableInfo(unsigned EntryKind) {
   return JumpTableInfo;
 }
 
 /// Should we be emitting segmented stack stuff for the function
 bool MachineFunction::shouldSplitStack() const {
   return getFunction()->hasFnAttribute("split-stack");
 }
 
+bool MachineFunction::shouldProbeStack() const {
+  return getFunction()->hasFnAttribute("probe-stack");
+}
+
 /// This discards all of the MachineBasicBlock numbers and recomputes them.
 /// This guarantees that the MBB numbers are sequential, dense, and match the
 /// ordering of the blocks within the function.  If a specific MachineBasicBlock
 /// is specified, only that block and those after it are renumbered.
 void MachineFunction::RenumberBlocks(MachineBasicBlock *MBB) {
   if (empty()) { MBBNumbering.clear(); return; }
   MachineFunction::iterator MBBI, E = end();
   if (MBB == nullptr)

lib/Target/X86/X86FrameLowering.h

public:
 
   /// True if the 64-bit frame or stack pointer should be used. True for most
   /// 64-bit targets with the exception of x32. If this is false, 32-bit
   /// instruction operands should be used to manipulate StackPtr and FramePtr.
   bool Uses64BitFramePtr;
 
   unsigned StackPtr;
 
+  void pushRegForStackProbeCall(MachineFunction &MF,
+                                MachineBasicBlock &MBB,
+                                MachineBasicBlock::iterator MBBI,
+                                DebugLoc DL,
+                                bool &IsAlive,
+                                unsigned RegType,
+                                uint64_t &NumBytes) const;
+  void popRegForStackProbeCall(MachineFunction &MF,
+                               MachineBasicBlock &MBB,
+                               MachineBasicBlock::iterator MBBI,
+                               DebugLoc DL,
+                               bool &IsAlive,
+                               unsigned RegType,
+                               uint64_t &NumBytes) const;
+
   /// Emit target stack probe code. This is required for all
   /// large stack allocations on Windows. The caller is required to materialize
   /// the number of bytes to probe in RAX/EAX. Returns instruction just
   /// after the expansion.
   MachineInstr *emitStackProbe(MachineFunction &MF, MachineBasicBlock &MBB,
                                MachineBasicBlock::iterator MBBI, DebugLoc DL,
                                bool InProlog) const;
 
   /// Replace a StackProbe inline-stub with the actual probe code inline.
   void inlineStackProbe(MachineFunction &MF,
                         MachineBasicBlock &PrologMBB) const override;
 
+  /// Emit a loop that probes the stack. The caller is required to materialize
+  /// the number of bytes to probe in RAX/EAX.
+  MachineInstr *emitStackProbes(MachineFunction &MF, MachineBasicBlock &MBB,
+                                MachineBasicBlock::iterator MBBI, DebugLoc DL,
+                                bool InProlog) const;
+
   void emitCalleeSavedFrameMoves(MachineBasicBlock &MBB,
                                  MachineBasicBlock::iterator MBBI,
                                  DebugLoc DL) const;
 
   /// emitProlog/emitEpilog - These methods insert prolog and epilog code into
   /// the function.
   void emitPrologue(MachineFunction &MF, MachineBasicBlock &MBB) const override;
   void emitEpilogue(MachineFunction &MF, MachineBasicBlock &MBB) const override;

lib/Target/X86/X86FrameLowering.cpp

for (auto CS : AvailableRegs)
       if (!Uses.count(CS) && CS != X86::RIP)
         return CS;
   }
   }
 
   return 0;
 }
 
-static bool isEAXLiveIn(MachineBasicBlock &MBB) {
+static bool isLiveIn(MachineBasicBlock &MBB, unsigned CheckReg) {
+  CheckReg = getX86SubSuperRegister(CheckReg, 32);
+
   for (MachineBasicBlock::RegisterMaskPair RegMask : MBB.liveins()) {
     unsigned Reg = RegMask.PhysReg;
 
-    if (Reg == X86::RAX || Reg == X86::EAX || Reg == X86::AX ||
-        Reg == X86::AH || Reg == X86::AL)
+    if (getX86SubSuperRegisterOrZero(Reg, 32) == CheckReg)
       return true;
   }
 
   return false;
 }
 
 /// Check if the flags need to be preserved before the terminators.
 /// This would be the case, if the eflags is live-in of the region
void X86FrameLowering::emitSPUpdate(MachineBasicBlock &MBB,
   DebugLoc DL = MBB.findDebugLoc(MBBI);
 
   while (Offset) {
     if (Offset > Chunk) {
       // Rather than emit a long series of instructions for large offsets,
       // load the offset into a register and do one sub/add
       unsigned Reg = 0;
 
-      if (isSub && !isEAXLiveIn(MBB))
+      if (isSub && !isLiveIn(MBB, X86::EAX))
         Reg = (unsigned)(Is64Bit ? X86::RAX : X86::EAX);
       else
         Reg = findDeadCallerSavedReg(MBB, MBBI, TRI, Is64Bit);
 
       if (Reg) {
         unsigned Opc = Is64Bit ? X86::MOV64ri : X86::MOV32ri;
         BuildMI(MBB, MBBI, DL, TII.get(Opc), Reg)
           .addImm(Offset);
MachineInstr *X86FrameLowering::emitStackProbe(MachineFunction &MF,
                                                bool InProlog) const {
   const X86Subtarget &STI = MF.getSubtarget<X86Subtarget>();
   if (STI.isTargetWindowsCoreCLR()) {
     if (InProlog) {
       return emitStackProbeInlineStub(MF, MBB, MBBI, DL, true);
     } else {
       return emitStackProbeInline(MF, MBB, MBBI, DL, false);
     }
+  } else if (!(STI.isOSWindows() && !STI.isTargetMachO())) {
+    if (InProlog) {
+      return emitStackProbeInlineStub(MF, MBB, MBBI, DL, true);
+    } else {
+      return emitStackProbes(MF, MBB, MBBI, DL, false);
+    }
   } else {
     return emitStackProbeCall(MF, MBB, MBBI, DL, InProlog);
   }
 }
 
 void X86FrameLowering::inlineStackProbe(MachineFunction &MF,
                                         MachineBasicBlock &PrologMBB) const {
   const StringRef ChkStkStubSymbol = "__chkstk_stub";
void X86FrameLowering::inlineStackProbe(MachineFunction &MF,
 
   if (ChkStkStub != nullptr) {
     assert(!ChkStkStub->isBundled() &&
            "Not expecting bundled instructions here");
     MachineBasicBlock::iterator MBBI = std::next(ChkStkStub->getIterator());
     assert(std::prev(MBBI).operator==(ChkStkStub) &&
       "MBBI expected after __chkstk_stub.");
     DebugLoc DL = PrologMBB.findDebugLoc(MBBI);
+    if (STI.isOSWindows() && !STI.isTargetMachO()) {
       emitStackProbeInline(MF, PrologMBB, MBBI, DL, true);
+    } else {
+      emitStackProbes(MF, PrologMBB, MBBI, DL, true);
+    }
     ChkStkStub->eraseFromParent();
   }
 }
 
 MachineInstr *X86FrameLowering::emitStackProbeInline(
   MachineFunction &MF, MachineBasicBlock &MBB,
   MachineBasicBlock::iterator MBBI, DebugLoc DL, bool InProlog) const {
   const X86Subtarget &STI = MF.getSubtarget<X86Subtarget>();
MachineInstr *X86FrameLowering::emitStackProbeInlineStub(
   assert(InProlog && "ChkStkStub called outside prolog!");
 
   BuildMI(MBB, MBBI, DL, TII.get(X86::CALLpcrel32))
       .addExternalSymbol("__chkstk_stub");
 
   return MBBI;
 }
 
+void X86FrameLowering::pushRegForStackProbeCall(MachineFunction &MF,
+                                                MachineBasicBlock &MBB,
+                                                MachineBasicBlock::iterator MBBI,
+                                                DebugLoc DL,
+                                                bool &IsAlive,
+                                                unsigned RegType,
+                                                uint64_t &NumBytes) const {
+  IsAlive = isLiveIn(MBB, RegType);
+
+  if (!IsAlive) {
+    return;
+  }
+
+  auto Reg = getX86SubSuperRegister(RegType, Is64Bit ? 64 : 32);
+
+  // Save the register on the stack.
+  BuildMI(MBB, MBBI, DL, TII.get(Is64Bit ? X86::PUSH64r : X86::PUSH32r))
+      .addReg(Reg, RegState::Kill)
+      .setMIFlag(MachineInstr::FrameSetup);
+
+  // Reuse the space from the spill as a stack allocation.
+  NumBytes -= SlotSize;
+}
+
+void X86FrameLowering::popRegForStackProbeCall(MachineFunction &MF,
+                                                MachineBasicBlock &MBB,
+                                                MachineBasicBlock::iterator MBBI,
+                                                DebugLoc DL,
+                                                bool &IsAlive,
+                                                unsigned RegType,
+                                                uint64_t &NumBytes) const {
+  if (!IsAlive) {
+    return;
+  }
+
+  // Restore the register from the stack slot.
+
+  auto Reg = getX86SubSuperRegister(RegType, Is64Bit ? 64 : 32);
+
+  auto MIB = BuildMI(MF, DL,
+                     TII.get(Is64Bit ? X86::MOV64rm : X86::MOV32rm),
+                     Reg);
+  MachineInstr *MI = addRegOffset(MIB, StackPtr, false, NumBytes);
+  MI->setFlag(MachineInstr::FrameSetup);
+  MBB.insert(MBBI, MI);
+
+  NumBytes += SlotSize;
+}
+
+MachineInstr *X86FrameLowering::emitStackProbes(MachineFunction &MF,
+                                                MachineBasicBlock &MBB,
+                                                MachineBasicBlock::iterator MBBI,
+                                                DebugLoc DL,
+                                                bool InProlog) const {
+  const X86Subtarget &STI = MF.getSubtarget<X86Subtarget>();
+  const TargetInstrInfo &TII = *STI.getInstrInfo();
+  const BasicBlock *LLVM_BB = MBB.getBasicBlock();
+
+  // RAX contains the number of bytes of desired stack adjustment.
+  // The handling here assumes this value has already been updated so as to
+  // maintain stack alignment.
+  //
+  // We need to exit with RSP modified by this amount and execute suitable
+  // page touches to notify the OS that we're growing the stack responsibly.
+  // All stack probing must be done without modifying RSP.
+
+  // Set up the new basic blocks
+  MachineBasicBlock *LoopMBB = MF.CreateMachineBasicBlock(LLVM_BB);
+  MachineBasicBlock *ContinueMBB = MF.CreateMachineBasicBlock(LLVM_BB);
+
+  MachineFunction::iterator MBBIter = std::next(MBB.getIterator());
+
+  MF.insert(MBBIter, LoopMBB);
+  MF.insert(MBBIter, ContinueMBB);
+
+  // Split MBB and move the tail portion down to ContinueMBB.
+  MachineBasicBlock::iterator BeforeMBBI = std::prev(MBBI);
+  ContinueMBB->splice(ContinueMBB->begin(), &MBB, MBBI, MBB.end());
+  ContinueMBB->transferSuccessorsAndUpdatePHIs(&MBB);
+
+  const int64_t GuardSize = 0x1000;
+
+  auto RType = Is64Bit ? 64 : 32;
+
+  auto InputReg = getX86SubSuperRegister(X86::RAX, RType);
+
+  // Registers we need. For the normal case we use virtual
+  // registers. For the prolog expansion we use RAX, RCX and RDX.
+  MachineRegisterInfo &MRI = MF.getRegInfo();
+  const TargetRegisterClass *RegClass = Is64Bit ? &X86::GR64RegClass : &X86::GR32RegClass;
+  unsigned
+      SizeReg = InProlog ? InputReg :
+                           MRI.createVirtualRegister(RegClass),
+      StackReg = InProlog ? getX86SubSuperRegister(X86::RCX, RType) :
+                            MRI.createVirtualRegister(RegClass),
+      CountReg = InProlog ? getX86SubSuperRegister(X86::RDX, RType) :
+                            MRI.createVirtualRegister(RegClass),
+      PStackReg = InProlog ? getX86SubSuperRegister(X86::RCX, RType) :
+                            MRI.createVirtualRegister(RegClass),
+      PCountReg = InProlog ? getX86SubSuperRegister(X86::RDX, RType) :
+                            MRI.createVirtualRegister(RegClass),
+      LStackReg = InProlog ? getX86SubSuperRegister(X86::RCX, RType) :
+                            MRI.createVirtualRegister(RegClass),
+      LCountReg = InProlog ? getX86SubSuperRegister(X86::RDX, RType) :
+                            MRI.createVirtualRegister(RegClass);
+
+  auto SPReg = getX86SubSuperRegister(X86::RSP, RType);
+
+  if (InProlog) {
+    ContinueMBB->addLiveIn((MCPhysReg)InputReg);
+    LoopMBB->addLiveIn((MCPhysReg)InputReg);
+    LoopMBB->addLiveIn((MCPhysReg)LCountReg);
+    LoopMBB->addLiveIn((MCPhysReg)LStackReg);
+
+    for (MachineBasicBlock::livein_iterator i = MBB.livein_begin(),
+                                            e = MBB.livein_end();
+         i != e; i++) {
+      if (!LoopMBB->isLiveIn(i->PhysReg, i->LaneMask))
+        LoopMBB->addLiveIn(*i);
+      if (!ContinueMBB->isLiveIn(i->PhysReg, i->LaneMask))
+        ContinueMBB->addLiveIn(*i);
+    }
+  } else {
+    BuildMI(&MBB, DL, TII.get(X86::COPY), SizeReg).addReg(InputReg);
+  }
+
+  BuildMI(&MBB, DL, TII.get(Is64Bit ? X86::MOV64rr : X86::MOV32rr), CountReg).addReg(SizeReg);
+  BuildMI(&MBB, DL, TII.get(Is64Bit ? X86::MOV64rr : X86::MOV32rr), StackReg).addReg(SPReg);
+
+  if (!InProlog) {
+    BuildMI(LoopMBB, DL, TII.get(X86::PHI), PCountReg)
+        .addReg(CountReg)
+        .addMBB(&MBB)
+        .addReg(LCountReg)
+        .addMBB(LoopMBB);
+    BuildMI(LoopMBB, DL, TII.get(X86::PHI), PStackReg)
+        .addReg(StackReg)
+        .addMBB(&MBB)
+        .addReg(LStackReg)
+        .addMBB(LoopMBB);
+  }
+
+  BuildMI(LoopMBB, DL, TII.get(Is64Bit ? X86::OR64mi8 : X86::OR32mi8))
+    .addReg(PStackReg)
+    .addImm(1)
+    .addReg(0)
+    .addImm(0)
+    .addReg(0)
+    .addImm(0);
+
+  BuildMI(LoopMBB, DL, TII.get(Is64Bit ? X86::SUB64ri32 : X86::SUB32ri),
+          LStackReg)
+      .addReg(PStackReg)
+      .addImm(GuardSize);
+
+  BuildMI(LoopMBB, DL, TII.get(Is64Bit ? X86::SUB64ri32 : X86::SUB32ri),
+          LCountReg)
+      .addReg(PCountReg)
+      .addImm(GuardSize);
+
+  BuildMI(LoopMBB, DL, TII.get(X86::JAE_1)).addMBB(LoopMBB);
+
+  MachineBasicBlock::iterator ContinueMBBI = ContinueMBB->getFirstNonPHI();
+
+  // Now that the probing is done, add code to continueMBB to update
+  // the stack pointer for real.
+  BuildMI(*ContinueMBB, ContinueMBBI, DL, TII.get(Is64Bit ? X86::SUB64rr : X86::SUB32rr), SPReg)
+      .addReg(SPReg)
+      .addReg(SizeReg);
+
+  // Add the control flow edges we need.
+  MBB.addSuccessor(LoopMBB);
+  LoopMBB->addSuccessor(ContinueMBB);
+  LoopMBB->addSuccessor(LoopMBB);
+
+  // Mark all the instructions added to the prolog as frame setup.
+  if (InProlog) {
+    for (++BeforeMBBI; BeforeMBBI != MBB.end(); ++BeforeMBBI) {
+      BeforeMBBI->setFlag(MachineInstr::FrameSetup);
+    }
+    for (MachineBasicBlock::iterator MI = LoopMBB->begin();
+         MI != LoopMBB->end(); ++MI) {
+      MI->setFlag(MachineInstr::FrameSetup);
+    }
+    for (MachineBasicBlock::iterator CMBBI = ContinueMBB->begin();
+         CMBBI != ContinueMBBI; ++CMBBI) {
+      CMBBI->setFlag(MachineInstr::FrameSetup);
+    }
+  }
+
+  return ContinueMBBI;
+}
+
 static unsigned calculateSetFPREG(uint64_t SPAdjust) {
   // Win64 ABI has a less restrictive limitation of 240; 128 works equally well
   // and might require smaller successive adjustments.
   const uint64_t Win64MaxSEHOffset = 128;
   uint64_t SEHFrameOffset = std::min(SPAdjust, Win64MaxSEHOffset);
   // Win64 ABI requires 16-byte alignment for the UWOP_SET_FPREG opcode.
   return SEHFrameOffset & -16;
 }
void X86FrameLowering::emitPrologue(MachineFunction &MF,
   int TailCallReturnAddrDelta = X86FI->getTCReturnAddrDelta();
   if (TailCallReturnAddrDelta && IsWin64Prologue)
     report_fatal_error("Can't handle guaranteed tail call under win64 yet");
 
   if (TailCallReturnAddrDelta < 0)
     X86FI->setCalleeSavedFrameSize(
       X86FI->getCalleeSavedFrameSize() - TailCallReturnAddrDelta);
 
-  bool UseStackProbe = (STI.isOSWindows() && !STI.isTargetMachO());
+  bool UseRedZone = false;
+  bool WindowsStackProbe = STI.isOSWindows() && !STI.isTargetMachO();
+  bool UseStackProbe = WindowsStackProbe || MF.shouldProbeStack();
 
   // The default stack probe size is 4096 if the function has no stackprobesize
   // attribute.
   unsigned StackProbeSize = 4096;
   if (Fn->hasFnAttribute("stack-probe-size"))
     Fn->getFnAttribute("stack-probe-size")
         .getValueAsString()
         .getAsInteger(0, StackProbeSize);
 
   // If this is x86-64 and the Red Zone is not disabled, if we are a leaf
   // function, and use up to 128 bytes of stack space, don't have a frame
   // pointer, calls, or dynamic alloca then we do not need to adjust the
   // stack pointer (we fit in the Red Zone). We also check that we don't
   // push and pop from the stack.
   if (Is64Bit && !Fn->hasFnAttribute(Attribute::NoRedZone) &&
       !TRI->needsStackRealignment(MF) &&
       !MFI->hasVarSizedObjects() &&             // No dynamic alloca.
       !MFI->adjustsStack() &&                   // No calls.
       !IsWin64CC &&                             // Win64 has no Red Zone
+
+      !(UseStackProbe && StackSize > 128) && // Only use the Red Zone if we can
+                                             // fit the whole stack in it
+                                             // and thus stack probes won't be
+                                             // needed
+
       !MFI->hasCopyImplyingStackAdjustment() && // Don't push and pop.
       !MF.shouldSplitStack()) {                 // Regular stack
     uint64_t MinSize = X86FI->getCalleeSavedFrameSize();
     if (HasFP) MinSize += SlotSize;
     X86FI->setUsesRedZone(MinSize > 0 || StackSize > 0);
     StackSize = std::max(MinSize, StackSize > 128 ? StackSize - 128 : 0);
     MFI->setStackSize(StackSize);
+    UseRedZone = true;
   }
 
   // Insert stack pointer adjustment for later moving of return addr.  Only
   // applies to tail call optimized functions where the callee argument stack
   // size is bigger than the callers.
   if (TailCallReturnAddrDelta < 0) {
     BuildStackAdjustment(MBB, MBBI, DL, TailCallReturnAddrDelta,
                          /*InEpilogue=*/false)
void X86FrameLowering::emitPrologue(MachineFunction &MF,
   // __chkstk is only responsible for probing the stack.  The 64-bit prologue is
   // responsible for adjusting the stack pointer.  Touching the stack at 4K
   // increments is necessary to ensure that the guard pages used by the OS
   // virtual memory manager are allocated in correct sequence.
   uint64_t AlignedNumBytes = NumBytes;
   if (IsWin64Prologue && !IsFunclet && TRI->needsStackRealignment(MF))
     AlignedNumBytes = alignTo(AlignedNumBytes, MaxAlign);
   if (AlignedNumBytes >= StackProbeSize && UseStackProbe) {
-    // Check whether EAX is livein for this block.
-    bool isEAXAlive = isEAXLiveIn(MBB);
+    assert(!UseRedZone && "The Red Zone is not accounted for in stack probes");
+
+    uint64_t PageSize = 0x1000;
 
-    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)
+    // If we only need to probe 5 pages or below, we just emit instructions to
+    // do that instead of using the loop. This is just what the loop in
+    // the called function would do. 5 probes is what GCC will do before using
+    // a loop.
+    if (!WindowsStackProbe && NumBytes <= 5 * PageSize) {
+      for (uint64_t i = 0; i < NumBytes / PageSize; ++i) {
+        BuildMI(MBB, MBBI, DL, TII.get(Is64Bit ? X86::OR64mi8 : X86::OR32mi8))
+          .addReg(StackPtr)
+          .addImm(1)
+          .addReg(0)
+          .addImm(- (i + 1) * PageSize)
+          .addReg(0)
+          .addImm(0)
           .setMIFlag(MachineInstr::FrameSetup);
       }
 
+      BuildMI(MBB, MBBI, DL, TII.get(Is64Bit ? X86::SUB64ri32 : X86::SUB32ri),
+              StackPtr)
+         .addReg(StackPtr)
+         .addImm(NumBytes)
+         .setMIFlag(MachineInstr::FrameSetup);
+
+    } else {
+      // In the large code model, we have to load the stack probe function into a scratch
+      // register to call it. R11 is used for that.
+      bool SpillR11 = WindowsStackProbe &&
+                      Is64Bit &&
+                      MF.getTarget().getCodeModel() == CodeModel::Large;
+
+      // We spill the registers we need for the stack probes.
+
+      bool RAXAlive, RCXAlive, RDXAlive, R11Alive;
+
+      pushRegForStackProbeCall(MF, MBB, MBBI, DL, RAXAlive, X86::RAX, NumBytes);
+      if (!WindowsStackProbe) {
+        pushRegForStackProbeCall(MF, MBB, MBBI, DL, RCXAlive, X86::RCX, NumBytes);
+        pushRegForStackProbeCall(MF, MBB, MBBI, DL, RDXAlive, X86::RDX, NumBytes);
+      }
+      if (SpillR11) {
+        pushRegForStackProbeCall(MF, MBB, MBBI, DL, R11Alive, X86::R11, NumBytes);
+      }
+
       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)) {
           BuildMI(MBB, MBBI, DL, TII.get(X86::MOV32ri), X86::EAX)
               .addImm(NumBytes)
               .setMIFlag(MachineInstr::FrameSetup);
         } else if (isInt<32>(NumBytes)) {
           BuildMI(MBB, MBBI, DL, TII.get(X86::MOV64ri32), X86::RAX)
               .addImm(NumBytes)
               .setMIFlag(MachineInstr::FrameSetup);
         } else {
           BuildMI(MBB, MBBI, DL, TII.get(X86::MOV64ri), X86::RAX)
               .addImm(NumBytes)
               .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)
+            .addImm(NumBytes)
             .setMIFlag(MachineInstr::FrameSetup);
       }
 
-    // Call __chkstk, __chkstk_ms, or __alloca.
+      // Ensure stack probes are done
       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);
-      MI->setFlag(MachineInstr::FrameSetup);
-      MBB.insert(MBBI, MI);
+      // Now we restore the spilled registers from the stack
+      if (SpillR11) {
+        popRegForStackProbeCall(MF, MBB, MBBI, DL, R11Alive, X86::R11, NumBytes);
+      }
+      if (!WindowsStackProbe) {
+        popRegForStackProbeCall(MF, MBB, MBBI, DL, RDXAlive, X86::RDX, NumBytes);
+        popRegForStackProbeCall(MF, MBB, MBBI, DL, RCXAlive, X86::RCX, NumBytes);
+      }
+      popRegForStackProbeCall(MF, MBB, MBBI, DL, RAXAlive, X86::RAX, NumBytes);
     }
   } else if (NumBytes) {
     emitSPUpdate(MBB, MBBI, -(int64_t)NumBytes, /*InEpilogue=*/false);
   }
 
   if (NeedsWinCFI && NumBytes)
     BuildMI(MBB, MBBI, DL, TII.get(X86::SEH_StackAlloc))
         .addImm(NumBytes)

lib/Target/X86/X86ISelLowering.cpp

 // that the guard pages used by the OS virtual memory manager are allocated in
 // correct sequence.
 SDValue
 X86TargetLowering::LowerDYNAMIC_STACKALLOC(SDValue Op,
                                            SelectionDAG &DAG) const {
   MachineFunction &MF = DAG.getMachineFunction();
   bool SplitStack = MF.shouldSplitStack();
   bool Lower = (Subtarget.isOSWindows() && !Subtarget.isTargetMachO()) ||
-               SplitStack;
+               SplitStack || MF.shouldProbeStack();
   SDLoc dl(Op);
 
   // Get the inputs.
   SDNode *Node = Op.getNode();
   SDValue Chain = Op.getOperand(0);
   SDValue Size  = Op.getOperand(1);
   unsigned Align = cast<ConstantSDNode>(Op.getOperand(2))->getZExtValue();
   EVT VT = Node->getValueType(0);

test/CodeGen/X86/movtopush.ll

 define void @test1() {
 entry:
   call void @good(i32 1, i32 2, i32 3, i32 4)
   ret void
 }
 
 ; If we have a reserved frame, we should have pushes
 ; NORMAL-LABEL: test2:
-; NORMAL-NOT: subl {{.*}} %esp
 ; NORMAL: pushl   $4
 ; NORMAL-NEXT: pushl   $3
 ; NORMAL-NEXT: pushl   $2
 ; NORMAL-NEXT: pushl   $1
 ; NORMAL-NEXT: call
 define void @test2(i32 %k) {
 entry:
   %a = alloca i32, i32 %k

test/CodeGen/X86/stack-probes.ll

+; RUN: llc -mtriple=i386-pc-linux-gnu < %s -o - | FileCheck --check-prefix=X86-LINUX %s
+; RUN: llc -mtriple=x86_64-pc-linux-gnu < %s -o - | FileCheck --check-prefix=X64-LINUX %s
+
+; Ensure calls to __probestack occur for large stack frames
+define void @test() "probe-stack" {
+  %array = alloca [40096 x i8]
+  ret void
+
+; X86-LINUX-LABEL: test:
+; X86-LINUX:       movl    $40096, %eax            # imm = 0x9CA0
+; X86-LINUX-NEXT:  movl    %eax, %edx
+; X86-LINUX-NEXT:  movl    %esp, %ecx
+; X86-LINUX-LABEL: .LBB0_1
+; X86-LINUX-NEXT:  orl     $0, (%ecx)
+; X86-LINUX-NEXT:  subl    $4096, %ecx
+; X86-LINUX-NEXT:  subl    $4096, %edx
+; X86-LINUX-NEXT:  jae     .LBB0_1
+; X86-LINUX:       subl    %eax, %esp
+; X86-LINUX:       addl    $40096, %esp            # imm = 0x9CA0
+
+; X64-LINUX-LABEL: test:
+; X64-LINUX:       movl    $40096, %eax            # imm = 0x9CA0
+; X64-LINUX-NEXT:  movq    %rax, %rdx
+; X64-LINUX-NEXT:  movq    %rsp, %rcx
+; X64-LINUX-LABEL: .LBB0_1
+; X64-LINUX-NEXT:  orq     $0, (%rcx)
+; X64-LINUX-NEXT:  subq    $4096, %rcx
+; X64-LINUX-NEXT:  subq    $4096, %rdx
+; X64-LINUX-NEXT:  jae     .LBB0_1
+; X64-LINUX:       subq    %rax, %rsp
+; X64-LINUX:       addq    $40096, %rsp            # imm = 0x9CA0
+
+}
+
+; Ensure the stack is probed for medium stack frames
+define void @testFast() "probe-stack" {
+  %array = alloca [4096 x i8]
+  ret void
+
+; X86-LINUX-LABEL: testFast:
+; X86-LINUX:       orl     $0, -4096(%esp)
+; X86-LINUX-NEXT:  subl    $4096, %esp             # imm = 0x1000
+
+; X64-LINUX-LABEL: testFast:
+; X64-LINUX:       orq     $0, -4096(%rsp)
+; X64-LINUX-NEXT:  subq    $4096, %rsp             # imm = 0x1000
+
+}
+
+; Ensure the stack is probed for dynamic allocations
+define void @testDynamic(i64 %n) "probe-stack" {
+  %buf = alloca i8, i64 %n
+  ret void
+
+; X86-LINUX-LABEL: testDynamic:
+; X86-LINUX:       movl    %eax, %ecx
+; X86-LINUX-NEXT:  movl    %esp, %edx
+; X86-LINUX-LABEL: .LBB2_1
+; X86-LINUX-NEXT:  orl     $0, (%edx)
+; X86-LINUX-NEXT:  subl    $4096, %edx
+; X86-LINUX-NEXT:  subl    $4096, %ecx
+; X86-LINUX-NEXT:  jae     .LBB2_1
+; X86-LINUX:       subl    %eax, %esp
+
+; X64-LINUX-LABEL: testDynamic:
+; X64-LINUX:       leaq 15(%rdi), %rax
+; X64-LINUX-NEXT:  andq $-16, %rax
+; X64-LINUX-NEXT:  movq    %rax, %rcx
+; X64-LINUX-NEXT:  movq    %rsp, %rdx
+; X64-LINUX-LABEL: .LBB2_1
+; X64-LINUX-NEXT:  orq     $0, (%rdx)
+; X64-LINUX-NEXT:  subq    $4096, %rdx
+; X64-LINUX-NEXT:  subq    $4096, %rcx
+; X64-LINUX-NEXT:  jae     .LBB2_1
+; X64-LINUX:       subq    %rax, %rsp
+
+}