[Coroutines] properly update CallGraph in CoroSplit

[yuxuanchen1997]

Fixes #107139.

We weren't updating the call graph properly in CoroSplit. This crash is due to the await_suspend() function calling the coroutine, forming a multi-node SCC. The issue bisected to #79712 but I think this is red herring. We haven't been properly updating the call graph.

Added an example of such code as a test case.

[llvmbot]

@llvm/pr-subscribers-coroutines

@llvm/pr-subscribers-llvm-transforms

Author: Yuxuan Chen (yuxuanchen1997)

Changes

Fixes #107139.

We weren't updating the call graph properly in CoroSplit. This crash is due to the split coroutine belongs to a multi-node SCC. The issue bisected to #79712 but I think this is red herring. We haven't been properly updating the call graph.

Added an example of a coroutine belong to an SCC as a test case.

---

Full diff: https://github.com/llvm/llvm-project/pull/107935.diff

2 Files Affected:

• (modified) llvm/lib/Transforms/Coroutines/CoroSplit.cpp (+12-6)
• (added) llvm/test/Transforms/Coroutines/gh107139-split-in-scc.ll (+38)

diff --git a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
index dc3829d7f28eb1..8ea460badaad5d 100644
--- a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
+++ b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
@@ -2080,12 +2080,13 @@ splitCoroutine(Function &F, SmallVectorImpl<Function *> &Clones,
   return Shape;
 }
 
