Add a "probe-stack" attribute

[Zoxc]

Summary

This attribute is used to ensure the guard page is triggered on stack overflow. Stack frames larger than the guard page size will generate a call to __probestack to touch each page so the guard page won't be skipped.

[silvas]

Has this been discussed on LLVMdev? Our development process requires
changes to the IR to have a design discussion on LLVMdev first.

• Sean Silva

[Zoxc]

No, this has not yet been discussed on LLVMdev.

[Zoxc]

This now uses the string attribute "probe-stack". Dynamic allocations will have stack probes added. I also added a fast path for smaller frames.

[Zoxc]

Fixed up the tests some.

[Zoxc]

Used the right scale value on the OR instruction. Change the relocation to the __probestack call on non-Windows x86-64 platforms.

[rnk]

This seems like a reasonable implementation of useful functionality.

lib/Target/X86/X86FrameLowering.cpp

675–692

This change seems like a separable optimization, correct? I'd like to have a separate patch for it.

752

Should be 2 space indentation

[Zoxc]

This is originally 4 commits: https://github.com/Zoxc/llvm/compare/llvm-mirror:master...stprobe

lib/Target/X86/X86FrameLowering.cpp

752

The original code has 4 space indentation here. Don't blame me.

[reames]

Overall, I like the direction of the change, but there's a number of details which need adjusted before it's ready for submission. The biggest is that this patch seems to combine a bunch of separate pieces and it would be better to factor this apart into a series of patches which each do one thing. This would make it much easier to review.

include/llvm/CodeGen/MachineFunction.h

268 ↗

