[X86] Elide references to _chkstk for dynamic allocas

[majnemer]

Summary

The _chkstk function is called by the compiler to probe the stack in an
order consistent with Windows' expectations. However, it is possible to
elide the call to _chkstk and manually adjust the stack pointer if we
can prove that the allocation is fixed size and smaller than the probe
size.

This shrinks chrome.dll, chrome_child.dll and chrome.exe by a
cummulative ~133 KB.

[hans]

Nice! lgtm

lib/Target/X86/X86ISelLowering.cpp

16377

Can we check whether the alloca is dynamic or not, or is that hard to figure out from here?

16379

s/it's/its/

[rnk]

lib/Target/X86/X86ISelLowering.cpp

16365

Can we name this something like CheckStack instead? I interpreted Lower to mean "lower this to modifying ESP directly".

16377

I don't think we need to think about the frame pointer at all for dynamic allocas. If we're concerned that there haven't been any accesses to top-of-stack since the prologue, we could instead emit an inline stack probe (orl $0, (%esp)) here.

[rnk]

lgtm

lib/Target/X86/X86ISelLowering.cpp

 // Calls to _alloca are needed to probe the stack when allocating more than 4k
 // 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();
+  const Function *F = MF.getFunction();
   bool SplitStack = MF.shouldSplitStack();
-  bool Lower = (Subtarget.isOSWindows() && !Subtarget.isTargetMachO()) ||
-               SplitStack;
   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);
 
   // Chain the dynamic stack allocation so that it doesn't modify the stack
   // pointer when other instructions are using the stack.
   Chain = DAG.getCALLSEQ_START(Chain, DAG.getIntPtrConstant(0, dl, true), dl);
 
+  const X86RegisterInfo *RegInfo = Subtarget.getRegisterInfo();
   bool Is64Bit = Subtarget.is64Bit();
   MVT SPTy = getPointerTy(DAG.getDataLayout());
 
