1 files changed, 186 insertions, 7 deletions
diff --git a/llvm/lib/Target/AArch64/AArch64CallLowering.cpp b/llvm/lib/Target/AArch64/AArch64CallLowering.cpp
index 9ace33b3985..e97b329985a 100644
--- a/llvm/lib/Target/AArch64/AArch64CallLowering.cpp
+++ b/llvm/lib/Target/AArch64/AArch64CallLowering.cpp
@@ -233,6 +233,17 @@ bool AArch64CallLowering::lowerReturn(MachineIRBuilder &MIRBuilder,
                                       const Value *Val,
                                       ArrayRef<Register> VRegs,
                                       Register SwiftErrorVReg) const {
+
+  // Check if a tail call was lowered in this block. If so, we already handled
+  // the terminator.
+  MachineFunction &MF = MIRBuilder.getMF();
+  if (MF.getFrameInfo().hasTailCall()) {
+    MachineBasicBlock &MBB = MIRBuilder.getMBB();
+    auto FirstTerm = MBB.getFirstTerminator();
+    if (FirstTerm != MBB.end() && FirstTerm->isCall())
+      return true;
+  }
+
   auto MIB = MIRBuilder.buildInstrNoInsert(AArch64::RET_ReallyLR);
   assert(((Val && !VRegs.empty()) || (!Val && VRegs.empty())) &&
          "Return value without a vreg");
@@ -403,6 +414,146 @@ bool AArch64CallLowering::lowerFormalArguments(
   return true;
 }
 
+/// Return true if the calling convention is one that we can guarantee TCO for.
+static bool canGuaranteeTCO(CallingConv::ID CC) {
+  return CC == CallingConv::Fast;
+}
+
+/// Return true if we might ever do TCO for calls with this calling convention.
+static bool mayTailCallThisCC(CallingConv::ID CC) {
+  switch (CC) {
+  case CallingConv::C:
+  case CallingConv::PreserveMost:
+  case CallingConv::Swift:
+    return true;
+  default:
+    return canGuaranteeTCO(CC);
+  }
+}
+
+bool AArch64CallLowering::isEligibleForTailCallOptimization(
+    MachineIRBuilder &MIRBuilder, CallLoweringInfo &Info) const {
+  CallingConv::ID CalleeCC = Info.CallConv;
+  MachineFunction &MF = MIRBuilder.getMF();
+  const Function &CallerF = MF.getFunction();
+  CallingConv::ID CallerCC = CallerF.getCallingConv();
+  bool CCMatch = CallerCC == CalleeCC;
+
+  LLVM_DEBUG(dbgs() << "Attempting to lower call as tail call\n");
+
+  if (Info.SwiftErrorVReg) {
+    // TODO: We should handle this.
+    // Note that this is also handled by the check for no outgoing arguments.
+    // Proactively disabling this though, because the swifterror handling in
+    // lowerCall inserts a COPY *after* the location of the call.
+    LLVM_DEBUG(dbgs() << "... Cannot handle tail calls with swifterror yet.\n");
+    return false;
+  }
+
+  if (!Info.OrigRet.Ty->isVoidTy()) {
+    // TODO: lowerCall will insert COPYs to handle the call's return value.
+    // This needs some refactoring to avoid this with tail call returns. For
+    // now, just don't handle that case.
+    LLVM_DEBUG(dbgs() << "... Cannot handle non-void return types yet.\n");
+    return false;
+  }
+
+  if (!mayTailCallThisCC(CalleeCC)) {
+    LLVM_DEBUG(dbgs() << "... Calling convention cannot be tail called.\n");
+    return false;
+  }
+
+  if (Info.IsVarArg) {
+    LLVM_DEBUG(dbgs() << "... Tail calling varargs not supported yet.\n");
+    return false;
+  }
+
+  // Byval parameters hand the function a pointer directly into the stack area
+  // we want to reuse during a tail call. Working around this *is* possible (see
+  // X86).
+  //
+  // FIXME: In AArch64ISelLowering, this isn't worked around. Can/should we try
+  // it?
+  //
+  // On Windows, "inreg" attributes signify non-aggregate indirect returns.
+  // In this case, it is necessary to save/restore X0 in the callee. Tail
+  // call opt interferes with this. So we disable tail call opt when the
+  // caller has an argument with "inreg" attribute.
+  //
+  // FIXME: Check whether the callee also has an "inreg" argument.
+  if (any_of(CallerF.args(), [](const Argument &A) {
+        return A.hasByValAttr() || A.hasInRegAttr();
+      })) {
+    LLVM_DEBUG(dbgs() << "... Cannot tail call from callers with byval or "
+                         "inreg arguments.\n");
+    return false;
+  }
+
+  // Externally-defined functions with weak linkage should not be
+  // tail-called on AArch64 when the OS does not support dynamic
+  // pre-emption of symbols, as the AAELF spec requires normal calls
+  // to undefined weak functions to be replaced with a NOP or jump to the
+  // next instruction. The behaviour of branch instructions in this
+  // situation (as used for tail calls) is implementation-defined, so we
+  // cannot rely on the linker replacing the tail call with a return.
+  if (Info.Callee.isGlobal()) {
+    const GlobalValue *GV = Info.Callee.getGlobal();
+    const Triple &TT = MF.getTarget().getTargetTriple();
+    if (GV->hasExternalWeakLinkage() &&
+        (!TT.isOSWindows() || TT.isOSBinFormatELF() ||
+         TT.isOSBinFormatMachO())) {
+      LLVM_DEBUG(dbgs() << "... Cannot tail call externally-defined function "
+                           "with weak linkage for this OS.\n");
+      return false;
+    }
+  }
+
+  // If we have -tailcallopt and matching CCs, at this point, we could return
+  // true. However, we don't have full tail call support yet. So, continue
+  // checking. We want to emit a sibling call.
+
+  // I want anyone implementing a new calling convention to think long and hard
+  // about this assert.
+  assert((!Info.IsVarArg || CalleeCC == CallingConv::C) &&
+         "Unexpected variadic calling convention");
+
+  // For now, only support the case where the calling conventions match.
+  if (!CCMatch) {
+    LLVM_DEBUG(
+        dbgs()
+        << "... Cannot tail call with mismatched calling conventions yet.\n");
+    return false;
+  }
+
+  // For now, only handle callees that take no arguments.
+  if (!Info.OrigArgs.empty()) {
+    LLVM_DEBUG(
+        dbgs()
+        << "... Cannot tail call callees with outgoing arguments yet.\n");
+    return false;
+  }
+
+  LLVM_DEBUG(
+      dbgs() << "... Call is eligible for tail call optimization.\n");
+  return true;
+}
+
+static unsigned getCallOpcode(const Function &CallerF, bool IsIndirect,
+                              bool IsTailCall) {
+  if (!IsTailCall)
+    return IsIndirect ? AArch64::BLR : AArch64::BL;
+
+  if (!IsIndirect)
+    return AArch64::TCRETURNdi;
+
+  // When BTI is enabled, we need to use TCRETURNriBTI to make sure that we use
+  // x16 or x17.
+  if (CallerF.hasFnAttribute("branch-target-enforcement"))
+    return AArch64::TCRETURNriBTI;
+
+  return AArch64::TCRETURNri;
+}
+
 bool AArch64CallLowering::lowerCall(MachineIRBuilder &MIRBuilder,
                                     CallLoweringInfo &Info) const {
   MachineFunction &MF = MIRBuilder.getMF();
@@ -411,6 +562,7 @@ bool AArch64CallLowering::lowerCall(MachineIRBuilder &MIRBuilder,
   auto &DL = F.getParent()->getDataLayout();
 
   if (Info.IsMustTailCall) {
+    // TODO: Until we lower all tail calls, we should fall back on this.
     LLVM_DEBUG(dbgs() << "Cannot lower musttail calls yet.\n");
     return false;
   }
@@ -423,6 +575,11 @@ bool AArch64CallLowering::lowerCall(MachineIRBuilder &MIRBuilder,
       SplitArgs.back().Flags[0].setZExt();
   }
 
+  bool IsSibCall =
+      Info.IsTailCall && isEligibleForTailCallOptimization(MIRBuilder, Info);
+  if (IsSibCall)
+    MF.getFrameInfo().setHasTailCall();
+
   // Find out which ABI gets to decide where things go.
   const AArch64TargetLowering &TLI = *getTLI<AArch64TargetLowering>();
   CCAssignFn *AssignFnFixed =
@@ -430,14 +587,33 @@ bool AArch64CallLowering::lowerCall(MachineIRBuilder &MIRBuilder,
   CCAssignFn *AssignFnVarArg =
       TLI.CCAssignFnForCall(Info.CallConv, /*IsVarArg=*/true);
 
-  auto CallSeqStart = MIRBuilder.buildInstr(AArch64::ADJCALLSTACKDOWN);
+  // If we have a sibling call, then we don't have to adjust the stack.
+  // Otherwise, we need to adjust it.
+  MachineInstrBuilder CallSeqStart;
+  if (!IsSibCall)
+    CallSeqStart = MIRBuilder.buildInstr(AArch64::ADJCALLSTACKDOWN);
 
   // Create a temporarily-floating call instruction so we can add the implicit
   // uses of arg registers.
-  auto MIB = MIRBuilder.buildInstrNoInsert(Info.Callee.isReg() ? AArch64::BLR
-                                                               : AArch64::BL);
+  unsigned Opc = getCallOpcode(F, Info.Callee.isReg(), IsSibCall);
+
+  // TODO: Right now, regbankselect doesn't know how to handle the rtcGPR64
+  // register class. Until we can do that, we should fall back here.
+  if (Opc == AArch64::TCRETURNriBTI) {
+    LLVM_DEBUG(
+        dbgs() << "Cannot lower indirect tail calls with BTI enabled yet.\n");
+    return false;
+  }
+
+  auto MIB = MIRBuilder.buildInstrNoInsert(Opc);
   MIB.add(Info.Callee);
 
+  // Add the byte offset for the tail call. We only have sibling calls, so this
+  // is always 0.
+  // TODO: Handle tail calls where we will have a different value here.
+  if (IsSibCall)
+    MIB.addImm(0);
+
   // Tell the call which registers are clobbered.
   auto TRI = MF.getSubtarget<AArch64Subtarget>().getRegisterInfo();
   const uint32_t *Mask = TRI->getCallPreservedMask(MF, F.getCallingConv());
@@ -486,10 +662,13 @@ bool AArch64CallLowering::lowerCall(MachineIRBuilder &MIRBuilder,
     MIRBuilder.buildCopy(Info.SwiftErrorVReg, Register(AArch64::X21));
   }
 
-  CallSeqStart.addImm(Handler.StackSize).addImm(0);
-  MIRBuilder.buildInstr(AArch64::ADJCALLSTACKUP)
-      .addImm(Handler.StackSize)
-      .addImm(0);
+  if (!IsSibCall) {
+    // If we aren't sibcalling, we need to move the stack.
+    CallSeqStart.addImm(Handler.StackSize).addImm(0);
+    MIRBuilder.buildInstr(AArch64::ADJCALLSTACKUP)
+        .addImm(Handler.StackSize)
+        .addImm(0);
+  }
 
   return true;
 }