(On Diff #12236)

Could you define the term "probing the stack" here? This is likely be where an LLVM developer might encounter the concept. One short sentence is fine.

lib/Target/ARM/ARMFrameLowering.cpp

303 ↗

(On Diff #12236)

Seems like you should rename WindowsRequiredStackProbe since this is no longer Windows specific.

lib/Target/X86/X86FrameLowering.cpp

500

What is this magic constant? At minimum, comment it.

This looks like it might be an independent bugfix? If so, please submit separately.

673

This section is a large diff that I'm having trouble understanding. If you can split this up into a set of smaller patches, that would help a lot.

675

Magic numbers? What do these mean?

708

In a separate change, exacting out this code out as a helper function (getStackProbeSymbol(..)) would make the functional diff easier to follow.

712

I don't understand this section at all. What's the motivation here?

Again, could this be separate patch?

lib/Target/X86/X86ISelLowering.cpp

18201 ↗

(On Diff #12236)

This condition "feels" wrong. It looks like you're only avoiding this case for 32 bit Windows? Why? The old code didn't.

18227 ↗

(On Diff #12236)

I'm not clear on what this change is doing. In addition to the changed symbol, it looks like you're adding 64 vs 32 bit case? Is this an unrelated bugfix?

lib/Transforms/IPO/Inliner.cpp

140

Add a comment to justify.

Would it be better to just prevent inlining when only one has probe-stack? (i.e. is this expected to ever happen?)

test/CodeGen/X86/mingw-alloca.ll

32 ↗

(On Diff #12236)

Why are you modifying the existing test rather than adding a new one? Is this test incorrect? It doesn't seem to be on first read.

test/CodeGen/X86/pr17631.ll

21 ↗

(On Diff #12236)

This looks like it goes with an optimization?

Can you add a test case for the unoptimized case?

Again, optimizations should be separate patches.

test/CodeGen/X86/stack-probes.ll

2

Why bother with a custom prefix if they are the same? You can just use the default CHECK prefix.

24

A comment describing what each test is checking for would be helpful.

test/CodeGen/X86/win64_alloca_dynalloca.ll

13 ↗

(On Diff #12236)

Why are you modify existing tests rather than adding new ones? Is this because of a new optimization?

[Zoxc]

A newer version was posted to the mailing list.

lib/Target/X86/X86ISelLowering.cpp

18201 ↗

(On Diff #12236)

The old code did also just avoid 32-bit Windows (the only other case was 64-bit Windows)

18227 ↗

(On Diff #12236)

This calls the x86-32 probestack function the same way as x86-64 chkstk/__probestack. I'll make this more explicit.

test/CodeGen/X86/stack-probes.ll

2

A custom prefix seems to be the pattern, and it makes sense if tests for other platforms will be added.

docs/LangRef.rst

 ``sanitize_thread``
     This attribute indicates that ThreadSanitizer checks
     (dynamic thread safety analysis) are enabled for this function.
+``stackprobe``
+    This attribute indicates that the function should emit stack
+    probes when the stack frame is larger than the guard page of the
+    target's stack. This ensures the function will be unable to overflow
+    the stack without it being detected.
+
+    If a function that has a ``stackprobe`` attribute is inlined into a
+    function that doesn't have a ``stackprobe`` attribute, then the resulting
+    function will have a ``stackprobe`` attribute.
 ``ssp``
     This attribute indicates that the function should emit a stack
     smashing protector. It is in the form of a "canary" --- a random value

include/llvm-c/Core.h

     LLVMNonNullAttribute = 1ULL << 37,
     LLVMJumpTableAttribute = 1ULL << 38,
     LLVMDereferenceableAttribute = 1ULL << 39,
+    LLVMStackProbe = 1ULL << 40,
     */
 } LLVMAttribute;
 

include/llvm/Bitcode/LLVMBitCodes.h

     ATTR_KIND_IN_ALLOCA = 38,
     ATTR_KIND_NON_NULL = 39,
     ATTR_KIND_JUMP_TABLE = 40,
-    ATTR_KIND_DEREFERENCEABLE = 41
+    ATTR_KIND_DEREFERENCEABLE = 41,
+    ATTR_KIND_STACK_PROBE = 42
   };
 
   enum ComdatSelectionKindCodes {

include/llvm/IR/Attributes.h

                            ///< stored as log2 of alignment with +1 bias 0
                            ///< means unaligned (different from
                            ///< alignstack=(1))
+    StackProbe,            ///< Insert stack probes.
     StackProtect,          ///< Stack protection.
     StackProtectReq,       ///< Stack protection required.
     StackProtectStrong,    ///< Strong Stack protection.

lib/AsmParser/LLLexer.cpp

   KEYWORD(returns_twice);
   KEYWORD(signext);
   KEYWORD(sret);
+  KEYWORD(stackprobe);
   KEYWORD(ssp);
   KEYWORD(sspreq);
   KEYWORD(sspstrong);

lib/AsmParser/LLParser.cpp

     case lltok::kw_readnone:          B.addAttribute(Attribute::ReadNone); break;
     case lltok::kw_readonly:          B.addAttribute(Attribute::ReadOnly); break;
     case lltok::kw_returns_twice:     B.addAttribute(Attribute::ReturnsTwice); break;
+    case lltok::kw_stackprobe:        B.addAttribute(Attribute::StackProbe); break;
     case lltok::kw_ssp:               B.addAttribute(Attribute::StackProtect); break;
     case lltok::kw_sspreq:            B.addAttribute(Attribute::StackProtectReq); break;
     case lltok::kw_sspstrong:         B.addAttribute(Attribute::StackProtectStrong); break;

     case lltok::kw_sanitize_address:
     case lltok::kw_sanitize_memory:
     case lltok::kw_sanitize_thread:
+    case lltok::kw_stackprobe:
     case lltok::kw_ssp:
     case lltok::kw_sspreq:
     case lltok::kw_sspstrong:

     case lltok::kw_sanitize_address:
     case lltok::kw_sanitize_memory:
     case lltok::kw_sanitize_thread:
+    case lltok::kw_stackprobe:
     case lltok::kw_ssp:
     case lltok::kw_sspreq:
     case lltok::kw_sspstrong:

lib/AsmParser/LLToken.h

     kw_returned,
     kw_returns_twice,
     kw_signext,
+    kw_stackprobe,
     kw_ssp,
     kw_sspreq,
     kw_sspstrong,

lib/Bitcode/Reader/BitcodeReader.cpp

     return Attribute::SExt;
   case bitc::ATTR_KIND_STACK_ALIGNMENT:
     return Attribute::StackAlignment;
+  case bitc::ATTR_KIND_STACK_PROBE:
+    return Attribute::StackProbe;
   case bitc::ATTR_KIND_STACK_PROTECT:
     return Attribute::StackProtect;
   case bitc::ATTR_KIND_STACK_PROTECT_REQ:

lib/Bitcode/Writer/BitcodeWriter.cpp

     return bitc::ATTR_KIND_S_EXT;
   case Attribute::StackAlignment:
     return bitc::ATTR_KIND_STACK_ALIGNMENT;
+  case Attribute::StackProbe:
+    return bitc::ATTR_KIND_STACK_PROBE;
   case Attribute::StackProtect:
     return bitc::ATTR_KIND_STACK_PROTECT;
   case Attribute::StackProtectReq:

lib/IR/Attributes.cpp

     return "returns_twice";
   if (hasAttribute(Attribute::SExt))
     return "signext";
+  if (hasAttribute(Attribute::StackProbe))
+    return "stackprobe";
   if (hasAttribute(Attribute::StackProtect))
     return "ssp";
   if (hasAttribute(Attribute::StackProtectReq))

   case Attribute::InAlloca:        return 1ULL << 43;
   case Attribute::NonNull:         return 1ULL << 44;
   case Attribute::JumpTable:       return 1ULL << 45;
+  case Attribute::StackProbe:      return 1ULL << 46;
   case Attribute::Dereferenceable:
     llvm_unreachable("dereferenceable attribute not supported in raw format");
   }

lib/IR/Verifier.cpp

         I->getKindAsEnum() == Attribute::NoInline ||
         I->getKindAsEnum() == Attribute::AlwaysInline ||
         I->getKindAsEnum() == Attribute::OptimizeForSize ||
+        I->getKindAsEnum() == Attribute::StackProbe ||
         I->getKindAsEnum() == Attribute::StackProtect ||
         I->getKindAsEnum() == Attribute::StackProtectReq ||
         I->getKindAsEnum() == Attribute::StackProtectStrong ||

lib/Target/CppBackend/CPPBackend.cpp

       HANDLE_ATTR(AlwaysInline);
       HANDLE_ATTR(OptimizeNone);
       HANDLE_ATTR(OptimizeForSize);
+      HANDLE_ATTR(StackProbe);
       HANDLE_ATTR(StackProtect);
       HANDLE_ATTR(StackProtectReq);
       HANDLE_ATTR(StackProtectStrong);

lib/Target/X86/X86FrameLowering.cpp

   // 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.
-  if (NumBytes >= 4096 && STI.isOSWindows() && !STI.isTargetMacho()) {
-    const char *StackProbeSymbol;
 
-    if (Is64Bit) {
-      if (STI.isTargetCygMing()) {
-        StackProbeSymbol = "___chkstk_ms";
-      } else {
-        StackProbeSymbol = "__chkstk";
-      }
-    } else if (STI.isTargetCygMing())
-      StackProbeSymbol = "_alloca";
-    else
-      StackProbeSymbol = "_chkstk";
+  if (NumBytes >= 4096 && ((STI.isOSWindows() && !STI.isTargetMacho()) ||
+                           Fn->getAttributes()
+                             .hasAttribute(AttributeSet::FunctionIndex,
+                                           Attribute::StackProbe))) {
+    const char *StackProbeSymbol;
 
-    // Check whether EAX is livein for this function.
-    bool isEAXAlive = isEAXLiveIn(MF);
+    if (STI.isOSWindows()) {
+      if (Is64Bit) {
+        if (STI.isTargetCygMing()) {
+          StackProbeSymbol = "___chkstk_ms";
+        } else {
+          StackProbeSymbol = "__chkstk";
+        }
+      } else if (STI.isTargetCygMing())
+        StackProbeSymbol = "_alloca";
+      else
+        StackProbeSymbol = "_chkstk";
+    } else {
+      StackProbeSymbol = "__probestack";
+    }
 
-    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!");
+    // Check whether the accumulator register is livein for this function.
+    bool isRegAccAlive = isEAXLiveIn(MF);
+    auto RegAcc = Is64Bit ? X86::RAX : X86::EAX;
 
-      // 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);
     }
 
-    if (Is64Bit) {
-      // Handle the 64-bit Windows ABI case where we need to call __chkstk.
-      // Function prologue is responsible for adjusting the stack pointer.
-      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)
-        .setMIFlag(MachineInstr::FrameSetup);
-    }
+    uint64_t NumBytesAdj = isRegAccAlive ? NumBytes - (Is64Bit ? 8 : 4) :
+                                         NumBytes;
+
+    // Allocate NumBytesAdj bytes on stack in case of isRegAAlive.
+    // We'll also use 8/4 already allocated bytes for EAX.
+    BuildMI(MBB, MBBI, DL, TII.get(Is64Bit ? X86::MOV64ri : X86::MOV32ri),
+            RegAcc)
+      .addImm(NumBytesAdj)
+      .setMIFlag(MachineInstr::FrameSetup);
 
     BuildMI(MBB, MBBI, DL,
             TII.get(Is64Bit ? X86::W64ALLOCA : X86::CALLpcrel32))

       .addReg(X86::EFLAGS, RegState::Define | RegState::Implicit)
       .setMIFlag(MachineInstr::FrameSetup);
 
-    if (Is64Bit) {
+    if (Is64Bit || !STI.isOSWindows()) {
       // MSVC x64's __chkstk and cygwin/mingw's ___chkstk_ms do not adjust %rsp
       // themself. It also does not clobber %rax so we can reuse it when
       // adjusting %rsp.
-      BuildMI(MBB, MBBI, DL, TII.get(X86::SUB64rr), StackPtr)
+      BuildMI(MBB, MBBI, DL, TII.get(Is64Bit ? X86::SUB64rr : X86::SUB32rr),
+              StackPtr)
         .addReg(StackPtr)
-        .addReg(X86::RAX)
+        .addReg(RegAcc)
         .setMIFlag(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);
     }

lib/Transforms/IPO/Inliner.cpp

 
   AdjustCallerSSPLevel(Caller, Callee);
 
+  if (Callee->getAttributes().hasAttribute(AttributeSet::FunctionIndex,
+                                           Attribute::StackProbe)) {
+    Caller->addFnAttr(Attribute::StackProbe);
+  }
+
   // Look at all of the allocas that we inlined through this call site.  If we
   // have already inlined other allocas through other calls into this function,
   // then we know that they have disjoint lifetimes and that we can merge them.

test/Bitcode/attributes.ll

 ; CHECK: define void @f34()
 {
         call void @nobuiltin() nobuiltin
-; CHECK: call void @nobuiltin() #25
+; CHECK: call void @nobuiltin() #26
         ret void;
 }
 

         ret void
 }
 
+define void @f42() stackprobe
+; CHECK: define void @f42() #25
+{
+        ret void;
+}
+
 ; CHECK: attributes #0 = { noreturn }
 ; CHECK: attributes #1 = { nounwind }
 ; CHECK: attributes #2 = { readnone }

 ; CHECK: attributes #22 = { minsize }
 ; CHECK: attributes #23 = { noinline optnone }
 ; CHECK: attributes #24 = { jumptable }
-; CHECK: attributes #25 = { nobuiltin }
+; CHECK: attributes #25 = { stackprobe }
+; CHECK: attributes #26 = { nobuiltin }

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([4096 x i8]*)
+
+define void @test() stackprobe {
+  %array = alloca [4096 x i8], align 16
+  call void @use([4096 x i8]* %array)
+  ret void
+
+; X86-Linux-LABEL: test:
+; X86-Linux:       movl $4124, %eax
+; X86-Linux-NEXT:  calll __probestack
+; X86-Linux-NEXT:  subl %eax, %esp
+; X86-Linux:       leal 16(%esp), %eax
+; X86-Linux-NEXT:  movl %eax, (%esp)
+; X86-Linux-NEXT:  calll use
+; X86-Linux-NEXT:  addl $4124, %esp
+; X86-Linux-NEXT:  retl
+
+; X64-Linux-LABEL: test:
+; X64-Linux:       movabsq $4104, %rax
+; X64-Linux-NEXT:  callq __probestack
+; X64-Linux-NEXT:  subq %rax, %rsp
+; X64-Linux:       leaq (%rsp), %rdi
+; X64-Linux-NEXT:  callq use
+; X64-Linux-NEXT:  addq $4104, %rsp
+; X64-Linux-NEXT:  retq
+  
+}

utils/kate/llvm.xml

       <item> optsize </item>
       <item> readnone </item>
       <item> readonly </item>
+      <item> stackprobe </item>
       <item> ssp </item>
       <item> sspreq </item>
       <item> sspstrong </item>

utils/vim/llvm.vim

 syn keyword llvmKeyword tail target thread_local to triple unnamed_addr
 syn keyword llvmKeyword unordered uwtable volatile weak weak_odr
 syn keyword llvmKeyword x86_fastcallcc x86_stdcallcc x86_thiscallcc x86_64_sysvcc
-syn keyword llvmKeyword x86_64_win64cc zeroext
+syn keyword llvmKeyword x86_64_win64cc zeroext stackprobe
 
 " Obsolete keywords.
 syn keyword llvmError  getresult begin end