[Coro][NewPM] Handle llvm.coro.prepare.retcon in NPM coro-split pass

[junparser]

what does this intrinsic use for?

[aeubanks]

In D87731#2275664, @junparser wrote:

what does this intrinsic use for?

I'm not actually familiar with coroutines, this is a blocker for the NPM switch that I took a look at and attempted to fix.

[rjmccall]

Thanks, this looks right.

llvm/lib/Transforms/Coroutines/CoroSplit.cpp

static void createDevirtTriggerFunc(CallGraph &CG, CallGraphSCC &SCC) {
   auto *Node = CG.getOrInsertFunction(DevirtFn);
 
   SmallVector<CallGraphNode *, 8> Nodes(SCC.begin(), SCC.end());
   Nodes.push_back(Node);
   SCC.initialize(Nodes);
 }
 
 /// Replace a call to llvm.coro.prepare.retcon.
+static void replacePrepare(CallInst *Prepare, LazyCallGraph &CG,
+                           LazyCallGraph::SCC &C) {
+  auto CastFn = Prepare->getArgOperand(0); // as an i8*
+  auto Fn = CastFn->stripPointerCasts();   // as its original type
+
+  // Attempt to peephole this pattern:
+  //    %0 = bitcast [[TYPE]] @some_function to i8*
+  //    %1 = call @llvm.coro.prepare.retcon(i8* %0)
+  //    %2 = bitcast %1 to [[TYPE]]
+  // ==>
+  //    %2 = @some_function
+  for (auto UI = Prepare->use_begin(), UE = Prepare->use_end(); UI != UE;) {
+    // Look for bitcasts back to the original function type.
+    auto *Cast = dyn_cast<BitCastInst>((UI++)->getUser());
+    if (!Cast || Cast->getType() != Fn->getType())
+      continue;
+
+    // Replace and remove the cast.
+    Cast->replaceAllUsesWith(Fn);
+    Cast->eraseFromParent();
+  }
+
+  // Replace any remaining uses with the function as an i8*.
+  // This can never directly be a callee, so we don't need to update CG.
+  Prepare->replaceAllUsesWith(CastFn);
+  Prepare->eraseFromParent();
+
+  // Kill dead bitcasts.
+  while (auto *Cast = dyn_cast<BitCastInst>(CastFn)) {
+    if (!Cast->use_empty())
+      break;
+    CastFn = Cast->getOperand(0);
+    Cast->eraseFromParent();
+  }
+}
+/// Replace a call to llvm.coro.prepare.retcon.
 static void replacePrepare(CallInst *Prepare, CallGraph &CG) {
   auto CastFn = Prepare->getArgOperand(0); // as an i8*
   auto Fn = CastFn->stripPointerCasts(); // as its original type
 
   // Find call graph nodes for the preparation.
   CallGraphNode *PrepareUserNode = nullptr, *FnNode = nullptr;
   if (auto ConcreteFn = dyn_cast<Function>(Fn)) {
     PrepareUserNode = CG[Prepare->getFunction()];
static void replacePrepare(CallInst *Prepare, CallGraph &CG) {
   // Kill dead bitcasts.
   while (auto *Cast = dyn_cast<BitCastInst>(CastFn)) {
     if (!Cast->use_empty()) break;
     CastFn = Cast->getOperand(0);
     Cast->eraseFromParent();
   }
 }
 
+static bool replaceAllPrepares(Function *PrepareFn, LazyCallGraph &CG,
+                               LazyCallGraph::SCC &C) {
+  bool Changed = false;
+  for (auto PI = PrepareFn->use_begin(), PE = PrepareFn->use_end(); PI != PE;) {
+    // Intrinsics can only be used in calls.
+    auto *Prepare = cast<CallInst>((PI++)->getUser());
+    replacePrepare(Prepare, CG, C);
+    Changed = true;
+  }
+
+  return Changed;
+}
+
 /// Remove calls to llvm.coro.prepare.retcon, a barrier meant to prevent
 /// IPO from operating on calls to a retcon coroutine before it's been
 /// split.  This is only safe to do after we've split all retcon
 /// coroutines in the module.  We can do that this in this pass because
 /// this pass does promise to split all retcon coroutines (as opposed to
 /// switch coroutines, which are lowered in multiple stages).
 static bool replaceAllPrepares(Function *PrepareFn, CallGraph &CG) {
   bool Changed = false;
PreservedAnalyses CoroSplitPass::run(LazyCallGraph::SCC &C,
   Module &M = *C.begin()->getFunction().getParent();
   auto &FAM =
       AM.getResult<FunctionAnalysisManagerCGSCCProxy>(C, CG).getManager();
 
   if (!declaresCoroSplitIntrinsics(M))
     return PreservedAnalyses::all();
 
   // Check for uses of llvm.coro.prepare.retcon.
-  const auto *PrepareFn = M.getFunction("llvm.coro.prepare.retcon");
+  auto *PrepareFn = M.getFunction("llvm.coro.prepare.retcon");
   if (PrepareFn && PrepareFn->use_empty())
     PrepareFn = nullptr;
 
   // Find coroutines for processing.
   SmallVector<LazyCallGraph::Node *, 4> Coroutines;
   for (LazyCallGraph::Node &N : C)
     if (N.getFunction().hasFnAttribute(CORO_PRESPLIT_ATTR))
       Coroutines.push_back(&N);
 
   if (Coroutines.empty() && !PrepareFn)
     return PreservedAnalyses::all();
 
   if (Coroutines.empty())
-    llvm_unreachable("new pass manager cannot yet handle "
-                     "'llvm.coro.prepare.retcon'");
+    replaceAllPrepares(PrepareFn, CG, C);
 
   // Split all the coroutines.
   for (LazyCallGraph::Node *N : Coroutines) {
     Function &F = N->getFunction();
     Attribute Attr = F.getFnAttribute(CORO_PRESPLIT_ATTR);
     StringRef Value = Attr.getValueAsString();
     LLVM_DEBUG(dbgs() << "CoroSplit: Processing coroutine '" << F.getName()
                       << "' state: " << Value << "\n");
for (LazyCallGraph::Node *N : Coroutines) {
     F.removeFnAttr(CORO_PRESPLIT_ATTR);
 
     SmallVector<Function *, 4> Clones;
     const coro::Shape Shape = splitCoroutine(F, Clones);
     updateCallGraphAfterCoroutineSplit(*N, Shape, Clones, C, CG, AM, UR, FAM);
   }
 
   if (PrepareFn)
-    llvm_unreachable("new pass manager cannot yet handle "
-                     "'llvm.coro.prepare.retcon'");
+    replaceAllPrepares(PrepareFn, CG, C);
 
   return PreservedAnalyses::none();
 }
 
 namespace {
 
 // We present a coroutine to LLVM as an ordinary function with suspension
 // points marked up with intrinsics. We let the optimizer party on the coroutine

llvm/test/Transforms/Coroutines/coro-retcon-frame.ll

 ; RUN: opt < %s -coro-split -S | FileCheck %s
+; RUN: opt < %s -passes=coro-split -S | FileCheck %s
 
 target datalayout = "p:64:64:64"
 
 declare void @prototype_f(i8*, i1)
 
 declare noalias i8* @allocate(i32 %size)
 declare void @deallocate(i8* %ptr)
 declare void @init(i64 *%ptr)