[llvm] [Inliner] Handle CGSCC changes from #94815 (PR #96274)

[Arthur Eubanks via llvm-commits]

https://github.com/aeubanks created https://github.com/llvm/llvm-project/pull/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.

>From a866a0ac487c5966e52122f5655d15d97af2d892 Mon Sep 17 00:00:00 2001
From: Arthur Eubanks <aeubanks at google.com>
Date: Fri, 21 Jun 2024 04:23:41 +0000
Subject: [PATCH] [Inliner] Handle CGSCC changes from #94815

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.
---
 llvm/lib/Analysis/MLInlineAdvisor.cpp | 43 +++++++++++++--------------
 llvm/lib/Transforms/IPO/Inliner.cpp   |  4 +--
 2 files changed, 23 insertions(+), 24 deletions(-)

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();