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

[Zoxc]

Summary

Adds support for "probe-stack" on x86

[majnemer]

lib/Target/X86/X86FrameLowering.cpp

495–499

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

525–527

Can you use getSUBrrOpcode here?

879

Please capitalize variable names.

892

(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

284

I think this part should be '\brief'

285–288

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

505–506

Please place braces around this else

908–909

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

205

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

436–446

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

439

Comments should be complete sentences and end with punctuation.

444

Please end this sentence with a period.

445

Please use SlotSize here.

456–468

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

876

I think this deserves a comment.

881

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();
 
+  /// \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() {
   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 a call to the target's stack probe function. 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 emitStackProbeCall(MachineFunction &MF, MachineBasicBlock &MBB,
                           MachineBasicBlock::iterator MBBI, DebugLoc DL) const;
 
   void emitCalleeSavedFrameMoves(MachineBasicBlock &MBB,
                                  MachineBasicBlock::iterator MBBI,

lib/Target/X86/X86FrameLowering.cpp

for (; *CS; ++CS)
       if (!Uses.count(*CS))
         return *CS;
   }
   }
 
   return 0;
 }
 
-static bool isEAXLiveIn(MachineFunction &MF) {
+static bool isLiveIn(MachineFunction &MF, unsigned CheckReg) {
   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)
+    if (getX86SubSuperRegister(Reg, MVT::i32) == getX86SubSuperRegister(CheckReg, MVT::i32))
       return true;
   }
 
   return false;
 }
 
 /// Check whether or not the terminators of \p MBB needs to read EFLAGS.
 static bool terminatorsNeedFlagsAsInput(const MachineBasicBlock &MBB) {
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.getParent()))
+      if (isSub && !isLiveIn(*MBB.getParent(), 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);
for (MachineRegisterInfo::reg_instr_iterator
        ri = MRI.reg_instr_begin(X86::EFLAGS), re = MRI.reg_instr_end();
        ri != re; ++ri)
     if (ri->isCopy())
       return true;
 
   return false;
 }
 
+void X86FrameLowering::pushRegForStackProbeCall(MachineFunction &MF,
+                                                MachineBasicBlock &MBB,
+                                                MachineBasicBlock::iterator MBBI,
+                                                DebugLoc DL,
+                                                bool &IsAlive,
+                                                unsigned RegType,
+                                                uint64_t &NumBytes) const {
+    IsAlive = isLiveIn(MF, RegType);
+
+    if (IsAlive) {
+      auto Reg = getX86SubSuperRegister(RegType, Is64Bit ? MVT::i64 : MVT::i32);
+
+      // Save Reg
+      BuildMI(MBB, MBBI, DL, TII.get(Is64Bit ? X86::PUSH64r : X86::PUSH32r))
+          .addReg(Reg, RegState::Kill)
+          .setMIFlag(MachineInstr::FrameSetup);
+
+      // Reuse the space as a stack allocation
+      NumBytes -= Is64Bit ? 8 : 4;
+    }
+}
+
+void X86FrameLowering::popRegForStackProbeCall(MachineFunction &MF,
+                                                MachineBasicBlock &MBB,
+                                                MachineBasicBlock::iterator MBBI,
+                                                DebugLoc DL,
+                                                bool &IsAlive,
+                                                unsigned RegType,
+                                                uint64_t &NumBytes) const {
+    if (IsAlive) {
+      // Restore Reg
+      auto Reg = getX86SubSuperRegister(RegType, Is64Bit ? MVT::i64 : MVT::i32);
+
+      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 += Is64Bit ? 8 : 4;
+    }
+}
+
 void X86FrameLowering::emitStackProbeCall(MachineFunction &MF,
                                           MachineBasicBlock &MBB,
                                           MachineBasicBlock::iterator MBBI,
                                           DebugLoc DL) 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 (STI.isOSWindows()) {
     if (Is64Bit) {
       if (STI.isTargetCygMing()) {
         Symbol = "___chkstk_ms";
       } else {
         Symbol = "__chkstk";
       }
-  } else if (STI.isTargetCygMing())
+    } else if (STI.isTargetCygMing()) {
       Symbol = "_alloca";
-  else
+    } 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 && 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.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) &&
       !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
