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

[Zoxc]

Summary

Adds support for "probe-stack" on x86

[majnemer]

lib/Target/X86/X86FrameLowering.cpp

411–426

Please combine this logic with the if/else-if chain.

452–454

Can you use getSUBrrOpcode here?

820

Please capitalize variable names.

831

(Is64Bit ? 8 : 4) is just SlotSize. Also, this looks formatted strangely; consider running clang-format on this.

test/CodeGen/X86/stack-probes.ll

2–3

FIleCheck prefixes are usually all caps.

[Zoxc]

Addressed comments

[majnemer]

include/llvm/CodeGen/MachineFunction.h

281

I think this part should be '\brief'

282–285

This bit should just be with the definition as per the coding standard: http://llvm.org/docs/CodingStandards.html#doxygen-use-in-documentation-comments

lib/Target/X86/X86FrameLowering.cpp

421–422

Please place braces around this else

850–851

This comment should probably be with the NumBytesAdj bit above. Please update the -4 bit as well.

[Zoxc]

Addressed comments again

[Zoxc]

Rebased

[Zoxc]

The guard region size is passed in RBX/EBX to __probestack. R11 and RBX/EBX is properly spilled if live.

[majnemer]

Can we have a test which exercises the [RE]BX case?

lib/Target/X86/X86FrameLowering.cpp

199

getX86SubSuperRegister(CheckReg, MVT::i32) can be hoisted out of the loop.

405–415

Please make this an early return, it would reduce indentation.

408

Comments should be complete sentences and end with punctuation.

413

Please end this sentence with a period.

414

Please use SlotSize here.

425–437

Please make this an early return, it would reduce indentation.

818–819

I think this deserves a comment.

822

I think this also deserves a comment.

[Zoxc]

Addressed comments.

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();
 
+  /// 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() {
   return getFunction()->hasFnAttribute("split-stack");
 }
 
+/// 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 MachineFunction::shouldProbeStack() const {
+  return getFunction()->hasFnAttribute("probe-stack");
+}
+
 /// 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
 /// it are renumbered.
 void MachineFunction::RenumberBlocks(MachineBasicBlock *MBB) {
   if (empty()) { MBBNumbering.clear(); return; }
   MachineFunction::iterator MBBI, E = end();

lib/Target/X86/X86FrameLowering.cpp

 }
 
 static bool isEAXLiveIn(MachineFunction &MF) {
   for (MachineRegisterInfo::livein_iterator II = MF.getRegInfo().livein_begin(),
        EE = MF.getRegInfo().livein_end(); II != EE; ++II) {
     unsigned Reg = II->first;
 
     if (Reg == X86::RAX || Reg == X86::EAX || Reg == X86::AX ||
         Reg == X86::AH || Reg == X86::AL)
       return true;
   }
 
   return false;
 }
 
 /// emitSPUpdate - Emit a series of instructions to increment / decrement the
 /// stack pointer by a constant value.
void X86FrameLowering::emitStackProbeCall(MachineFunction &MF,
                                           MachineBasicBlock &MBB,
                                           MachineBasicBlock::iterator MBBI,
                                           DebugLoc DL) {
   const X86Subtarget &STI = MF.getSubtarget<X86Subtarget>();
   const TargetInstrInfo &TII = *STI.getInstrInfo();
   bool Is64Bit = STI.is64Bit();
   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 (STI.isOSWindows()) {
     if (Is64Bit) {
       if (STI.isTargetCygMing()) {
         Symbol = "___chkstk_ms";
       } else {
         Symbol = "__chkstk";
       }
     } else if (STI.isTargetCygMing())
       Symbol = "_alloca";
     else
       Symbol = "_chkstk";
-
+  } else {
+    Symbol = "__probestack";
+  }
   MachineInstrBuilder CI;
 
   // 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);
     CI = BuildMI(MBB, MBBI, DL, TII.get(CallOp)).addReg(X86::R11);
   } else {
     CI = BuildMI(MBB, MBBI, DL, TII.get(CallOp)).addExternalSymbol(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);
   }
 }
 
 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);
void X86FrameLowering::emitPrologue(MachineFunction &MF,
   int TailCallReturnAddrDelta = X86FI->getTCReturnAddrDelta();
   if (TailCallReturnAddrDelta && IsWinEH)
     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.isOSWindows() && !STI.isTargetMachO()) || 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) &&
       !RegInfo->needsStackRealignment(MF) &&
       !MFI->hasVarSizedObjects() && // No dynamic alloca.
       !MFI->adjustsStack() &&       // No calls.
       !IsWin64 &&                   // 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
