[AArch64] Rewrite stack frame handling for win64 vararg functions

[mstorsjo]

Summary

The previous attempt, which made do with a single offset in computeCalleeSaveRegisterPairs, wasn't quite enough. The previous attempt only worked as long as CombineSPBump == true (since the offset would be adjusted later in fixupCalleeSaveRestoreStackOffset).

Instead include the size for the fixed stack area used for win64 varargs in calculations in emitPrologue/emitEpilogue. The stack consists of mainly three parts;

• AFI->getLocalStackSize()
• AFI->getCalleeSavedStackSize()
• FixedObject

Most of the places in the code which previously used the CSStackSize now use PrologueSaveSize instead, which is the sum of the latter two, while some cases which need exactly the middle one use AFI->getCalleeSavedStackSize() explicitly instead of a local variable.

In addition to moving the offsetting into emitPrologue/emitEpilogue (which fixes functions with CombineSPBump == false), also set the frame pointer to point to the right location, where the frame pointer and link register actually are stored. In addition to the prologue/epilogue, this also requires changes to resolveFrameIndexReference.

Add tests for a function that keeps a frame pointer and that uses a VLA.

[mstorsjo]

Hmm, I think this still breaks some part of my setup - I need to check it more. But does this look like a good direction?

[t.p.northover]

There's far too much ad-hoc hackery going on here for my tastes. This is the 3rd (soon to be 4th) patch trying to get something working, I think it's time to step back and really look at the code that's being modified. Or at least come up with a thorough set of tests.

[mstorsjo]

In D35919#822150, @t.p.northover wrote:

There's far too much ad-hoc hackery going on here for my tastes. This is the 3rd (soon to be 4th) patch trying to get something working, I think it's time to step back and really look at the code that's being modified. Or at least come up with a thorough set of tests.

Yeah. The first approach I had (which redefined the canonical frame address) was ugly but probably didn't have any of these other issues.

I'll try to get all codepaths thoroughly tested here before trying to get the next version landed.

[mgrang]

@mstorsjo Could you please add me as subscriber to your COFF ARM64 patches? We are trying to build spec2k/spec2k6 for COFF and running into issues with vararg functions. Some of your patches have fixed our issues :)

[mstorsjo]

In D35919#824351, @mgrang wrote:

@mstorsjo Could you please add me as subscriber to your COFF ARM64 patches? We are trying to build spec2k/spec2k6 for COFF and running into issues with vararg functions. Some of your patches have fixed our issues :)

Sorry, I forgot to add you on this one :-) I've got a revision of this one coming up pretty soon, I'm trying to get all the corner cases tested properly this time. A partial fix is to skip the change in line 536/539, but that change is also needed (with some other bits) for completeness.

[mstorsjo]

Tim: Updated after a whole lot of more testing, testing all codepaths through the emit prologue/epilogue as far as I can see (tested both with and without a frame pointer), testing unwinding for C++ exceptions through a win64cc function, etc. Hopefully this instills a bit more confidence this time. I've included a few open questions about how to rename variables in the patch message and inline that I hope you can give your opinion on.

[mstorsjo]

@mgrang does this version fix all your issues with vararg functions, or are there still issues that can be traced back to this? With this in place, all of the test code I've tried to build seems to work. (Also, have you tried using lld-link for linking? It seems complete enough for coff/arm64 for me now.)

[mstorsjo]

In D35919#825156, @mstorsjo wrote:

@mgrang does this version fix all your issues with vararg functions, or are there still issues that can be traced back to this? With this in place, all of the test code I've tried to build seems to work. (Also, have you tried using lld-link for linking? It seems complete enough for coff/arm64 for me now.)

Ping @mgrang - did this solve your issues with vararg functions, or is there still some case that I've missed?

[mgrang]

In D35919#827821, @mstorsjo wrote:

In D35919#825156, @mstorsjo wrote:

@mgrang does this version fix all your issues with vararg functions, or are there still issues that can be traced back to this? With this in place, all of the test code I've tried to build seems to work. (Also, have you tried using lld-link for linking? It seems complete enough for coff/arm64 for me now.)

Ping @mgrang - did this solve your issues with vararg functions, or is there still some case that I've missed?

Yes, so far we went past all issues related to varargs for spec2000. I think this patch LGTM. Also, we don't use lld-link for linking (we have an in-house ARM linker).

[mgrang]

LGTM.

[t.p.northover]

lib/Target/AArch64/AArch64FrameLowering.cpp

512–514 ↗

(On Diff #108715)

I'm marginally in favour of all the TODOs, but either way the decision should be made and resolved before committing. FWIW my preferred bikeshed colour would be something like PrologueSaveSize or PrologueStackSize.

[mstorsjo]

Applied the TODOs as suggested by Tim, renamed CSStackSize to PrologueSaveSize (and renamed FixedStack to FixedObject), wrapped some long lines with clang-format-diff.

[t.p.northover]

Thanks Martin, I think this looks reasonable now too.

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

else {
           MCCFIInstruction::createDefCfaOffset(FrameLabel, -NumBytes));
       BuildMI(MBB, MBBI, DL, TII->get(TargetOpcode::CFI_INSTRUCTION))
           .addCFIIndex(CFIIndex)
           .setMIFlags(MachineInstr::FrameSetup);
     }
     return;
   }
 
-  auto CSStackSize = AFI->getCalleeSavedStackSize();
+  bool IsWin64 =
+      Subtarget.isCallingConvWin64(MF.getFunction()->getCallingConv());
+  unsigned FixedObject = IsWin64 ? alignTo(AFI->getVarArgsGPRSize(), 16) : 0;
+
+  auto PrologueSaveSize = AFI->getCalleeSavedStackSize() + FixedObject;
   // All of the remaining stack allocations are for locals.
-  AFI->setLocalStackSize(NumBytes - CSStackSize);
+  AFI->setLocalStackSize(NumBytes - PrologueSaveSize);
 
   bool CombineSPBump = shouldCombineCSRLocalStackBump(MF, NumBytes);
   if (CombineSPBump) {
     emitFrameOffset(MBB, MBBI, DL, AArch64::SP, AArch64::SP, -NumBytes, TII,
                     MachineInstr::FrameSetup);
     NumBytes = 0;
-  } else if (CSStackSize != 0) {
+  } else if (PrologueSaveSize != 0) {
     MBBI = convertCalleeSaveRestoreToSPPrePostIncDec(MBB, MBBI, DL, TII,
-                                                     -CSStackSize);
-    NumBytes -= CSStackSize;
+                                                     -PrologueSaveSize);
+    NumBytes -= PrologueSaveSize;
   }
   assert(NumBytes >= 0 && "Negative stack allocation size!?");
 
   // Move past the saves of the callee-saved registers, fixing up the offsets
   // and pre-inc if we decided to combine the callee-save and local stack
   // pointer bump above.
   MachineBasicBlock::iterator End = MBB.end();
   while (MBBI != End && MBBI->getFlag(MachineInstr::FrameSetup)) {
     if (CombineSPBump)
       fixupCalleeSaveRestoreStackOffset(*MBBI, AFI->getLocalStackSize());
     ++MBBI;
   }
   if (HasFP) {
-    // Only set up FP if we actually need to. Frame pointer is fp = sp - 16.
-    int FPOffset = CSStackSize - 16;
+    // Only set up FP if we actually need to. Frame pointer is fp =
+    // sp - fixedobject - 16.
+    int FPOffset = AFI->getCalleeSavedStackSize() - 16;
     if (CombineSPBump)
       FPOffset += AFI->getLocalStackSize();
 
     // Issue    sub fp, sp, FPOffset or
     //          mov fp,sp          when FPOffset is zero.
     // Note: All stores of callee-saved registers are marked as "FrameSetup".
     // This code marks the instruction(s) that set the FP also.
     emitFrameOffset(MBB, MBBI, DL, AArch64::FP, AArch64::SP, FPOffset, TII,
if (needsFrameMoves) {
     //  Ltmp4:
     //     .cfi_offset w27, -24
     //  Ltmp5:
     //     .cfi_offset w28, -32
 
     if (HasFP) {
       // Define the current CFA rule to use the provided FP.
       unsigned Reg = RegInfo->getDwarfRegNum(FramePtr, true);
-      unsigned CFIIndex = MF.addFrameInst(
-          MCCFIInstruction::createDefCfa(nullptr, Reg, 2 * StackGrowth));
+      unsigned CFIIndex = MF.addFrameInst(MCCFIInstruction::createDefCfa(
+          nullptr, Reg, 2 * StackGrowth - FixedObject));
       BuildMI(MBB, MBBI, DL, TII->get(TargetOpcode::CFI_INSTRUCTION))
           .addCFIIndex(CFIIndex)
           .setMIFlags(MachineInstr::FrameSetup);
     } else {
       // Encode the stack size of the leaf function.
       unsigned CFIIndex = MF.addFrameInst(
           MCCFIInstruction::createDefCfaOffset(nullptr, -MFI.getStackSize()));
       BuildMI(MBB, MBBI, DL, TII->get(TargetOpcode::CFI_INSTRUCTION))
void AArch64FrameLowering::emitEpilogue(MachineFunction &MF,
   //      ----------------------        ---          ---
   //
   // So NumBytes = StackSize + BytesInStackArgArea - CalleeArgStackSize
   //             = StackSize + ArgumentPopSize
   //
   // AArch64TargetLowering::LowerCall figures out ArgumentPopSize and keeps
   // it as the 2nd argument of AArch64ISD::TC_RETURN.
 
-  auto CSStackSize = AFI->getCalleeSavedStackSize();
+  bool IsWin64 =
+      Subtarget.isCallingConvWin64(MF.getFunction()->getCallingConv());
+  unsigned FixedObject = IsWin64 ? alignTo(AFI->getVarArgsGPRSize(), 16) : 0;
+
+  auto PrologueSaveSize = AFI->getCalleeSavedStackSize() + FixedObject;
   bool CombineSPBump = shouldCombineCSRLocalStackBump(MF, NumBytes);
 
-  if (!CombineSPBump && CSStackSize != 0)
+  if (!CombineSPBump && PrologueSaveSize != 0)
     convertCalleeSaveRestoreToSPPrePostIncDec(
-        MBB, std::prev(MBB.getFirstTerminator()), DL, TII, CSStackSize);
+        MBB, std::prev(MBB.getFirstTerminator()), DL, TII, PrologueSaveSize);
 
   // Move past the restores of the callee-saved registers.
   MachineBasicBlock::iterator LastPopI = MBB.getFirstTerminator();
   MachineBasicBlock::iterator Begin = MBB.begin();
   while (LastPopI != Begin) {
     --LastPopI;
     if (!LastPopI->getFlag(MachineInstr::FrameDestroy)) {
       ++LastPopI;
       break;
     } else if (CombineSPBump)
       fixupCalleeSaveRestoreStackOffset(*LastPopI, AFI->getLocalStackSize());
   }
 
   // If there is a single SP update, insert it before the ret and we're done.
   if (CombineSPBump) {
     emitFrameOffset(MBB, MBB.getFirstTerminator(), DL, AArch64::SP, AArch64::SP,
                     NumBytes + ArgumentPopSize, TII,
                     MachineInstr::FrameDestroy);
     return;
   }
 
-  NumBytes -= CSStackSize;
+  NumBytes -= PrologueSaveSize;
   assert(NumBytes >= 0 && "Negative stack allocation size!?");
 
   if (!hasFP(MF)) {
     bool RedZone = canUseRedZone(MF);
     // If this was a redzone leaf function, we don't need to restore the
     // stack pointer (but we may need to pop stack args for fastcc).
     if (RedZone && ArgumentPopSize == 0)
       return;
 
-    bool NoCalleeSaveRestore = CSStackSize == 0;
+    bool NoCalleeSaveRestore = PrologueSaveSize == 0;
     int StackRestoreBytes = RedZone ? 0 : NumBytes;
     if (NoCalleeSaveRestore)
       StackRestoreBytes += ArgumentPopSize;
     emitFrameOffset(MBB, LastPopI, DL, AArch64::SP, AArch64::SP,
                     StackRestoreBytes, TII, MachineInstr::FrameDestroy);
     // If we were able to combine the local stack pop with the argument pop,
     // then we're done.
     if (NoCalleeSaveRestore || ArgumentPopSize == 0)
       return;
     NumBytes = 0;
   }
 
   // Restore the original stack pointer.
   // FIXME: Rather than doing the math here, we should instead just use
   // non-post-indexed loads for the restores if we aren't actually going to
   // be able to save any instructions.
   if (MFI.hasVarSizedObjects() || AFI->isStackRealigned())
     emitFrameOffset(MBB, LastPopI, DL, AArch64::SP, AArch64::FP,
-                    -CSStackSize + 16, TII, MachineInstr::FrameDestroy);
+                    -AFI->getCalleeSavedStackSize() + 16, TII,
+                    MachineInstr::FrameDestroy);
   else if (NumBytes)
     emitFrameOffset(MBB, LastPopI, DL, AArch64::SP, AArch64::SP, NumBytes, TII,
                     MachineInstr::FrameDestroy);
 
   // This must be placed after the callee-save restore code because that code
   // assumes the SP is at the same location as it was after the callee-save save
   // code in the prologue.
   if (ArgumentPopSize)

 
 int AArch64FrameLowering::resolveFrameIndexReference(const MachineFunction &MF,
                                                      int FI, unsigned &FrameReg,
                                                      bool PreferFP) const {
   const MachineFrameInfo &MFI = MF.getFrameInfo();
   const AArch64RegisterInfo *RegInfo = static_cast<const AArch64RegisterInfo *>(
       MF.getSubtarget().getRegisterInfo());
   const AArch64FunctionInfo *AFI = MF.getInfo<AArch64FunctionInfo>();
-  int FPOffset = MFI.getObjectOffset(FI) + 16;
+  const AArch64Subtarget &Subtarget = MF.getSubtarget<AArch64Subtarget>();
+  bool IsWin64 =
+      Subtarget.isCallingConvWin64(MF.getFunction()->getCallingConv());
+  unsigned FixedObject = IsWin64 ? alignTo(AFI->getVarArgsGPRSize(), 16) : 0;
+  int FPOffset = MFI.getObjectOffset(FI) + FixedObject + 16;
   int Offset = MFI.getObjectOffset(FI) + MFI.getStackSize();
   bool isFixed = MFI.isFixedObjectIndex(FI);
 
   // Use frame pointer to reference fixed objects. Use it for locals if
   // there are VLAs or a dynamically realigned SP (and thus the SP isn't
   // reliable as a base). Make sure useFPForScavengingIndex() does the
   // right thing for the emergency spill slot.
   bool UseFP = false;
static void computeCalleeSaveRegisterPairs(
   // MachO's compact unwind format relies on all registers being stored in
   // pairs.
   assert((!produceCompactUnwindFrame(MF) ||
           CC == CallingConv::PreserveMost ||
           (Count & 1) == 0) &&
          "Odd number of callee-saved regs to spill!");
   int Offset = AFI->getCalleeSavedStackSize();
 
-  unsigned GPRSaveSize = AFI->getVarArgsGPRSize();
-  const AArch64Subtarget &Subtarget = MF.getSubtarget<AArch64Subtarget>();
-  bool IsWin64 = Subtarget.isCallingConvWin64(MF.getFunction()->getCallingConv());
-  if (IsWin64)
-    Offset -= alignTo(GPRSaveSize, 16);
-
   for (unsigned i = 0; i < Count; ++i) {
     RegPairInfo RPI;
     RPI.Reg1 = CSI[i].getReg();
 
     assert(AArch64::GPR64RegClass.contains(RPI.Reg1) ||
            AArch64::FPR64RegClass.contains(RPI.Reg1));
     RPI.IsGPR = AArch64::GPR64RegClass.contains(RPI.Reg1);
 

llvm/trunk/test/CodeGen/AArch64/aarch64_win64cc_vararg.ll

 ; RUN: llc < %s -mtriple=aarch64-linux-gnu | FileCheck %s
 
 define win64cc void @pass_va(i32 %count, ...) nounwind {
 entry:
-; CHECK: sub     sp, sp, #80
+; CHECK: str     x30, [sp, #-80]!
 ; CHECK: add     x8, sp, #24
 ; CHECK: add     x0, sp, #24
 ; CHECK: stp     x6, x7, [sp, #64]
 ; CHECK: stp     x4, x5, [sp, #48]
 ; CHECK: stp     x2, x3, [sp, #32]
 ; CHECK: str     x1, [sp, #24]
-; CHECK: stp     x30, x8, [sp]
+; CHECK: str     x8, [sp, #8]
 ; CHECK: bl      other_func
 ; CHECK: ldr     x30, [sp], #80
 ; CHECK: ret
   %ap = alloca i8*, align 8
   %ap1 = bitcast i8** %ap to i8*
   call void @llvm.va_start(i8* %ap1)
   %ap2 = load i8*, i8** %ap, align 8
   call void @other_func(i8* %ap2)

llvm/trunk/test/CodeGen/AArch64/win64_vararg.ll

 ; RUN: llc < %s -mtriple=aarch64-pc-win32 | FileCheck %s
 
 define void @pass_va(i32 %count, ...) nounwind {
 entry:
-; CHECK: sub     sp, sp, #80
+; CHECK: str     x30, [sp, #-80]!
 ; CHECK: add     x8, sp, #24
 ; CHECK: add     x0, sp, #24
 ; CHECK: stp     x6, x7, [sp, #64]
 ; CHECK: stp     x4, x5, [sp, #48]
 ; CHECK: stp     x2, x3, [sp, #32]
 ; CHECK: str     x1, [sp, #24]
-; CHECK: stp     x30, x8, [sp]
+; CHECK: str     x8, [sp, #8]
 ; CHECK: bl      other_func
 ; CHECK: ldr     x30, [sp], #80
 ; CHECK: ret
   %ap = alloca i8*, align 8
   %ap1 = bitcast i8** %ap to i8*
   call void @llvm.va_start(i8* %ap1)
   %ap2 = load i8*, i8** %ap, align 8
   call void @other_func(i8* %ap2)

 
 declare void @llvm.va_end(i8*)
 declare void @llvm.lifetime.start.p0i8(i64, i8* nocapture) #1
 declare void @llvm.lifetime.end.p0i8(i64, i8* nocapture) #1
 
 declare i32 @__stdio_common_vsprintf(i64, i8*, i64, i8*, i8*, i8*) local_unnamed_addr #3
 declare i64* @__local_stdio_printf_options() local_unnamed_addr #4
 
+; CHECK-LABEL: fp
+; CHECK: str     x21, [sp, #-96]!
+; CHECK: stp     x20, x19, [sp, #16]
+; CHECK: stp     x29, x30, [sp, #32]
+; CHECK: add     x29, sp, #32
+; CHECK: add     x8, x29, #24
+; CHECK: mov     x19, x2
+; CHECK: mov     x20, x1
+; CHECK: mov     x21, x0
+; CHECK: stp     x6, x7, [x29, #48]
+; CHECK: stp     x4, x5, [x29, #32]
+; CHECK: str     x3, [x29, #24]
+; CHECK: str     x8, [sp, #8]
+; CHECK: bl      __local_stdio_printf_options
+; CHECK: ldr     x8, [x0]
+; CHECK: add     x5, x29, #24
+; CHECK: mov     x1, x21
+; CHECK: mov     x2, x20
+; CHECK: orr     x0, x8, #0x2
+; CHECK: mov     x3, x19
+; CHECK: mov     x4, xzr
+; CHECK: bl      __stdio_common_vsprintf
+; CHECK: ldp     x29, x30, [sp, #32]
+; CHECK: ldp     x20, x19, [sp, #16]
+; CHECK: cmp     w0, #0
+; CHECK: csinv   w0, w0, wzr, ge
+; CHECK: ldr     x21, [sp], #96
+; CHECK: ret
+define i32 @fp(i8*, i64, i8*, ...) local_unnamed_addr #6 {
+  %4 = alloca i8*, align 8
+  %5 = bitcast i8** %4 to i8*
+  call void @llvm.lifetime.start.p0i8(i64 8, i8* nonnull %5) #2
+  call void @llvm.va_start(i8* nonnull %5)
+  %6 = load i8*, i8** %4, align 8
+  %7 = call i64* @__local_stdio_printf_options() #2
+  %8 = load i64, i64* %7, align 8
+  %9 = or i64 %8, 2
+  %10 = call i32 @__stdio_common_vsprintf(i64 %9, i8* %0, i64 %1, i8* %2, i8* null, i8* %6) #2
+  %11 = icmp sgt i32 %10, -1
+  %12 = select i1 %11, i32 %10, i32 -1
+  call void @llvm.va_end(i8* nonnull %5)
+  call void @llvm.lifetime.end.p0i8(i64 8, i8* nonnull %5) #2
+  ret i32 %12
+}
+
+attributes #6 = { "no-frame-pointer-elim"="true" }
+
+; CHECK-LABEL: vla
+; CHECK: str     x23, [sp, #-112]!
+; 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: 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: 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
+; CHECK: mov     x4, xzr
+; CHECK: mov     x5, x21
+; CHECK: bl      __stdio_common_vsprintf
+; CHECK: mov     sp, x23
+; CHECK: sub     sp, x29, #48
+; CHECK: ldp     x29, x30, [sp, #48]
+; CHECK: ldp     x20, x19, [sp, #32]
+; CHECK: ldp     x22, x21, [sp, #16]
+; CHECK: ldr     x23, [sp], #112
+; CHECK: ret
+define void @vla(i32, i8*, ...) local_unnamed_addr {
+  %3 = alloca i8*, align 8
+  %4 = bitcast i8** %3 to i8*
+  call void @llvm.lifetime.start.p0i8(i64 8, i8* nonnull %4) #5
+  call void @llvm.va_start(i8* nonnull %4)
+  %5 = zext i32 %0 to i64
+  %6 = call i8* @llvm.stacksave()
+  %7 = alloca i8, i64 %5, align 1
+  %8 = load i8*, i8** %3, align 8
+  %9 = sext i32 %0 to i64
+  %10 = call i64* @__local_stdio_printf_options()
+  %11 = load i64, i64* %10, align 8
+  %12 = or i64 %11, 2
+  %13 = call i32 @__stdio_common_vsprintf(i64 %12, i8* nonnull %7, i64 %9, i8* %1, i8* null, i8* %8)
+  call void @llvm.va_end(i8* nonnull %4)
+  call void @llvm.stackrestore(i8* %6)
+  call void @llvm.lifetime.end.p0i8(i64 8, i8* nonnull %4) #5
+  ret void
+}
+
+declare i8* @llvm.stacksave()
+declare void @llvm.stackrestore(i8*)
+
 ; CHECK-LABEL: snprintf
 ; CHECK: sub     sp,  sp, #96
 ; CHECK: stp     x21, x20, [sp, #16]
 ; CHECK: stp     x19, x30, [sp, #32]
 ; CHECK: add     x8, sp, #56
 ; CHECK: mov     x19, x2
 ; CHECK: mov     x20, x1
 ; CHECK: mov     x21, x0