[AArch64][GlobalISel] Support sibling calls with mismatched calling conventions

[paquette]

Summary

Add support for sib/tail calling calls whose calling convention differs from the caller's.

• Port over CCState::resultsCombatible from CallingConvLower.cpp into CallLowering. This is used to verify that the way the caller and callee CC handle incoming arguments matches up.

• Add CallLowering::analyzeCallResult. This is basically a port of CCState::AnalyzeCallResult, but using ArgInfo rather than ISD::InputArg.

• Add AArch64CallLowering::doCallerAndCalleePassArgsTheSameWay. This checks that the calling conventions are compatible, and that the caller and callee preserve the same registers.

For testing:

• Update call-translator-tail-call.ll to show that we can now handle mismatched calling conventions.

• Add a GISel line to tailcall-ccmismatch.ll to show that we will not tail call when the regmasks don't line up.

[aemerson]

LGTM.

llvm/include/llvm/CodeGen/GlobalISel/CallLowering.h

bool handleAssignments(MachineIRBuilder &MIRBuilder,
                          SmallVectorImpl<ArgInfo> &Args,
                          ValueHandler &Handler) const;
   bool handleAssignments(CCState &CCState,
                          SmallVectorImpl<CCValAssign> &ArgLocs,
                          MachineIRBuilder &MIRBuilder,
                          SmallVectorImpl<ArgInfo> &Args,
                          ValueHandler &Handler) const;
 
+  /// Analyze the return values of a call, incorporating info about the passed
+  /// values into \p CCState.
+  bool analyzeCallResult(CCState &CCState, SmallVectorImpl<ArgInfo> &Args,
+                         CCAssignFn &Fn) const;
+
+  /// \returns True if the calling convention for a callee and its caller pass
+  /// results in the same way. Typically used for tail call eligibility checks.
+  ///
+  /// \p Info is the CallLoweringInfo for the call.
+  /// \p MF is the MachineFunction for the caller.
+  /// \p InArgs contains the results of the call.
+  /// \p CalleeAssignFn is the CCAssignFn to be used for the callee.
+  /// \p CallerAssignFn is the CCAssignFn to be used for the caller.
+  bool resultsCompatible(CallLoweringInfo &Info, MachineFunction &MF,
+                         SmallVectorImpl<ArgInfo> &InArgs,
+                         CCAssignFn &CalleeAssignFn,
+                         CCAssignFn &CallerAssignFn) const;
+
 public:
   CallLowering(const TargetLowering *TLI) : TLI(TLI) {}
   virtual ~CallLowering() = default;
 
   /// \return true if the target is capable of handling swifterror values that
   /// have been promoted to a specified register. The extended versions of
   /// lowerReturn and lowerCall should be implemented.
   virtual bool supportSwiftError() const {

llvm/lib/CodeGen/GlobalISel/CallLowering.cpp

for (unsigned i = 0, e = Args.size(), j = 0; i != e; ++i, ++j) {
     } else {
       // FIXME: Support byvals and other weirdness
       return false;
     }
   }
   return true;
 }
 