+
       !usesTheStack(MF) &&          // Don't push and pop.
       !MF.shouldSplitStack()) {     // Regular stack
     uint64_t MinSize = X86FI->getCalleeSavedFrameSize();
     if (HasFP) MinSize += SlotSize;
     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) {
     MachineInstr *MI =
       BuildMI(MBB, MBBI, DL,
void X86FrameLowering::emitPrologue(MachineFunction &MF,
   // stack and adjust the stack pointer in one go.  The 64-bit version of
   // __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 (IsWinEH && RegInfo->needsStackRealignment(MF))
     AlignedNumBytes = RoundUpToAlignment(AlignedNumBytes, MaxAlign);
   if (AlignedNumBytes >= StackProbeSize && UseStackProbe) {
-    // Check whether EAX is livein for this function.
-    bool isEAXAlive = isEAXLiveIn(MF);
+    assert(!UseRedZone && "The Red Zone is not accounted for in stack probes");
+    // Check whether the accumulator register is livein for this function.
+    bool IsRegAccAlive = isEAXLiveIn(MF);
+    auto RegAcc = Is64Bit ? X86::RAX : X86::EAX;
 
-    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 (IsRegAccAlive) {
+      // Save RegAcc
+      BuildMI(MBB, MBBI, DL, TII.get(Is64Bit ? X86::PUSH64r : X86::PUSH32r))
+          .addReg(RegAcc, RegState::Kill)
           .setMIFlag(MachineInstr::FrameSetup);
     }
 
+    uint64_t NumBytesAdj = IsRegAccAlive ? NumBytes - SlotSize : 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)) {
+      if (isUInt<32>(NumBytesAdj)) {
         BuildMI(MBB, MBBI, DL, TII.get(X86::MOV32ri), X86::EAX)
-            .addImm(NumBytes)
+            .addImm(NumBytesAdj)
             .setMIFlag(MachineInstr::FrameSetup);
-      } else if (isInt<32>(NumBytes)) {
+      } else if (isInt<32>(NumBytesAdj)) {
         BuildMI(MBB, MBBI, DL, TII.get(X86::MOV64ri32), X86::RAX)
-            .addImm(NumBytes)
+            .addImm(NumBytesAdj)
             .setMIFlag(MachineInstr::FrameSetup);
       } else {
         BuildMI(MBB, MBBI, DL, TII.get(X86::MOV64ri), X86::RAX)
-            .addImm(NumBytes)
+            .addImm(NumBytesAdj)
             .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(NumBytesAdj)
           .setMIFlag(MachineInstr::FrameSetup);
     }
 
     // Save a pointer to the MI where we set AX.
     MachineBasicBlock::iterator SetRAX = MBBI;
     --SetRAX;
 
     // Call __chkstk, __chkstk_ms, or __alloca.
     emitStackProbeCall(MF, MBB, MBBI, DL);
 
     // Apply the frame setup flag to all inserted instrs.
     for (; SetRAX != MBBI; ++SetRAX)
       SetRAX->setFlag(MachineInstr::FrameSetup);
 
-    if (isEAXAlive) {
-      // Restore EAX
-      MachineInstr *MI = addRegOffset(BuildMI(MF, DL, TII.get(X86::MOV32rm),
-                                              X86::EAX),
-                                      StackPtr, false, NumBytes - 4);
+    if (IsRegAccAlive) {
+      // Restore RegAcc
+      auto MIB = BuildMI(MF, DL,
+                         TII.get(Is64Bit ? X86::MOV64rm : X86::MOV32rm),
+                         RegAcc);
+      MachineInstr *MI = addRegOffset(MIB, StackPtr, false, NumBytesAdj);
       MI->setFlag(MachineInstr::FrameSetup);
       MBB.insert(MBBI, MI);
     }
   } else if (NumBytes) {
     emitSPUpdate(MBB, MBBI, StackPtr, -(int64_t)NumBytes, Is64Bit, Uses64BitFramePtr,
                  UseLEA, TII, *RegInfo);
   }
 

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);
 
   if (!Lower) {
     const TargetLowering &TLI = DAG.getTargetLoweringInfo();
     SDNode* Node = Op.getNode();
 
     unsigned SPReg = TLI.getStackPointerRegisterToSaveRestore();
     assert(SPReg && "Target cannot require DYNAMIC_STACKALLOC expansion and"

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
+
+declare void @use([40096 x i8]*)
+
+; Ensure calls to __probestack occur for large stack frames
+define void @test() "probe-stack" {
+  %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
+  
+}