[PowerPC][Coroutines] Add tail-call check with context information for coroutines

[tingwang]

Summary

This is an attempt to solve https://github.com/llvm/llvm-project/issues/56679 by rejecting musttail in CoroSplitPass on ppc conditionally based on CallInst and target configuration.

TailCallOptimization is conditionally eligible on PPC as details explained in PPCTargetLowering::IsEligibleForTailCallOptimization_64SVR4(). The check happens during instruction selection.

Coroutines need to know the eligibility during transformation. From the CallInst provided, ppc can decide if indirect call can be turned into tail call.

[tingwang]

Test case clang/test/CodeGenCoroutines/pr56329.cpp with this patch will not generate musttail on ppc.

[ChuanqiXu]

In D131953#3725578, @tingwang wrote:

Test case clang/test/CodeGenCoroutines/pr56329.cpp with this patch will not generate musttail on ppc.

So will pr56329.cpp fail after this patch? I guess we need REQUIRE/XFAIL clause to pr56329.cpp for pcc.

llvm/include/llvm/Analysis/TargetTransformInfo.h

769

llvm/lib/Transforms/Coroutines/CoroSplit.cpp

1378–1379

[tingwang]

Update according to comments.

[ChuanqiXu]

LGTM with comment addressed. Thanks!

llvm/include/llvm/Analysis/TargetTransformInfoImpl.h

346

The default implementation should return supportsTailCalls().

[tingwang]

For default implementation of supportsTailCallFor, return supportsTailCalls() as suggested in comments.

[tingwang]

In D131953#3727900, @ChuanqiXu wrote:

LGTM with comment addressed. Thanks!

Thank you! I'm looking forward to comments from ppc.

[shchenz]

This LGTM. Although supportsTailCallFor is not completely equal to how PPC backend checks tail call support, we can continuously improve this when we have more caller of supportsTailCallFor() in future.

clang/test/CodeGenCoroutines/pr56329.cpp

 // Test for PR56919. Tests the we won't contain the resumption of final suspend point.
 //
 // RUN: %clang_cc1 -triple %itanium_abi_triple -std=c++20 %s -O3 -S -emit-llvm -o - | FileCheck %s
+// This test is expected to fail on PowerPC.
+// XFAIL: powerpc
 
 #include "Inputs/coroutine.h"
 
 void _exit(int status) __attribute__ ((__noreturn__));
 
 class Promise;
 
 // An object that can be co_awaited, but we always resume immediately from

llvm/include/llvm/Analysis/TargetTransformInfo.h

public:
   /// If target has efficient vector element load/store instructions, it can
   /// return true here so that insertion/extraction costs are not added to
   /// the scalarization cost of a load/store.
   bool supportsEfficientVectorElementLoadStore() const;
 
   /// If the target supports tail calls.
   bool supportsTailCalls() const;
 
