[AArch64] Implement dynamic stack probing for windows

[mstorsjo]

Summary

This makes sure that alloca() function calls properly probe the stack as needed.

This is partially based on D40863 by @aemerson.

[aemerson]

For GlobalISel we haven't had to have very target specific code in the IRTranslator for this before. @qcolombet any opinion on whether we want to include target/ABI specific code into the IRTranslator? My feeling is that we want to avoid this, the alternative I can see is to introduce a new opcode for G_DYNALLOCA and lower that at isel.

lib/Target/AArch64/AArch64ISelLowering.cpp

7452

Unfortunately it seems we need to pass a const char* to this because of lifetime issues at codegen. Do you have any ideas on how to do this better?

7458

Can you separate this out into some kind of helper for the Windows specific ABI?

[mstorsjo]

In D42356#986541, @aemerson wrote:

For GlobalISel we haven't had to have very target specific code in the IRTranslator for this before. @qcolombet any opinion on whether we want to include target/ABI specific code into the IRTranslator? My feeling is that we want to avoid this, the alternative I can see is to introduce a new opcode for G_DYNALLOCA and lower that at isel.

Using a new opcode that maps down to this case here probably sounds like the best option.

lib/Target/AArch64/AArch64ISelLowering.cpp

7452

I don't think I'm following what the issue is here? Right now I'm passing "__chkstk" as a const char* here. Do you mean further, if refactoring this into a form that also suits darwin, with a different function name? In that case, the method providing the function name could probably just return const char* as well, instead of e.g. a StringRef.

7458

I can try to make it easier fit in the darwin version from D40863, by more clearly separating out the windows specific bits.

test/CodeGen/AArch64/win64_vararg.ll

171

This form produces the pattern of "mov x8, sp; sub x20, x8, x15, lsl #4; mov sp x20", when it could just as well do "sub sp, sp, x15, lsl #4; mov x20, sp" - however I didn't manage to make the code in AArch64ISelLowering do that - any hints?

The version I did in D41131 for the fixed stack parts use BuildMI(MBB, MBBI, DL, TII->get(AArch64::SUBXrx64), AArch64::SP).addReg(AArch64::SP, RegState::Kill).addReg(AArch64::X15, RegState::Kill).addImm(AArch64_AM::getArithExtendImm(AArch64_AM::UXTX, 4)) - which achieves exactly what I want.

[aemerson]

lib/Target/AArch64/AArch64ISelLowering.cpp

7452

Sorry I didn't word that right. I meant that when I tried to use the getStackProbingSymbol hook from TargetLowering it returns a StringRef, so I had to hard code the string in this function as a short term solution. It's ok if we only need it for one platform, but it would be nice to have it configurable by the specific platform.

test/CodeGen/AArch64/win64_vararg.ll

171

I don't know of any way to do that without using getCopyFromReg, not to say there definitely isn't a better way.

[mstorsjo]

Structured the code to easier allow adding a darwin specific version of stack probing without conflicting with the windows specific parts here. Added a "return false" in GlobalISel for windows for alloca, to make sure we invoke the proper codegen for this.

On ARM and X86, the implementation of stack probing uses a custom opcode like WIN__CHKSTK - that would allow to more directly emit exactly the intended machine code for this call. Doing it via SelectionDAG is less code but ends up less elegant at -O0 - a compromise I think is ok.

[mstorsjo]

Ping @aemerson

[aemerson]

Thanks for refactoring it. LGTM with a minor comment addition.

lib/CodeGen/GlobalISel/IRTranslator.cpp

1007

Please add a comment here explaining why Windows isn't supported.

lib/CodeGen/GlobalISel/IRTranslator.cpp

 
   if (AI.isStaticAlloca()) {
     unsigned Res = getOrCreateVReg(AI);
     int FI = getOrCreateFrameIndex(AI);
     MIRBuilder.buildFrameIndex(Res, FI);
     return true;
   }
 
