Make sure to forward arguments from vararg to musttail vararg

[yinma]

Summary

Thunk functions in Windows are varag functions that call a musttail function
to pass the arguments after the fixup is done.  We need to make sure that we
forward the arguments from the caller vararg to the callee vararg function.
This is the same mechanism that is used for Windows on X86.

[rnk]

lib/Target/AArch64/AArch64ISelLowering.cpp

3152

Please do this for all OSs, not just windows. This is a general purpose feature that is supposed to work for all targets, it's just only implemented for one right now.

3158

Neat, I guess the code I wrote to find all the remaining regparms was generic enough to work for any convention.

3622

Please remove the IsWin64 test so we just do this all the time. One day, someone will come asking to use this on AArch64, probably.

test/CodeGen/AArch64/vararg-tallcall.ll

2

If it works, test some other targets, especially if the default calling convention changes based on OS.

[efriedma]

Testcase should also check that the vector registers are correctly saved/restored. (Something like call void asm sideeffect "", "~{d0}"() should force a spill.) It looks like the code handles this correctly already.

Otherwise LGTM.

[efriedma]

LGTM

lib/Target/AArch64/AArch64ISelLowering.cpp

if (!Subtarget->isTargetDarwin() || IsWin64) {
       saveVarArgRegisters(CCInfo, DAG, DL, Chain);
     }
 
     // This will point to the next argument passed via stack.
     unsigned StackOffset = CCInfo.getNextStackOffset();
     // We currently pass all varargs at 8-byte alignment.
     StackOffset = ((StackOffset + 7) & ~7);
     FuncInfo->setVarArgsStackIndex(MFI.CreateFixedObject(4, StackOffset, true));