+bool CallLowering::analyzeCallResult(CCState &CCState,
+                                     SmallVectorImpl<ArgInfo> &Args,
+                                     CCAssignFn &Fn) const {
+  for (unsigned i = 0, e = Args.size(); i < e; ++i) {
+    MVT VT = MVT::getVT(Args[i].Ty);
+    if (Fn(i, VT, VT, CCValAssign::Full, Args[i].Flags[0], CCState)) {
+      // Bail out on anything we can't handle.
+      LLVM_DEBUG(dbgs() << "Cannot analyze " << EVT(VT).getEVTString()
+                        << " (arg number = " << i << "\n");
+      return false;
+    }
+  }
+  return true;
+}
+
+bool CallLowering::resultsCompatible(CallLoweringInfo &Info,
+                                     MachineFunction &MF,
+                                     SmallVectorImpl<ArgInfo> &InArgs,
+                                     CCAssignFn &CalleeAssignFn,
+                                     CCAssignFn &CallerAssignFn) const {
+  const Function &F = MF.getFunction();
+  CallingConv::ID CalleeCC = Info.CallConv;
+  CallingConv::ID CallerCC = F.getCallingConv();
+
+  if (CallerCC == CalleeCC)
+    return true;
+
+  SmallVector<CCValAssign, 16> ArgLocs1;
+  CCState CCInfo1(CalleeCC, false, MF, ArgLocs1, F.getContext());
+  if (!analyzeCallResult(CCInfo1, InArgs, CalleeAssignFn))
+    return false;
+
+  SmallVector<CCValAssign, 16> ArgLocs2;
+  CCState CCInfo2(CallerCC, false, MF, ArgLocs2, F.getContext());
+  if (!analyzeCallResult(CCInfo2, InArgs, CallerAssignFn))
+    return false;
+
+  // We need the argument locations to match up exactly. If there's more in
+  // one than the other, then we are done.
+  if (ArgLocs1.size() != ArgLocs2.size())
+    return false;
+
+  // Make sure that each location is passed in exactly the same way.
+  for (unsigned i = 0, e = ArgLocs1.size(); i < e; ++i) {
+    const CCValAssign &Loc1 = ArgLocs1[i];
+    const CCValAssign &Loc2 = ArgLocs2[i];
+
+    // We need both of them to be the same. So if one is a register and one
+    // isn't, we're done.
+    if (Loc1.isRegLoc() != Loc2.isRegLoc())
+      return false;
+
+    if (Loc1.isRegLoc()) {
+      // If they don't have the same register location, we're done.
+      if (Loc1.getLocReg() != Loc2.getLocReg())
+        return false;
+
+      // They matched, so we can move to the next ArgLoc.
+      continue;
+    }
+
+    // Loc1 wasn't a RegLoc, so they both must be MemLocs. Check if they match.
+    if (Loc1.getLocMemOffset() != Loc2.getLocMemOffset())
+      return false;
+  }
+
+  return true;
+}
+
 Register CallLowering::ValueHandler::extendRegister(Register ValReg,
                                                     CCValAssign &VA) {
   LLT LocTy{VA.getLocVT()};
   if (LocTy.getSizeInBits() == MRI.getType(ValReg).getSizeInBits())
     return ValReg;
   switch (VA.getLocInfo()) {
   default: break;
   case CCValAssign::Full:

llvm/lib/Target/AArch64/AArch64CallLowering.h

public:
 
   bool lowerFormalArguments(MachineIRBuilder &MIRBuilder, const Function &F,
                             ArrayRef<ArrayRef<Register>> VRegs) const override;
 
   bool lowerCall(MachineIRBuilder &MIRBuilder,
                  CallLoweringInfo &Info) const override;
 
   /// Returns true if the call can be lowered as a tail call.
-  bool isEligibleForTailCallOptimization(MachineIRBuilder &MIRBuilder,
-                                         CallLoweringInfo &Info) const;
+  bool
+  isEligibleForTailCallOptimization(MachineIRBuilder &MIRBuilder,
+                                    CallLoweringInfo &Info,
+                                    SmallVectorImpl<ArgInfo> &InArgs) const;
 
   bool supportSwiftError() const override { return true; }
 
 private:
   using RegHandler = std::function<void(MachineIRBuilder &, Type *, unsigned,
                                         CCValAssign &)>;
 
   using MemHandler =
       std::function<void(MachineIRBuilder &, int, CCValAssign &)>;
 
   void splitToValueTypes(const ArgInfo &OrigArgInfo,
                          SmallVectorImpl<ArgInfo> &SplitArgs,
                          const DataLayout &DL, MachineRegisterInfo &MRI,
                          CallingConv::ID CallConv) const;
+
+  bool
+  doCallerAndCalleePassArgsTheSameWay(CallLoweringInfo &Info,
+                                      MachineFunction &MF,
+                                      SmallVectorImpl<ArgInfo> &InArgs) const;
 };
 
 } // end namespace llvm
 
 #endif // LLVM_LIB_TARGET_AARCH64_AARCH64CALLLOWERING_H

llvm/lib/Target/AArch64/AArch64CallLowering.cpp

static bool mayTailCallThisCC(CallingConv::ID CC) {
   case CallingConv::PreserveMost:
   case CallingConv::Swift:
     return true;
   default:
     return canGuaranteeTCO(CC);
   }
 }
 
