[AArch64] Fix fastcc -tailcallopt epilog code generation.

[gberry]

Summary

Fix a bug in epilog generation where the incoming stack arguments were
not being popped for fastcc functions when -tailcallopt was passed.

[jmolloy]

Tim is a better reviewer than me for this.

[gberry]

Ping

[t.p.northover]

Hi Geoff,

Sorry this took so long to get around to. I think it looks OK.

Tim.

test/CodeGen/AArch64/fastcc.ll

237 ↗

(On Diff #46940)

Might want to fix that up before commit.

[gberry]

Tim,

No problem, thanks for the review.

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

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.
-  NumBytes += ArgumentPopSize;
 
   // 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;
     }
   }
   NumBytes -= AFI->getCalleeSavedStackSize();
   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.
-    if (!canUseRedZone(MF))
-      emitFrameOffset(MBB, LastPopI, DL, AArch64::SP, AArch64::SP, NumBytes,
-                      TII, MachineInstr::FrameDestroy);
+    // stack pointer (but we may need to pop stack args for fastcc).
+    if (RedZone && ArgumentPopSize == 0)
       return;
+
+    bool NoCalleeSaveRestore = AFI->getCalleeSavedStackSize() == 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 (NumBytes || MFI->hasVarSizedObjects())
     emitFrameOffset(MBB, LastPopI, DL, AArch64::SP, AArch64::FP,
                     -AFI->getCalleeSavedStackSize() + 16, 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)