+  bool Lower = SplitStack;
+  if (!Lower && Subtarget.isOSWindows() && !Subtarget.isTargetMachO()) {
+    // The Windows ABI requires us to probe the stack for allocations beyond
+    // the probe size.
+    if (auto *SizeC = dyn_cast<ConstantSDNode>(Size)) {
+      // Try to elide the probe if we can prove that this dynamic allocation is
+      // smaller than the probe size.
+      unsigned StackProbeSize = 4096;
+      if (F->hasFnAttribute("stack-probe-size"))
+        F->getFnAttribute("stack-probe-size")
+            .getValueAsString()
+            .getAsInteger(0, StackProbeSize);
+      // Dynamic allocas force a frame pointer, take that into account.
+      unsigned AlignedAlloc = SizeC->getZExtValue() + RegInfo->getSlotSize();
+      // Round the dynamic alloca's size up to it's alignment.
+      if (Align)
+        AlignedAlloc = alignTo(AlignedAlloc, Align);
+
+      // If the aligned allocation is smaller than the probe size, then we don't
+      // need to probe the stack.
+      Lower = AlignedAlloc >= StackProbeSize;
+    } else {
+      // We cannot tell how big this dynamic alloca will be, probe the stack.
+      Lower = true;
+    }
+  }
+
   SDValue Result;
   if (!Lower) {
     const TargetLowering &TLI = DAG.getTargetLoweringInfo();
     unsigned SPReg = TLI.getStackPointerRegisterToSaveRestore();
     assert(SPReg && "Target cannot require DYNAMIC_STACKALLOC expansion and"
                     " not tell us which reg is the stack pointer!");
-    EVT VT = Node->getValueType(0);
-    SDValue Tmp3 = Node->getOperand(2);
 
     SDValue SP = DAG.getCopyFromReg(Chain, dl, SPReg, VT);
     Chain = SP.getValue(1);
-    unsigned Align = cast<ConstantSDNode>(Tmp3)->getZExtValue();
     const TargetFrameLowering &TFI = *Subtarget.getFrameLowering();
     unsigned StackAlign = TFI.getStackAlignment();
     Result = DAG.getNode(ISD::SUB, dl, VT, SP, Size); // Value
     if (Align > StackAlign)
       Result = DAG.getNode(ISD::AND, dl, VT, Result,
                          DAG.getConstant(-(uint64_t)Align, dl, VT));
     Chain = DAG.getCopyToReg(Chain, dl, SPReg, Result); // Output chain
   } else if (SplitStack) {
     MachineRegisterInfo &MRI = MF.getRegInfo();
 
     if (Is64Bit) {
       // The 64 bit implementation of segmented stacks needs to clobber both r10
       // r11. This makes it impossible to use it along with nested parameters.
-      const Function *F = MF.getFunction();
-
       for (Function::const_arg_iterator I = F->arg_begin(), E = F->arg_end();
            I != E; ++I)
         if (I->hasNestAttr())
           report_fatal_error("Cannot use segmented stacks with functions that "
                              "have nested arguments.");
     }
 
     const TargetRegisterClass *AddrRegClass = getRegClassFor(SPTy);
     unsigned Vreg = MRI.createVirtualRegister(AddrRegClass);
     Chain = DAG.getCopyToReg(Chain, dl, Vreg, Size);
     Result = DAG.getNode(X86ISD::SEG_ALLOCA, dl, SPTy, Chain,
                                 DAG.getRegister(Vreg, SPTy));
   } else {
     SDValue Flag;
     const unsigned Reg = (Subtarget.isTarget64BitLP64() ? X86::RAX : X86::EAX);
 
     Chain = DAG.getCopyToReg(Chain, dl, Reg, Size, 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);
 
     if (Align) {
       SP = DAG.getNode(ISD::AND, dl, VT, SP.getValue(0),
                        DAG.getConstant(-(uint64_t)Align, dl, VT));
       Chain = DAG.getCopyToReg(Chain, dl, SPReg, SP);

test/CodeGen/X86/cleanuppad-inalloca.ll

ehcleanup:                                        ; preds = %entry
   call x86_thiscallcc void @"\01??1A@@QAE@XZ"(%struct.A* %0) [ "funclet"(token %2) ]
   cleanupret from %2 unwind to caller
 }
 
 ; CHECK: _passes_two:
 ; CHECK: pushl %ebp
 ; CHECK: movl %esp, %ebp
 ; CHECK: subl ${{[0-9]+}}, %esp
-; CHECK: movl $8, %eax
-; CHECK: calll __chkstk
+; CHECK: movl %esp, %[[tmp_sp1:.*]]
+; CHECK: leal -8(%[[tmp_sp1]]), %[[tmp_sp2:.*]]
+; CHECK: %[[tmp_sp2]], %esp
 ; CHECK: calll "??0A@@QAE@XZ"
 ; CHECK: calll "??0A@@QAE@XZ"
 ; CHECK: calll _takes_two
 ;   ESP must be restored via EBP due to "dynamic" alloca.
 ; CHECK: leal -{{[0-9]+}}(%ebp), %esp
 ; CHECK: popl %ebp
 ; CHECK: retl
 

test/CodeGen/X86/dynamic-alloca-in-entry.ll

 ; CHECK: retl
 
 ; Use of inalloca implies that that the alloca is not static.
 define void @bar() {
   %m = alloca inalloca i32
   ret void
 }
 ; CHECK-LABEL: _bar:
-; CHECK: calll __chkstk
+; CHECK: movl %esp, %ebp
+; CHECK: movl %esp, %[[sp_tmp:.*]]
+; CHECK: addl $-4, %[[sp_tmp]]
+; CHECK: movl %[[sp_tmp]], %esp
 ; CHECK: retl

test/CodeGen/X86/inalloca-ctor.ll

 ; RUN: llc < %s -mtriple=i686-pc-win32 | FileCheck %s
 
 %Foo = type { i32, i32 }
 
 %frame = type { %Foo, i32, %Foo }
 
 declare void @f(%frame* inalloca %a)
 
 declare void @Foo_ctor(%Foo* %this)
 
 define void @g() {
 entry:
+; CHECK: movl    %esp, %ebp
   %args = alloca inalloca %frame
   %c = getelementptr %frame, %frame* %args, i32 0, i32 2
-; CHECK: movl    $20, %eax
-; CHECK: calll   __chkstk
-; CHECK: movl %esp,
+; CHECK: movl    %esp, %[[tmp_sp1:.*]]
+; CHECK: leal    -20(%[[tmp_sp1]]), %[[tmp_sp2:.*]]
+; CHECK: movl    %[[tmp_sp2]], %esp
   call void @Foo_ctor(%Foo* %c)
-; CHECK: leal 12(%{{.*}}),
+; CHECK: leal -8(%[[tmp_sp1]]),
 ; CHECK-NEXT: pushl
 ; CHECK-NEXT: calll _Foo_ctor
 ; CHECK: addl $4, %esp
   %b = getelementptr %frame, %frame* %args, i32 0, i32 1
   store i32 42, i32* %b
 ; CHECK: movl $42,
   %a = getelementptr %frame, %frame* %args, i32 0, i32 0
   call void @Foo_ctor(%Foo* %a)
 ; CHECK-NEXT: pushl
 ; CHECK-NEXT: calll _Foo_ctor
 ; CHECK: addl $4, %esp
   call void @f(%frame* inalloca %args)
 ; CHECK: calll   _f
   ret void
 }

test/CodeGen/X86/inalloca-invoke.ll

 ; RUN: llc < %s -mtriple=i686-pc-win32 | FileCheck %s
 
 %Iter = type { i32, i32, i32 }
 
 %frame.reverse = type { %Iter, %Iter }
 
 declare i32 @pers(...)
 declare void @llvm.stackrestore(i8*)
 declare i8* @llvm.stacksave()
 declare void @begin(%Iter* sret)
 declare void @plus(%Iter* sret, %Iter*, i32)
 declare void @reverse(%frame.reverse* inalloca align 4)
 
 define i32 @main() personality i32 (...)* @pers {
+; CHECK: movl    %esp, %ebp
   %temp.lvalue = alloca %Iter
   br label %blah
 
 blah:
   %inalloca.save = call i8* @llvm.stacksave()
   %rev_args = alloca inalloca %frame.reverse, align 4
   %beg = getelementptr %frame.reverse, %frame.reverse* %rev_args, i32 0, i32 0
   %end = getelementptr %frame.reverse, %frame.reverse* %rev_args, i32 0, i32 1
 
-; CHECK:  calll   __chkstk
-; CHECK:  movl %esp, %[[beg:[^ ]*]]
-; CHECK:  leal 12(%[[beg]]), %[[end:[^ ]*]]
+; CHECK:  movl %esp, %[[end:.*]]
+; CHECK:  leal -24(%[[end]]), %[[beg:.*]]
+; CHECK:  movl %[[beg]], %esp
+; CHECK:  addl $-12, %[[end]]
 
   call void @begin(%Iter* sret %temp.lvalue)
 ; CHECK:  calll _begin
 
   invoke void @plus(%Iter* sret %end, %Iter* %temp.lvalue, i32 4)
           to label %invoke.cont unwind label %lpad
 
 ;  Uses end as sret param.

test/CodeGen/X86/inalloca-stdcall.ll

 ; RUN: llc < %s -mtriple=i686-pc-win32 | FileCheck %s
 
 %Foo = type { i32, i32 }
 
 declare x86_stdcallcc void @f(%Foo* inalloca %a)
 declare x86_stdcallcc void @i(i32 %a)
 
 define void @g() {
 ; CHECK-LABEL: _g:
+; CHECK: movl    %esp, %ebp
   %b = alloca inalloca %Foo
-; CHECK: movl    $8, %eax
-; CHECK: calll   __chkstk
+; CHECK: movl    %esp, %[[tmp_sp:.*]]
+; CHECK: leal    -8(%[[tmp_sp]]), %esp
   %f1 = getelementptr %Foo, %Foo* %b, i32 0, i32 0
   %f2 = getelementptr %Foo, %Foo* %b, i32 0, i32 1
   store i32 13, i32* %f1
   store i32 42, i32* %f2
-; CHECK: movl %esp, %eax
-; CHECK: movl    $13, (%eax)
-; CHECK: movl    $42, 4(%eax)
+; CHECK: movl    $13, -8(%[[tmp_sp]])
+; CHECK: movl    $42, -4(%[[tmp_sp]])
   call x86_stdcallcc void @f(%Foo* inalloca %b)
 ; CHECK: calll   _f@8
 ; CHECK-NOT: %esp
 ; CHECK: pushl
 ; CHECK: calll   _i@4
   call x86_stdcallcc void @i(i32 0)
   ret void
 }

test/CodeGen/X86/inalloca.ll

 ; RUN: llc < %s -mtriple=i686-pc-win32 | FileCheck %s
 
 %Foo = type { i32, i32 }
 
 declare void @f(%Foo* inalloca %b)
 
 define void @a() {
 ; CHECK-LABEL: _a:
 entry:
+; CHECK: movl    %esp, %ebp
   %b = alloca inalloca %Foo
-; CHECK: movl    $8, %eax
-; CHECK: calll   __chkstk
+; CHECK: movl    %esp, %[[tmp_sp:.*]]
+; CHECK: leal    -8(%[[tmp_sp]]), %esp
   %f1 = getelementptr %Foo, %Foo* %b, i32 0, i32 0
   %f2 = getelementptr %Foo, %Foo* %b, i32 0, i32 1
   store i32 13, i32* %f1
   store i32 42, i32* %f2
-; CHECK: movl %esp, %eax
-; CHECK: movl    $13, (%eax)
-; CHECK: movl    $42, 4(%eax)
+; CHECK: movl    $13, -8(%[[tmp_sp]])
+; CHECK: movl    $42, -4(%[[tmp_sp]])
   call void @f(%Foo* inalloca %b)
 ; CHECK: calll   _f
   ret void
 }
 
 declare void @inreg_with_inalloca(i32 inreg %a, %Foo* inalloca %b)
 
 define void @b() {
 ; CHECK-LABEL: _b:
 entry:
+; CHECK: movl    %esp, %ebp
   %b = alloca inalloca %Foo
-; CHECK: movl    $8, %eax
-; CHECK: calll   __chkstk
+; CHECK: movl    %esp, %[[tmp_sp:.*]]
+; CHECK: leal    -8(%[[tmp_sp]]), %esp
   %f1 = getelementptr %Foo, %Foo* %b, i32 0, i32 0
   %f2 = getelementptr %Foo, %Foo* %b, i32 0, i32 1
   store i32 13, i32* %f1
   store i32 42, i32* %f2
-; CHECK: movl %esp, %eax
-; CHECK: movl    $13, (%eax)
-; CHECK: movl    $42, 4(%eax)
+; CHECK: movl    $13, -8(%[[tmp_sp]])
+; CHECK: movl    $42, -4(%[[tmp_sp]])
   call void @inreg_with_inalloca(i32 inreg 1, %Foo* inalloca %b)
 ; CHECK: movl    $1, %eax
 ; CHECK: calll   _inreg_with_inalloca
   ret void
 }
 
 declare x86_thiscallcc void @thiscall_with_inalloca(i8* %a, %Foo* inalloca %b)
 
 define void @c() {
 ; CHECK-LABEL: _c:
 entry:
+; CHECK: movl    %esp, %ebp
   %b = alloca inalloca %Foo
-; CHECK: movl    $8, %eax
-; CHECK: calll   __chkstk
+; CHECK: movl    %esp, %[[tmp_sp:.*]]
+; CHECK: leal    -8(%[[tmp_sp]]), %esp
   %f1 = getelementptr %Foo, %Foo* %b, i32 0, i32 0
   %f2 = getelementptr %Foo, %Foo* %b, i32 0, i32 1
   store i32 13, i32* %f1
   store i32 42, i32* %f2
-; CHECK: movl %esp, %eax
-; CHECK-DAG: movl    $13, (%eax)
-; CHECK-DAG: movl    $42, 4(%eax)
+; CHECK-DAG: movl    $13, -8(%[[tmp_sp]])
+; CHECK-DAG: movl    $42, -4(%[[tmp_sp]])
   call x86_thiscallcc void @thiscall_with_inalloca(i8* null, %Foo* inalloca %b)
 ; CHECK-DAG: xorl    %ecx, %ecx
 ; CHECK: calll   _thiscall_with_inalloca
   ret void
 }

test/CodeGen/X86/shrink-wrap-chkstk.ll

 ; RUN: llc < %s -enable-shrink-wrap=true | FileCheck %s
 
 ; chkstk cannot come before the usual prologue, since it adjusts ESP.
 ; If chkstk is used in the prologue, we also have to be careful about preserving
 ; EAX if it is used.
 
 target datalayout = "e-m:x-p:32:32-i64:64-f80:32-n8:16:32-a:0:32-S32"
 target triple = "i686-pc-windows-msvc18.0.0"
 
 %struct.S = type { [12 x i8] }
 
-define x86_thiscallcc void @call_inalloca(i1 %x) {
+define x86_thiscallcc void @call_inalloca(i1 %x) "stack-probe-size"="12" {
 entry:
   %argmem = alloca inalloca <{ %struct.S }>, align 4
   %argidx1 = getelementptr inbounds <{ %struct.S }>, <{ %struct.S }>* %argmem, i32 0, i32 0, i32 0, i32 0
   %argidx2 = getelementptr inbounds <{ %struct.S }>, <{ %struct.S }>* %argmem, i32 0, i32 0, i32 0, i32 1
   store i8 42, i8* %argidx2, align 4
   br i1 %x, label %bb1, label %bb2
 
 bb1: