[Codegen] Ensure stack is properly aligned for call argument initialization

[jketema]

Summary

Arguments spilled on the stack before a function call may have
alignment requirements, for example in the case of vectors.
These requirements are exploited by the code generator by using
move instructions that have similar alignment requirements, e.g.,
movaps on x86.

Although the code generator properly aligns the arguments with
respect to the displacement of the stack pointer it computes,
the displacement itself may cause misalignment. For example if
we have

%3 = load <16 x float>, <16 x float>* %1, align 64
call void @bar(<16 x float> %3, i32 0)

The x86 back-end emits:

movaps  32(%ecx), %xmm2
movaps  (%ecx), %xmm0
movaps  16(%ecx), %xmm1
movaps  48(%ecx), %xmm3
subl    $20, %esp       <-- if %esp was 16-byte aligned before this instruction, it no longer will be afterwards 
movaps  %xmm3, (%esp)   <-- movaps requires 16-byte alignment, while %esp is not aligned as such.
movl    $0, 16(%esp)
calll   __bar

To solve this, we need to make sure that the computed value with which
the stack pointer is changed is a multiple af the maximal alignment seen
during its computation. With this change we get proper alignment:

subl    $32, %esp
movaps  %xmm3, (%esp)

[rnk]

include/llvm/CodeGen/CallingConvLower.h

204

I think this is really more like MaxStackArgAlign. It's the maximum alignment of arguments that ended up on the stack, right?

273

We should probably put some Doxygen here about what NextStackOffset really means.

274–276

Are you sure this is the right place to change? Isn't this method used to figure out where variadic argument packs start also? Shouldn't we *not* align StackOffset for that use case?

I think it also gets used when we have a reserved stack frame (i.e. no dynamic alloca). As written, your code will overallocate some extra padding that isn't necessary.

275

Use RoundUpToAlignment.

406

RoundUpToAlignment here would be a nice cleanup.

409

This looks like std::max(Align, MinStackAlign) to me.

[jketema]

Thanks for your feedback! I'll address the easy bits, and have another look at the use of getNextStackOffset. I might need some help with that though, as I'm not very familiar with the code generation parts of llvm (still learning).

include/llvm/CodeGen/CallingConvLower.h

204

It was meant to be the minimum stack alignment I'll eventually need. I like your suggestion though.

274–276

I'll have another look at the use cases.

[jketema]

Addressed easy comments. getNextStackOffset concerns still to investigate (including adding Doxygen comment)

[jketema]

include/llvm/CodeGen/CallingConvLower.h

274–276

In the case of variadic arguments: can't we run into the same problems? If one of the arguments in the pack is a <4 x float> vector, then wouldn't there be cases where the X86 code generator uses and an movaps instruction to load the vector into one of the XMM registers? Or will it always used movups in that case?

And similar for reserved stack frames?

[jketema]

Add Doxygen comment

[jketema]

Hi Reid,

Could you have another look at this? After having had a renewed look at the code, getNextStackOffset still seems the correct place for this change to me for the reasons I provided in the inline comments. However, if you think this really belongs somewhere else, would you be able to point me in the right direction?

Thanks,

Jeroen

[jketema]

Bump

[rnk]

include/llvm/CodeGen/CallingConvLower.h

274–277

Yes, but here's the 32-bit test case I'm imagining:

void f(__m128 v, const char *f, ...) {
  va_list ap;
  va_start(ap, f);
  int d = va_arg(ap, int);
}

Do we have padding between 'f' and the variadic int parameter? I think the answer should be know, because if they were not variadic, they would not have padding between them.

[jketema]

bump

[jketema]

include/llvm/CodeGen/CallingConvLower.h

274–277

Sorry, I missed this before I bumped. Somehow the system didn't sent me a notification email.

I don't see why this is a problem, my patch only changes the alignment of the whole frame, not the padding between the individual members; the computation for that stays unchanged. Hence, there will only be additional padding after the very last of the variadic parameters supplied.

[jketema]

include/llvm/CodeGen/CallingConvLower.h

274–277

Reading back all the comments, I think I see what you're getting at: What you're saying is that getNextStackOffset is also used to compute the offset of the variadic arguments with respect to the beginning of the stack frame?

If that's the case, then the patch would indeed be wrong, and I think we would then need both the old variant and the new one I'm proposing and adapt every use accordingly. Or do you see a better solution?