+  if (MF->getTarget().getTargetTriple().isOSWindows())
+    return false;
+
   // Now we're in the harder dynamic case.
   Type *Ty = AI.getAllocatedType();
   unsigned Align =
       std::max((unsigned)DL->getPrefTypeAlignment(Ty), AI.getAlignment());
 
   unsigned NumElts = getOrCreateVReg(*AI.getArraySize());
 
   Type *IntPtrIRTy = DL->getIntPtrType(AI.getType());

lib/Target/AArch64/AArch64CallingConvention.td

: CalleeSavedRegs<(add (sequence "W%u", 0, 30), WSP,
                            (sequence "S%u", 0, 31), (sequence "D%u", 0, 31),
                            (sequence "Q%u", 0, 31))>;
 
 def CSR_AArch64_NoRegs : CalleeSavedRegs<(add)>;
 
 def CSR_AArch64_RT_MostRegs :  CalleeSavedRegs<(add CSR_AArch64_AAPCS,
                                                 (sequence "X%u", 9, 15))>;
 
+def CSR_AArch64_StackProbe_Windows
+    : CalleeSavedRegs<(add (sequence "X%u", 0, 15),
+                           (sequence "X%u", 18, 28), FP, SP,
+                           (sequence "Q%u", 0, 31))>;

lib/Target/AArch64/AArch64ISelLowering.h