+    emitFrameOffset(MBB, MBB.getFirstTerminator(), DL, AArch64::SP, AArch64::SP,
+                    ArgumentPopSize, TII, MachineInstr::FrameDestroy);
 }
 
 /// getFrameIndexReference - Provide a base+offset reference to an FI slot for
 /// debug info.  It's the same as what we use for resolving the code-gen
 /// references for now.  FIXME: This can go wrong when references are
 /// SP-relative and simple call frames aren't used.
 int AArch64FrameLowering::getFrameIndexReference(const MachineFunction &MF,
                                                  int FI,

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

 ; RUN: llc -verify-machineinstrs < %s -mtriple=aarch64-none-linux-gnu -tailcallopt | FileCheck %s -check-prefix CHECK-TAIL
 ; RUN: llc -verify-machineinstrs < %s -mtriple=aarch64-none-linux-gnu | FileCheck %s
+; RUN: llc -verify-machineinstrs < %s -mtriple=aarch64-none-linux-gnu -tailcallopt -aarch64-redzone | FileCheck %s -check-prefix CHECK-TAIL-RZ
 
 ; Without tailcallopt fastcc still means the caller cleans up the
 ; stack, so try to make sure this is respected.
 
 define fastcc void @func_stack0() {
 ; CHECK-LABEL: func_stack0:
 ; CHECK: mov x29, sp
 ; CHECK: str w{{[0-9]+}}, [sp, #-32]!
; CHECK-TAIL-NOT: sub sp, sp
   ret void
 ; CHECK: mov sp, x29
 ; CHECK-NEXT: ldp     x29, x30, [sp], #16
 ; CHECK-NEXT: ret
 
 
 ; CHECK-TAIL: mov sp, x29
 ; CHECK-TAIL-NEXT: ldp     x29, x30, [sp], #16
+; CHECK-TAIL-NEXT: add     sp, sp, #16
 ; CHECK-TAIL-NEXT: ret
 }
 
 define fastcc void @func_stack32([8 x i32], i128 %stacked0, i128 %stacked1) {
 ; CHECK-LABEL: func_stack32:
 ; CHECK: mov x29, sp
 
 ; CHECK-TAIL-LABEL: func_stack32:
; CHECK-TAIL-NOT: sub sp, sp
 
   ret void
 ; CHECK: mov sp, x29
 ; CHECK-NEXT: ldp     x29, x30, [sp], #16
 ; CHECK-NEXT: ret
 
 ; CHECK-TAIL: mov sp, x29
 ; CHECK-TAIL-NEXT: ldp     x29, x30, [sp], #16
+; CHECK-TAIL-NEXT: add     sp, sp, #32
 ; CHECK-TAIL-NEXT: ret
 }
+
+; Check that arg stack pop is done after callee-save restore when no frame pointer is used.
+define fastcc void @func_stack32_leaf([8 x i32], i128 %stacked0, i128 %stacked1) {
+; CHECK-LABEL: func_stack32_leaf:
+; CHECK: str     x20, [sp, #-16]!
+; CHECK: nop
+; CHECK-NEXT: //NO_APP
+; CHECK-NEXT: ldr     x20, [sp], #16
+; CHECK-NEXT: ret
+
+; CHECK-TAIL-LABEL: func_stack32_leaf:
+; CHECK-TAIL: str     x20, [sp, #-16]!
+; CHECK-TAIL: nop
+; CHECK-TAIL-NEXT: //NO_APP
+; CHECK-TAIL-NEXT: ldr     x20, [sp], #16
+; CHECK-TAIL-NEXT: add     sp, sp, #32
+; CHECK-TAIL-NEXT: ret
+
+; CHECK-TAIL-RZ-LABEL: func_stack32_leaf:
+; CHECK-TAIL-RZ: str     x20, [sp, #-16]!
+; CHECK-TAIL-RZ-NOT: sub     sp, sp
+; CHECK-TAIL-RZ: nop
+; CHECK-TAIL-RZ-NEXT: //NO_APP
+; CHECK-TAIL-RZ-NEXT: ldr     x20, [sp], #16
+; CHECK-TAIL-RZ-NEXT: add     sp, sp, #32
+; CHECK-TAIL-RZ-NEXT: ret
+
+  ; Make sure there is a callee-save register to save/restore.
+  call void asm sideeffect "nop", "~{x20}"() nounwind
+  ret void
+}
+
+; Check that arg stack pop is done after callee-save restore when no frame pointer is used.
+define fastcc void @func_stack32_leaf_local([8 x i32], i128 %stacked0, i128 %stacked1) {
+; CHECK-LABEL: func_stack32_leaf_local:
+; CHECK: str     x20, [sp, #-16]!
+; CHECK-NEXT: sub     sp, sp, #16
+; CHECK: nop
+; CHECK-NEXT: //NO_APP
+; CHECK-NEXT: add     sp, sp, #16
+; CHECK-NEXT: ldr     x20, [sp], #16
+; CHECK-NEXT: ret
+
+; CHECK-TAIL-LABEL: func_stack32_leaf_local:
+; CHECK-TAIL: str     x20, [sp, #-16]!
+; CHECK-TAIL-NEXT: sub     sp, sp, #16
+; CHECK-TAIL: nop
+; CHECK-TAIL-NEXT: //NO_APP
+; CHECK-TAIL-NEXT: add     sp, sp, #16
+; CHECK-TAIL-NEXT: ldr     x20, [sp], #16
+; CHECK-TAIL-NEXT: add     sp, sp, #32
+; CHECK-TAIL-NEXT: ret
+
+; CHECK-TAIL-RZ-LABEL: func_stack32_leaf_local:
+; CHECK-TAIL-RZ: str     x20, [sp, #-16]!
+; CHECK-TAIL-RZ-NOT: sub     sp, sp
+; CHECK-TAIL-RZ: nop
+; CHECK-TAIL-RZ-NEXT: //NO_APP
+; CHECK-TAIL-RZ-NEXT: ldr     x20, [sp], #16
+; CHECK-TAIL-RZ-NEXT: add     sp, sp, #32
+; CHECK-TAIL-RZ-NEXT: ret
+
+  %val0 = alloca [2 x i64], align 8
+
+  ; Make sure there is a callee-save register to save/restore.
+  call void asm sideeffect "nop", "~{x20}"() nounwind
+  ret void
+}
+
+; Check that arg stack pop is done after callee-save restore when no frame pointer is used.
+define fastcc void @func_stack32_leaf_local_nocs([8 x i32], i128 %stacked0, i128 %stacked1) {
+; CHECK-LABEL: func_stack32_leaf_local_nocs:
+; CHECK: sub     sp, sp, #16
+; CHECK: add     sp, sp, #16
+; CHECK-NEXT: ret
+
+; CHECK-TAIL-LABEL: func_stack32_leaf_local_nocs:
+; CHECK-TAIL: sub     sp, sp, #16
+; CHECK-TAIL: add     sp, sp, #48
+; CHECK-TAIL-NEXT: ret
+
+; CHECK-TAIL-RZ-LABEL: func_stack32_leaf_local_nocs:
+; CHECK-TAIL-RZ: add     sp, sp, #32
+; CHECK-TAIL-RZ-NEXT: ret
+
+  %val0 = alloca [2 x i64], align 8
+
+  ret void
+}