[ARM, AArch64] Check the no-stack-arg-probe attribute for dynamic stack probes

This extends the use of this attribute on ARM and AArch64 from
SVN r325900 (where it was only checked for fixed stack
allocations on ARM/AArch64, but for all stack allocations on X86).

This also adds a testcase for the existing use of disabling the
fixed stack probe with the attribute on ARM and AArch64.

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

llvm-svn: 327897

6 files changed, 89 insertions, 0 deletions

diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
index 43b2b1cc46b..b566810520a 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -7520,6 +7520,19 @@ AArch64TargetLowering::LowerDYNAMIC_STACKALLOC(SDValue Op,
   unsigned Align = cast<ConstantSDNode>(Op.getOperand(2))->getZExtValue();
   EVT VT = Node->getValueType(0);
 
+  if (DAG.getMachineFunction().getFunction().hasFnAttribute(
+          "no-stack-arg-probe")) {
+    SDValue SP = DAG.getCopyFromReg(Chain, dl, AArch64::SP, MVT::i64);
+    Chain = SP.getValue(1);
+    SP = DAG.getNode(ISD::SUB, dl, MVT::i64, SP, Size);
+    if (Align)
+      SP = DAG.getNode(ISD::AND, dl, VT, SP.getValue(0),
+                       DAG.getConstant(-(uint64_t)Align, dl, VT));
+    Chain = DAG.getCopyToReg(Chain, dl, AArch64::SP, SP);
+    SDValue Ops[2] = {SP, Chain};
+    return DAG.getMergeValues(Ops, dl);
+  }
+
   Chain = DAG.getCALLSEQ_START(Chain, 0, 0, dl);
 
   Chain = LowerWindowsDYNAMIC_STACKALLOC(Op, Chain, Size, DAG);
diff --git a/llvm/lib/Target/ARM/ARMISelLowering.cpp b/llvm/lib/Target/ARM/ARMISelLowering.cpp
index dc4185752ea..762da041244 100644
--- a/llvm/lib/Target/ARM/ARMISelLowering.cpp
+++ b/llvm/lib/Target/ARM/ARMISelLowering.cpp
@@ -13956,6 +13956,20 @@ ARMTargetLowering::LowerDYNAMIC_STACKALLOC(SDValue Op, SelectionDAG &DAG) const
   SDValue Chain = Op.getOperand(0);
   SDValue Size  = Op.getOperand(1);
 
+  if (DAG.getMachineFunction().getFunction().hasFnAttribute(
+          "no-stack-arg-probe")) {
+    unsigned Align = cast<ConstantSDNode>(Op.getOperand(2))->getZExtValue();
+    SDValue SP = DAG.getCopyFromReg(Chain, DL, ARM::SP, MVT::i32);
+    Chain = SP.getValue(1);
+    SP = DAG.getNode(ISD::SUB, DL, MVT::i32, SP, Size);
+    if (Align)
+      SP = DAG.getNode(ISD::AND, DL, MVT::i32, SP.getValue(0),
+                       DAG.getConstant(-(uint64_t)Align, DL, MVT::i32));
+    Chain = DAG.getCopyToReg(Chain, DL, ARM::SP, SP);
+    SDValue Ops[2] = { SP, Chain };
+    return DAG.getMergeValues(Ops, DL);
+  }
+
   SDValue Words = DAG.getNode(ISD::SRL, DL, MVT::i32, Size,
                               DAG.getConstant(2, DL, MVT::i32));
 
diff --git a/llvm/test/CodeGen/AArch64/no-stack-arg-probe.ll b/llvm/test/CodeGen/AArch64/no-stack-arg-probe.ll
new file mode 100644
index 00000000000..0a4097f4d38
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/no-stack-arg-probe.ll
@@ -0,0 +1,12 @@
+; RUN: llc -mtriple=aarch64-windows -verify-machineinstrs %s -o - \
+; RUN:  | FileCheck %s
+
+define void @check_watermark() "no-stack-arg-probe" {
+entry:
+  %buffer = alloca [4096 x i8], align 1
+  ret void
+}
+
+; CHECK: check_watermark:
+; CHECK: sub sp, sp, #1, lsl #12
+; CHECK-NOT: bl __chkstk
diff --git a/llvm/test/CodeGen/AArch64/win-alloca-no-stack-probe.ll b/llvm/test/CodeGen/AArch64/win-alloca-no-stack-probe.ll
new file mode 100644
index 00000000000..0ab161f8f27
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/win-alloca-no-stack-probe.ll
@@ -0,0 +1,17 @@
+; RUN: llc -mtriple aarch64-windows -verify-machineinstrs -filetype asm -o - %s | FileCheck %s
+
+define void @func(i64 %a) "no-stack-arg-probe" {
+entry:
+  %0 = alloca i8, i64 %a, align 16
+  call void @func2(i8* nonnull %0)
+  ret void
+}
+
+declare void @func2(i8*)
+
+; CHECK: add [[REG1:x[0-9]+]], x0, #15
+; CHECK-NOT: bl __chkstk
+; CHECK: mov [[REG2:x[0-9]+]], sp
+; CHECK: and [[REG1]], [[REG1]], #0xfffffffffffffff0
+; CHECK: sub [[REG3:x[0-9]+]], [[REG2]], [[REG1]]
+; CHECK: mov sp, [[REG3]]
diff --git a/llvm/test/CodeGen/ARM/Windows/alloca-no-stack-arg-probe.ll b/llvm/test/CodeGen/ARM/Windows/alloca-no-stack-arg-probe.ll
new file mode 100644
index 00000000000..a1dce7fd4a3
--- /dev/null
+++ b/llvm/test/CodeGen/ARM/Windows/alloca-no-stack-arg-probe.ll
@@ -0,0 +1,21 @@
+; RUN: llc -mtriple thumbv7-windows -filetype asm -o - %s | FileCheck %s
+
+declare arm_aapcs_vfpcc i32 @num_entries()
+
+define arm_aapcs_vfpcc void @test___builtin_alloca() "no-stack-arg-probe" {
+entry:
+  %array = alloca i8*, align 4
+  %call = call arm_aapcs_vfpcc i32 @num_entries()
+  %mul = mul i32 4, %call
+  %0 = alloca i8, i32 %mul
+  store i8* %0, i8** %array, align 4
+  ret void
+}
+
+; CHECK: bl num_entries
+; CHECK: movs [[R1:r[0-9]+]], #7
+; CHECK: add.w [[R0:r[0-9]+]], [[R1]], [[R0]], lsl #2
+; CHECK: bic [[R0]], [[R0]], #7
+; CHECK-NOT: bl __chkstk
+; CHECK: sub.w [[R0]], sp, [[R0]]
+; CHECK: mov sp, [[R0]]
diff --git a/llvm/test/CodeGen/ARM/Windows/chkstk-no-stack-arg-probe.ll b/llvm/test/CodeGen/ARM/Windows/chkstk-no-stack-arg-probe.ll
new file mode 100644
index 00000000000..a2a3bf0df88
--- /dev/null
+++ b/llvm/test/CodeGen/ARM/Windows/chkstk-no-stack-arg-probe.ll
@@ -0,0 +1,12 @@
+; RUN: llc -mtriple=thumbv7-windows -verify-machineinstrs %s -o - \
+; RUN:  | FileCheck %s
+
+define arm_aapcs_vfpcc void @check_watermark() "no-stack-arg-probe" {
+entry:
+  %buffer = alloca [4096 x i8], align 1
+  ret void
+}
+
+; CHECK: check_watermark:
+; CHECK-NOT: bl __chkstk
+; CHECK: sub.w sp, sp, #4096