Inliner optimization fixups

llvm/include/llvm/Analysis/CGSCCPassManager.h

@@ -230,17 +230,6 @@ using ModuleAnalysisManagerCGSCCProxy =
 /// Passes which do not change the call graph structure in any way can just
 /// ignore this argument to their run method.
 struct CGSCCUpdateResult {
-  /// Worklist of the RefSCCs queued for processing.
-  ///
-  /// When a pass refines the graph and creates new RefSCCs or causes them to
-  /// have a different shape or set of component SCCs it should add the RefSCCs
-  /// to this worklist so that we visit them in the refined form.
-  ///
-  /// This worklist is in reverse post-order, as we pop off the back in order
-  /// to observe RefSCCs in post-order. When adding RefSCCs, clients should add
-  /// them in reverse post-order.
-  SmallPriorityWorklist<LazyCallGraph::RefSCC *, 1> &RCWorklist;
-
   /// Worklist of the SCCs queued for processing.
   ///
   /// When a pass refines the graph and creates new SCCs or causes them to have
@@ -256,14 +245,6 @@ struct CGSCCUpdateResult {
   /// in reverse post-order.
   SmallPriorityWorklist<LazyCallGraph::SCC *, 1> &CWorklist;
 
-  /// The set of invalidated RefSCCs which should be skipped if they are found
-  /// in \c RCWorklist.
-  ///
-  /// This is used to quickly prune out RefSCCs when they get deleted and
-  /// happen to already be on the worklist. We use this primarily to avoid
-  /// scanning the list and removing entries from it.
-  SmallPtrSetImpl<LazyCallGraph::RefSCC *> &InvalidatedRefSCCs;
-
   /// The set of invalidated SCCs which should be skipped if they are found
   /// in \c CWorklist.
   ///

llvm/include/llvm/Analysis/MLInlineAdvisor.h

@@ -86,6 +86,7 @@ class MLInlineAdvisor : public InlineAdvisor {
   int32_t CurrentIRSize = 0;
   llvm::SmallPtrSet<const LazyCallGraph::Node *, 1> NodesInLastSCC;
   DenseSet<const LazyCallGraph::Node *> AllNodes;
+  DenseSet<Function *> DeadFunctions;
   bool ForceStop = false;
 };
 

llvm/lib/Analysis/CGSCCPassManager.cpp

@@ -152,20 +152,22 @@ ModuleToPostOrderCGSCCPassAdaptor::run(Module &M, ModuleAnalysisManager &AM) {
   SmallPriorityWorklist<LazyCallGraph::RefSCC *, 1> RCWorklist;
   SmallPriorityWorklist<LazyCallGraph::SCC *, 1> CWorklist;
 
-  // Keep sets for invalidated SCCs and RefSCCs that should be skipped when
+  // Keep sets for invalidated SCCs that should be skipped when
   // iterating off the worklists.
-  SmallPtrSet<LazyCallGraph::RefSCC *, 4> InvalidRefSCCSet;
   SmallPtrSet<LazyCallGraph::SCC *, 4> InvalidSCCSet;
 
   SmallDenseSet<std::pair<LazyCallGraph::Node *, LazyCallGraph::SCC *>, 4>
       InlinedInternalEdges;
 
   SmallVector<Function *, 4> DeadFunctions;
 
-  CGSCCUpdateResult UR = {
-      RCWorklist,           CWorklist,     InvalidRefSCCSet,
-      InvalidSCCSet,        nullptr,       PreservedAnalyses::all(),
-      InlinedInternalEdges, DeadFunctions, {}};
+  CGSCCUpdateResult UR = {CWorklist,
+                          InvalidSCCSet,
+                          nullptr,
+                          PreservedAnalyses::all(),
+                          InlinedInternalEdges,
+                          DeadFunctions,
+                          {}};
 
   // Request PassInstrumentation from analysis manager, will use it to run
   // instrumenting callbacks for the passes later.
@@ -192,11 +194,6 @@ ModuleToPostOrderCGSCCPassAdaptor::run(Module &M, ModuleAnalysisManager &AM) {
 
     do {
       LazyCallGraph::RefSCC *RC = RCWorklist.pop_back_val();
-      if (InvalidRefSCCSet.count(RC)) {
-        LLVM_DEBUG(dbgs() << "Skipping an invalid RefSCC...\n");
-        continue;
-      }
-
       assert(CWorklist.empty() &&
              "Should always start with an empty SCC worklist");
 
@@ -1175,7 +1172,6 @@ static LazyCallGraph::SCC &updateCGAndAnalysisManagerForPass(
   }
 
   assert(!UR.InvalidatedSCCs.count(C) && "Invalidated the current SCC!");
-  assert(!UR.InvalidatedRefSCCs.count(RC) && "Invalidated the current RefSCC!");
   assert(&C->getOuterRefSCC() == RC && "Current SCC not in current RefSCC!");
 
   // Record the current SCC for higher layers of the CGSCC pass manager now that

llvm/lib/Analysis/MLInlineAdvisor.cpp

@@ -139,8 +139,8 @@ unsigned MLInlineAdvisor::getInitialFunctionLevel(const Function &F) const {
   return CG.lookup(F) ? FunctionLevels.at(CG.lookup(F)) : 0;
 }
 
-void MLInlineAdvisor::onPassEntry(LazyCallGraph::SCC *LastSCC) {
-  if (!LastSCC || ForceStop)
+void MLInlineAdvisor::onPassEntry(LazyCallGraph::SCC *CurSCC) {
+  if (!CurSCC || ForceStop)
     return;
   FPICache.clear();
   // Function passes executed between InlinerPass runs may have changed the
@@ -156,23 +156,22 @@ void MLInlineAdvisor::onPassEntry(LazyCallGraph::SCC *LastSCC) {
   // they'd be adjacent to Nodes in the last SCC. So we just need to check the
   // boundary of Nodes in NodesInLastSCC for Nodes we haven't seen. We don't
   // care about the nature of the Edge (call or ref).
-  NodeCount -= static_cast<int64_t>(NodesInLastSCC.size());
+  // - nodes are only deleted at the end of a call graph walk where they are
+  // batch deleted, so we shouldn't see any dead nodes here.
   while (!NodesInLastSCC.empty()) {
     const auto *N = *NodesInLastSCC.begin();
+    assert(!N->isDead());
     NodesInLastSCC.erase(N);
-    // The Function wrapped by N could have been deleted since we last saw it.
-    if (N->isDead()) {
-      assert(!N->getFunction().isDeclaration());
-      continue;
-    }
-    ++NodeCount;
     EdgeCount += getLocalCalls(N->getFunction());
     for (const auto &E : *(*N)) {
       const auto *AdjNode = &E.getNode();
       assert(!AdjNode->isDead() && !AdjNode->getFunction().isDeclaration());
       auto I = AllNodes.insert(AdjNode);
-      if (I.second)
+      // We've discovered a new function.
+      if (I.second) {
+        ++NodeCount;
         NodesInLastSCC.insert(AdjNode);
+      }
     }
   }
 
@@ -182,31 +181,29 @@ void MLInlineAdvisor::onPassEntry(LazyCallGraph::SCC *LastSCC) {
   // (Re)use NodesInLastSCC to remember the nodes in the SCC right now,
   // in case the SCC is split before onPassExit and some nodes are split out
   assert(NodesInLastSCC.empty());
-  for (const auto &N : *LastSCC)
+  for (const auto &N : *CurSCC)
     NodesInLastSCC.insert(&N);
 }
 
-void MLInlineAdvisor::onPassExit(LazyCallGraph::SCC *LastSCC) {
+void MLInlineAdvisor::onPassExit(LazyCallGraph::SCC *CurSCC) {
   // No need to keep this around - function passes will invalidate it.
   if (!KeepFPICache)
     FPICache.clear();
-  if (!LastSCC || ForceStop)
+  if (!CurSCC || ForceStop)
     return;
   // Keep track of the nodes and edges we last saw. Then, in onPassEntry,
   // we update the node count and edge count from the subset of these nodes that
   // survived.
   EdgesOfLastSeenNodes = 0;
 
   // Check on nodes that were in SCC onPassEntry
-  for (auto I = NodesInLastSCC.begin(); I != NodesInLastSCC.end();) {
-    if ((*I)->isDead())
-      NodesInLastSCC.erase(*I++);
-    else
-      EdgesOfLastSeenNodes += getLocalCalls((*I++)->getFunction());
+  for (const LazyCallGraph::Node *N : NodesInLastSCC) {
+    assert(!N->isDead());
+    EdgesOfLastSeenNodes += getLocalCalls(N->getFunction());
   }
 
   // Check on nodes that may have got added to SCC
-  for (const auto &N : *LastSCC) {
+  for (const auto &N : *CurSCC) {
     assert(!N.isDead());
     auto I = NodesInLastSCC.insert(&N);
     if (I.second)
@@ -253,11 +250,17 @@ void MLInlineAdvisor::onSuccessfulInlining(const MLInlineAdvice &Advice,
   int64_t NewCallerAndCalleeEdges =
       getCachedFPI(*Caller).DirectCallsToDefinedFunctions;
 
-  if (CalleeWasDeleted)
+  // A dead function's node is not actually removed from the call graph until
+  // the end of the call graph walk, but the node no longer belongs to any valid
+  // SCC.
+  if (CalleeWasDeleted) {
     --NodeCount;
-  else
+    NodesInLastSCC.erase(CG.lookup(*Callee));
+    DeadFunctions.insert(Callee);
+  } else {
     NewCallerAndCalleeEdges +=
         getCachedFPI(*Callee).DirectCallsToDefinedFunctions;
+  }
   EdgeCount += (NewCallerAndCalleeEdges - Advice.CallerAndCalleeEdges);
   assert(CurrentIRSize >= 0 && EdgeCount >= 0 && NodeCount >= 0);
 }

llvm/lib/Transforms/Utils/CallGraphUpdater.cpp

@@ -59,20 +59,14 @@ bool CallGraphUpdater::finalize() {
         auto *DeadSCC = LCG->lookupSCC(N);
         assert(DeadSCC && DeadSCC->size() == 1 &&
                &DeadSCC->begin()->getFunction() == DeadFn);
-        auto &DeadRC = DeadSCC->getOuterRefSCC();
 
-        FunctionAnalysisManager &FAM =
-            AM->getResult<FunctionAnalysisManagerCGSCCProxy>(*DeadSCC, *LCG)
-                .getManager();
-
-        FAM.clear(*DeadFn, DeadFn->getName());
+        FAM->clear(*DeadFn, DeadFn->getName());
         AM->clear(*DeadSCC, DeadSCC->getName());
         LCG->markDeadFunction(*DeadFn);
 
         // Mark the relevant parts of the call graph as invalid so we don't
         // visit them.
-        UR->InvalidatedSCCs.insert(DeadSCC);
-        UR->InvalidatedRefSCCs.insert(&DeadRC);
+        UR->InvalidatedSCCs.insert(LCG->lookupSCC(N));
         UR->DeadFunctions.push_back(DeadFn);
       } else {
         // The CGSCC infrastructure batch deletes functions at the end of the

llvm/test/Transforms/Inline/ML/dead-callee.ll

@@ -0,0 +1,17 @@
+; REQUIRES: llvm_inliner_model_autogenerated
+; RUN: opt -passes=inliner-ml-advisor-release -S < %s | FileCheck %s
+
+; Check that our accounting works when a function in a non-trivial SCC is dead.
+
+; CHECK: define void @f
+; CHECK-NOT: @g
+
+define void @f() {
+    call void @g()
+    ret void
+}
+
+define internal void @g() {
+    call void @f()
+    ret void
+}

llvm/unittests/Analysis/CGSCCPassManagerTest.cpp

@@ -1875,6 +1875,61 @@ TEST_F(CGSCCPassManagerTest, TestInsertionOfNewFunctions2) {
   ASSERT_TRUE(Ran);
 }
 
+TEST_F(CGSCCPassManagerTest, TestDeletionOfFunctionInNonTrivialRefSCC) {
+  std::unique_ptr<Module> M = parseIR("define void @f1() {\n"
+                                      "entry:\n"
+                                      "  call void @f2()\n"
+                                      "  ret void\n"
+                                      "}\n"
+                                      "define void @f2() {\n"
+                                      "entry:\n"
+                                      "  call void @f1()\n"
+                                      "  ret void\n"
+                                      "}\n");
+
+  bool Ran = false;
+  CGSCCPassManager CGPM;
+  CGPM.addPass(LambdaSCCPassNoPreserve(
+      [&](LazyCallGraph::SCC &C, CGSCCAnalysisManager &AM, LazyCallGraph &CG,
+          CGSCCUpdateResult &UR) {
+        if (Ran)
+          return;
+
+        LazyCallGraph::Node *N1 = nullptr;
+
+        for (LazyCallGraph::Node *N : SCCNodes(C)) {
+          Function &F = N->getFunction();
+          if (F.getName() != "f1")
+            continue;
+          N1 = N;
+
+          Function &F2 = *F.getParent()->getFunction("f2");
+
+          // Remove f1 <-> f2 references
+          F.getEntryBlock().front().eraseFromParent();
+          F2.getEntryBlock().front().eraseFromParent();
+
+          CallGraphUpdater CGU;
+          CGU.initialize(CG, C, AM, UR);
+          CGU.removeFunction(F2);
+          CGU.reanalyzeFunction(F);
+
+          Ran = true;
+        }
+
+        // Check that updateCGAndAnalysisManagerForCGSCCPass() after
+        // CallGraphUpdater::removeFunction() succeeds.
+        updateCGAndAnalysisManagerForCGSCCPass(CG, *CG.lookupSCC(*N1), *N1, AM,
+                                               UR, FAM);
+      }));
+
+  ModulePassManager MPM;
+  MPM.addPass(createModuleToPostOrderCGSCCPassAdaptor(std::move(CGPM)));
+  MPM.run(*M, MAM);
+
+  ASSERT_TRUE(Ran);
+}
+
 TEST_F(CGSCCPassManagerTest, TestInsertionOfNewNonTrivialCallEdge) {
   std::unique_ptr<Module> M = parseIR("define void @f1() {\n"
                                       "entry:\n"