+bool AArch64CallLowering::doCallerAndCalleePassArgsTheSameWay(
+    CallLoweringInfo &Info, MachineFunction &MF,
+    SmallVectorImpl<ArgInfo> &InArgs) const {
+  const Function &CallerF = MF.getFunction();
+  CallingConv::ID CalleeCC = Info.CallConv;
+  CallingConv::ID CallerCC = CallerF.getCallingConv();
+
+  // If the calling conventions match, then everything must be the same.
+  if (CalleeCC == CallerCC)
+    return true;
+
+  // Check if the caller and callee will handle arguments in the same way.
+  const AArch64TargetLowering &TLI = *getTLI<AArch64TargetLowering>();
+  CCAssignFn *CalleeAssignFn = TLI.CCAssignFnForCall(CalleeCC, Info.IsVarArg);
+  CCAssignFn *CallerAssignFn =
+      TLI.CCAssignFnForCall(CallerCC, CallerF.isVarArg());
+
+  if (!resultsCompatible(Info, MF, InArgs, *CalleeAssignFn, *CallerAssignFn))
+    return false;
+
+  // Make sure that the caller and callee preserve all of the same registers.
+  auto TRI = MF.getSubtarget<AArch64Subtarget>().getRegisterInfo();
+  const uint32_t *CallerPreserved = TRI->getCallPreservedMask(MF, CallerCC);
+  const uint32_t *CalleePreserved = TRI->getCallPreservedMask(MF, CalleeCC);
+  if (MF.getSubtarget<AArch64Subtarget>().hasCustomCallingConv()) {
+    TRI->UpdateCustomCallPreservedMask(MF, &CallerPreserved);
+    TRI->UpdateCustomCallPreservedMask(MF, &CalleePreserved);
+  }
+
+  return TRI->regmaskSubsetEqual(CallerPreserved, CalleePreserved);
+}
+
 bool AArch64CallLowering::isEligibleForTailCallOptimization(
-    MachineIRBuilder &MIRBuilder, CallLoweringInfo &Info) const {
+    MachineIRBuilder &MIRBuilder, CallLoweringInfo &Info,
+    SmallVectorImpl<ArgInfo> &InArgs) 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.
bool AArch64CallLowering::isEligibleForTailCallOptimization(
   // 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) {
+  // Look at the incoming values.
+  if (!doCallerAndCalleePassArgsTheSameWay(Info, MF, InArgs)) {
     LLVM_DEBUG(
         dbgs()
-        << "... Cannot tail call with mismatched calling conventions yet.\n");
+        << "... Caller and callee have incompatible calling conventions.\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");

 }
 
 bool AArch64CallLowering::lowerCall(MachineIRBuilder &MIRBuilder,
                                     CallLoweringInfo &Info) const {
   MachineFunction &MF = MIRBuilder.getMF();
   const Function &F = MF.getFunction();
   MachineRegisterInfo &MRI = MF.getRegInfo();
   auto &DL = F.getParent()->getDataLayout();
+  const AArch64TargetLowering &TLI = *getTLI<AArch64TargetLowering>();
 
   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;
   }
 
   if (Info.IsTailCall && MF.getTarget().Options.GuaranteedTailCallOpt) {
     // TODO: Until we lower all tail calls, we should fall back on this.
     LLVM_DEBUG(dbgs() << "Cannot handle -tailcallopt yet.\n");
     return false;
   }
 
   SmallVector<ArgInfo, 8> SplitArgs;
   for (auto &OrigArg : Info.OrigArgs) {
     splitToValueTypes(OrigArg, SplitArgs, DL, MRI, Info.CallConv);
     // AAPCS requires that we zero-extend i1 to 8 bits by the caller.
     if (OrigArg.Ty->isIntegerTy(1))
       SplitArgs.back().Flags[0].setZExt();
   }
 
-  bool IsSibCall =
-      Info.IsTailCall && isEligibleForTailCallOptimization(MIRBuilder, Info);
+  SmallVector<ArgInfo, 8> InArgs;
+  if (!Info.OrigRet.Ty->isVoidTy())
+    splitToValueTypes(Info.OrigRet, InArgs, DL, MRI, F.getCallingConv());
+
+  bool IsSibCall = Info.IsTailCall &&
+                   isEligibleForTailCallOptimization(MIRBuilder, Info, InArgs);
   if (IsSibCall)
     MF.getFrameInfo().setHasTailCall();
 
   // Find out which ABI gets to decide where things go.