+
+    if (MFI.hasMustTailInVarArgFunc()) {
+      SmallVector<MVT, 2> RegParmTypes;
+      RegParmTypes.push_back(MVT::i64);
+      RegParmTypes.push_back(MVT::f128);
+      // Compute the set of forwarded registers. The rest are scratch.
+      SmallVectorImpl<ForwardedRegister> &Forwards =
+                                       FuncInfo->getForwardedMustTailRegParms();
+      CCInfo.analyzeMustTailForwardedRegisters(Forwards, RegParmTypes,
+                                               CC_AArch64_AAPCS);
+    }
   }
 
   unsigned StackArgSize = CCInfo.getNextStackOffset();
   bool TailCallOpt = MF.getTarget().Options.GuaranteedTailCallOpt;
   if (DoesCalleeRestoreStack(CallConv, TailCallOpt)) {
     // This is a non-standard ABI so by fiat I say we're allowed to make full
     // use of the stack area to be popped, which must be aligned to 16 bytes in
     // any case:
AArch64TargetLowering::LowerCall(CallLoweringInfo &CLI,
 
   SDValue StackPtr = DAG.getCopyFromReg(Chain, DL, AArch64::SP,
                                         getPointerTy(DAG.getDataLayout()));
 
   SmallVector<std::pair<unsigned, SDValue>, 8> RegsToPass;
   SmallVector<SDValue, 8> MemOpChains;
   auto PtrVT = getPointerTy(DAG.getDataLayout());
 
+  if (IsVarArg && CLI.CS && CLI.CS.isMustTailCall()) {
+    const auto &Forwards = FuncInfo->getForwardedMustTailRegParms();
+    for (const auto &F : Forwards) {
+      SDValue Val = DAG.getCopyFromReg(Chain, DL, F.VReg, F.VT);
+       RegsToPass.push_back(std::make_pair(unsigned(F.PReg), Val));
+    }
+  }
+
   // Walk the register/memloc assignments, inserting copies/loads.
   for (unsigned i = 0, realArgIdx = 0, e = ArgLocs.size(); i != e;
        ++i, ++realArgIdx) {
     CCValAssign &VA = ArgLocs[i];
     SDValue Arg = OutVals[realArgIdx];
     ISD::ArgFlagsTy Flags = Outs[realArgIdx].Flags;
 
     // Promote the value if needed.

lib/Target/AArch64/AArch64MachineFunctionInfo.h

 
 #ifndef LLVM_LIB_TARGET_AARCH64_AARCH64MACHINEFUNCTIONINFO_H
 #define LLVM_LIB_TARGET_AARCH64_AARCH64MACHINEFUNCTIONINFO_H
 
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/CodeGen/CallingConvLower.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/MC/MCLinkerOptimizationHint.h"
 #include <cassert>
 
 namespace llvm {
 
 class MachineInstr;
 
class AArch64FunctionInfo final : public MachineFunctionInfo {
   bool CalleeSaveStackHasFreeSpace = false;
 
   /// Has a value when it is known whether or not the function uses a
   /// redzone, and no value otherwise.
   /// Initialized during frame lowering, unless the function has the noredzone
   /// attribute, in which case it is set to false at construction.
   Optional<bool> HasRedZone;
 
+  /// ForwardedMustTailRegParms - A list of virtual and physical registers
+  /// that must be forwarded to every musttail call.
+  SmallVector<ForwardedRegister, 1> ForwardedMustTailRegParms;
 public:
   AArch64FunctionInfo() = default;
 
   explicit AArch64FunctionInfo(MachineFunction &MF) {
     (void)MF;
 
     // If we already know that the function doesn't have a redzone, set
     // HasRedZone here.
public:
   const MILOHContainer &getLOHContainer() const { return LOHContainerSet; }
 
   /// Add a LOH directive of this @p Kind and this @p Args.
   void addLOHDirective(MCLOHType Kind, MILOHArgs Args) {
     LOHContainerSet.push_back(MILOHDirective(Kind, Args));
     LOHRelated.insert(Args.begin(), Args.end());
   }
 
+  SmallVectorImpl<ForwardedRegister> &getForwardedMustTailRegParms() {
+    return ForwardedMustTailRegParms;
+  }
+
 private:
   // Hold the lists of LOHs.
   MILOHContainer LOHContainerSet;
   SetOfInstructions LOHRelated;
 
   DenseMap<int, std::pair<unsigned, MCSymbol *>> JumpTableEntryInfo;
 };
 
 } // end namespace llvm
 
 #endif // LLVM_LIB_TARGET_AARCH64_AARCH64MACHINEFUNCTIONINFO_H

test/CodeGen/AArch64/vararg-tallcall.ll

+; RUN: llc -mtriple=aarch64-windows-msvc %s -o - | FileCheck %s
+; RUN: llc -mtriple=aarch64-linux-gnu %s -o - | FileCheck %s
+
+target datalayout = "e-m:w-p:64:64-i32:32-i64:64-i128:128-n32:64-S128"
+
+%class.X = type { i8 }
+%struct.B = type { i32 (...)** }
+
+$"??_9B@@$BA@AA" = comdat any
+
+; Function Attrs: noinline optnone
+define linkonce_odr void @"??_9B@@$BA@AA"(%struct.B* %this, ...) #1 comdat align 2  {
+entry:
+  %this.addr = alloca %struct.B*, align 8
+  store %struct.B* %this, %struct.B** %this.addr, align 8
+  %this1 = load %struct.B*, %struct.B** %this.addr, align 8
+  call void asm sideeffect "", "~{d0}"()
+  %0 = bitcast %struct.B* %this1 to void (%struct.B*, ...)***
+  %vtable = load void (%struct.B*, ...)**, void (%struct.B*, ...)*** %0, align 8
+  %vfn = getelementptr inbounds void (%struct.B*, ...)*, void (%struct.B*, ...)** %vtable, i64 0
+  %1 = load void (%struct.B*, ...)*, void (%struct.B*, ...)** %vfn, align 8
+  musttail call void (%struct.B*, ...) %1(%struct.B* %this1, ...)
+  ret void
+                                                  ; No predecessors!
+  ret void
+}
+
+attributes #1 = { noinline optnone "thunk" }
+
+; CHECK: mov     v16.16b, v0.16b
+; CHECK: ldr     x8, [x0]
+; CHECK: ldr     x8, [x8]
+; CHECK: mov     v0.16b, v16.16b
+; CHECK: br      x8