[PATCH 2/2] Implement "probe-stack" on x86

[pcwalton]

Summary

I've fixed up https://reviews.llvm.org/D9654. It's a follow-up to https://reviews.llvm.org/D34386. Please let me know if you have any further suggestions. Thanks!

[majnemer]

How does this review differ from D34386?

[fhahn]

llvm-commits should be added as a subscriber when the review is created to mirror the review on the mailing list

[fhahn]

I'll repost some of my comments from D34386. I think it would be good to avoid unrelated whitespace changes to reduce the noise in the diff. Also, could you add more context to the diff? That makes reviewing easier :)

I think you should also include a test for your change. You could probably extend test/CodeGen/X86/stack-probe-size.ll, by adding an additional RUN line with a triple where stack probes aren't generated at the moment and add probe-stack attributes to some functions.

[pcwalton]

Added more context, added missing tests, removed whitespace-only changes.

[majnemer]

lib/CodeGen/MachineFunction.cpp

216–224 ↗

(On Diff #103236)

I think this should be renamed or change its definition to reflect its name, we will probe the stack on x86 Windows regardless of the whether or not this attribute was supplied.

Perhaps hasProbeStackAttr?

lib/Target/X86/X86FrameLowering.cpp

800–801

Could this be a StringRef?

1023–1024

You could make a function which gets the stack probe symbol name, then UseStackProbe could be defined in terms of that name being non-empty.

1243–1244

Test case for this?

1248–1249

Saving RBX seems new. I'd have a comment describing both of these saves.

1256

Magic number should have a comment or stuck in a const variable with a nice name. Even better would be both.

[pcwalton]

lib/Target/X86/X86FrameLowering.cpp

800–801

I originally wanted to do that, but the docs say that StringRefs don't have null termination, which the addExternalSymbol API needs. If you'd prefer to work around that in another way, let me know.

[majnemer]

lib/Target/X86/X86FrameLowering.cpp

800–801

Hmm. Isn't there a liveness problem then? Symbol.c_str() is only valid until this functoin returns. addExternalSymbol doesn't make a copy of its input unless I am mistaken.

[majnemer]

lib/Target/X86/X86FrameLowering.cpp

800–801

If I am correct, I think you will want to use createExternalSymbolName.

[pcwalton]

lib/Target/X86/X86FrameLowering.cpp

1248–1249

I don't know why the original patch author chose to spill RBX. Seems not needed. I'll remove it.

1256

As before, I don't know what this magic number is for. Seems unneeded. I'll remove it.

[pcwalton]

lib/Target/X86/X86FrameLowering.cpp

1256

Oh, I guess this is the page size. I don't see why user-requested stack probes should need this to be supplied explicitly; it just bloats the size of every function prolog for something that is going to be constant on the architecture. Will remove.

[whitequark]

@pcwalton No, both the spill and the "magic" number are necessary. See https://reviews.llvm.org/D9858. ebx is the size of the frame, and 0x1000 is the probe interval to be performed (in other words, page size).

[whitequark]

@pcwalton But I agree that it's better to change both sides of the __probestack ABI to eliminate the argument and the spill (on x86_64), of course.

[pcwalton]

Addressed review comments.

[majnemer]

lib/Target/X86/X86FrameLowering.cpp

510

Why is this by ref?

801

I think it is more natural to see: StringRef Symbol = STI.getTargetLowering()->getStackProbeSymbolName(MF);

1243–1244

Am I correct in thinking that we will save and restore R11 when we previously would not in the large code model for Win64? If so, was that a bug in the implementation of the large code model for Win64?

lib/Target/X86/X86ISelLowering.cpp

36342

const StringRef looks weird, StringRefs are immutable.

36344–36347

I'd kill the braces here.

36355–36357

Ditto.

lib/Target/X86/X86ISelLowering.h

1059

Ditto regarding const StringRef.

[pcwalton]

In D34387#785802, @whitequark wrote:

@pcwalton No, both the spill and the "magic" number are necessary. See https://reviews.llvm.org/D9858. ebx is the size of the frame, and 0x1000 is the probe interval to be performed (in other words, page size).

That patch never landed, and the current approach, with customizable symbol names, is to move the responsibility out of compiler-rt. So I would argue that https://reviews.llvm.org/D9858 is irrelevant.

lib/Target/X86/X86FrameLowering.cpp

1243–1244

I'm not familiar enough with the Windows 64-bit ABI to be able to answer that question. I would assume that it would have to have been a bug in the previous implementation, as I don't see any other way to make the call to __chkstk work…

[pcwalton]

Addressed comments.

[majnemer]

lib/Target/X86/X86FrameLowering.cpp

1250

Hmm, I believe a testcase for this situation hasn't been added.

[pcwalton]

lib/Target/X86/X86FrameLowering.cpp

1250

Well, I'm not sure how R11 can actually be live-in in the function prolog right now. As I understand things, the reason why EAX is spilled if necessary is to handle the edge case on 32-bit x86 where inreg is specified for a function parameter, which causes EAX to be live-in at the very start of the function. But there's no corresponding inreg attribute for x86-64 that causes R11 to become live-in that early. Checking for that case even for R11 seemed prudent to me though, because (a) the corresponding situation can happen for EAX and it seemed cleaner to use the same logic for both registers instead of duplicating some of it; (b) in the future, if somehow R11 can become live-in at function start, this logic will be needed.

Am I interpreting the situation correctly, and if so, what would you prefer I do regarding testing?

[whitequark]

lib/Target/X86/X86FrameLowering.cpp

1250

I am quite certain that R11 should *not* be spilled this way. LLVM uses it as a scratch register in quite a few places. E.g. the comment above:

// For the large code model, we have to call through a register. Use R11,
// as it is scratch in all supported calling conventions.

the GetScratchRegister function in X86FrameLowering.cpp, the X86TargetLowering::getScratchRegisters, and the X86TargetLowering::LowerINIT_TRAMPOLINE function, which needs it for essentially the same job.

You could perhaps replace the call with one of the getScratchRegister functions but it seems fine to hardcode it while not spilling the way it is, too.

[pcwalton]

Addressed review comments.

I removed the R11 spill code as requested. I left in the refactoring that moves the existing EAX spill and reload code into separate functions, as X86FrameLowering::emitPrologue() is rather large and that seemed like a worthwhile refactoring.

[whitequark]

LGTM, but I'd still like to hear from @majnemer

[majnemer]

In D34387#787200, @pcwalton wrote:

Addressed review comments.

I removed the R11 spill code as requested. I left in the refactoring that moves the existing EAX spill and reload code into separate functions, as X86FrameLowering::emitPrologue() is rather large and that seemed like a worthwhile refactoring.

I am not entirely sure that the refactored code is identical to what was left behind. Moreover, the refactoring is now incidental with respect to the bulk of the change. I'd request that if you were to pursue it, to please factor it out into another review.

[pcwalton]

Addressed comments. Reverted the EAX save refactoring.

[majnemer]

Looking pretty good, I think we are almost there :)