-  const AArch64TargetLowering &TLI = *getTLI<AArch64TargetLowering>();
   CCAssignFn *AssignFnFixed =
       TLI.CCAssignFnForCall(Info.CallConv, /*IsVarArg=*/false);
   CCAssignFn *AssignFnVarArg =
       TLI.CCAssignFnForCall(Info.CallConv, /*IsVarArg=*/true);
 
   // If we have a sibling call, then we don't have to adjust the stack.
   // Otherwise, we need to adjust it.
   MachineInstrBuilder CallSeqStart;
bool AArch64CallLowering::lowerCall(MachineIRBuilder &MIRBuilder,
   // If we're tail calling, then we're the return from the block. So, we don't
   // want to copy anything.
   if (IsSibCall)
     return true;
 
   // Finally we can copy the returned value back into its virtual-register. In
   // symmetry with the arugments, the physical register must be an
   // implicit-define of the call instruction.
-  CCAssignFn *RetAssignFn = TLI.CCAssignFnForReturn(F.getCallingConv());
   if (!Info.OrigRet.Ty->isVoidTy()) {
-    SplitArgs.clear();
-
-    splitToValueTypes(Info.OrigRet, SplitArgs, DL, MRI, F.getCallingConv());
-
+    CCAssignFn *RetAssignFn = TLI.CCAssignFnForReturn(F.getCallingConv());
     CallReturnHandler Handler(MIRBuilder, MRI, MIB, RetAssignFn);
-    if (!handleAssignments(MIRBuilder, SplitArgs, Handler))
+    if (!handleAssignments(MIRBuilder, InArgs, Handler))
       return false;
   }
 
   if (Info.SwiftErrorVReg) {
     MIB.addDef(AArch64::X21, RegState::Implicit);
     MIRBuilder.buildCopy(Info.SwiftErrorVReg, Register(AArch64::X21));
   }
 
   CallSeqStart.addImm(Handler.StackSize).addImm(0);
   MIRBuilder.buildInstr(AArch64::ADJCALLSTACKUP)
       .addImm(Handler.StackSize)
       .addImm(0);
 
   return true;
 }

llvm/test/CodeGen/AArch64/GlobalISel/call-translator-tail-call.ll

define void @test_extern_weak() {
   ; DARWIN:   RET_ReallyLR
   ; WINDOWS-LABEL: name: test_extern_weak
   ; WINDOWS: bb.1 (%ir-block.0):
   ; WINDOWS:   TCRETURNdi @extern_weak_fn, 0, csr_aarch64_aapcs, implicit $sp
   tail call void @extern_weak_fn()
   ret void
 }
 
-; Right now, mismatched calling conventions should not be tail called.
-; TODO: Support this.
 declare fastcc void @fast_fn()
 define void @test_mismatched_caller() {
   ; COMMON-LABEL: name: test_mismatched_caller
   ; COMMON: bb.1 (%ir-block.0):
-  ; COMMON:   ADJCALLSTACKDOWN 0, 0, implicit-def $sp, implicit $sp
-  ; COMMON:   BL @fast_fn, csr_aarch64_aapcs, implicit-def $lr, implicit $sp
-  ; COMMON:   ADJCALLSTACKUP 0, 0, implicit-def $sp, implicit $sp
-  ; COMMON:   RET_ReallyLR
+  ; COMMON:   TCRETURNdi @fast_fn, 0, csr_aarch64_aapcs, implicit $sp
   tail call fastcc void @fast_fn()
   ret void
 }

llvm/test/CodeGen/AArch64/tailcall-ccmismatch.ll

 ; RUN: llc -o - %s | FileCheck %s
+; RUN: llc -global-isel -verify-machineinstrs -o - %s | FileCheck %s
 target triple="aarch64--"
 
 declare void @somefunc()
 define preserve_mostcc void @test_ccmismatch_notail() {
 ; Ensure that no tail call is used here, as the called function somefunc does
 ; not preserve enough registers for preserve_mostcc.
 ; CHECK-LABEL: test_ccmismatch_notail:
 ; CHECK-NOT: b somefunc