[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 ↗

(On Diff #130821)

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 ↗

(On Diff #130821)

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 ↗

(On Diff #130821)

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 ↗

(On Diff #130821)

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 ↗

(On Diff #130821)

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 ↗

(On Diff #130821)

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 ↗

(On Diff #130821)

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 ↗

(On Diff #133838)

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

llvm/trunk/lib/CodeGen/GlobalISel/IRTranslator.cpp

bool IRTranslator::translateAlloca(const User &U,
 
   if (AI.isStaticAlloca()) {
     unsigned Res = getOrCreateVReg(AI);
     int FI = getOrCreateFrameIndex(AI);
     MIRBuilder.buildFrameIndex(Res, FI);
     return true;
   }
 
+  // FIXME: support stack probing for Windows.
+  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());

llvm/trunk/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))>;

llvm/trunk/lib/Target/AArch64/AArch64ISelLowering.h

private:
   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 LowerATOMIC_LOAD_SUB(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerATOMIC_LOAD_AND(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;

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

AArch64TargetLowering::AArch64TargetLowering(const TargetMachine &TM,
   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);
+
+  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.
SDValue AArch64TargetLowering::LowerOperation(SDValue Op,
   case ISD::VECREDUCE_UMIN:
   case ISD::VECREDUCE_FMAX:
   case ISD::VECREDUCE_FMIN:
     return LowerVECREDUCE(Op, DAG);
   case ISD::ATOMIC_LOAD_SUB:
     return LowerATOMIC_LOAD_SUB(Op, DAG);
   case ISD::ATOMIC_LOAD_AND:
     return LowerATOMIC_LOAD_AND(Op, DAG);
+  case ISD::DYNAMIC_STACKALLOC:
+    return LowerDYNAMIC_STACKALLOC(Op, DAG);
   }
 }
 
 //===----------------------------------------------------------------------===//
 //                      Calling Convention Implementation
 //===----------------------------------------------------------------------===//
 
 #include "AArch64GenCallingConv.inc"
SDValue AArch64TargetLowering::LowerATOMIC_LOAD_AND(SDValue Op,
   SDValue RHS = Op.getOperand(2);
   AtomicSDNode *AN = cast<AtomicSDNode>(Op.getNode());
   RHS = DAG.getNode(ISD::XOR, dl, VT, DAG.getConstant(-1ULL, dl, VT), RHS);
   return DAG.getAtomic(ISD::ATOMIC_LOAD_CLR, dl, AN->getMemoryVT(),
                        Op.getOperand(0), Op.getOperand(1), RHS,
                        AN->getMemOperand());
 }
 
+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();

llvm/trunk/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;
 

llvm/trunk/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);

llvm/trunk/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

llvm/trunk/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     [[REG:x[0-9]+]], x9, x8
+; CHECK: lsr     x15, x8, #4
 ; CHECK: mov     x19, x1
 ; CHECK: mov     [[REG2:x[0-9]+]], 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     [[REG:x[0-9]+]], x8, x15, lsl #4
 ; CHECK: mov     sp, [[REG]]
 ; CHECK: ldur    [[REG3:x[0-9]+]], [x29, #-40]
 ; CHECK: sxtw    [[REG4:x[0-9]+]], w0
 ; CHECK: bl      __local_stdio_printf_options
 ; CHECK: ldr     x8, [x0]
 ; CHECK: mov     x1, [[REG]]
 ; CHECK: mov     x2, [[REG4]]
 ; CHECK: mov     x3, x19