lib/Target/X86/X86FrameLowering.cpp

1053–1055

We can use the redzone and stack probes at the same time? This surprises me. Also, there is no test for this.

[pcwalton]

Addressed review comments.

[majnemer]

LGTM!

[joerg]

It would be nice to make the stackprobesize a proper TLI attribute, so that OSes can decide how much guarding they are willing to guarantee. I think I would also like to have an option to inline the probing without the complexity of the function call. But both can be done as a separate step.

include/llvm/Target/TargetLowering.h

 protected:
   Value *getDefaultSafeStackPointerLocation(IRBuilder<> &IRB,
                                             bool UseTLS) const;
 
 public:
   /// Returns the target-specific address of the unsafe stack pointer.
   virtual Value *getSafeStackPointerLocation(IRBuilder<> &IRB) const;
 
+  /// Returns the name of the symbol used to emit stack probes or the empty
+  /// string if not applicable.
+  virtual StringRef getStackProbeSymbolName(MachineFunction &MF) const {
+    return "";
+  }
+
   /// Returns true if a cast between SrcAS and DestAS is a noop.
   virtual bool isNoopAddrSpaceCast(unsigned SrcAS, unsigned DestAS) const {
     return false;
   }
 
   /// Returns true if a cast from SrcAS to DestAS is "cheap", such that e.g. we
   /// are happy to sink it into basic blocks.
   virtual bool isCheapAddrSpaceCast(unsigned SrcAS, unsigned DestAS) const {

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;
 
+  /// Emits code to push a caller-save register before a call to the stack
+  /// probing function.
+  void pushRegForStackProbeCall(MachineFunction &MF,
+                                MachineBasicBlock &MBB,
+                                MachineBasicBlock::iterator MBBI,
+                                DebugLoc DL,
+                                bool &IsAlive,
+                                unsigned RegType,
+                                uint64_t &NumBytes) const;
+
+  /// Emits code to pop a caller-save register after a call to the stack
+  /// probing function.
+  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.
   void emitStackProbe(MachineFunction &MF, MachineBasicBlock &MBB,
                       MachineBasicBlock::iterator MBBI, const DebugLoc &DL,
                       bool InProlog) const;
 
   /// Replace a StackProbe inline-stub with the actual probe code inline.

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);
 
   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;
     unsigned Rax = (unsigned)(Is64Bit ? X86::RAX : X86::EAX);
 