+  /// If target supports tail call on \p CB
+  bool supportsTailCallFor(const CallBase *CB) const;
+
   /// Don't restrict interleaved unrolling to small loops.
   bool enableAggressiveInterleaving(bool LoopHasReductions) const;
 
   /// Returns options for expansion of memcmp. IsZeroCmp is
   // true if this is the expansion of memcmp(p1, p2, s) == 0.
   struct MemCmpExpansionOptions {
     // Return true if memcmp expansion is enabled.
     operator bool() const { return MaxNumLoads > 0; }
virtual InstructionCost getScalarizationOverhead(VectorType *Ty,
                                                    const APInt &DemandedElts,
                                                    bool Insert,
                                                    bool Extract) = 0;
   virtual InstructionCost
   getOperandsScalarizationOverhead(ArrayRef<const Value *> Args,
                                    ArrayRef<Type *> Tys) = 0;
   virtual bool supportsEfficientVectorElementLoadStore() = 0;
   virtual bool supportsTailCalls() = 0;
+  virtual bool supportsTailCallFor(const CallBase *CB) = 0;
   virtual bool enableAggressiveInterleaving(bool LoopHasReductions) = 0;
   virtual MemCmpExpansionOptions
   enableMemCmpExpansion(bool OptSize, bool IsZeroCmp) const = 0;
   virtual bool enableInterleavedAccessVectorization() = 0;
   virtual bool enableMaskedInterleavedAccessVectorization() = 0;
   virtual bool isFPVectorizationPotentiallyUnsafe() = 0;
   virtual bool allowsMisalignedMemoryAccesses(LLVMContext &Context,
                                               unsigned BitWidth,
getOperandsScalarizationOverhead(ArrayRef<const Value *> Args,
     return Impl.getOperandsScalarizationOverhead(Args, Tys);
   }
 
   bool supportsEfficientVectorElementLoadStore() override {
     return Impl.supportsEfficientVectorElementLoadStore();
   }
 
   bool supportsTailCalls() override { return Impl.supportsTailCalls(); }
+  bool supportsTailCallFor(const CallBase *CB) override {
+    return Impl.supportsTailCallFor(CB);
+  }
 
   bool enableAggressiveInterleaving(bool LoopHasReductions) override {
     return Impl.enableAggressiveInterleaving(LoopHasReductions);
   }
   MemCmpExpansionOptions enableMemCmpExpansion(bool OptSize,
                                                bool IsZeroCmp) const override {
     return Impl.enableMemCmpExpansion(OptSize, IsZeroCmp);
   }

llvm/include/llvm/Analysis/TargetTransformInfoImpl.h

InstructionCost getOperandsScalarizationOverhead(ArrayRef<const Value *> Args,
                                                    ArrayRef<Type *> Tys) const {
     return 0;
   }
 
   bool supportsEfficientVectorElementLoadStore() const { return false; }
 
   bool supportsTailCalls() const { return true; }
 
+  bool supportsTailCallFor(const CallBase *CB) const {
+    return supportsTailCalls();
+  }
+
   bool enableAggressiveInterleaving(bool LoopHasReductions) const {
     return false;
   }
 
   TTI::MemCmpExpansionOptions enableMemCmpExpansion(bool OptSize,
                                                     bool IsZeroCmp) const {
     return {};
   }

llvm/lib/Analysis/TargetTransformInfo.cpp

 bool TargetTransformInfo::supportsEfficientVectorElementLoadStore() const {
   return TTIImpl->supportsEfficientVectorElementLoadStore();
 }
 
 bool TargetTransformInfo::supportsTailCalls() const {
   return TTIImpl->supportsTailCalls();
 }
 
+bool TargetTransformInfo::supportsTailCallFor(const CallBase *CB) const {
+  return TTIImpl->supportsTailCallFor(CB);
+}
+
 bool TargetTransformInfo::enableAggressiveInterleaving(
     bool LoopHasReductions) const {
   return TTIImpl->enableAggressiveInterleaving(LoopHasReductions);
 }
 
 TargetTransformInfo::MemCmpExpansionOptions
 TargetTransformInfo::enableMemCmpExpansion(bool OptSize, bool IsZeroCmp) const {
   return TTIImpl->enableMemCmpExpansion(OptSize, IsZeroCmp);

llvm/lib/Target/PowerPC/PPCTargetTransformInfo.h

public:
   bool areTypesABICompatible(const Function *Caller, const Function *Callee,
                              const ArrayRef<Type *> &Types) const;
   bool hasActiveVectorLength(unsigned Opcode, Type *DataType,
                              Align Alignment) const;
   InstructionCost getVPMemoryOpCost(unsigned Opcode, Type *Src, Align Alignment,
                                     unsigned AddressSpace,
                                     TTI::TargetCostKind CostKind,
                                     const Instruction *I = nullptr);
+  bool supportsTailCallFor(const CallBase *CB) const;
 
 private:
   // The following constant is used for estimating costs on power9.
   static const InstructionCost::CostType P9PipelineFlushEstimate = 80;
 
   /// @}
 };
 
 } // end namespace llvm
 
 #endif

llvm/lib/Target/PowerPC/PPCTargetTransformInfo.cpp

return (MisalignmentProb * P9PipelineFlushEstimate) +
            (AlignmentProb * *Cost.getValue());
   }
 
   // Usually we should not get to this point, but the following is an attempt to
   // model the cost of legalization. Currently we can only lower intrinsics with
   // evl but no mask, on Power 9/10. Otherwise, we must scalarize.
   return getMaskedMemoryOpCost(Opcode, Src, Alignment, AddressSpace, CostKind);
 }