[rnk]

include/llvm/CodeGen/CallingConvLower.h

274–277

Exactly, getNextStackOffset currently has more than one use. You should take the example C code above and turn it into a test case for lit to show that at least this use isn't changing.

Given that there is more than one use, the I think we should introduce a new method called getAlignedCallFrameSize, update the appropriate call sites that use this in PEI, and declare victory. :)

[jketema]

Hi Reid,

I've introduced a getAlignedCallFrameSize method as suggested and changed call lowering in the X86 target appropriately. I've also introduced a test involving var args, as you suggested.

The X86 target also contains a call to getNextStackOffset in the code that determines whether tail call optimization can be performed. I'm not sure if that call should also be changed to getAlignedCallFrameSize. I'm also wondering whether getAlignedCallFrameSize should be introduced in other targets too.

[jketema]

Bump

[rnk]

lgtm

test/CodeGen/X86/win32-spill-xmm.ll

9

; CHECK: calll bar

[jketema]

Thanks for all your feedback Reid!

include/llvm/CodeGen/CallingConvLower.h

private:
   CallingConv::ID CallingConv;
   bool IsVarArg;
   MachineFunction &MF;
   const TargetRegisterInfo &TRI;
   SmallVectorImpl<CCValAssign> &Locs;
   LLVMContext &Context;
 
   unsigned StackOffset;
+  unsigned MaxStackArgAlign;
   SmallVector<uint32_t, 16> UsedRegs;
   SmallVector<CCValAssign, 4> PendingLocs;
 
   // ByValInfo and SmallVector<ByValInfo, 4> ByValRegs:
   //
   // Vector of ByValInfo instances (ByValRegs) is introduced for byval registers
   // tracking.
   // Or, in another words it tracks byval parameters that are stored in
public:
   void addLoc(const CCValAssign &V) {
     Locs.push_back(V);
   }
 
   LLVMContext &getContext() const { return Context; }
   MachineFunction &getMachineFunction() const { return MF; }
   CallingConv::ID getCallingConv() const { return CallingConv; }
   bool isVarArg() const { return IsVarArg; }
 
-  unsigned getNextStackOffset() const { return StackOffset; }
+  unsigned getNextStackOffset() const {
+    return RoundUpToAlignment(StackOffset, MaxStackArgAlign);
+  }
 
   /// isAllocated - Return true if the specified register (or an alias) is
   /// allocated.
   bool isAllocated(unsigned Reg) const {
     return UsedRegs[Reg/32] & (1 << (Reg&31));
   }
 
   /// AnalyzeFormalArguments - Analyze an array of argument values,
   /// incorporating info about the formals into this state.
