[MLInliner] Handle CGSCC changes from #94815 (#96274)

With #94815, the nodes belonging to dead functions are no longer
invalidated, but kept around to batch delete at the end of the call
graph walk.

The ML inliner needs to be updated to handle this. This fixes some
asserts getting hit, e.g. https://crbug.com/348376263.

2 files changed, 27 insertions, 10 deletions

diff --git a/llvm/lib/Analysis/MLInlineAdvisor.cpp b/llvm/lib/Analysis/MLInlineAdvisor.cpp
index 8131acb..b59aa48 100644
--- a/llvm/lib/Analysis/MLInlineAdvisor.cpp
+++ b/llvm/lib/Analysis/MLInlineAdvisor.cpp
@@ -211,14 +211,12 @@ void MLInlineAdvisor::onPassEntry(LazyCallGraph::SCC *CurSCC) {
   // care about the nature of the Edge (call or ref). `FunctionLevels`-wise, we
   // record them at the same level as the original node (this is a choice, may
   // need revisiting).
+  // - 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;
-    }
     EdgeCount += getLocalCalls(N->getFunction());
     const auto NLevel = FunctionLevels.at(N);
     for (const auto &E : *(*N)) {
@@ -256,11 +254,9 @@ void MLInlineAdvisor::onPassExit(LazyCallGraph::SCC *CurSCC) {
   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
@@ -311,8 +307,12 @@ void MLInlineAdvisor::onSuccessfulInlining(const MLInlineAdvice &Advice,
   int64_t NewCallerAndCalleeEdges =
       getCachedFPI(*Caller).DirectCallsToDefinedFunctions;
 
+  // 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;
+    NodesInLastSCC.erase(CG.lookup(*Callee));
     DeadFunctions.insert(Callee);
   } else {
     NewCallerAndCalleeEdges +=
diff --git a/llvm/test/Transforms/Inline/ML/dead-callee.ll b/llvm/test/Transforms/Inline/ML/dead-callee.ll
new file mode 100644
index 0000000..a886557
--- /dev/null
+++ b/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
+}