-static void updateCallGraphAfterCoroutineSplit(
+static LazyCallGraph::SCC &updateCallGraphAfterCoroutineSplit(
     LazyCallGraph::Node &N, const coro::Shape &Shape,
     const SmallVectorImpl<Function *> &Clones, LazyCallGraph::SCC &C,
     LazyCallGraph &CG, CGSCCAnalysisManager &AM, CGSCCUpdateResult &UR,
     FunctionAnalysisManager &FAM) {
 
+  auto *CurrentSCC = &C;
   if (!Clones.empty()) {
     switch (Shape.ABI) {
     case coro::ABI::Switch:
@@ -2105,13 +2106,16 @@ static void updateCallGraphAfterCoroutineSplit(
     }
 
     // Let the CGSCC infra handle the changes to the original function.
-    updateCGAndAnalysisManagerForCGSCCPass(CG, C, N, AM, UR, FAM);
+    CurrentSCC = &updateCGAndAnalysisManagerForCGSCCPass(CG, *CurrentSCC, N, AM,
+                                                         UR, FAM);
   }
 
   // Do some cleanup and let the CGSCC infra see if we've cleaned up any edges
   // to the split functions.
   postSplitCleanup(N.getFunction());
-  updateCGAndAnalysisManagerForFunctionPass(CG, C, N, AM, UR, FAM);
+  CurrentSCC = &updateCGAndAnalysisManagerForFunctionPass(CG, *CurrentSCC, N,
+                                                          AM, UR, FAM);
+  return *CurrentSCC;
 }
 
 /// Replace a call to llvm.coro.prepare.retcon.
@@ -2200,6 +2204,7 @@ PreservedAnalyses CoroSplitPass::run(LazyCallGraph::SCC &C,
   if (Coroutines.empty() && PrepareFns.empty())
     return PreservedAnalyses::all();
 
+  auto *CurrentSCC = &C;
   // Split all the coroutines.
   for (LazyCallGraph::Node *N : Coroutines) {
     Function &F = N->getFunction();
@@ -2211,7 +2216,8 @@ PreservedAnalyses CoroSplitPass::run(LazyCallGraph::SCC &C,
     coro::Shape Shape =
         splitCoroutine(F, Clones, FAM.getResult<TargetIRAnalysis>(F),
                        OptimizeFrame, MaterializableCallback);
-    updateCallGraphAfterCoroutineSplit(*N, Shape, Clones, C, CG, AM, UR, FAM);
+    CurrentSCC = &updateCallGraphAfterCoroutineSplit(
+        *N, Shape, Clones, *CurrentSCC, CG, AM, UR, FAM);
 
     auto &ORE = FAM.getResult<OptimizationRemarkEmitterAnalysis>(F);
     ORE.emit([&]() {
@@ -2223,14 +2229,14 @@ PreservedAnalyses CoroSplitPass::run(LazyCallGraph::SCC &C,
 
     if (!Shape.CoroSuspends.empty()) {
       // Run the CGSCC pipeline on the original and newly split functions.
-      UR.CWorklist.insert(&C);
+      UR.CWorklist.insert(CurrentSCC);
       for (Function *Clone : Clones)
         UR.CWorklist.insert(CG.lookupSCC(CG.get(*Clone)));
     }
   }
 
   for (auto *PrepareFn : PrepareFns) {
-    replaceAllPrepares(PrepareFn, CG, C);
+    replaceAllPrepares(PrepareFn, CG, *CurrentSCC);
   }
 
   return PreservedAnalyses::none();
diff --git a/llvm/test/Transforms/Coroutines/gh107139-split-in-scc.ll b/llvm/test/Transforms/Coroutines/gh107139-split-in-scc.ll
new file mode 100644
index 00000000000000..1fae4d2c21e80f
--- /dev/null
+++ b/llvm/test/Transforms/Coroutines/gh107139-split-in-scc.ll
@@ -0,0 +1,38 @@
+; Verify that we don't crash on mutually recursive coroutines
+; RUN: opt < %s -passes='cgscc(coro-split)' -S | FileCheck %s
+
+target triple = "x86_64-redhat-linux-gnu"
+
+; CHECK-LABEL: define void @foo
+define void @foo() presplitcoroutine personality ptr null {
+entry:
+
+  %0 = call token @llvm.coro.id(i32 0, ptr null, ptr null, ptr null)
+  %1 = call ptr @llvm.coro.begin(token %0, ptr null)
+  %2 = call token @llvm.coro.save(ptr null)
+  %3 = call i8 @llvm.coro.suspend(token none, i1 false)
+  %4 = call token @llvm.coro.save(ptr null)
+  ; CHECK: call void @bar(ptr null, ptr null)
+  call void @llvm.coro.await.suspend.void(ptr null, ptr null, ptr @bar)
+  ret void
+}
+
+; CHECK-LABEL: define void @bar({{.*}})
+define void @bar(ptr %0, ptr %1) {
+entry:
+  ; CHECK: call void @foo()
+  call void @foo()
+  ret void
+}
+
+; CHECK-LABEL: @foo.resume({{.*}})
+; CHECK-LABEL: @foo.destroy({{.*}})
+; CHECK-LABEL: @foo.cleanup({{.*}})
+
+declare token @llvm.coro.id(i32, ptr readnone, ptr nocapture readonly, ptr) #0
+declare ptr @llvm.coro.begin(token, ptr writeonly) nounwind
+declare token @llvm.coro.save(ptr) nomerge nounwind
+declare void @llvm.coro.await.suspend.void(ptr, ptr, ptr)
+declare i8 @llvm.coro.suspend(token, i1) nounwind
+
+attributes #0 = { nocallback nofree nosync nounwind willreturn memory(argmem: read) }

[ChuanqiXu9]

nit: the signature is slightly odd to me. It looks better to make it the below CurrentSCC to be an output parameter of the function. So that we can avoid the odd the pattern

[yuxuanchen1997]

Actually, I can't. It's not just to keep it consistent with the semantics of updateCGAndAnalysisManagerForCGSCCPass. The LazyCallGraph::SCC &C is an lvalue ref and cannot be rebound. If I do C = update... it will invoke the Move Assignment operator for LazyCallGraph::SCC. Not sure the consequence will be intended.

My other option is to make this input a LazyCallGraph::SCC ** or LazyCallGraph::SCC *&, which is also not very obvious.

[ChuanqiXu9]

Can't we make the type of C to LazyCallGraph::SCC *&?

[TylerNowicki]

I noted that the signature is actually very similar to LLVM's updateCGAndAnalysisManagerForPass() signature. Perhaps it is better to be consistent with that? https://llvm.org/doxygen/CGSCCPassManager_8cpp.html#a490117b63072462d035a6933fdb94c1f

[yuxuanchen1997]

I noted that the signature is actually very similar to LLVM's updateCGAndAnalysisManagerForPass() signature. Perhaps it is better to be consistent with that? https://llvm.org/doxygen/CGSCCPassManager_8cpp.html#a490117b63072462d035a6933fdb94c1f

Perhaps. Or perhaps the answer is to rewrite updateCGAndAnalysisManagerForPass. I really don't like the current semantics. I will see what I can do in another patch.

[yuxuanchen1997]

@ChuanqiXu9 I am reverting this suggestion for now. We can do a refactor to change all updateCGAndAnalysisManagerFor(Function|CGSCC)?Pass usage once we decide on a new design.

[yuxuanchen1997]

$ rg "updateCGAndAnalysisManagerFor(Function|CGSCC)?Pass"
llvm/unittests/Analysis/CGSCCPassManagerTest.cpp
1212:            updateCGAndAnalysisManagerForFunctionPass(CG, C, H2N, AM, UR, FAM);
1258:            updateCGAndAnalysisManagerForFunctionPass(CG, C, H2N, AM, UR, FAM);
1370:            updateCGAndAnalysisManagerForCGSCCPass(CG, C, H2N, AM, UR, FAM));
1407:        updateCGAndAnalysisManagerForFunctionPass(CG, C, H2N, AM, UR, FAM),
1438:            updateCGAndAnalysisManagerForCGSCCPass(CG, C, FN, AM, UR, FAM));
1469:        updateCGAndAnalysisManagerForFunctionPass(CG, C, FN, AM, UR, FAM),
1508:            updateCGAndAnalysisManagerForCGSCCPass(CG, C, FN, AM, UR, FAM));
1545:    ASSERT_DEATH(updateCGAndAnalysisManagerForCGSCCPass(CG, C, FN, AM, UR, FAM),
1777:              updateCGAndAnalysisManagerForCGSCCPass(CG, C, *N, AM, UR, FAM))
1867:          updateCGAndAnalysisManagerForCGSCCPass(CG, C, *N, AM, UR, FAM))
1923:        // Check that updateCGAndAnalysisManagerForCGSCCPass() after
1925:        updateCGAndAnalysisManagerForCGSCCPass(CG, *CG.lookupSCC(*N1), *N1, AM,
1987:          updateCGAndAnalysisManagerForCGSCCPass(CG, C, *N, AM, UR, FAM))

llvm/lib/Transforms/Utils/CallGraphUpdater.cpp
68:    updateCGAndAnalysisManagerForCGSCCPass(*LCG, *C, N, *AM, *UR, *FAM);

llvm/lib/Transforms/IPO/Inliner.cpp
497:    C = &updateCGAndAnalysisManagerForCGSCCPass(CG, *C, N, AM, UR, FAM);

llvm/lib/Transforms/Coroutines/CoroSplit.cpp
2108:    updateCGAndAnalysisManagerForCGSCCPass(CG, C, N, AM, UR, FAM);
2114:  updateCGAndAnalysisManagerForFunctionPass(CG, C, N, AM, UR, FAM);

llvm/lib/Analysis/CGSCCPassManager.cpp
561:      CurrentC = &updateCGAndAnalysisManagerForFunctionPass(CG, *CurrentC, *N,
876:static LazyCallGraph::SCC &updateCGAndAnalysisManagerForPass(
925:        // promoted before updateCGAndAnalysisManagerForPass runs.
1178:LazyCallGraph::SCC &llvm::updateCGAndAnalysisManagerForFunctionPass(
1182:  return updateCGAndAnalysisManagerForPass(G, InitialC, N, AM, UR, FAM,
1185:LazyCallGraph::SCC &llvm::updateCGAndAnalysisManagerForCGSCCPass(
1189:  return updateCGAndAnalysisManagerForPass(G, InitialC, N, AM, UR, FAM,

llvm/include/llvm/Analysis/CGSCCPassManager.h
419:LazyCallGraph::SCC &updateCGAndAnalysisManagerForFunctionPass(
430:LazyCallGraph::SCC &updateCGAndAnalysisManagerForCGSCCPass(

This is manageable refactoring I'd say. Let's do this in a separate patch.