+
+bool PPCTTIImpl::supportsTailCallFor(const CallBase *CB) const {
+  // Subtargets using PC-Relative addressing supported.
+  if (ST->isUsingPCRelativeCalls())
+    return true;
+
+  const Function *Callee = CB->getCalledFunction();
+  // Indirect calls and variadic argument functions not supported.
+  if (!Callee || Callee->isVarArg())
+    return false;
+
+  const Function *Caller = CB->getCaller();
+  // Support if we can share TOC base.
+  return ST->getTargetMachine().shouldAssumeDSOLocal(*Caller->getParent(),
+                                                     Callee);
+}

llvm/lib/Transforms/Coroutines/CoroSplit.cpp

static bool shouldBeMustTail(const CallInst &CI, const Function &F) {
   return true;
 }
 
 // Add musttail to any resume instructions that is immediately followed by a
 // suspend (i.e. ret). We do this even in -O0 to support guaranteed tail call
 // for symmetrical coroutine control transfer (C++ Coroutines TS extension).
 // This transformation is done only in the resume part of the coroutine that has
 // identical signature and calling convention as the coro.resume call.
-static void addMustTailToCoroResumes(Function &F) {
+static void addMustTailToCoroResumes(Function &F, TargetTransformInfo &TTI) {
   bool changed = false;
 
   // Collect potential resume instructions.
   SmallVector<CallInst *, 4> Resumes;
   for (auto &I : instructions(F))
     if (auto *Call = dyn_cast<CallInst>(&I))
       if (shouldBeMustTail(*Call, F))
         Resumes.push_back(Call);
 
   // Set musttail on those that are followed by a ret instruction.
   for (CallInst *Call : Resumes)
-    if (simplifyTerminatorLeadingToRet(Call->getNextNode())) {
+    // Skip targets which don't support tail call on the specific case.
+    if (TTI.supportsTailCallFor(Call) &&
+        simplifyTerminatorLeadingToRet(Call->getNextNode())) {
       Call->setTailCallKind(CallInst::TCK_MustTail);
       changed = true;
     }
 
   if (changed)
     removeUnreachableBlocks(F);
 }
 
static void splitSwitchCoroutine(Function &F, coro::Shape &Shape,
   postSplitCleanup(*CleanupClone);
 
   // Adding musttail call to support symmetric transfer.
   // Skip targets which don't support tail call.
   //
   // FIXME: Could we support symmetric transfer effectively without musttail
   // call?
   if (TTI.supportsTailCalls())
-    addMustTailToCoroResumes(*ResumeClone);
+    addMustTailToCoroResumes(*ResumeClone, TTI);
 
   // Store addresses resume/destroy/cleanup functions in the coroutine frame.
   updateCoroFrame(Shape, ResumeClone, DestroyClone, CleanupClone);
 
   assert(Clones.empty());
   Clones.push_back(ResumeClone);
   Clones.push_back(DestroyClone);
   Clones.push_back(CleanupClone);

llvm/test/Transforms/Coroutines/coro-split-musttail-ppc64le.ll

+; Tests that some target (e.g. ppc) can support tail call under condition.
+; RUN: opt < %s -passes='cgscc(coro-split),simplifycfg,early-cse' -S \
+; RUN:     -mtriple=powerpc64le-unknown-linux-gnu -mcpu=pwr9 | FileCheck %s
+; RUN: opt < %s -passes='cgscc(coro-split),simplifycfg,early-cse' -S \
+; RUN:     -mtriple=powerpc64le-unknown-linux-gnu -mcpu=pwr10 --code-model=medium \
+; RUN:     | FileCheck %s --check-prefix=CHECK-PCREL
+
+define void @f() #0 {
+entry:
+  %id = call token @llvm.coro.id(i32 0, i8* null, i8* null, i8* null)
+  %alloc = call i8* @malloc(i64 16) #3
+  %vFrame = call noalias nonnull i8* @llvm.coro.begin(token %id, i8* %alloc)
+
+  %save = call token @llvm.coro.save(i8* null)
+  %addr1 = call i8* @llvm.coro.subfn.addr(i8* null, i8 0)
+  %pv1 = bitcast i8* %addr1 to void (i8*)*
+  call fastcc void %pv1(i8* null)
+
+  %suspend = call i8 @llvm.coro.suspend(token %save, i1 false)
+  switch i8 %suspend, label %exit [
+    i8 0, label %await.ready
+    i8 1, label %exit
+  ]
+await.ready:
+  %save2 = call token @llvm.coro.save(i8* null)
+  %addr2 = call i8* @llvm.coro.subfn.addr(i8* null, i8 0)
+  %pv2 = bitcast i8* %addr2 to void (i8*)*
+  call fastcc void %pv2(i8* null)
+
+  %suspend2 = call i8 @llvm.coro.suspend(token %save2, i1 false)
+  switch i8 %suspend2, label %exit [
+    i8 0, label %exit
+    i8 1, label %exit
+  ]
+exit:
+  call i1 @llvm.coro.end(i8* null, i1 false)
+  ret void
+}
+
+; Verify that in the initial function resume is not marked with musttail.
+; CHECK-LABEL: @f(
+; CHECK: %[[addr1:.+]] = call i8* @llvm.coro.subfn.addr(i8* null, i8 0)
+; CHECK-NEXT: %[[pv1:.+]] = bitcast i8* %[[addr1]] to void (i8*)*
+; CHECK-NOT: musttail call fastcc void %[[pv1]](i8* null)
+
+; Verify that ppc target not using PC-Relative addressing in the resume part resume call is not marked with musttail.
+; CHECK-LABEL: @f.resume(
+; CHECK: %[[addr2:.+]] = call i8* @llvm.coro.subfn.addr(i8* null, i8 0)
+; CHECK-NEXT: %[[pv2:.+]] = bitcast i8* %[[addr2]] to void (i8*)*
+; CHECK-NEXT: call fastcc void %[[pv2]](i8* null)
+
+; Verify that ppc target using PC-Relative addressing in the resume part resume call is marked with musttail.
+; CHECK-PCREL-LABEL: @f.resume(
+; CHECK-PCREL: %[[addr2:.+]] = call i8* @llvm.coro.subfn.addr(i8* null, i8 0)
+; CHECK-PCREL-NEXT: %[[pv2:.+]] = bitcast i8* %[[addr2]] to void (i8*)*
+; CHECK-PCREL-NEXT: musttail call fastcc void %[[pv2]](i8* null)
+; CHECK-PCREL-NEXT: ret void
+
+declare token @llvm.coro.id(i32, i8* readnone, i8* nocapture readonly, i8*) #1
+declare i1 @llvm.coro.alloc(token) #2
+declare i64 @llvm.coro.size.i64() #3
+declare i8* @llvm.coro.begin(token, i8* writeonly) #2
+declare token @llvm.coro.save(i8*) #2
+declare i8* @llvm.coro.frame() #3
+declare i8 @llvm.coro.suspend(token, i1) #2
+declare i8* @llvm.coro.free(token, i8* nocapture readonly) #1
+declare i1 @llvm.coro.end(i8*, i1) #2
+declare i8* @llvm.coro.subfn.addr(i8* nocapture readonly, i8) #1
+declare i8* @malloc(i64)
+
+attributes #0 = { presplitcoroutine }
+attributes #1 = { argmemonly nounwind readonly }
+attributes #2 = { nounwind }
+attributes #3 = { nounwind readnone }