Merge pull request #86 from alexcrichton/probestack

Add x86 probestack support

Author: alexcrichton

docs/LangRef.rst

@@ -1467,6 +1467,19 @@ example:
     This attribute by itself does not imply restrictions on
     inter-procedural optimizations.  All of the semantic effects the
     patching may have to be separately conveyed via the linkage type.
+``"probe-stack"``
+    This attribute indicates that the function will trigger a guard region
+    in the end of the stack. It ensures that accesses to the stack must be
+    no further apart than the size of the guard region to a previous
+    access of the stack. It takes one required string value, the name of
+    the stack probing function that will be called.
+
+    If a function that has a ``"probe-stack"`` attribute is inlined into
+    a function with another ``"probe-stack"`` attribute, the resulting
+    function has the ``"probe-stack"`` attribute of the caller. If a
+    function that has a ``"probe-stack"`` attribute is inlined into a
+    function that has no ``"probe-stack"`` attribute at all, the resulting
+    function has the ``"probe-stack"`` attribute of the callee.
 ``readnone``
     On a function, this attribute indicates that the function computes its
     result (or decides to unwind an exception) based strictly on its arguments,

include/llvm/IR/Attributes.td

@@ -211,3 +211,4 @@ def : MergeRule<"setAND<UnsafeFPMathAttr>">;
 def : MergeRule<"setOR<NoImplicitFloatAttr>">;
 def : MergeRule<"setOR<NoJumpTablesAttr>">;
 def : MergeRule<"adjustCallerSSPLevel">;
+def : MergeRule<"adjustCallerStackProbes">;

include/llvm/Target/TargetLowering.h

@@ -1179,6 +1179,12 @@ class TargetLoweringBase {
   /// 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;

lib/IR/Attributes.cpp

@@ -1502,6 +1502,13 @@ static void adjustCallerSSPLevel(Function &Caller, const Function &Callee) {
     Caller.addFnAttr(Attribute::StackProtect);
 }
 
+/// \brief If the inlined function required stack probes, then ensure that
+/// the calling function has those too.
+static void adjustCallerStackProbes(Function &Caller, const Function &Callee) {
+  if (!Caller.hasFnAttribute("probe-stack") && Callee.hasFnAttribute("probe-stack"))
+    Caller.addFnAttr("probe-stack", Callee.getFnAttribute("probe-stack").getValueAsString());
+}
+
 #define GET_ATTR_COMPAT_FUNC
 #include "AttributesCompatFunc.inc"
 

lib/Target/X86/X86FrameLowering.cpp

@@ -719,17 +719,7 @@ void X86FrameLowering::emitStackProbeCall(MachineFunction &MF,
   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);
@@ -740,10 +730,11 @@ void X86FrameLowering::emitStackProbeCall(MachineFunction &MF,
     // 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;
@@ -754,13 +745,13 @@ void X86FrameLowering::emitStackProbeCall(MachineFunction &MF,
       .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) {
@@ -949,7 +940,8 @@ void X86FrameLowering::emitPrologue(MachineFunction &MF,
     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.
@@ -968,6 +960,7 @@ void X86FrameLowering::emitPrologue(MachineFunction &MF,
       !TRI->needsStackRealignment(MF) &&
       !MFI.hasVarSizedObjects() &&             // No dynamic alloca.
       !MFI.adjustsStack() &&                   // No calls.
+      !UseStackProbe &&                        // No stack probes.
       !IsWin64CC &&                            // Win64 has no Red Zone
       !MFI.hasCopyImplyingStackAdjustment() && // Don't push and pop.
       !MF.shouldSplitStack()) {                // Regular stack
@@ -976,6 +969,7 @@ void X86FrameLowering::emitPrologue(MachineFunction &MF,
     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
@@ -1158,6 +1152,8 @@ void X86FrameLowering::emitPrologue(MachineFunction &MF,
   if (IsWin64Prologue && !IsFunclet && TRI->needsStackRealignment(MF))
     AlignedNumBytes = alignTo(AlignedNumBytes, MaxAlign);
   if (AlignedNumBytes >= StackProbeSize && UseStackProbe) {
+    assert(!UseRedZone && "The Red Zone is not accounted for in stack probes");
+
     // Check whether EAX is livein for this block.
     bool isEAXAlive = isEAXLiveIn(MBB);
 

lib/Target/X86/X86ISelLowering.cpp

@@ -18272,8 +18272,9 @@ 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.
@@ -34918,3 +34919,22 @@ void X86TargetLowering::insertCopiesSplitCSR(
 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";
+}

lib/Target/X86/X86ISelLowering.h

@@ -1042,6 +1042,8 @@ namespace llvm {
 
     bool supportSwiftError() const override;
 
+    StringRef getStackProbeSymbolName(MachineFunction &MF) const override;
+
     unsigned getMaxSupportedInterleaveFactor() const override { return 4; }
 
     /// \brief Lower interleaved load(s) into target specific

test/CodeGen/X86/stack-probe-red-zone.ll

@@ -0,0 +1,21 @@
+; RUN: llc -mtriple=x86_64-pc-linux-gnu < %s -o - | FileCheck %s
+
+; Ensure that red zone usage occurs.
+define void @testStackProbesOff() {
+  %array = alloca [40096 x i8], align 16
+  ret void
+
+; CHECK-LABEL:  testStackProbesOff:
+; CHECK:        subq $39976, %rsp # imm = 0x9C28
+}
+
+; Ensure stack probes do not result in red zone usage.
+define void @testStackProbesOn() "probe-stack"="__probestack" {
+  %array = alloca [40096 x i8], align 16
+  ret void
+
+; CHECK-LABEL:  testStackProbesOn:
+; CHECK:        movl $40104, %eax # imm = 0x9CA8
+; CHECK-NEXT:   callq __probestack
+; CHECK-NEXT:   subq %rax, %rsp
+}

test/CodeGen/X86/stack-probes.ll

@@ -0,0 +1,29 @@
+; 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
+
+}

test/Transforms/Inline/inline-probe-stack.ll

@@ -0,0 +1,20 @@
+; RUN: opt %s -inline -S | FileCheck %s
+
+define internal void @inner() "probe-stack"="__probestackinner" {
+  ret void
+}
+
+define void @outerNoAttribute() {
+  call void @inner()
+  ret void
+}
+
+define void @outerConflictingAttribute() "probe-stack"="__probestackouter" {
+  call void @inner()
+  ret void
+}
+
+; CHECK: define void @outerNoAttribute() #0
+; CHECK: define void @outerConflictingAttribute() #1
+; CHECK: attributes #0 = { "probe-stack"="__probestackinner" }
+; CHECK: attributes #1 = { "probe-stack"="__probestackouter" }