[LLVM][Coroutines] Switch CoroAnnotationElidePass to a FunctionPass (#107897)

After landing #99285 we found
that the call graph update was causing the following crash when
expensive checks are turned on
```
llvm-project/llvm/lib/Analysis/CGSCCPassManager.cpp:982: LazyCallGraph::SCC &updateCGAndAnalysisManagerForPass(LazyCallGraph &, LazyCallGraph::SCC &, LazyCallGraph::Node &, CGSCCAnalysisManager &, CGSCCUpdateResult &, FunctionAnalysisManager &, bool): Assertion `(RC == &TargetRC || RC->isAncestorOf(Targe
tRC)) && "New call edge is not trivial!"' failed.                                                                                                                                                                                                                                                                               
```
I have to admit I believe that the call graph update process I did for
that patch could be wrong.

After reading the code in `CGSCCToFunctionPassAdaptor`, I am convinced
that `CoroAnnotationElidePass` can be a FunctionPass and rely on the
adaptor to update the call graph for us, so long as we properly
invalidate the caller's analyses.

After this patch,
`llvm/test/Transforms/Coroutines/coro-transform-must-elide.ll` no longer
fails under expensive checks.

Author: yuxuanchen1997

llvm/include/llvm/Transforms/Coroutines/CoroAnnotationElide.h

@@ -17,18 +17,14 @@
 #ifndef LLVM_TRANSFORMS_COROUTINES_COROANNOTATIONELIDE_H
 #define LLVM_TRANSFORMS_COROUTINES_COROANNOTATIONELIDE_H
 
-#include "llvm/Analysis/CGSCCPassManager.h"
-#include "llvm/Analysis/LazyCallGraph.h"
 #include "llvm/IR/PassManager.h"
 
 namespace llvm {
 
-struct CoroAnnotationElidePass : PassInfoMixin<CoroAnnotationElidePass> {
-  CoroAnnotationElidePass() {}
-
-  PreservedAnalyses run(LazyCallGraph::SCC &C, CGSCCAnalysisManager &AM,
-                        LazyCallGraph &CG, CGSCCUpdateResult &UR);
+class Function;
 
+struct CoroAnnotationElidePass : PassInfoMixin<CoroAnnotationElidePass> {
+  PreservedAnalyses run(Function &F, FunctionAnalysisManager &FAM);
   static bool isRequired() { return false; }
 };
 } // end namespace llvm

llvm/lib/Passes/PassBuilderPipelines.cpp

@@ -977,7 +977,8 @@ PassBuilder::buildInlinerPipeline(OptimizationLevel Level,
 
   if (Phase != ThinOrFullLTOPhase::ThinLTOPreLink) {
     MainCGPipeline.addPass(CoroSplitPass(Level != OptimizationLevel::O0));
-    MainCGPipeline.addPass(CoroAnnotationElidePass());
+    MainCGPipeline.addPass(
+        createCGSCCToFunctionPassAdaptor(CoroAnnotationElidePass()));
   }
 
   // Make sure we don't affect potential future NoRerun CGSCC adaptors.
@@ -1028,8 +1029,7 @@ PassBuilder::buildModuleInlinerPipeline(OptimizationLevel Level,
   if (Phase != ThinOrFullLTOPhase::ThinLTOPreLink) {
     MPM.addPass(createModuleToPostOrderCGSCCPassAdaptor(
         CoroSplitPass(Level != OptimizationLevel::O0)));
-    MPM.addPass(
-        createModuleToPostOrderCGSCCPassAdaptor(CoroAnnotationElidePass()));
+    MPM.addPass(createModuleToFunctionPassAdaptor(CoroAnnotationElidePass()));
   }
 
   return MPM;

llvm/lib/Passes/PassRegistry.def

@@ -243,7 +243,6 @@ CGSCC_PASS("attributor-light-cgscc", AttributorLightCGSCCPass())
 CGSCC_PASS("invalidate<all>", InvalidateAllAnalysesPass())
 CGSCC_PASS("no-op-cgscc", NoOpCGSCCPass())
 CGSCC_PASS("openmp-opt-cgscc", OpenMPOptCGSCCPass())
-CGSCC_PASS("coro-annotation-elide", CoroAnnotationElidePass())
 #undef CGSCC_PASS
 
 #ifndef CGSCC_PASS_WITH_PARAMS
@@ -343,6 +342,7 @@ FUNCTION_PASS("codegenprepare", CodeGenPreparePass(TM))
 FUNCTION_PASS("consthoist", ConstantHoistingPass())
 FUNCTION_PASS("constraint-elimination", ConstraintEliminationPass())
 FUNCTION_PASS("coro-elide", CoroElidePass())
+FUNCTION_PASS("coro-annotation-elide", CoroAnnotationElidePass())
 FUNCTION_PASS("correlated-propagation", CorrelatedValuePropagationPass())
 FUNCTION_PASS("count-visits", CountVisitsPass())
 FUNCTION_PASS("dce", DCEPass())

llvm/lib/Transforms/Coroutines/CoroAnnotationElide.cpp

@@ -16,7 +16,6 @@
 
 #include "llvm/Transforms/Coroutines/CoroAnnotationElide.h"
 
-#include "llvm/Analysis/CGSCCPassManager.h"
 #include "llvm/Analysis/LazyCallGraph.h"
 #include "llvm/Analysis/OptimizationRemarkEmitter.h"
 #include "llvm/IR/Analysis.h"
@@ -96,24 +95,15 @@ static void processCall(CallBase *CB, Function *Caller, Function *NewCallee,
   CB->eraseFromParent();
 }
 
-PreservedAnalyses CoroAnnotationElidePass::run(LazyCallGraph::SCC &C,
-                                               CGSCCAnalysisManager &AM,
-                                               LazyCallGraph &CG,
-                                               CGSCCUpdateResult &UR) {
+PreservedAnalyses CoroAnnotationElidePass::run(Function &F,
+                                               FunctionAnalysisManager &FAM) {
   bool Changed = false;
-  CallGraphUpdater CGUpdater;
-  CGUpdater.initialize(CG, C, AM, UR);
-
-  auto &FAM =
-      AM.getResult<FunctionAnalysisManagerCGSCCProxy>(C, CG).getManager();
-
-  for (LazyCallGraph::Node &N : C) {
-    Function *Callee = &N.getFunction();
-    Function *NewCallee = Callee->getParent()->getFunction(
-        (Callee->getName() + ".noalloc").str());
-    if (!NewCallee) {
-      continue;
-    }
+
+  Function *NewCallee =
+      F.getParent()->getFunction((F.getName() + ".noalloc").str());
+  if (!NewCallee) {
+    return PreservedAnalyses::all();
+  }
 
     auto FramePtrArgPosition = NewCallee->arg_size() - 1;
     auto FrameSize =
@@ -122,34 +112,31 @@ PreservedAnalyses CoroAnnotationElidePass::run(LazyCallGraph::SCC &C,
         NewCallee->getParamAlign(FramePtrArgPosition).valueOrOne();
 
     SmallVector<CallBase *, 4> Users;
-    for (auto *U : Callee->users()) {
+    for (auto *U : F.users()) {
       if (auto *CB = dyn_cast<CallBase>(U)) {
-        if (CB->getCalledFunction() == Callee)
+        if (CB->getCalledFunction() == &F)
           Users.push_back(CB);
       }
     }
 
-    auto &ORE = FAM.getResult<OptimizationRemarkEmitterAnalysis>(*Callee);
+    auto &ORE = FAM.getResult<OptimizationRemarkEmitterAnalysis>(F);
 
     for (auto *CB : Users) {
       auto *Caller = CB->getFunction();
       if (Caller && Caller->isPresplitCoroutine() &&
           CB->hasFnAttr(llvm::Attribute::CoroElideSafe)) {
 
-        auto *CallerN = CG.lookup(*Caller);
-        auto *CallerC = CG.lookupSCC(*CallerN);
         processCall(CB, Caller, NewCallee, FrameSize, FrameAlign);
 
         ORE.emit([&]() {
           return OptimizationRemark(DEBUG_TYPE, "CoroAnnotationElide", Caller)
-                 << "'" << ore::NV("callee", Callee->getName())
-                 << "' elided in '" << ore::NV("caller", Caller->getName());
+                 << "'" << ore::NV("callee", F.getName()) << "' elided in '"
+                 << ore::NV("caller", Caller->getName());
         });
+        FAM.invalidate(*Caller, PreservedAnalyses::none());
         Changed = true;
-        updateCGAndAnalysisManagerForCGSCCPass(CG, *CallerC, *CallerN, AM, UR,
-                                               FAM);
       }
     }
-  }
+
   return Changed ? PreservedAnalyses::none() : PreservedAnalyses::all();
 }