[Coroutines 1/2] Improve symmetric control transfer feature

[junparser]

Summary

As @rsmith find in https://bugs.llvm.org/show_bug.cgi?id=45130 that musttail attribute is not added to some of the resume intrinsics. These are some optimizations need to do here:

1. We only check llvm.coro.subfn.addr for now, however, the intrinsic can be replaced in CoroElide pass. Furthermore, we can support other function calls as long as they obey rules of musttail.

2. in simplifyTerminatorLeadingToRet, we only check switch and unconditional branch instructions while suspend switch instruction can be reduced as conditional branch. This also break the optimization.

This patch implements part 1, all of the rules in shouldBeMustTail are same as verifyMustTailCall

TestPlan: check-llvm

[modocache]

Looks good to me! Initially I was wondering about whether shouldBeMustTail ought to check whether the next instruction is a ret, or a bitcast that's used by a ret, but I see simplifyTerminatorLeadingToRet takes care of that check. Thanks for the patch!

llvm/lib/Transforms/Coroutines/CoroSplit.cpp

if (auto *BR = dyn_cast<BranchInst>(I)) {
         continue;
       }
     }
     return false;
   }
   return false;
 }
 
+// Check whether CI obeys the rules of musttail attribute.
+static bool shouldBeMustTail(const CallInst &CI, const Function &F) {
+  if (CI.isInlineAsm())
+    return false;
+
+  // Match prototypes and calling conventions of resume function.
+  FunctionType *CalleeTy = CI.getFunctionType();
+  if (!CalleeTy->getReturnType()->isVoidTy() || (CalleeTy->getNumParams() != 1))
+    return false;
+
+  Type *CalleeParmTy = CalleeTy->getParamType(0);
+  if (!CalleeParmTy->isPointerTy() ||
+      (CalleeParmTy->getPointerAddressSpace() != 0))
+    return false;
+
+  if (CI.getCallingConv() != F.getCallingConv())
+    return false;
+
+  // CI should not has any ABI-impacting function attributes.
+  static const Attribute::AttrKind ABIAttrs[] = {
+      Attribute::StructRet,  Attribute::ByVal,    Attribute::InAlloca,
+      Attribute::InReg,      Attribute::Returned, Attribute::SwiftSelf,
+      Attribute::SwiftError, Attribute::Alignment};
+  AttributeList Attrs = CI.getAttributes();
+  for (auto AK : ABIAttrs)
+    if (Attrs.hasParamAttribute(0, AK))
+      return false;
+
+  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) {
   bool changed = false;
 
   // Collect potential resume instructions.
   SmallVector<CallInst *, 4> Resumes;
   for (auto &I : instructions(F))
     if (auto *Call = dyn_cast<CallInst>(&I))
-      if (auto *CalledValue = Call->getCalledValue())
-        // CoroEarly pass replaced coro resumes with indirect calls to an
-        // address return by CoroSubFnInst intrinsic. See if it is one of those.
-        if (isa<CoroSubFnInst>(CalledValue->stripPointerCasts()))
-          Resumes.push_back(Call);
+      if (shouldBeMustTail(*Call, F))
 
   // Set musttail on those that are followed by a ret instruction.
   for (CallInst *Call : Resumes)
     if (simplifyTerminatorLeadingToRet(Call->getNextNode())) {
       Call->setTailCallKind(CallInst::TCK_MustTail);
       changed = true;
     }
 

llvm/test/Transforms/Coroutines/coro-split-musttail2.ll

+; Tests that coro-split will convert coro.resume followed by a suspend to a
+; musttail call.
+; RUN: opt < %s -coro-split -S | FileCheck %s
+; RUN: opt < %s -passes=coro-split -S | FileCheck %s
+
+define void @fakeresume1(i8*)  {
+entry:
+  ret void;
+}
+
+define void @fakeresume2(i64*)  {
+entry:
+  ret void;
+}
+
+define void @g() #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)
+  call fastcc void @fakeresume1(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)
+  call fastcc void @fakeresume2(i64* 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: @g(
+; CHECK-NOT: musttail call fastcc void @fakeresume1(i8* null)
+
+; Verify that in the resume part resume call is marked with musttail.
+; CHECK-LABEL: @g.resume(
+; CHECK: musttail call fastcc void @fakeresume2(i64* null)
+; CHECK-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 = { "coroutine.presplit"="1" }
+attributes #1 = { argmemonly nounwind readonly }
+attributes #2 = { nounwind }
+attributes #3 = { nounwind readnone }