Support for the mno-stack-arg-probe flag

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

By Ruslan Nikolaev!

Differential Revision: https://reviews.llvm.org/D43107

llvm-svn: 325900

6 files changed, 65 insertions, 10 deletions

diff --git a/llvm/docs/LangRef.rst b/llvm/docs/LangRef.rst
index 6999e8f55f8..97a8471c9e6 100644
--- a/llvm/docs/LangRef.rst
+++ b/llvm/docs/LangRef.rst
@@ -1561,6 +1561,8 @@ example:
     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.
diff --git a/llvm/lib/Target/AArch64/AArch64FrameLowering.cpp b/llvm/lib/Target/AArch64/AArch64FrameLowering.cpp
index ea4bfe7e8d9..f49ca11d98f 100644
--- a/llvm/lib/Target/AArch64/AArch64FrameLowering.cpp
+++ b/llvm/lib/Target/AArch64/AArch64FrameLowering.cpp
@@ -369,7 +369,8 @@ static bool windowsRequiresStackProbe(MachineFunction &MF,
     F.getFnAttribute("stack-probe-size")
         .getValueAsString()
         .getAsInteger(0, StackProbeSize);
-  return StackSizeInBytes >= StackProbeSize;
+  return (StackSizeInBytes >= StackProbeSize) &&
+         !F.hasFnAttribute("no-stack-arg-probe");
 }
 
 bool AArch64FrameLowering::shouldCombineCSRLocalStackBump(
diff --git a/llvm/lib/Target/ARM/ARMFrameLowering.cpp b/llvm/lib/Target/ARM/ARMFrameLowering.cpp
index 4ff864ac6cc..f438469acb2 100644
--- a/llvm/lib/Target/ARM/ARMFrameLowering.cpp
+++ b/llvm/lib/Target/ARM/ARMFrameLowering.cpp
@@ -209,7 +209,8 @@ static bool WindowsRequiresStackProbe(const MachineFunction &MF,
     F.getFnAttribute("stack-probe-size")
         .getValueAsString()
         .getAsInteger(0, StackProbeSize);
-  return StackSizeInBytes >= StackProbeSize;
+  return (StackSizeInBytes >= StackProbeSize) &&
+         !F.hasFnAttribute("no-stack-arg-probe");
 }
 
 namespace {
diff --git a/llvm/lib/Target/X86/X86ISelLowering.cpp b/llvm/lib/Target/X86/X86ISelLowering.cpp
index 0f9f9bf1b83..e0767c856f2 100644
--- a/llvm/lib/Target/X86/X86ISelLowering.cpp
+++ b/llvm/lib/Target/X86/X86ISelLowering.cpp
@@ -39186,7 +39186,8 @@ StringRef X86TargetLowering::getStackProbeSymbolName(MachineFunction &MF) const
 
   // 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
diff --git a/llvm/lib/Target/X86/X86WinAllocaExpander.cpp b/llvm/lib/Target/X86/X86WinAllocaExpander.cpp
index 1046696587d..d298aaa97ec 100644
--- a/llvm/lib/Target/X86/X86WinAllocaExpander.cpp
+++ b/llvm/lib/Target/X86/X86WinAllocaExpander.cpp
@@ -62,6 +62,7 @@ private:
   unsigned StackPtr;
   unsigned SlotSize;
   int64_t StackProbeSize;
+  bool NoStackArgProbe;
 
   StringRef getPassName() const override { return "X86 WinAlloca Expander"; }
   static char ID;
@@ -240,13 +241,21 @@ void X86WinAllocaExpander::lower(MachineInstr* MI, Lowering L) {
     }
     break;
   case Probe:
-    // 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);
+    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;
   }
 
@@ -285,6 +294,9 @@ bool X86WinAllocaExpander::runOnMachineFunction(MachineFunction &MF) {
         .getValueAsString()
         .getAsInteger(0, StackProbeSize);
   }
+  NoStackArgProbe = MF.getFunction().hasFnAttribute("no-stack-arg-probe");
+  if (NoStackArgProbe)
+    StackProbeSize = INT64_MAX;
 
   LoweringMap Lowerings;
   computeLowerings(MF, Lowerings);
diff --git a/llvm/test/CodeGen/X86/no-stack-arg-probe.ll b/llvm/test/CodeGen/X86/no-stack-arg-probe.ll
new file mode 100644
index 00000000000..7639996636d
--- /dev/null
+++ b/llvm/test/CodeGen/X86/no-stack-arg-probe.ll
@@ -0,0 +1,38 @@
+; 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
+}