-    if (isSub && !isEAXLiveIn(MBB))
+    if (isSub && !isLiveIn(MBB, X86::EAX))
       Reg = Rax;
     else
       Reg = findDeadCallerSavedReg(MBB, MBBI, TRI, Is64Bit);
 
     unsigned MovRIOpc = Is64Bit ? X86::MOV64ri : X86::MOV32ri;
     unsigned AddSubRROpc =
         isSub ? getSUBrrOpcode(Is64Bit) : getADDrrOpcode(Is64Bit);
     if (Reg) {
for (std::vector<CalleeSavedInfo>::const_iterator
     unsigned Reg = I->getReg();
 
     unsigned DwarfReg = MRI->getDwarfRegNum(Reg, true);
     BuildCFI(MBB, MBBI, DL,
              MCCFIInstruction::createOffset(nullptr, DwarfReg, Offset));
   }
 }
 
+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;
+}
+
 void X86FrameLowering::emitStackProbe(MachineFunction &MF,
                                       MachineBasicBlock &MBB,
                                       MachineBasicBlock::iterator MBBI,
                                       const DebugLoc &DL, bool InProlog) const {
   const X86Subtarget &STI = MF.getSubtarget<X86Subtarget>();
   if (STI.isTargetWindowsCoreCLR()) {
     if (InProlog) {
       emitStackProbeInlineStub(MF, MBB, MBBI, DL, true);
void X86FrameLowering::emitStackProbeCall(MachineFunction &MF,
                                           bool InProlog) const {
   bool IsLargeCodeModel = MF.getTarget().getCodeModel() == CodeModel::Large;
 
   unsigned CallOp;
   if (Is64Bit)
     CallOp = IsLargeCodeModel ? X86::CALL64r : X86::CALL64pcrel32;
   else
     CallOp = X86::CALLpcrel32;
 
-  const char *Symbol;
-  if (Is64Bit) {
-    if (STI.isTargetCygMing()) {
-      Symbol = "___chkstk_ms";
-    } else {
-      Symbol = "__chkstk";
-    }
-  } else if (STI.isTargetCygMing())
-    Symbol = "_alloca";
-  else
-    Symbol = "_chkstk";
+  StringRef Symbol = STI.getTargetLowering()->getStackProbeSymbolName(MF);
 
   MachineInstrBuilder CI;
   MachineBasicBlock::iterator ExpansionMBBI = std::prev(MBBI);
 
   // All current stack probes take AX and SP as input, clobber flags, and
   // preserve all registers. x86_64 probes leave RSP unmodified.
   if (Is64Bit && MF.getTarget().getCodeModel() == CodeModel::Large) {
     // For the large code model, we have to call through a register. Use R11,
     // as it is scratch in all supported calling conventions.
     BuildMI(MBB, MBBI, DL, TII.get(X86::MOV64ri), X86::R11)
-        .addExternalSymbol(Symbol);
+        .addExternalSymbol(MF.createExternalSymbolName(Symbol));
     CI = BuildMI(MBB, MBBI, DL, TII.get(CallOp)).addReg(X86::R11);
   } else {
-    CI = BuildMI(MBB, MBBI, DL, TII.get(CallOp)).addExternalSymbol(Symbol);
+    CI = BuildMI(MBB, MBBI, DL, TII.get(CallOp))
+        .addExternalSymbol(MF.createExternalSymbolName(Symbol));
   }
 
   unsigned AX = Is64Bit ? X86::RAX : X86::EAX;
   unsigned SP = Is64Bit ? X86::RSP : X86::ESP;
   CI.addReg(AX, RegState::Implicit)
       .addReg(SP, RegState::Implicit)
       .addReg(AX, RegState::Define | RegState::Implicit)
       .addReg(SP, RegState::Define | RegState::Implicit)
       .addReg(X86::EFLAGS, RegState::Define | RegState::Implicit);
 
-  if (Is64Bit) {
+  if (!STI.isTargetWin32()) {
     // MSVC x64's __chkstk and cygwin/mingw's ___chkstk_ms do not adjust %rsp
     // themselves. It also does not clobber %rax so we can reuse it when
     // adjusting %rsp.
-    BuildMI(MBB, MBBI, DL, TII.get(X86::SUB64rr), X86::RSP)
-        .addReg(X86::RSP)
-        .addReg(X86::RAX);
+    BuildMI(MBB, MBBI, DL, TII.get(getSUBrrOpcode(Is64Bit)), SP)
+        .addReg(SP)
+        .addReg(AX);
   }
 
   if (InProlog) {
     // Apply the frame setup flag to all inserted instrs.
     for (++ExpansionMBBI; ExpansionMBBI != MBBI; ++ExpansionMBBI)
       ExpansionMBBI->setFlag(MachineInstr::FrameSetup);
   }
 }
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 UseStackProbe = !STI.getTargetLowering()->getStackProbeSymbolName(MF).empty();
 
   // 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);
 
void X86FrameLowering::emitPrologue(MachineFunction &MF,
   // 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 avoid stack probes
       !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");
 
-    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);
-    }
+    // We spill the registers we need to call the stack probe function.
+
+    bool RAXAlive;
+    pushRegForStackProbeCall(MF, MBB, MBBI, DL, RAXAlive, X86::RAX, 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.
     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);
-    }
+    // Restore the spilled registers from the stack.
+    popRegForStackProbeCall(MF, MBB, MBBI, DL, RAXAlive, X86::RAX, NumBytes);
   } else if (NumBytes) {
     emitSPUpdate(MBB, MBBI, -(int64_t)NumBytes, /*InEpilogue=*/false);
   }
 
   if (NeedsWinCFI && NumBytes) {
     HasWinCFI = true;
     BuildMI(MBB, MBBI, DL, TII.get(X86::SEH_StackAlloc))
         .addImm(NumBytes)

lib/Target/X86/X86ISelLowering.h

public:
 
     /// \brief Customize the preferred legalization strategy for certain types.
     LegalizeTypeAction getPreferredVectorAction(EVT VT) const override;
 
     bool isIntDivCheap(EVT VT, AttributeList Attr) const override;
 
     bool supportSwiftError() const override;
 
+    StringRef getStackProbeSymbolName(MachineFunction &MF) const override;
+
     unsigned getMaxSupportedInterleaveFactor() const override { return 4; }
 
     /// \brief Lower interleaved load(s) into target specific
     /// instructions/intrinsics.
     bool lowerInterleavedLoad(LoadInst *LI,
                               ArrayRef<ShuffleVectorInst *> Shuffles,
                               ArrayRef<unsigned> Indices,
                               unsigned Factor) const override;

lib/Target/X86/X86ISelLowering.cpp

 // bytes in one go. 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.
 SDValue
 X86TargetLowering::LowerDYNAMIC_STACKALLOC(SDValue Op,
                                            SelectionDAG &DAG) const {
   MachineFunction &MF = DAG.getMachineFunction();
   bool SplitStack = MF.shouldSplitStack();
+  bool EmitStackProbe = !getStackProbeSymbolName(MF).empty();
   bool Lower = (Subtarget.isOSWindows() && !Subtarget.isTargetMachO()) ||
-               SplitStack;
+               SplitStack || EmitStackProbe;
   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);
for (auto *Exit : Exits)
               TII->get(TargetOpcode::COPY), *I)
           .addReg(NewVR);
   }
 }
 
 bool X86TargetLowering::supportSwiftError() const {
   return Subtarget.is64Bit();
 }
