ARM, AArch64, X86: Check preserved registers for tail calls.

[MatzeB]

Summary

We can only perform a tail call to a callee that preserves all the
registers that the caller needs to preserve.

This situation happens with calling conventions like preserver_mostcc or
cxx_fast_tls. It was explicitely handled for the fast_tls case and failing for preserve_most. This patch generalizes
the check to any calling convention.

Related to rdar://24207743

[qcolombet]

lib/CodeGen/TargetRegisterInfo.cpp

397 ↗

(On Diff #52323)

No braces.

lib/Target/AArch64/AArch64ISelLowering.cpp

2882 ↗

(On Diff #52323)

Isn’t sufficient to be inclusive?
I.e., CallerCC C= CalleeCC

lib/Target/ARM/ARMISelLowering.cpp

2155 ↗

(On Diff #52323)

Ditto

test/CodeGen/AArch64/tailcall-ccmismatch.ll

5 ↗

(On Diff #52323)

Add a comment on what you are testing here.

test/CodeGen/ARM/cxx-tlscc.ll

130 ↗

(On Diff #52323)

Ditto

[MatzeB]

Thanks for the review! Addressed the comments as discussed.

[MatzeB]

upload the correct patch this time.

[rengolin]

Overall, I think the change makes sense in that it should keep the original intention intact and make it more generic and not less safe (but I'm not an expert).

lib/CodeGen/TargetRegisterInfo.cpp

395 ↗

(On Diff #52454)

IIUC, you need to check that all reserved registers on the left are also on the right, but not necessarily all on the right to be on the left, right? In that case, this check would be too restrictive, mostly matching recursive tail calls only, no?

[MatzeB]

lib/CodeGen/TargetRegisterInfo.cpp

395 ↗

(On Diff #52454)

That's what the check is doing, isn't it?

Bit0	Bit1	Bit0&Bit1 == Bit0
0	0/1	1
1	0	0
1	1	1

Is there a clearer way to write the code?

[rengolin]

lib/CodeGen/TargetRegisterInfo.cpp

395 ↗

(On Diff #52454)

Dang! It is, sorry, brain fart.

[qcolombet]

LGTM.

llvm/trunk/include/llvm/Target/TargetRegisterInfo.h

virtual const uint32_t *getCallPreservedMask(const MachineFunction &MF,
     return nullptr;
   }
 
   /// Return a register mask that clobbers everything.
   virtual const uint32_t *getNoPreservedMask() const {
     llvm_unreachable("target does not provide no preserved mask");
   }
 
+  /// Return true if all bits that are set in mask \p mask0 are also set in
+  /// \p mask1.
+  bool regmaskSubsetEqual(const uint32_t *mask0, const uint32_t *mask1) const;
+
   /// Return all the call-preserved register masks defined for this target.
   virtual ArrayRef<const uint32_t *> getRegMasks() const = 0;
   virtual ArrayRef<const char *> getRegMaskNames() const = 0;
 
   /// Returns a bitset indexed by physical register number indicating if a
   /// register is a special register that has particular uses and should be
   /// considered unavailable at all times, e.g. SP, RA. This is
   /// used by register scavenger to determine what registers are free.

llvm/trunk/lib/CodeGen/TargetRegisterInfo.cpp

bool TargetRegisterInfo::needsStackRealignment(
   if (MF.getFunction()->hasFnAttribute("stackrealign") || requiresRealignment) {
     if (canRealignStack(MF))
       return true;
     DEBUG(dbgs() << "Can't realign function's stack: " << F->getName() << "\n");
   }
   return false;
 }
 
+bool TargetRegisterInfo::regmaskSubsetEqual(const uint32_t *mask0,
+                                            const uint32_t *mask1) const {
+  unsigned N = (getNumRegs()+31) / 32;
+  for (unsigned I = 0; I < N; ++I)
+    if ((mask0[I] & mask1[I]) != mask0[I])
+      return false;
+  return true;
+}
+
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
 void
 TargetRegisterInfo::dumpReg(unsigned Reg, unsigned SubRegIndex,
                             const TargetRegisterInfo *TRI) {
   dbgs() << PrintReg(Reg, TRI, SubRegIndex) << "\n";
 }
 #endif

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

bool AArch64TargetLowering::isEligibleForTailCallOptimization(
   if (!IsTailCallConvention(CalleeCC) && CalleeCC != CallingConv::C)
     return false;
 
   MachineFunction &MF = DAG.getMachineFunction();
   const Function *CallerF = MF.getFunction();
   CallingConv::ID CallerCC = CallerF->getCallingConv();
   bool CCMatch = CallerCC == CalleeCC;
 
-  // Disable tailcall for CXX_FAST_TLS when callee and caller have different
-  // calling conventions, given that CXX_FAST_TLS has a bigger CSR set.
-  if (!CCMatch &&
-      (CallerCC == CallingConv::CXX_FAST_TLS ||
-       CalleeCC == CallingConv::CXX_FAST_TLS))
-    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) but less efficient and uglier in LowerCall.
   for (Function::const_arg_iterator i = CallerF->arg_begin(),
                                     e = CallerF->arg_end();
        i != e; ++i)
     if (i->hasByValAttr())
       return false;
for (const CCValAssign &ArgLoc : ArgLocs)
         return false;
   }
 
   // Check that the call results are passed in the same way.
   if (!CCState::resultsCompatible(CalleeCC, CallerCC, MF, C, Ins,
                                   CCAssignFnForCall(CalleeCC, isVarArg),
                                   CCAssignFnForCall(CallerCC, isVarArg)))
     return false;
+  // The callee has to preserve all registers the caller needs to preserve.
+  if (!CCMatch) {
+    const AArch64RegisterInfo *TRI = Subtarget->getRegisterInfo();
+    if (!TRI->regmaskSubsetEqual(TRI->getCallPreservedMask(MF, CallerCC),
+                                 TRI->getCallPreservedMask(MF, CalleeCC)))
+      return false;
+  }
 
   // Nothing more to check if the callee is taking no arguments
   if (Outs.empty())
     return true;
 
   SmallVector<CCValAssign, 16> ArgLocs;
   CCState CCInfo(CalleeCC, isVarArg, MF, ArgLocs, C);
 

llvm/trunk/lib/Target/ARM/ARMISelLowering.cpp

ARMTargetLowering::IsEligibleForTailCallOptimization(SDValue Callee,
                                                      bool isCallerStructRet,
                                     const SmallVectorImpl<ISD::OutputArg> &Outs,
                                     const SmallVectorImpl<SDValue> &OutVals,
                                     const SmallVectorImpl<ISD::InputArg> &Ins,
                                                      SelectionDAG& DAG) const {
   MachineFunction &MF = DAG.getMachineFunction();
   const Function *CallerF = MF.getFunction();
   CallingConv::ID CallerCC = CallerF->getCallingConv();
-  bool CCMatch = CallerCC == CalleeCC;
-
-  // Disable tailcall for CXX_FAST_TLS when callee and caller have different
-  // calling conventions, given that CXX_FAST_TLS has a bigger CSR set.
-  if (!CCMatch &&
-      (CallerCC == CallingConv::CXX_FAST_TLS ||
-       CalleeCC == CallingConv::CXX_FAST_TLS))
-    return false;
 
   assert(Subtarget->supportsTailCall());
 
   // Look for obvious safe cases to perform tail call optimization that do not
   // require ABI changes. This is what gcc calls sibcall.
 
   // Do not sibcall optimize vararg calls unless the call site is not passing
   // any arguments.
ARMTargetLowering::IsEligibleForTailCallOptimization(SDValue Callee,
   }
 
   // Check that the call results are passed in the same way.
   LLVMContext &C = *DAG.getContext();
   if (!CCState::resultsCompatible(CalleeCC, CallerCC, MF, C, Ins,
                                   CCAssignFnForNode(CalleeCC, true, isVarArg),
                                   CCAssignFnForNode(CallerCC, true, isVarArg)))
     return false;
+  // The callee has to preserve all registers the caller needs to preserve.
+  if (CalleeCC != CallerCC) {
+    const ARMBaseRegisterInfo *TRI = Subtarget->getRegisterInfo();
+    if (!TRI->regmaskSubsetEqual(TRI->getCallPreservedMask(MF, CallerCC),
+                                 TRI->getCallPreservedMask(MF, CalleeCC)))
+      return false;
+  }
 
   // If Caller's vararg or byval argument has been split between registers and
   // stack, do not perform tail call, since part of the argument is in caller's
   // local frame.
   const ARMFunctionInfo *AFI_Caller = MF.getInfo<ARMFunctionInfo>();
   if (AFI_Caller->getArgRegsSaveSize())
     return false;
 

llvm/trunk/lib/Target/X86/X86ISelLowering.cpp

bool X86TargetLowering::IsEligibleForTailCallOptimization(
   bool IsCallerWin64 = Subtarget.isCallingConvWin64(CallerCC);
 
   // Win64 functions have extra shadow space for argument homing. Don't do the
   // sibcall if the caller and callee have mismatched expectations for this
   // space.
   if (IsCalleeWin64 != IsCallerWin64)
     return false;
 
-  // Disable tailcall for CXX_FAST_TLS when callee and caller have different
-  // calling conventions, given that CXX_FAST_TLS has a bigger CSR set.
-  if (!CCMatch &&
-      (CallerCC == CallingConv::CXX_FAST_TLS ||
-       CalleeCC == CallingConv::CXX_FAST_TLS))
-    return false;
-
   if (DAG.getTarget().Options.GuaranteedTailCallOpt) {
     if (canGuaranteeTCO(CalleeCC) && CCMatch)
       return true;
     return false;
   }
 
   // Look for obvious safe cases to perform tail call optimization that do not
   // require ABI changes. This is what gcc calls sibcall.
for (unsigned i = 0, e = RVLocs.size(); i != e; ++i) {
         return false;
     }
   }
 
   // Check that the call results are passed in the same way.
   if (!CCState::resultsCompatible(CalleeCC, CallerCC, MF, C, Ins,
                                   RetCC_X86, RetCC_X86))
     return false;
+  // The callee has to preserve all registers the caller needs to preserve.
+  if (!CCMatch) {
+    const X86RegisterInfo *TRI = Subtarget.getRegisterInfo();
+    if (!TRI->regmaskSubsetEqual(TRI->getCallPreservedMask(MF, CallerCC),
+                                 TRI->getCallPreservedMask(MF, CalleeCC)))
+      return false;
+  }
 
   unsigned StackArgsSize = 0;
 
   // If the callee takes no arguments then go on to check the results of the
   // call.
   if (!Outs.empty()) {
     // Check if stack adjustment is needed. For now, do not do this if any
     // argument is passed on the stack.

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

+; RUN: llc -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
+; CHECK: bl somefunc
+  tail call void @somefunc()
+  ret void
+}
+
+declare preserve_mostcc void @some_preserve_most_func()
+define void @test_ccmismatch_tail() {
+; We can perform a tail call here, because some_preserve_most_func preserves
+; all registers necessary for test_ccmismatch_tail.
+; CHECK-LABEL: test_ccmismatch_tail:
+; CHECK-NOT: bl some_preserve_most_func
+; CHECK: b some_preserve_most_func
+  tail call preserve_mostcc void @some_preserve_most_func()
+  ret void
+}

llvm/trunk/test/CodeGen/ARM/cxx-tlscc.ll

 ; CHECK: bl __tlv_atexit
 define cxx_fast_tlscc void @__tls_test() {
 entry:
   store i32 0, i32* getelementptr inbounds (%class.C, %class.C* @tC, i64 0, i32 0), align 4
   %0 = tail call i32 @_tlv_atexit(void (i8*)* bitcast (%class.C* (%class.C*)* @_ZN1CD1Ev to void (i8*)*), i8* bitcast (%class.C* @tC to i8*), i8* nonnull @__dso_handle) #1
   ret void
 }
 
+declare void @somefunc()
+define cxx_fast_tlscc void @test_ccmismatch_notail() {
+; A tail call is not possible here because somefunc does not preserve enough
+; registers.
+; CHECK-LABEL: test_ccmismatch_notail:
+; CHECK-NOT: b _somefunc
+; CHECK: bl _somefunc
+  tail call void @somefunc()
+  ret void
+}
+
+declare cxx_fast_tlscc void @some_fast_tls_func()
+define void @test_ccmismatch_tail() {
+; We can perform a tail call here because some_fast_tls_func preserves all
+; necessary registers (and more).
+; CHECK-LABEL: test_ccmismatch_tail:
+; CHECK-NOT: bl _some_fast_tls_func
+; CHECK: b _some_fast_tls_func
+  tail call cxx_fast_tlscc void @some_fast_tls_func()
+  ret void
+}
+
 attributes #0 = { nounwind "no-frame-pointer-elim"="true" }
 attributes #1 = { nounwind }