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

Arthur Eubanks via llvm-commits llvm-commits at lists.llvm.org
Tue Jul 2 14:30:42 PDT 2024


https://github.com/aeubanks updated https://github.com/llvm/llvm-project/pull/96274

>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 1/3] [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();
 

>From c395ebb3084d21b708a1c3ee7f0d120523279630 Mon Sep 17 00:00:00 2001
From: Arthur Eubanks <aeubanks at google.com>
Date: Tue, 2 Jul 2024 20:49:22 +0000
Subject: [PATCH 2/3] add comment

---
 llvm/lib/Analysis/MLInlineAdvisor.cpp | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/llvm/lib/Analysis/MLInlineAdvisor.cpp b/llvm/lib/Analysis/MLInlineAdvisor.cpp
index e6d30be8f7073..a8fcaf5cd289e 100644
--- a/llvm/lib/Analysis/MLInlineAdvisor.cpp
+++ b/llvm/lib/Analysis/MLInlineAdvisor.cpp
@@ -211,6 +211,8 @@ 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());

>From 47a696827e0460f81c386bf6f45c989f0088a133 Mon Sep 17 00:00:00 2001
From: Arthur Eubanks <aeubanks at google.com>
Date: Tue, 2 Jul 2024 21:28:46 +0000
Subject: [PATCH 3/3] fixes

add test
remove node from from NodesInLastSCC
---
 llvm/lib/Analysis/MLInlineAdvisor.cpp         |  1 +
 llvm/test/Transforms/Inline/ML/dead-callee.ll | 17 +++++++++++++++++
 2 files changed, 18 insertions(+)
 create mode 100644 llvm/test/Transforms/Inline/ML/dead-callee.ll

diff --git a/llvm/lib/Analysis/MLInlineAdvisor.cpp b/llvm/lib/Analysis/MLInlineAdvisor.cpp
index a8fcaf5cd289e..b59aa4810005b 100644
--- a/llvm/lib/Analysis/MLInlineAdvisor.cpp
+++ b/llvm/lib/Analysis/MLInlineAdvisor.cpp
@@ -312,6 +312,7 @@ void MLInlineAdvisor::onSuccessfulInlining(const MLInlineAdvice &Advice,
   // 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 0000000000000..a88655777e6f9
--- /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
+}



More information about the llvm-commits mailing list