unsigned AllocateReg(ArrayRef<MCPhysReg> Regs, const MCPhysReg *ShadowRegs) {
     MarkAllocated(ShadowReg);
     return Reg;
   }
 
   /// AllocateStack - Allocate a chunk of stack space with the specified size
   /// and alignment.
   unsigned AllocateStack(unsigned Size, unsigned Align) {
     assert(Align && ((Align - 1) & Align) == 0); // Align is power of 2.
-    StackOffset = ((StackOffset + Align - 1) & ~(Align - 1));
+    StackOffset = RoundUpToAlignment(StackOffset, Align);
     unsigned Result = StackOffset;
     StackOffset += Size;
+    MaxStackArgAlign = std::max(Align, MaxStackArgAlign);
     MF.getFrameInfo()->ensureMaxAlignment(Align);
     return Result;
   }
 
   /// Version of AllocateStack with extra register to be shadowed.
   unsigned AllocateStack(unsigned Size, unsigned Align, unsigned ShadowReg) {
     MarkAllocated(ShadowReg);
     return AllocateStack(Size, Align);

lib/CodeGen/CallingConvLower.cpp

 
 CCState::CCState(CallingConv::ID CC, bool isVarArg, MachineFunction &mf,
                  SmallVectorImpl<CCValAssign> &locs, LLVMContext &C)
     : CallingConv(CC), IsVarArg(isVarArg), MF(mf),
       TRI(*MF.getSubtarget().getRegisterInfo()), Locs(locs), Context(C),
       CallOrPrologue(Unknown) {
   // No stack is used.
   StackOffset = 0;
+  MaxStackArgAlign = 1;
 
   clearByValRegsInfo();
   UsedRegs.resize((TRI.getNumRegs()+31)/32);
 }
 
 /// Allocate space on the stack large enough to pass an argument by value.
 /// The size and alignment information of the argument is encoded in
 /// its parameter attribute.
static bool isValueTypeInRegForCC(CallingConv::ID CC, MVT VT) {
   if (CC == CallingConv::X86_VectorCall || CC == CallingConv::X86_FastCall)
     return true;
   return false;
 }
 
 void CCState::getRemainingRegParmsForType(SmallVectorImpl<MCPhysReg> &Regs,
                                           MVT VT, CCAssignFn Fn) {
   unsigned SavedStackOffset = StackOffset;
+  unsigned SavedMaxStackArgAlign = MaxStackArgAlign;
   unsigned NumLocs = Locs.size();
 
   // Set the 'inreg' flag if it is used for this calling convention.
   ISD::ArgFlagsTy Flags;
   if (isValueTypeInRegForCC(CallingConv, VT))
     Flags.setInReg();
 
   // Allocate something of this value type repeatedly until we get assigned a
#endif
   for (unsigned I = NumLocs, E = Locs.size(); I != E; ++I)
     if (Locs[I].isRegLoc())
       Regs.push_back(MCPhysReg(Locs[I].getLocReg()));
 
   // Clear the assigned values and stack memory. We leave the registers marked
   // as allocated so that future queries don't return the same registers, i.e.
   // when i64 and f64 are both passed in GPRs.
   StackOffset = SavedStackOffset;
+  MaxStackArgAlign = SavedMaxStackArgAlign;
   Locs.resize(NumLocs);
 }
 
 void CCState::analyzeMustTailForwardedRegisters(
     SmallVectorImpl<ForwardedRegister> &Forwards, ArrayRef<MVT> RegParmTypes,
     CCAssignFn Fn) {
   // Oftentimes calling conventions will not user register parameters for
   // variadic functions, so we need to assume we're not variadic so that we get

test/CodeGen/X86/aligned-variadic.ll

entry:
   %va = alloca [1 x %struct.__va_list_tag], align 16
   %arraydecay = getelementptr inbounds [1 x %struct.__va_list_tag], [1 x %struct.__va_list_tag]* %va, i64 0, i64 0
   %arraydecay1 = bitcast [1 x %struct.__va_list_tag]* %va to i8*
   call void @llvm.va_start(i8* %arraydecay1)
   %overflow_arg_area_p = getelementptr inbounds [1 x %struct.__va_list_tag], [1 x %struct.__va_list_tag]* %va, i64 0, i64 0, i32 2
   %overflow_arg_area = load i8*, i8** %overflow_arg_area_p, align 8
   %overflow_arg_area.next = getelementptr i8, i8* %overflow_arg_area, i64 24
   store i8* %overflow_arg_area.next, i8** %overflow_arg_area_p, align 8
-; X32: leal    68(%esp), [[REG:%.*]]
+; X32: leal    72(%esp), [[REG:%.*]]
 ; X32: movl    [[REG]], 16(%esp)
 ; X64: leaq    232(%rsp), [[REG:%.*]]
 ; X64: movq    [[REG]], 184(%rsp)
 ; X64: leaq    176(%rsp), %rdi
   call void @qux(%struct.__va_list_tag* %arraydecay)
   ret void
 }
 
 ; Function Attrs: nounwind
 declare void @llvm.va_start(i8*)
 
 declare void @qux(%struct.__va_list_tag*)

test/CodeGen/X86/win32-spill-xmm.ll

+; RUN: llc -mcpu=generic -mtriple=i686-pc-windows-msvc -mattr=+sse < %s | FileCheck %s
+; CHECK: subl    $32, %esp
+; CHECK: movaps  %xmm3, (%esp)
+; CHECK: movl    $0, 16(%esp)
+
+declare void @bar(<16 x float> %a, i32 %b) nounwind
+
+define void @foo(i32, <16 x float> * nocapture readonly) nounwind {
+entry:
+  %2 = alloca i32, i32 %0
+  %3 = load <16 x float>, <16 x float> * %1, align 64
+  tail call void @bar(<16 x float> %3, i32 0) nounwind
+  ret void
+}