Bug type: crash
Component: Coroutines
Bug location: llvm/lib/Transforms/Coroutines/CoroSplit.cpp: CoroSplitPass::run
replacePrepare
splitCoroutine
updateCallGraphAfterCoroutineSplit
Issue title: "Invalidated the current SCC" assertion failure after #79712
Labels: crash
coroutines
llvm:transforms
Reported by: smeenai

For the source file in https://gist.github.com/smeenai/fcd14380467e813a28f1ef8c5484b775 (reduced from a large example that uses libunifex; it's probably possible to reduce this further but I'm not at all familiar with coroutines), the following command succeeds with Clang 18 but fails with Clang 19 and trunk:

```
$ clang -c -o /dev/null coroutine_crash.cpp -std=c++20 -w
clang: llvm/lib/Analysis/CGSCCPassManager.cpp:1167: LazyCallGraph::SCC &updateCGAndAnalysisManagerForPass(llvm::LazyCallGraph &, LazyCallGraph::SCC &, LazyCallGraph::Node &, llvm::CGSCCAnalysisManager &, llvm::CGSCCUpdateResult &, llvm::FunctionAnalysisManager &, bool): Assertion `!UR.InvalidatedSCCs.count(C) && "Invalidated the current SCC!"' failed.
```

I bisected this to #79712. CC @fpasserby, @ChuanqiXu9, and @yuxuanchen1997 

Comments: author: yuxuanchen1997
body: I took a look at this and I think it was due to incorrect call graph updates. `updateCGAndAnalysisManagerFor(CGSCC|Function)Pass` seems to return an updated SCC but we discarded the updated SCC.

See example usage: https://github.com/llvm/llvm-project/blob/main/llvm/lib/Analysis/CGSCCPassManager.cpp#L561-L562
See implementation: https://github.com/llvm/llvm-project/blob/main/llvm/lib/Analysis/CGSCCPassManager.cpp#L1170-L1175

I have a patch https://github.com/llvm/llvm-project/pull/107935 but still need some time for a test case. 
author: yuxuanchen1997
body: This happens when foo has is not a single node SCC:
```
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)
  call void @llvm.coro.await.suspend.void(ptr null, ptr null, ptr @bar)
  ret void
}

define void @bar(ptr %0, ptr %1) {
entry:
  call void @foo()
  ret void
}

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) }
```

Repro the crash with `clang -O0 bad.ll`.
author: yuxuanchen1997
body: The 3rd argument to `@llvm.coro.await.suspend` is the `await_suspend` function in C++. This happens when `await_suspend()` calls the same coroutine from its body. 

Fix patch:
commit 853bff2122e1f42398587f76234c05d56f16318e
Author: Yuxuan Chen <ych@fb.com>
Date:   Thu Sep 12 10:45:20 2024 -0700

    [Coroutines] properly update CallGraph in CoroSplit (#107935)
    
    Fixes https://github.com/llvm/llvm-project/issues/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
    https://github.com/llvm/llvm-project/pull/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.

diff --git a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
index dc3829d7f28e..8ea460badaad 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();


Regression tests:
file: llvm/test/Transforms/Coroutines/gh107139-split-in-scc.ll
commands: opt < %s -passes='cgscc(coro-split)' -S
tests: test_name: <module>
test_body: ; Verify that we don't crash on mutually recursive coroutines

target triple = "x86_64-redhat-linux-gnu"

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)
  call void @llvm.coro.await.suspend.void(ptr null, ptr null, ptr @bar)
  ret void
}

define void @bar(ptr %0, ptr %1) {
entry:
  call void @foo()
  ret void
}


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) }