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

[mstorsjo]

Summary

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 both fixed and dynamic 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.

The manual SelectionDAG code for subtracting from the SP register feels very clunky; is there any way to make this fall back to the default behaviour of Expand in this particular case, instead of handcoding it?

[mstorsjo]

Removed a leftover out of place comment in one of the tests.

[hans]

Sorry for not being responsive here. I think someone more familiar with the ARM target needs to take a look.

[mstorsjo]

Ping @aemerson

[aemerson]

lib/Target/AArch64/AArch64ISelLowering.cpp

7543

This is unfortunate, ideally we could fall back to Expand here.

7550

What about the alignment?

lib/Target/ARM/ARMISelLowering.cpp

13959

Alignment again?

[mstorsjo]

lib/Target/AArch64/AArch64ISelLowering.cpp

7543

Yes, I'd like to do that - do you have any suggestions on how to?

7550

Right, will fix.

[aemerson]

lib/Target/AArch64/AArch64ISelLowering.cpp

7543

I don't, unless we could make this a codegen option.

[mstorsjo]

Take alignment into account

lib/Target/AArch64/AArch64ISelLowering.cpp

   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);
 
+  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);
 
   SDValue SP = DAG.getCopyFromReg(Chain, dl, AArch64::SP, MVT::i64);
   Chain = SP.getValue(1);
   SP = DAG.getNode(ISD::SUB, dl, MVT::i64, SP, Size);
   Chain = DAG.getCopyToReg(Chain, dl, AArch64::SP, SP);
 
   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);
   }
 
   Chain = DAG.getCALLSEQ_END(Chain, DAG.getIntPtrConstant(0, dl, true),
                              DAG.getIntPtrConstant(0, dl, true), SDValue(), dl);
 

lib/Target/ARM/ARMISelLowering.cpp

 ARMTargetLowering::LowerDYNAMIC_STACKALLOC(SDValue Op, SelectionDAG &DAG) const {
   assert(Subtarget->isTargetWindows() && "unsupported target platform");
   SDLoc DL(Op);
 
   // Get the inputs.
   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));
 
   SDValue Flag;
   Chain = DAG.getCopyToReg(Chain, DL, ARM::R4, Words, Flag);
   Flag = Chain.getValue(1);
 
   SDVTList NodeTys = DAG.getVTList(MVT::Other, MVT::Glue);

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

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

test/CodeGen/AArch64/win-alloca-no-stack-probe.ll

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

test/CodeGen/ARM/Windows/alloca-no-stack-arg-probe.ll

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

test/CodeGen/ARM/Windows/chkstk-no-stack-arg-probe.ll

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