[PATCH 2/2] Implement "probe-stack" on x86

[pcwalton]

Summary

I've fixed up https://reviews.llvm.org/D9654. It's a follow-up to https://reviews.llvm.org/D34386. Please let me know if you have any further suggestions. Thanks!

[majnemer]

How does this review differ from D34386?

[fhahn]

llvm-commits should be added as a subscriber when the review is created to mirror the review on the mailing list

[fhahn]

I'll repost some of my comments from D34386. I think it would be good to avoid unrelated whitespace changes to reduce the noise in the diff. Also, could you add more context to the diff? That makes reviewing easier :)

I think you should also include a test for your change. You could probably extend test/CodeGen/X86/stack-probe-size.ll, by adding an additional RUN line with a triple where stack probes aren't generated at the moment and add probe-stack attributes to some functions.

[pcwalton]

Added more context, added missing tests, removed whitespace-only changes.

[majnemer]

lib/CodeGen/MachineFunction.cpp

216–224 ↗

(On Diff #103236)

I think this should be renamed or change its definition to reflect its name, we will probe the stack on x86 Windows regardless of the whether or not this attribute was supplied.

Perhaps hasProbeStackAttr?

lib/Target/X86/X86FrameLowering.cpp

801 ↗

(On Diff #103236)

Could this be a StringRef?

1034–1035 ↗

(On Diff #103236)

You could make a function which gets the stack probe symbol name, then UseStackProbe could be defined in terms of that name being non-empty.

1255–1257 ↗

(On Diff #103236)

Test case for this?

1263–1264 ↗

(On Diff #103236)

Saving RBX seems new. I'd have a comment describing both of these saves.

1271 ↗

(On Diff #103236)

Magic number should have a comment or stuck in a const variable with a nice name. Even better would be both.

[pcwalton]

lib/Target/X86/X86FrameLowering.cpp

801 ↗

(On Diff #103236)

I originally wanted to do that, but the docs say that StringRefs don't have null termination, which the addExternalSymbol API needs. If you'd prefer to work around that in another way, let me know.

[majnemer]

lib/Target/X86/X86FrameLowering.cpp

801 ↗

(On Diff #103236)

Hmm. Isn't there a liveness problem then? Symbol.c_str() is only valid until this functoin returns. addExternalSymbol doesn't make a copy of its input unless I am mistaken.

[majnemer]

lib/Target/X86/X86FrameLowering.cpp

801 ↗

(On Diff #103236)

If I am correct, I think you will want to use createExternalSymbolName.

[pcwalton]

lib/Target/X86/X86FrameLowering.cpp

1263–1264 ↗

(On Diff #103236)

I don't know why the original patch author chose to spill RBX. Seems not needed. I'll remove it.

1271 ↗

(On Diff #103236)

As before, I don't know what this magic number is for. Seems unneeded. I'll remove it.

[pcwalton]

lib/Target/X86/X86FrameLowering.cpp

1271 ↗

(On Diff #103236)

Oh, I guess this is the page size. I don't see why user-requested stack probes should need this to be supplied explicitly; it just bloats the size of every function prolog for something that is going to be constant on the architecture. Will remove.

[whitequark]

@pcwalton No, both the spill and the "magic" number are necessary. See https://reviews.llvm.org/D9858. ebx is the size of the frame, and 0x1000 is the probe interval to be performed (in other words, page size).

[whitequark]

@pcwalton But I agree that it's better to change both sides of the __probestack ABI to eliminate the argument and the spill (on x86_64), of course.

[pcwalton]

Addressed review comments.

[majnemer]

lib/Target/X86/X86FrameLowering.cpp

510 ↗

(On Diff #103293)

Why is this by ref?

801 ↗

(On Diff #103293)

I think it is more natural to see: StringRef Symbol = STI.getTargetLowering()->getStackProbeSymbolName(MF);

1243–1245 ↗

(On Diff #103293)

Am I correct in thinking that we will save and restore R11 when we previously would not in the large code model for Win64? If so, was that a bug in the implementation of the large code model for Win64?

lib/Target/X86/X86ISelLowering.cpp

36342 ↗

(On Diff #103293)

const StringRef looks weird, StringRefs are immutable.

36344–36347 ↗

(On Diff #103293)

I'd kill the braces here.

36355–36357 ↗

(On Diff #103293)

Ditto.

lib/Target/X86/X86ISelLowering.h

1059 ↗

(On Diff #103293)

Ditto regarding const StringRef.

[pcwalton]

In D34387#785802, @whitequark wrote:

@pcwalton No, both the spill and the "magic" number are necessary. See https://reviews.llvm.org/D9858. ebx is the size of the frame, and 0x1000 is the probe interval to be performed (in other words, page size).

That patch never landed, and the current approach, with customizable symbol names, is to move the responsibility out of compiler-rt. So I would argue that https://reviews.llvm.org/D9858 is irrelevant.

lib/Target/X86/X86FrameLowering.cpp

1243–1245 ↗

(On Diff #103293)

I'm not familiar enough with the Windows 64-bit ABI to be able to answer that question. I would assume that it would have to have been a bug in the previous implementation, as I don't see any other way to make the call to __chkstk work…

[pcwalton]

Addressed comments.

[majnemer]

lib/Target/X86/X86FrameLowering.cpp

1253 ↗

(On Diff #103308)

Hmm, I believe a testcase for this situation hasn't been added.

[pcwalton]

lib/Target/X86/X86FrameLowering.cpp

1253 ↗

(On Diff #103308)

Well, I'm not sure how R11 can actually be live-in in the function prolog right now. As I understand things, the reason why EAX is spilled if necessary is to handle the edge case on 32-bit x86 where inreg is specified for a function parameter, which causes EAX to be live-in at the very start of the function. But there's no corresponding inreg attribute for x86-64 that causes R11 to become live-in that early. Checking for that case even for R11 seemed prudent to me though, because (a) the corresponding situation can happen for EAX and it seemed cleaner to use the same logic for both registers instead of duplicating some of it; (b) in the future, if somehow R11 can become live-in at function start, this logic will be needed.

Am I interpreting the situation correctly, and if so, what would you prefer I do regarding testing?

[whitequark]

lib/Target/X86/X86FrameLowering.cpp

1253 ↗

(On Diff #103308)

I am quite certain that R11 should *not* be spilled this way. LLVM uses it as a scratch register in quite a few places. E.g. the comment above:

// For the large code model, we have to call through a register. Use R11,
// as it is scratch in all supported calling conventions.

the GetScratchRegister function in X86FrameLowering.cpp, the X86TargetLowering::getScratchRegisters, and the X86TargetLowering::LowerINIT_TRAMPOLINE function, which needs it for essentially the same job.

You could perhaps replace the call with one of the getScratchRegister functions but it seems fine to hardcode it while not spilling the way it is, too.

[pcwalton]

Addressed review comments.

I removed the R11 spill code as requested. I left in the refactoring that moves the existing EAX spill and reload code into separate functions, as X86FrameLowering::emitPrologue() is rather large and that seemed like a worthwhile refactoring.

[whitequark]

LGTM, but I'd still like to hear from @majnemer

[majnemer]

In D34387#787200, @pcwalton wrote:

Addressed review comments.

I removed the R11 spill code as requested. I left in the refactoring that moves the existing EAX spill and reload code into separate functions, as X86FrameLowering::emitPrologue() is rather large and that seemed like a worthwhile refactoring.

I am not entirely sure that the refactored code is identical to what was left behind. Moreover, the refactoring is now incidental with respect to the bulk of the change. I'd request that if you were to pursue it, to please factor it out into another review.

[pcwalton]

Addressed comments. Reverted the EAX save refactoring.

[majnemer]

Looking pretty good, I think we are almost there :)

lib/Target/X86/X86FrameLowering.cpp

1003–1005 ↗

(On Diff #103477)

We can use the redzone and stack probes at the same time? This surprises me. Also, there is no test for this.

[pcwalton]

Addressed review comments.

[majnemer]

LGTM!

[joerg]

It would be nice to make the stackprobesize a proper TLI attribute, so that OSes can decide how much guarding they are willing to guarantee. I think I would also like to have an option to inline the probing without the complexity of the function call. But both can be done as a separate step.

llvm/trunk/include/llvm/Target/TargetLowering.h

 protected:
   Value *getDefaultSafeStackPointerLocation(IRBuilder<> &IRB,
                                             bool UseTLS) const;
 
 public:
   /// 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;
   }
 
   /// Returns true if a cast from SrcAS to DestAS is "cheap", such that e.g. we
   /// are happy to sink it into basic blocks.
   virtual bool isCheapAddrSpaceCast(unsigned SrcAS, unsigned DestAS) const {

llvm/trunk/lib/Target/X86/X86FrameLowering.cpp

void X86FrameLowering::emitStackProbeCall(MachineFunction &MF,
   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 (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);
 
   // 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);
+        .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;
   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);
   }
 
   if (InProlog) {
     // Apply the frame setup flag to all inserted instrs.
     for (++ExpansionMBBI; ExpansionMBBI != MBBI; ++ExpansionMBBI)
       ExpansionMBBI->setFlag(MachineInstr::FrameSetup);
   }
 }
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.getTargetLowering()->getStackProbeSymbolName(MF).empty();
 
   // 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);
void X86FrameLowering::emitPrologue(MachineFunction &MF,
   // 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.
+      !UseStackProbe &&                        // No stack probes.
       !IsWin64CC &&                            // Win64 has no Red Zone
       !MFI.hasCopyImplyingStackAdjustment() && // Don't push and pop.
       !MF.shouldSplitStack()) {                // Regular stack
     uint64_t MinSize = X86FI->getCalleeSavedFrameSize();
     if (HasFP) MinSize += SlotSize;
     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
   // 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 && !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);
 
     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!");
 

llvm/trunk/lib/Target/X86/X86ISelLowering.h

public:
 
     /// \brief Customize the preferred legalization strategy for certain types.
     LegalizeTypeAction getPreferredVectorAction(EVT VT) const override;
 
     bool isIntDivCheap(EVT VT, AttributeList Attr) const override;
 
     bool supportSwiftError() const override;
 
+    StringRef getStackProbeSymbolName(MachineFunction &MF) const override;
+
     unsigned getMaxSupportedInterleaveFactor() const override { return 4; }
 
     /// \brief Lower interleaved load(s) into target specific
     /// instructions/intrinsics.
     bool lowerInterleavedLoad(LoadInst *LI,
                               ArrayRef<ShuffleVectorInst *> Shuffles,
                               ArrayRef<unsigned> Indices,
                               unsigned Factor) const override;

llvm/trunk/lib/Target/X86/X86ISelLowering.cpp

 // 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();
   bool SplitStack = MF.shouldSplitStack();
+  bool EmitStackProbe = !getStackProbeSymbolName(MF).empty();
   bool Lower = (Subtarget.isOSWindows() && !Subtarget.isTargetMachO()) ||
-               SplitStack;
+               SplitStack || EmitStackProbe;
   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);
for (auto *Exit : Exits)
               TII->get(TargetOpcode::COPY), *I)
           .addReg(NewVR);
   }
 }
 
 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";
+}

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

+; 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
+}

llvm/trunk/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
+; 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
+
+}