+
       !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) {
     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 && TRI->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");
+
+    bool SpillR11 = Is64Bit && MF.getTarget().getCodeModel() == CodeModel::Large;
 
-    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)
+    bool RAXAlive, RBXAlive, R11Alive;
+
+    pushRegForStackProbeCall(MF, MBB, MBBI, DL, RAXAlive, X86::RAX, NumBytes);
+    pushRegForStackProbeCall(MF, MBB, MBBI, DL, RBXAlive, X86::RBX, NumBytes);
+    if (SpillR11) {
+      pushRegForStackProbeCall(MF, MBB, MBBI, DL, R11Alive, X86::R11, NumBytes);
+    }
+
+    if (!STI.isOSWindows()) {
+      BuildMI(MBB, MBBI, DL, TII.get(X86::MOV32ri), X86::EDX)
+           .addImm(0x1000)
            .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)) {
         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);
     }
 
     // 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);
-      MI->setFlag(MachineInstr::FrameSetup);
-      MBB.insert(MBBI, MI);
+    if (SpillR11) {
+      popRegForStackProbeCall(MF, MBB, MBBI, DL, R11Alive, X86::R11, NumBytes);
     }
+    popRegForStackProbeCall(MF, MBB, MBBI, DL, RBXAlive, X86::RBX, 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)
         .setMIFlag(MachineInstr::FrameSetup);

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"
if (SplitStack) {
     SDValue Ops1[2] = { Value, Chain };
     return DAG.getMergeValues(Ops1, dl);
   } else {
     SDValue Flag;
     const unsigned Reg = (Subtarget->isTarget64BitLP64() ? X86::RAX : X86::EAX);
 
     Chain = DAG.getCopyToReg(Chain, dl, Reg, Size, Flag);
     Flag = Chain.getValue(1);
+
+    if (!Subtarget->isOSWindows()) {
+        const unsigned Reg = (Subtarget->isTarget64BitLP64() ? X86::RBX : X86::EBX);
+
+        Chain = DAG.getCopyToReg(Chain, dl, Reg, DAG.getConstant(0x1000, dl, SPTy), Flag);
+        Flag = Chain.getValue(1);
+    }
+
     SDVTList NodeTys = DAG.getVTList(MVT::Other, MVT::Glue);
 
     Chain = DAG.getNode(X86ISD::WIN_ALLOCA, dl, NodeTys, Chain, Flag);
 
     const X86RegisterInfo *RegInfo = Subtarget->getRegisterInfo();
     unsigned SPReg = RegInfo->getStackRegister();
     SDValue SP = DAG.getCopyFromReg(Chain, dl, SPReg, SPTy);
     Chain = SP.getValue(1);
X86TargetLowering::EmitLoweredSegAlloca(MachineInstr *MI,
   return continueMBB;
 }
 
 MachineBasicBlock *
 X86TargetLowering::EmitLoweredWinAlloca(MachineInstr *MI,
                                         MachineBasicBlock *BB) const {
   DebugLoc DL = MI->getDebugLoc();
 
-  assert(!Subtarget->isTargetMachO());
-
   Subtarget->getFrameLowering()->emitStackProbeCall(*BB->getParent(), *BB, MI,
                                                     DL);
 
   MI->eraseFromParent();   // The pseudo instruction is gone now.
   return BB;
 }
 
 MachineBasicBlock *

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 $4096, %edx # imm = 0x1000
+; 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 $4096, %edx # imm = 0x1000
+; X64-LINUX:       movl $40104, %eax # imm = 0x9CA8
+; X64-LINUX-NEXT:  callq __probestack
+; X64-LINUX-NEXT:  subq %rax, %rsp
+  
+}