private:
   SDValue LowerFP_ROUND(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerFP_TO_INT(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerINT_TO_FP(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerVectorAND(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerVectorOR(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerCONCAT_VECTORS(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerFSINCOS(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerVECREDUCE(SDValue Op, SelectionDAG &DAG) const;
+  SDValue LowerDYNAMIC_STACKALLOC(SDValue Op, SelectionDAG &DAG) const;
+  SDValue LowerWindowsDYNAMIC_STACKALLOC(SDValue Op, SDValue Chain,
+                                         SDValue &Size,
+                                         SelectionDAG &DAG) const;
 
   SDValue BuildSDIVPow2(SDNode *N, const APInt &Divisor, SelectionDAG &DAG,
                         std::vector<SDNode *> *Created) const override;
   SDValue getSqrtEstimate(SDValue Operand, SelectionDAG &DAG, int Enabled,
                           int &ExtraSteps, bool &UseOneConst,
                           bool Reciprocal) const override;
   SDValue getRecipEstimate(SDValue Operand, SelectionDAG &DAG, int Enabled,
                            int &ExtraSteps) const override;

lib/Target/AArch64/AArch64ISelLowering.cpp

   setOperationAction(ISD::VASTART, MVT::Other, Custom);
   setOperationAction(ISD::VAARG, MVT::Other, Custom);
   setOperationAction(ISD::VACOPY, MVT::Other, Custom);
   setOperationAction(ISD::VAEND, MVT::Other, Expand);
 
   // Variable-sized objects.
   setOperationAction(ISD::STACKSAVE, MVT::Other, Expand);
   setOperationAction(ISD::STACKRESTORE, MVT::Other, Expand);
-  setOperationAction(ISD::DYNAMIC_STACKALLOC, MVT::i64, Expand);
+
+  if (Subtarget->isTargetWindows())
+    setOperationAction(ISD::DYNAMIC_STACKALLOC, MVT::i64, Custom);
+  else
+    setOperationAction(ISD::DYNAMIC_STACKALLOC, MVT::i64, Expand);
 
   // Constant pool entries
   setOperationAction(ISD::ConstantPool, MVT::i64, Custom);
 
   // BlockAddress
   setOperationAction(ISD::BlockAddress, MVT::i64, Custom);
 
   // Add/Sub overflow ops with MVT::Glues are lowered to NZCV dependences.

   case ISD::VECREDUCE_ADD:
   case ISD::VECREDUCE_SMAX:
   case ISD::VECREDUCE_SMIN:
   case ISD::VECREDUCE_UMAX:
   case ISD::VECREDUCE_UMIN:
   case ISD::VECREDUCE_FMAX:
   case ISD::VECREDUCE_FMIN:
     return LowerVECREDUCE(Op, DAG);
+  case ISD::DYNAMIC_STACKALLOC:
+    return LowerDYNAMIC_STACKALLOC(Op, DAG);
   }
 }
 
 //===----------------------------------------------------------------------===//
 //                      Calling Convention Implementation
 //===----------------------------------------------------------------------===//
 
 #include "AArch64GenCallingConv.inc"

         DAG.getConstant(Intrinsic::aarch64_neon_fminnmv, dl, MVT::i32),
         Op.getOperand(0));
   }
   default:
     llvm_unreachable("Unhandled reduction");
   }
 }
 
+SDValue AArch64TargetLowering::LowerWindowsDYNAMIC_STACKALLOC(
+    SDValue Op, SDValue Chain, SDValue &Size, SelectionDAG &DAG) const {
+  SDLoc dl(Op);
+  EVT PtrVT = getPointerTy(DAG.getDataLayout());
+  SDValue Callee = DAG.getTargetExternalSymbol("__chkstk", PtrVT, 0);
+
+  const uint32_t *Mask =
+      Subtarget->getRegisterInfo()->getWindowsStackProbePreservedMask();
+
+  Size = DAG.getNode(ISD::SRL, dl, MVT::i64, Size,
+                     DAG.getConstant(4, dl, MVT::i64));
+  Chain = DAG.getCopyToReg(Chain, dl, AArch64::X15, Size, SDValue());
+  Chain =
+      DAG.getNode(AArch64ISD::CALL, dl, DAG.getVTList(MVT::Other, MVT::Glue),
+                  Chain, Callee, DAG.getRegister(AArch64::X15, MVT::i64),
+                  DAG.getRegisterMask(Mask), Chain.getValue(1));
+  // To match the actual intent better, we should read the output from X15 here
+  // again (instead of potentially spilling it to the stack), but rereading Size
+  // from X15 here doesn't work at -O0, since it thinks that X15 is undefined
+  // here.
+
+  Size = DAG.getNode(ISD::SHL, dl, MVT::i64, Size,
+                     DAG.getConstant(4, dl, MVT::i64));
+  return Chain;
+}
+
+SDValue
+AArch64TargetLowering::LowerDYNAMIC_STACKALLOC(SDValue Op,
+                                               SelectionDAG &DAG) const {
+  assert(Subtarget->isTargetWindows() &&
+         "Only Windows alloca probing supported");
+  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);
+
+  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);
+
+  SDValue Ops[2] = {SP, Chain};
+  return DAG.getMergeValues(Ops, dl);
+}
+
 /// getTgtMemIntrinsic - Represent NEON load and store intrinsics as
 /// MemIntrinsicNodes.  The associated MachineMemOperands record the alignment
 /// specified in the intrinsic calls.
 bool AArch64TargetLowering::getTgtMemIntrinsic(IntrinsicInfo &Info,
                                                const CallInst &I,
                                                MachineFunction &MF,
                                                unsigned Intrinsic) const {
   auto &DL = I.getModule()->getDataLayout();

lib/Target/AArch64/AArch64RegisterInfo.h

public:
   /// (i.e. it is a calling convention that uses the same register for the first
   /// i64 argument and an i64 return value)
   ///
   /// Should return NULL in the case that the calling convention does not have
   /// this property
   const uint32_t *getThisReturnPreservedMask(const MachineFunction &MF,
                                              CallingConv::ID) const;
 
+  /// Stack probing calls preserve different CSRs to the normal CC.
+  const uint32_t *getWindowsStackProbePreservedMask() const;
+
   BitVector getReservedRegs(const MachineFunction &MF) const override;
   bool isConstantPhysReg(unsigned PhysReg) const override;
   const TargetRegisterClass *
   getPointerRegClass(const MachineFunction &MF,
                      unsigned Kind = 0) const override;
   const TargetRegisterClass *
   getCrossCopyRegClass(const TargetRegisterClass *RC) const override;
 

lib/Target/AArch64/AArch64RegisterInfo.cpp

AArch64RegisterInfo::getThisReturnPreservedMask(const MachineFunction &MF,
   // single i64 return value)
   //
   // In case that the calling convention does not use the same register for
   // both, the function should return NULL (does not currently apply)
   assert(CC != CallingConv::GHC && "should not be GHC calling convention.");
   return CSR_AArch64_AAPCS_ThisReturn_RegMask;
 }
 
+const uint32_t *AArch64RegisterInfo::getWindowsStackProbePreservedMask() const {
+  return CSR_AArch64_StackProbe_Windows_RegMask;
+}
+
 BitVector
 AArch64RegisterInfo::getReservedRegs(const MachineFunction &MF) const {
   const AArch64FrameLowering *TFI = getFrameLowering(MF);
 
   // FIXME: avoid re-calculating this every time.
   BitVector Reserved(getNumRegs());
   markSuperRegs(Reserved, AArch64::WSP);
   markSuperRegs(Reserved, AArch64::WZR);

test/CodeGen/AArch64/win-alloca.ll

+; RUN: llc -mtriple aarch64-windows -verify-machineinstrs -filetype asm -o - %s | FileCheck %s -check-prefixes=CHECK,CHECK-OPT
+; RUN: llc -mtriple aarch64-windows -verify-machineinstrs -filetype asm -o - %s -O0 | FileCheck %s
+
+define void @func(i64 %a) {
+entry:
+  %0 = alloca i8, i64 %a, align 16
+  call void @func2(i8* nonnull %0)
+  ret void
+}
+
+declare void @func2(i8*)
+
+; The -O0 version here ends up much less elegant, so just check the
+; details of the optimized form, but check that -O0 at least emits the
+; call to __chkstk.
+
+; CHECK: add [[REG1:x[0-9]+]], x0, #15
+; CHECK-OPT: lsr x15, [[REG1]], #4
+; CHECK: bl __chkstk
+; CHECK: mov [[REG2:x[0-9]+]], sp
+; CHECK-OPT: sub [[REG3:x[0-9]+]], [[REG2]], x15, lsl #4
+; CHECK-OPT: mov sp, [[REG3:x[0-9]+]]
+; CHECK: bl func2

test/CodeGen/AArch64/win64_vararg.ll

 ; CHECK: stp     x22, x21, [sp, #16]
 ; CHECK: stp     x20, x19, [sp, #32]
 ; CHECK: stp     x29, x30, [sp, #48]
 ; CHECK: add     x29, sp, #48
 ; CHECK: add     x8, x29, #16
 ; CHECK: stur    x8, [x29, #-40]
 ; CHECK: mov     w8, w0
 ; CHECK: add     x8, x8, #15
-; CHECK: mov     x9, sp
-; CHECK: and     x8, x8, #0x1fffffff0
-; CHECK: sub     x20, x9, x8
+; CHECK: lsr     x15, x8, #4
 ; CHECK: mov     x19, x1
 ; CHECK: mov     x23, sp
 ; CHECK: stp     x6, x7, [x29, #48]
 ; CHECK: stp     x4, x5, [x29, #32]
 ; CHECK: stp     x2, x3, [x29, #16]
+; CHECK: bl      __chkstk
+; CHECK: mov     x8, sp
+; CHECK: sub     x20, x8, x15, lsl #4
 ; CHECK: mov     sp, x20
 ; CHECK: ldur    x21, [x29, #-40]
 ; CHECK: sxtw    x22, w0
 ; CHECK: bl      __local_stdio_printf_options
 ; CHECK: ldr     x8, [x0]
 ; CHECK: mov     x1, x20
 ; CHECK: mov     x2, x22
 ; CHECK: mov     x3, x19
 ; CHECK: orr     x0, x8, #0x2