[PATCH] Ensure that variable length allocas with large alignments gets proper stack probes

[Zoxc]

Summary

Previously allocas was aligned after allocation, which means that there were up to <alignment - 1> bytes potentially unchecked.

[mkuper]

Thanks for patch.

As we've discussed on IRC though - I'm not sure this works.
What guarantees the initial alignment of %rsp is high enough that subtracting %rax from it results in a properly-aligned variable?

lib/Target/X86/X86ISelLowering.cpp

 X86TargetLowering::LowerDYNAMIC_STACKALLOC(SDValue Op,
                                            SelectionDAG &DAG) const {
   MachineFunction &MF = DAG.getMachineFunction();
   bool SplitStack = MF.shouldSplitStack();
   bool Lower = (Subtarget->isOSWindows() && !Subtarget->isTargetMachO()) ||
                SplitStack || MF.shouldProbeStack();
   SDLoc dl(Op);
 
+  const TargetFrameLowering &TFI = *Subtarget->getFrameLowering();
+
   if (!Lower) {
     const TargetLowering &TLI = DAG.getTargetLoweringInfo();
     SDNode* Node = Op.getNode();
 
     unsigned SPReg = TLI.getStackPointerRegisterToSaveRestore();
     assert(SPReg && "Target cannot require DYNAMIC_STACKALLOC expansion and"
         " not tell us which reg is the stack pointer!");
     EVT VT = Node->getValueType(0);
     SDValue Tmp1 = SDValue(Node, 0);
     SDValue Tmp2 = SDValue(Node, 1);
     SDValue Tmp3 = Node->getOperand(2);
     SDValue Chain = Tmp1.getOperand(0);
 
     // Chain the dynamic stack allocation so that it doesn't modify the stack
     // pointer when other instructions are using the stack.
     Chain = DAG.getCALLSEQ_START(Chain, DAG.getIntPtrConstant(0, dl, true),
         SDLoc(Node));
 
     SDValue Size = Tmp2.getOperand(1);
     SDValue SP = DAG.getCopyFromReg(Chain, dl, SPReg, VT);
     Chain = SP.getValue(1);
     unsigned Align = cast<ConstantSDNode>(Tmp3)->getZExtValue();
-    const TargetFrameLowering &TFI = *Subtarget->getFrameLowering();
     unsigned StackAlign = TFI.getStackAlignment();
     Tmp1 = DAG.getNode(ISD::SUB, dl, VT, SP, Size); // Value
     if (Align > StackAlign)
       Tmp1 = DAG.getNode(ISD::AND, dl, VT, Tmp1,
           DAG.getConstant(-(uint64_t)Align, dl, VT));
     Chain = DAG.getCopyToReg(Chain, dl, SPReg, Tmp1); // Output chain
 
     Tmp2 = DAG.getCALLSEQ_END(Chain, DAG.getIntPtrConstant(0, dl, true),
         DAG.getIntPtrConstant(0, dl, true), SDValue(),
         SDLoc(Node));
 
     SDValue Ops[2] = { Tmp1, Tmp2 };
     return DAG.getMergeValues(Ops, dl);
   }
 
   // Get the inputs.
   SDValue Chain = Op.getOperand(0);
   SDValue Size  = Op.getOperand(1);
   unsigned Align = cast<ConstantSDNode>(Op.getOperand(2))->getZExtValue();
-  EVT VT = Op.getNode()->getValueType(0);
 
   bool Is64Bit = Subtarget->is64Bit();
   MVT SPTy = getPointerTy(DAG.getDataLayout());
 
   if (SplitStack) {
     MachineRegisterInfo &MRI = MF.getRegInfo();
 
     if (Is64Bit) {
if (SplitStack) {
     const TargetRegisterClass *AddrRegClass = getRegClassFor(SPTy);
     unsigned Vreg = MRI.createVirtualRegister(AddrRegClass);
     Chain = DAG.getCopyToReg(Chain, dl, Vreg, Size);
     SDValue Value = DAG.getNode(X86ISD::SEG_ALLOCA, dl, SPTy, Chain,
                                 DAG.getRegister(Vreg, SPTy));
     SDValue Ops1[2] = { Value, Chain };
     return DAG.getMergeValues(Ops1, dl);
   } else {
+    unsigned StackAlign = TFI.getStackAlignment();
     SDValue Flag;
+    if (Align > StackAlign) {
+      Size = DAG.getNode(ISD::ADD, dl, SPTy, Size, DAG.getIntPtrConstant(Align - 1, dl));
+      Size = DAG.getNode(ISD::AND, dl, SPTy, Size,
+                         DAG.getIntPtrConstant(~(uint64_t)(Align - 1), dl));
+    }
+
     const unsigned Reg = (Subtarget->isTarget64BitLP64() ? X86::RAX : X86::EAX);
 
     Chain = DAG.getCopyToReg(Chain, dl, Reg, Size, Flag);
     Flag = Chain.getValue(1);
 
     if (!Subtarget->isOSWindows()) {
         const unsigned Reg = (Subtarget->isTarget64BitLP64() ? X86::RBX : X86::EBX);
 
         Chain = DAG.getCopyToReg(Chain, dl, Reg, DAG.getConstant(0x1000, dl, SPTy), Flag);
         Flag = Chain.getValue(1);
     }
 
     SDVTList NodeTys = DAG.getVTList(MVT::Other, MVT::Glue);
 
     Chain = DAG.getNode(X86ISD::WIN_ALLOCA, dl, NodeTys, Chain, Flag);
 
     const X86RegisterInfo *RegInfo = Subtarget->getRegisterInfo();
     unsigned SPReg = RegInfo->getStackRegister();
     SDValue SP = DAG.getCopyFromReg(Chain, dl, SPReg, SPTy);
     Chain = SP.getValue(1);
 
-    if (Align) {
-      SP = DAG.getNode(ISD::AND, dl, VT, SP.getValue(0),
-                       DAG.getConstant(-(uint64_t)Align, dl, VT));
-      Chain = DAG.getCopyToReg(Chain, dl, SPReg, SP);
-    }
-
     SDValue Ops1[2] = { SP, Chain };
     return DAG.getMergeValues(Ops1, dl);
   }
 }
 
 SDValue X86TargetLowering::LowerVASTART(SDValue Op, SelectionDAG &DAG) const {
   MachineFunction &MF = DAG.getMachineFunction();
   auto PtrVT = getPointerTy(MF.getDataLayout());

test/CodeGen/X86/win64_alloca_dynalloca.ll

 define i64 @aligned(i64 %n, i64 %x) nounwind {
 ; M64-LABEL: aligned:
 ; W64-LABEL: aligned:
 ; EFI-LABEL: aligned:
 entry:
 
   %buf1 = alloca i8, i64 %n, align 128
 
-; M64: leaq  15(%{{.*}}), %rax
-; M64: andq  $-16, %rax
+; M64: addq  $127, %rax
+; M64: andq  $-128, %rax
 ; M64: callq ___chkstk_ms
 ; M64: subq  %rax, %rsp
 ; M64: movq  %rsp, [[R2:%r.*]]
-; M64: andq  $-128, [[R2]]
-; M64: movq  [[R2]], %rsp
 
-; W64: leaq  15(%{{.*}}), %rax
-; W64: andq  $-16, %rax
+; W64: addq  $127, %rax
+; W64: andq  $-128, %rax
 ; W64: callq __chkstk
 ; W64: subq  %rax, %rsp
 ; W64: movq  %rsp, [[R2:%r.*]]
-; W64: andq  $-128, [[R2]]
-; W64: movq  [[R2]], %rsp
 
 ; EFI: leaq  15(%{{.*}}), [[R1:%r.*]]
 ; EFI: andq  $-16, [[R1]]
 ; EFI: movq  %rsp, [[R64:%r.*]]
 ; EFI: subq  [[R1]], [[R64]]
 ; EFI: andq  $-128, [[R64]]
 ; EFI: movq  [[R64]], %rsp