Support for the mno-stack-arg-probe flag

[nruslan]

Summary

Adds support for this flag. There is also another piece for clang (separate review). More info: https://bugs.llvm.org/show_bug.cgi?id=36221

Related: D43108

[nruslan]

@craig.topper, @MatzeB, @hans: Can someone take a look at this change?

[craig.topper]

Please upload patch with full context as described here https://llvm.org/docs/Phabricator.html#requesting-a-review-via-the-web-interface

[hans]

docs/LangRef.rst

1564 ↗

(On Diff #133551)

Is the "arg" part of the name significant? It seems to me this applies to stack probing in general, and that's not just for arguments?

lib/Target/X86/X86WinAllocaExpander.cpp

243 ↗

(On Diff #133551)

It seems odd that we end up in this case when we're not actually going to probe. Could the code be changed so that we end up in the Sub case instead, since that seems to do what we want here?

[nruslan]

docs/LangRef.rst

1564 ↗

(On Diff #133551)

It is just to be consistent and compatible with MinGW

lib/Target/X86/X86WinAllocaExpander.cpp

243 ↗

(On Diff #133551)

I am not sure what would be better here without major code changes. I tried initially to put it under Sub, but the problem is that it assumes specific fixed value which is -1 if unknown, so offsets will not be calculated correctly. So, I figured, it may be better to do directly under Probe. BTW -- It is the only place Probe is used, so possibly it can be renamed to something else?

[nruslan]

@craig.topper: Updated the diff.

@hans: I responded to the comments and also updated diff for clang with new changes.

[hans]

test/CodeGen/X86/no-stack-arg-probe.ll

28 ↗

(On Diff #134378)

This test passes even if I comment out the code in X86WinAllocaExpander::lower. It probably needs a test with a dynamic alloca.

[MatzeB]

• Have you considered using stack-probe-size=0 instead?
• Doing it as a separate no-stack-arg-probe probably needs an IR verifier rule that forbids setting stack-probe-size and no-stack-arg-probe at the same time. Maybe also some attribute compatibility rules for inlining?

[nruslan]

@MatzeB : mstack-probe-size=0 is already used for a different purpose. It basically, if I remember correctly, always forces calls to chkstk regardless of the actual stack usage. Normally that would only happen after 4K. Regarding your second question -- do we really need to have these flags to be mutually-exclusive? They are related but serve a different purpose. stack-probe-size merely changes the default probe size from 4K to the specified value. mstack-arg-probe/mno-stack-arg-probe enables/disables stack probes. If we disable stack probes, their size no longer matters, but we do not necessarily have to prohibit stack-probe-size. So, in fact, I can imagine somebody may even define system-wide alias if they want to override 4K (for whatever reason). Then some Makefile (which is unaware of this alias) may, for example, disable stack probes completely (like in UEFI code case), and we do not necessarily want to prohibit this combination.

[nruslan]

added a test case for variable-sized alloca

[nruslan]

@hans, @MatzeB : I have updated the diff (added the test case for variable-sized alloca). Please let me know if it looks good now.

[nruslan]

Added a small change to set max value for StackProbeSize in WinAlloca. This is to (potentially) avoid generating suboptimal code if -mno-stack-arg-probe is specified

[hans]

Seems fine to me, but MatzeB should probably sign off too.

[nruslan]

@MatzeB: do you need more time for the review, or should we just go ahead and merge this and clang changes?

@hans: I probably cannot commit the change myself, right? Who should I ask to do that?

[MatzeB]

In D43107#1011163, @nruslan wrote:

@MatzeB : mstack-probe-size=0 is already used for a different purpose. It basically, if I remember correctly, always forces calls to chkstk regardless of the actual stack usage. Normally that would only happen after 4K. Regarding your second question -- do we really need to have these flags to be mutually-exclusive? They are related but serve a different purpose. stack-probe-size merely changes the default probe size from 4K to the specified value. mstack-arg-probe/mno-stack-arg-probe enables/disables stack probes. If we disable stack probes, their size no longer matters, but we do not necessarily have to prohibit stack-probe-size. So, in fact, I can imagine somebody may even define system-wide alias if they want to override 4K (for whatever reason). Then some Makefile (which is unaware of this alias) may, for example, disable stack probes completely (like in UEFI code case), and we do not necessarily want to prohibit this combination.

• I don't really know what the current stack-probe-size=0 would be useful for; but then again I'm not actively working on it either.
• There's a difference in what flags clang allows and what IR allows. While I see your point for what clang should allow, for llvm IR I think it would be more desirable to be have things stricter and generally disallow non-useful/ambiguous combinations.

Anyway I don't care deeply about this and don't do that much middle-end work anyway so don't wait for my approval and go with what you think is best.

[whitequark]

I don't really know what the current stack-probe-size=0 would be useful for; but then again I'm not actively working on it either.

@MatzeB It is useful in cases where you want to perform manual stack probing, i.e. if you don't have an MMU or MPU and compare the requested stack size to some thread-local stack limit on every chkstk call. This is useful for microcontrollers.

[nruslan]

@hans: I can modify the clang patch to make sure it does not set stack-probe-size for IR if no-stack-arg-probe is set. Do you think, it is necessary? Also, if we want to further have some changes in the IR verifier, please let me know what example(s) I can take a look at.

[nruslan]

@MatzeB, @hans, @whitequark: Should we go ahead and merge both clang and llvm changes as is? (since they are both approved already). We can add IR checks in a separate commit if you think it is necessary.

[hans]

In D43107#1012948, @nruslan wrote:

@hans: I can modify the clang patch to make sure it does not set stack-probe-size for IR if no-stack-arg-probe is set. Do you think, it is necessary? Also, if we want to further have some changes in the IR verifier, please let me know what example(s) I can take a look at.

I don't know whether it's necessary to enforce that the attributes can't be used together, but I don't have a lot of experience in what the IR allows in this regard. I don't really see that setting both flags would cause problems, though it would of course be a little bit redundant.

I'm fine with the change as it is.

[nruslan]

@hans, @craig.topper, @MatzeB: Can someone commit D43108 and D43107 (this) on my behalf? I do not think, I have commit access.

[nruslan]

@hans: Thanks! Is it already too late to cherry-pick these changes to the upcoming 6.0 version?

[hans]

In D43107#1018277, @nruslan wrote:

@hans: Thanks! Is it already too late to cherry-pick these changes to the upcoming 6.0 version?

Yes, it's too late for 6.0.

llvm/trunk/docs/LangRef.rst

``"stack-probe-size"``
 
     If a function that has a ``"stack-probe-size"`` attribute is inlined into
     a function with another ``"stack-probe-size"`` attribute, the resulting
     function has the ``"stack-probe-size"`` attribute that has the lower
     numeric value. If a function that has a ``"stack-probe-size"`` attribute is
     inlined into a function that has no ``"stack-probe-size"`` attribute
     at all, the resulting function has the ``"stack-probe-size"`` attribute
     of the callee.
+``"no-stack-arg-probe"``
+    This attribute disables ABI-required stack probes, if any.
 ``writeonly``
     On a function, this attribute indicates that the function may write to but
     does not read from memory.
 
     On an argument, this attribute indicates that the function may write to but
     does not read through this pointer argument (even though it may read from
     the memory that the pointer points to).
 ``argmemonly``

llvm/trunk/lib/Target/AArch64/AArch64FrameLowering.cpp

static bool windowsRequiresStackProbe(MachineFunction &MF,
   const Function &F = MF.getFunction();
   // TODO: When implementing stack protectors, take that into account
   // for the probe threshold.
   unsigned StackProbeSize = 4096;
   if (F.hasFnAttribute("stack-probe-size"))
     F.getFnAttribute("stack-probe-size")
         .getValueAsString()
         .getAsInteger(0, StackProbeSize);
-  return StackSizeInBytes >= StackProbeSize;
+  return (StackSizeInBytes >= StackProbeSize) &&
+         !F.hasFnAttribute("no-stack-arg-probe");
 }
 
 bool AArch64FrameLowering::shouldCombineCSRLocalStackBump(
     MachineFunction &MF, unsigned StackBumpBytes) const {
   AArch64FunctionInfo *AFI = MF.getInfo<AArch64FunctionInfo>();
   const MachineFrameInfo &MFI = MF.getFrameInfo();
   const AArch64Subtarget &Subtarget = MF.getSubtarget<AArch64Subtarget>();
   const AArch64RegisterInfo *RegInfo = Subtarget.getRegisterInfo();

llvm/trunk/lib/Target/ARM/ARMFrameLowering.cpp

static bool WindowsRequiresStackProbe(const MachineFunction &MF,
                                       size_t StackSizeInBytes) {
   const MachineFrameInfo &MFI = MF.getFrameInfo();
   const Function &F = MF.getFunction();
   unsigned StackProbeSize = (MFI.getStackProtectorIndex() > 0) ? 4080 : 4096;
   if (F.hasFnAttribute("stack-probe-size"))
     F.getFnAttribute("stack-probe-size")
         .getValueAsString()
         .getAsInteger(0, StackProbeSize);
-  return StackSizeInBytes >= StackProbeSize;
+  return (StackSizeInBytes >= StackProbeSize) &&
+         !F.hasFnAttribute("no-stack-arg-probe");
 }
 
 namespace {
 
 struct StackAdjustingInsts {
   struct InstInfo {
     MachineBasicBlock::iterator I;
     unsigned SPAdjust;

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

 /// 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())
+  if (!Subtarget.isOSWindows() || Subtarget.isTargetMachO() ||
+      MF.getFunction().hasFnAttribute("no-stack-arg-probe"))
     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/lib/Target/X86/X86WinAllocaExpander.cpp

private:
 
   MachineRegisterInfo *MRI;
   const X86Subtarget *STI;
   const TargetInstrInfo *TII;
   const X86RegisterInfo *TRI;
   unsigned StackPtr;
   unsigned SlotSize;
   int64_t StackProbeSize;
+  bool NoStackArgProbe;
 
   StringRef getPassName() const override { return "X86 WinAlloca Expander"; }
   static char ID;
 };
 
 char X86WinAllocaExpander::ID = 0;
 
 } // end anonymous namespace
case Sub:
     } else {
       // Sub.
       BuildMI(*MBB, I, DL, TII->get(getSubOpcode(Is64Bit, Amount)), StackPtr)
           .addReg(StackPtr)
           .addImm(Amount);
     }
     break;
   case Probe:
+    if (!NoStackArgProbe) {
       // The probe lowering expects the amount in RAX/EAX.
       BuildMI(*MBB, MI, DL, TII->get(TargetOpcode::COPY), RegA)
           .addReg(MI->getOperand(0).getReg());
 
       // Do the probe.
       STI->getFrameLowering()->emitStackProbe(*MBB->getParent(), *MBB, MI, DL,
                                               /*InPrologue=*/false);
+    } else {
+      // Sub
+      BuildMI(*MBB, I, DL, TII->get(Is64Bit ? X86::SUB64rr : X86::SUB32rr),
+              StackPtr)
+          .addReg(StackPtr)
+          .addReg(MI->getOperand(0).getReg());
+    }
     break;
   }
 
   unsigned AmountReg = MI->getOperand(0).getReg();
   MI->eraseFromParent();
 
   // Delete the definition of AmountReg, possibly walking a chain of copies.
   for (;;) {
bool X86WinAllocaExpander::runOnMachineFunction(MachineFunction &MF) {
 
   StackProbeSize = 4096;
   if (MF.getFunction().hasFnAttribute("stack-probe-size")) {
     MF.getFunction()
         .getFnAttribute("stack-probe-size")
         .getValueAsString()
         .getAsInteger(0, StackProbeSize);
   }
+  NoStackArgProbe = MF.getFunction().hasFnAttribute("no-stack-arg-probe");
+  if (NoStackArgProbe)
+    StackProbeSize = INT64_MAX;
 
   LoweringMap Lowerings;
   computeLowerings(MF, Lowerings);
   for (auto &P : Lowerings)
     lower(P.first, P.second);
 
   return true;
 }

llvm/trunk/test/CodeGen/X86/no-stack-arg-probe.ll

+; This test is attempting to detect that the compiler disables stack
+; probe calls when the corresponding option is specified.
+;
+; RUN: llc -mtriple=i686-windows-msvc < %s | FileCheck %s
+
+target datalayout = "e-m:w-p:32:32-i64:64-f80:32-n8:16:32-S32"
+
+define i32 @test1() "no-stack-arg-probe" {
+  %buffer = alloca [4095 x i8]
+
+  ret i32 0
+
+; CHECK-LABEL: _test1:
+; CHECK-NOT: movl $4095, %eax
+; CHECK: subl $4095, %esp
+; CHECK-NOT: calll __chkstk
+}
+
+define i32 @test2() "no-stack-arg-probe" {
+  %buffer = alloca [4096 x i8]
+
+  ret i32 0
+
+; CHECK-LABEL: _test2:
+; CHECK-NOT: movl $4096, %eax
+; CHECK: subl $4096, %esp
+; CHECK-NOT: calll __chkstk
+}
+
+define i32 @test3(i32 %size) "no-stack-arg-probe" {
+  %buffer = alloca i8, i32 %size
+
+  ret i32 0
+
+; CHECK-LABEL: _test3:
+; CHECK: subl {{.*}}, %esp
+; CHECK-NOT: calll __chkstk
+}