AMDGPU: Fix indirect tail calls

[arsenm]

Summary

Fix a selection error on uniform callees, and use a regular call if
divergent.

[rampitec]

LGTM

[arsenm]

70ab76a81b982354fea7f0cee85f7001ad30b2e6

[madhur13490]

llvm/test/CodeGen/AMDGPU/sibling-call.ll

282

Sorry for late comments. Nit - %c is unused. Do you need attribute on the signature?

301

Just wondering, why is this a divergent call? Because the pointer is incoming as a parameter?

[arsenm]

llvm/test/CodeGen/AMDGPU/sibling-call.ll

282

The attribute adds noinline in case we were ever using the force inline mode for some reason

301

Yes

llvm/lib/Target/AMDGPU/SIISelLowering.cpp

 bool SITargetLowering::isEligibleForTailCallOptimization(
     SDValue Callee, CallingConv::ID CalleeCC, bool IsVarArg,
     const SmallVectorImpl<ISD::OutputArg> &Outs,
     const SmallVectorImpl<SDValue> &OutVals,
     const SmallVectorImpl<ISD::InputArg> &Ins, SelectionDAG &DAG) const {
   if (!mayTailCallThisCC(CalleeCC))
     return false;
 
+  // For a divergent call target, we need to do a waterfall loop over the
+  // possible callees which precludes us from using a simple jump.
+  if (Callee->isDivergent())
+    return false;
+
   MachineFunction &MF = DAG.getMachineFunction();
   const Function &CallerF = MF.getFunction();
   CallingConv::ID CallerCC = CallerF.getCallingConv();
   const SIRegisterInfo *TRI = getSubtarget()->getRegisterInfo();
   const uint32_t *CallerPreserved = TRI->getCallPreservedMask(MF, CallerCC);
 
   // Kernels aren't callable, and don't have a live in return address so it
   // doesn't make sense to do a tail call with entry functions.

llvm/lib/Target/AMDGPU/SIInstructions.td

def SI_TCRETURN : SPseudoInstSI <(outs),
   let isReturn = 1;
   let isBarrier = 1;
   let UseNamedOperandTable = 1;
   let SchedRW = [WriteBranch];
   // TODO: Should really base this on the call target
   let isConvergent = 1;
 }
 
+// Handle selecting indirect tail calls
+def : GCNPat<
+  (AMDGPUtc_return i64:$src0, (i64 0), (i32 timm:$fpdiff)),
+  (SI_TCRETURN SReg_64:$src0, (i64 0), i32imm:$fpdiff)
+>;
 
 def ADJCALLSTACKUP : SPseudoInstSI<
   (outs), (ins i32imm:$amt0, i32imm:$amt1),
   [(callseq_start timm:$amt0, timm:$amt1)],
   "; adjcallstackup $amt0 $amt1"> {
   let Size = 8; // Worst case. (s_add_u32 + constant)
   let FixedSize = 1;
   let hasSideEffects = 1;

llvm/test/CodeGen/AMDGPU/sibling-call.ll

 entry:
   %alloca = alloca [16 x i32], align 4, addrspace(5)
   %gep = getelementptr inbounds [16 x i32], [16 x i32] addrspace(5)* %alloca, i32 0, i32 5
   store volatile i32 9, i32 addrspace(5)* %gep
   %ret = tail call fastcc i32 @i32_fastcc_i32_i32_a32i32(i32 %a, i32 %b, [32 x i32] zeroinitializer)
   ret i32 %ret
 }
 
+@func_ptr_gv = external unnamed_addr addrspace(4) constant i32(i32, i32)*, align 4
+
+; Do support tail calls with a uniform, but unknown, callee.
+; GCN-LABEL: {{^}}indirect_uniform_sibling_call_i32_fastcc_i32_i32:
+; GCN: s_load_dwordx2 [[GV_ADDR:s\[[0-9]+:[0-9]+\]]]
+; GCN: s_load_dwordx2 [[FUNC_PTR:s\[[0-9]+:[0-9]+\]]], [[GV_ADDR]]
+; GCN: s_setpc_b64 [[FUNC_PTR]]
+define hidden fastcc i32 @indirect_uniform_sibling_call_i32_fastcc_i32_i32(i32 %a, i32 %b, i32 %c) #1 {
+entry:
+  %func.ptr.load = load i32(i32, i32)*, i32(i32, i32)* addrspace(4)* @func_ptr_gv
+  %ret = tail call fastcc i32 %func.ptr.load(i32 %a, i32 %b)
+  ret i32 %ret
+}
+
+; We can't support a tail call to a divergent target. Use a waterfall
+; loop around a regular call
+; GCN-LABEL: {{^}}indirect_divergent_sibling_call_i32_fastcc_i32_i32:
+; GCN: v_readfirstlane_b32
+; GCN: v_readfirstlane_b32
+; GCN: s_and_saveexec_b64
+; GCN: s_swappc_b64
+; GCN: s_cbranch_execnz
+; GCN: s_setpc_b64
+define hidden fastcc i32 @indirect_divergent_sibling_call_i32_fastcc_i32_i32(i32(i32, i32)* %func.ptr, i32 %a, i32 %b, i32 %c) #1 {
+entry:
+  %add = add i32 %b, %c
+  %ret = tail call fastcc i32 %func.ptr(i32 %a, i32 %add)
+  ret i32 %ret
+}
+
 attributes #0 = { nounwind }
 attributes #1 = { nounwind noinline }