[MLInliner] Handle CGSCC changes from #94815

[aeubanks]

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.

[aeubanks]

Still needs tests, installing tflite seems annoying

[llvmbot]

@llvm/pr-subscribers-mlgo

@llvm/pr-subscribers-llvm-transforms

Author: Arthur Eubanks (aeubanks)

Changes

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.

---

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

2 Files Affected:

• (modified) llvm/lib/Analysis/MLInlineAdvisor.cpp (+21-22)
• (modified) llvm/lib/Transforms/IPO/Inliner.cpp (+2-2)

diff --git a/llvm/lib/Analysis/MLInlineAdvisor.cpp b/llvm/lib/Analysis/MLInlineAdvisor.cpp
index 21946572339b9..57c7801053555 100644
--- a/llvm/lib/Analysis/MLInlineAdvisor.cpp
+++ b/llvm/lib/Analysis/MLInlineAdvisor.cpp
@@ -187,8 +187,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
@@ -206,16 +206,10 @@ void MLInlineAdvisor::onPassEntry(LazyCallGraph::SCC *LastSCC) {
   // 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).
-  NodeCount -= static_cast<int64_t>(NodesInLastSCC.size());
   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());
     const auto NLevel = FunctionLevels.at(N);
     for (const auto &E : *(*N)) {
@@ -223,6 +217,7 @@ void MLInlineAdvisor::onPassEntry(LazyCallGraph::SCC *LastSCC) {
       assert(!AdjNode->isDead() && !AdjNode->getFunction().isDeclaration());
       auto I = AllNodes.insert(AdjNode);
       if (I.second) {
+        ++NodeCount;
         NodesInLastSCC.insert(AdjNode);
         FunctionLevels[AdjNode] = NLevel;
       }
@@ -235,15 +230,15 @@ 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
@@ -251,15 +246,13 @@ void MLInlineAdvisor::onPassExit(LazyCallGraph::SCC *LastSCC) {
   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)
@@ -306,11 +299,18 @@ 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
+    LazyCallGraph::Node *CalleeN = CG.lookup(*Callee);
+    NodesInLastSCC.erase(CalleeN);
+    AllNodes.erase(CalleeN);
+  } else {
     NewCallerAndCalleeEdges +=
         getCachedFPI(*Callee).DirectCallsToDefinedFunctions;
+  }
   EdgeCount += (NewCallerAndCalleeEdges - Advice.CallerAndCalleeEdges);
   assert(CurrentIRSize >= 0 && EdgeCount >= 0 && NodeCount >= 0);
 }
@@ -484,8 +484,7 @@ void MLInlineAdvisor::print(raw_ostream &OS) const {
   OS << "\n";
   OS << "[MLInlineAdvisor] FuncLevels:\n";
   for (auto I : FunctionLevels)
-    OS << (I.first->isDead() ? "<deleted>" : I.first->getFunction().getName())
-       << " : " << I.second << "\n";
+    OS << I.first->getFunction().getName() << " : " << I.second << "\n";
 
   OS << "\n";
 }
diff --git a/llvm/lib/Transforms/IPO/Inliner.cpp b/llvm/lib/Transforms/IPO/Inliner.cpp
index 1a7b9bc8e3e77..6f0146142d494 100644
--- a/llvm/lib/Transforms/IPO/Inliner.cpp
+++ b/llvm/lib/Transforms/IPO/Inliner.cpp
@@ -223,8 +223,6 @@ PreservedAnalyses InlinerPass::run(LazyCallGraph::SCC &InitialC,
   InlineAdvisor &Advisor = getAdvisor(MAMProxy, FAM, M);
   Advisor.onPassEntry(&InitialC);
 
-  auto AdvisorOnExit = make_scope_exit([&] { Advisor.onPassExit(&InitialC); });
-
   // We use a single common worklist for calls across the entire SCC. We
   // process these in-order and append new calls introduced during inlining to
   // the end. The PriorityInlineOrder is optional here, in which the smaller
@@ -564,6 +562,8 @@ PreservedAnalyses InlinerPass::run(LazyCallGraph::SCC &InitialC,
     ++NumDeleted;
   }
 
+  Advisor.onPassExit(C);
+
   if (!Changed)
     return PreservedAnalyses::all();
 

[mtrofin]

could you add a note in the comment blurb above the while, to the "The cgscc pass manager rules", saying "Nodes are never deleted"? (IIUC that's now the case, right?)

[aeubanks]

added comment

[aeubanks]

I've split out a bunch of the changes into other PRs that have been merge, PTAL again