+
+/// Returns the name of the symbol used to emit stack probes or the empty
+/// string if not applicable.
+StringRef X86TargetLowering::getStackProbeSymbolName(MachineFunction &MF) const {
+  // If the function specifically requests stack probes, emit them.
+  if (MF.getFunction()->hasFnAttribute("probe-stack"))
+    return MF.getFunction()->getFnAttribute("probe-stack").getValueAsString();
+
+  // Generally, if we aren't on Windows, the platform ABI does not include
+  // support for stack probes, so don't emit them.
+  if (!Subtarget.isOSWindows() || Subtarget.isTargetMachO())
+    return "";
+
+  // We need a stack probe to conform to the Windows ABI. Choose the right
+  // symbol.
+  if (Subtarget.is64Bit())
+    return Subtarget.isTargetCygMing() ? "___chkstk_ms" : "__chkstk";
+  return Subtarget.isTargetCygMing() ? "_alloca" : "_chkstk";
+}

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
+; RUN: llc -mtriple=x86_64-pc-linux-gnu -code-model=large < %s -o - | FileCheck --check-prefix=X64-LINUX-LARGE %s
+
+declare void @use([40096 x i8]*)
+
+; Ensure calls to __probestack occur for large stack frames
+define void @test() "probe-stack"="__probestack" {
+  %array = alloca [40096 x i8], align 16
+  call void @use([40096 x i8]* %array)
+  ret void
+
+; X86-LINUX-LABEL:       test:
+; X86-LINUX:             movl $40124, %eax # imm = 0x9CBC
+; X86-LINUX-NEXT:        calll __probestack
+; X86-LINUX-NEXT:        subl %eax, %esp
+
+; X64-LINUX-LABEL:       test:
+; X64-LINUX:             movl $40104, %eax # imm = 0x9CA8
+; X64-LINUX-NEXT:        callq __probestack
+; X64-LINUX-NEXT:        subq %rax, %rsp
+
+; X64-LINUX-LARGE-LABEL: test:
+; X64-LINUX-LARGE:       movl $40104, %eax # imm = 0x9CA8
+; X64-LINUX-LARGE-NEXT:  movabsq $__probestack, %r11
+; X64-LINUX-LARGE-NEXT:  callq *%r11
+; X64-LINUX-LARGE-NEXT:  subq %rax, %rsp
+
+}