[CGSCC] Fix compile time blowup with large RefSCCs (llvm#94815)

aeubanks · SvyatoslavScherbina · commit 176f3d57617a · 2024-06-12T16:24:37.000+02:00
In some modules, e.g. Kotlin-generated IR, we end up with a huge RefSCC and the call graph updates done as a result of the inliner take a long time. This is due to RefSCC::removeInternalRefEdges() getting called many times, each time removing one function from the RefSCC, but each call to removeInternalRefEdges() is proportional to the size of the RefSCC. There are two places that call removeInternalRefEdges(), in updateCGAndAnalysisManagerForPass() and LazyCallGraph::removeDeadFunction(). 1) Since LazyCallGraph can deal with spurious (edges that exist in the graph but not in the IR) ref edges, we can simply not call removeInternalRefEdges() in updateCGAndAnalysisManagerForPass(). 2) LazyCallGraph::removeDeadFunction() still ends up taking the brunt of compile time with the above change for the original reason. So instead we batch all the dead function removals so we can call removeInternalRefEdges() just once. This requires some changes to callers of removeDeadFunction() to not actually erase the function from the module, but defer it to when we batch delete dead functions at the end of the CGSCC run, leaving the function body as "unreachable" in the meantime. We still need to ensure that call edges are accurate. I had also tried deleting dead functions after visiting a RefSCC, but deleting them all at once at the end was simpler. Many test changes are due to not performing unnecessary revisits of an SCC (the CGSCC infrastructure deems ref edge refinements as unimportant when it comes to revisiting SCCs, although that seems to not be consistently true given these changes) because we don't remove some ref edges. Specifically for devirt-invalidated.ll this seems to expose an inlining order issue with the inliner. Probably unimportant for this type of intentionally weird call graph. Compile time: https://llvm-compile-time-tracker.com/compare.php?from=6f2c61071c274a1b5e212e6ad4114641ec7c7fc3&to=b08c90d05e290dd065755ea776ceaf1420680224&stat=instructions:u (cherry picked from commit 71497cc)
diff --git a/llvm/include/llvm/Analysis/CGSCCPassManager.h b/llvm/include/llvm/Analysis/CGSCCPassManager.h
@@ -306,6 +306,10 @@ struct CGSCCUpdateResult {
   SmallDenseSet<std::pair<LazyCallGraph::Node *, LazyCallGraph::SCC *>, 4>
       &InlinedInternalEdges;
 
+  /// Functions that a pass has considered to be dead to be removed at the end
+  /// of the call graph walk in batch.
+  SmallVector<Function *, 4> &DeadFunctions;
+
   /// Weak VHs to keep track of indirect calls for the purposes of detecting
   /// devirtualization.
   ///
diff --git a/llvm/include/llvm/Analysis/LazyCallGraph.h b/llvm/include/llvm/Analysis/LazyCallGraph.h
@@ -832,7 +832,7 @@ class LazyCallGraph {
     /// self-edges and edge removals which result in a spanning tree with no
     /// more cycles.
     [[nodiscard]] SmallVector<RefSCC *, 1>
-    removeInternalRefEdge(Node &SourceN, ArrayRef<Node *> TargetNs);
+    removeInternalRefEdges(ArrayRef<std::pair<Node *, Node *>> Edges);
 
     /// A convenience wrapper around the above to handle trivial cases of
     /// inserting a new call edge.
@@ -1056,18 +1056,18 @@ class LazyCallGraph {
   /// once SCCs have started to be formed. These routines have strict contracts
   /// but may be called at any point.
 
-  /// Remove a dead function from the call graph (typically to delete it).
+  /// Remove dead functions from the call graph.
   ///
-  /// Note that the function must have an empty use list, and the call graph
-  /// must be up-to-date prior to calling this. That means it is by itself in
-  /// a maximal SCC which is by itself in a maximal RefSCC, etc. No structural
-  /// changes result from calling this routine other than potentially removing
-  /// entry points into the call graph.
+  /// These functions should have already been passed to markDeadFunction().
+  /// This is done as a batch to prevent compile time blowup as a result of
+  /// handling a single function at a time.
+  void removeDeadFunctions(ArrayRef<Function *> DeadFs);
+
+  /// Mark a function as dead to be removed later by removeDeadFunctions().
   ///
-  /// If SCC formation has begun, this function must not be part of the current
-  /// DFS in order to call this safely. Typically, the function will have been
-  /// fully visited by the DFS prior to calling this routine.
-  void removeDeadFunction(Function &F);
+  /// The function body should have no incoming or outgoing call or ref edges.
+  /// For example, a function with a single "unreachable" instruction.
+  void markDeadFunction(Function &F);
 
   /// Add a new function split/outlined from an existing function.
   ///
diff --git a/llvm/lib/Analysis/CGSCCPassManager.cpp b/llvm/lib/Analysis/CGSCCPassManager.cpp
@@ -160,10 +160,12 @@ ModuleToPostOrderCGSCCPassAdaptor::run(Module &M, ModuleAnalysisManager &AM) {
   SmallDenseSet<std::pair<LazyCallGraph::Node *, LazyCallGraph::SCC *>, 4>
       InlinedInternalEdges;
 
+  SmallVector<Function *, 4> DeadFunctions;
+
   CGSCCUpdateResult UR = {
-      RCWorklist,           CWorklist, InvalidRefSCCSet,
-      InvalidSCCSet,        nullptr,   PreservedAnalyses::all(),
-      InlinedInternalEdges, {}};
+      RCWorklist,           CWorklist,     InvalidRefSCCSet,
+      InvalidSCCSet,        nullptr,       PreservedAnalyses::all(),
+      InlinedInternalEdges, DeadFunctions, {}};
 
   // Request PassInstrumentation from analysis manager, will use it to run
   // instrumenting callbacks for the passes later.
@@ -344,6 +346,10 @@ ModuleToPostOrderCGSCCPassAdaptor::run(Module &M, ModuleAnalysisManager &AM) {
     } while (!RCWorklist.empty());
   }
 
+  CG.removeDeadFunctions(DeadFunctions);
+  for (Function *DeadF : DeadFunctions)
+    DeadF->eraseFromParent();
+
 #if defined(EXPENSIVE_CHECKS)
   // Verify that the call graph is still valid.
   CG.verify();
@@ -1037,36 +1043,6 @@ static LazyCallGraph::SCC &updateCGAndAnalysisManagerForPass(
     return true;
   });
 
-  // Now do a batch removal of the internal ref edges left.
-  auto NewRefSCCs = RC->removeInternalRefEdge(N, DeadTargets);
-  if (!NewRefSCCs.empty()) {
-    // The old RefSCC is dead, mark it as such.
-    UR.InvalidatedRefSCCs.insert(RC);
-
-    // Note that we don't bother to invalidate analyses as ref-edge
-    // connectivity is not really observable in any way and is intended
-    // exclusively to be used for ordering of transforms rather than for
-    // analysis conclusions.
-
-    // Update RC to the "bottom".
-    assert(G.lookupSCC(N) == C && "Changed the SCC when splitting RefSCCs!");
-    RC = &C->getOuterRefSCC();
-    assert(G.lookupRefSCC(N) == RC && "Failed to update current RefSCC!");
-
-    // The RC worklist is in reverse postorder, so we enqueue the new ones in
-    // RPO except for the one which contains the source node as that is the
-    // "bottom" we will continue processing in the bottom-up walk.
-    assert(NewRefSCCs.front() == RC &&
-           "New current RefSCC not first in the returned list!");
-    for (RefSCC *NewRC : llvm::reverse(llvm::drop_begin(NewRefSCCs))) {
-      assert(NewRC != RC && "Should not encounter the current RefSCC further "
-                            "in the postorder list of new RefSCCs.");
-      UR.RCWorklist.insert(NewRC);
-      LLVM_DEBUG(dbgs() << "Enqueuing a new RefSCC in the update worklist: "
-                        << *NewRC << "\n");
-    }
-  }
-
   // Next demote all the call edges that are now ref edges. This helps make
   // the SCCs small which should minimize the work below as we don't want to
   // form cycles that this would break.
diff --git a/llvm/lib/Analysis/LazyCallGraph.cpp b/llvm/lib/Analysis/LazyCallGraph.cpp
@@ -1160,8 +1160,8 @@ void LazyCallGraph::RefSCC::removeOutgoingEdge(Node &SourceN, Node &TargetN) {
 }
 
 SmallVector<LazyCallGraph::RefSCC *, 1>
-LazyCallGraph::RefSCC::removeInternalRefEdge(Node &SourceN,
-                                             ArrayRef<Node *> TargetNs) {
+LazyCallGraph::RefSCC::removeInternalRefEdges(
+    ArrayRef<std::pair<Node *, Node *>> Edges) {
   // We return a list of the resulting *new* RefSCCs in post-order.
   SmallVector<RefSCC *, 1> Result;
 
@@ -1179,25 +1179,21 @@ LazyCallGraph::RefSCC::removeInternalRefEdge(Node &SourceN,
 #endif
 
   // First remove the actual edges.
-  for (Node *TargetN : TargetNs) {
-    assert(!(*SourceN)[*TargetN].isCall() &&
+  for (auto [SourceN, TargetN] : Edges) {
+    assert(!(**SourceN)[*TargetN].isCall() &&
            "Cannot remove a call edge, it must first be made a ref edge");
 
-    bool Removed = SourceN->removeEdgeInternal(*TargetN);
+    bool Removed = (*SourceN)->removeEdgeInternal(*TargetN);
     (void)Removed;
     assert(Removed && "Target not in the edge set for this caller?");
   }
 
   // Direct self references don't impact the ref graph at all.
-  if (llvm::all_of(TargetNs,
-                   [&](Node *TargetN) { return &SourceN == TargetN; }))
-    return Result;
-
   // If all targets are in the same SCC as the source, because no call edges
   // were removed there is no RefSCC structure change.
-  SCC &SourceC = *G->lookupSCC(SourceN);
-  if (llvm::all_of(TargetNs, [&](Node *TargetN) {
-        return G->lookupSCC(*TargetN) == &SourceC;
+  if (llvm::all_of(Edges, [&](std::pair<Node *, Node *> E) {
+        return E.first == E.second ||
+               G->lookupSCC(*E.first) == G->lookupSCC(*E.second);
       }))
     return Result;
 
@@ -1499,7 +1495,7 @@ void LazyCallGraph::removeEdge(Node &SourceN, Node &TargetN) {
   assert(Removed && "Target not in the edge set for this caller?");
 }
 
-void LazyCallGraph::removeDeadFunction(Function &F) {
+void LazyCallGraph::markDeadFunction(Function &F) {
   // FIXME: This is unnecessarily restrictive. We should be able to remove
   // functions which recursively call themselves.
   assert(F.hasZeroLiveUses() &&
@@ -1517,57 +1513,66 @@ void LazyCallGraph::removeDeadFunction(Function &F) {
 
   Node &N = *NI->second;
 
-  // Cannot remove a function which has yet to be visited in the DFS walk, so
-  // if we have a node at all then we must have an SCC and RefSCC.
-  auto CI = SCCMap.find(&N);
-  assert(CI != SCCMap.end() &&
-         "Tried to remove a node without an SCC after DFS walk started!");
-  SCC &C = *CI->second;
-  RefSCC *RC = &C.getOuterRefSCC();
-
-  // In extremely rare cases, we can delete a dead function which is still in a
-  // non-trivial RefSCC. This can happen due to spurious ref edges sticking
-  // around after an IR function reference is removed.
-  if (RC->size() != 1) {
-    SmallVector<Node *, 0> NodesInRC;
-    for (SCC &OtherC : *RC) {
-      for (Node &OtherN : OtherC)
-        NodesInRC.push_back(&OtherN);
+  // Remove all call edges out of dead function.
+  for (Edge E : *N) {
+    if (E.isCall())
+      N->setEdgeKind(E.getNode(), Edge::Ref);
+  }
+}
+
+void LazyCallGraph::removeDeadFunctions(ArrayRef<Function *> DeadFs) {
+  if (DeadFs.empty())
+    return;
+
+  // Group dead functions by the RefSCC they're in.
+  DenseMap<RefSCC *, SmallVector<Node *, 1>> RCs;
+  for (Function *DeadF : DeadFs) {
+    Node *N = lookup(*DeadF);
+#ifndef NDEBUG
+    for (Edge &E : **N) {
+      assert(!E.isCall() &&
+             "dead function shouldn't have any outgoing call edges");
     }
-    for (Node *OtherN : NodesInRC) {
-      if ((*OtherN)->lookup(N)) {
-        auto NewRefSCCs =
-            RC->removeInternalRefEdge(*OtherN, ArrayRef<Node *>(&N));
-        // If we've split into multiple RefSCCs, RC is now invalid and the
-        // RefSCC containing C will be different.
-        if (!NewRefSCCs.empty())
-          RC = &C.getOuterRefSCC();
+#endif
+    RefSCC *RC = lookupRefSCC(*N);
+    RCs[RC].push_back(N);
+  }
+  // Remove outgoing edges from all dead functions. Dead functions should
+  // already have had their call edges removed in markDeadFunction(), so we only
+  // need to worry about spurious ref edges.
+  for (auto [RC, DeadNs] : RCs) {
+    SmallVector<std::pair<Node *, Node *>> InternalEdgesToRemove;
+    for (Node *DeadN : DeadNs) {
+      for (Edge &E : **DeadN) {
+        if (lookupRefSCC(E.getNode()) == RC)
+          InternalEdgesToRemove.push_back({DeadN, &E.getNode()});
+        else
+          RC->removeOutgoingEdge(*DeadN, E.getNode());
       }
     }
+    // We ignore the returned RefSCCs since at this point we're done with CGSCC
+    // iteration and don't need to add it to any worklists.
+    (void)RC->removeInternalRefEdges(InternalEdgesToRemove);
+    for (Node *DeadN : DeadNs) {
+      RefSCC *DeadRC = lookupRefSCC(*DeadN);
+      assert(DeadRC->size() == 1);
+      assert(DeadRC->begin()->size() == 1);
+      DeadRC->clear();
+      DeadRC->G = nullptr;
+    }
   }
+  // Clean up data structures.
+  for (Function *DeadF : DeadFs) {
+    Node &N = *lookup(*DeadF);
+
+    EntryEdges.removeEdgeInternal(N);
+    SCCMap.erase(SCCMap.find(&N));
+    NodeMap.erase(NodeMap.find(DeadF));
 
-  NodeMap.erase(NI);
-  EntryEdges.removeEdgeInternal(N);
-  SCCMap.erase(CI);
-
-  // This node must be the only member of its SCC as it has no callers, and
-  // that SCC must be the only member of a RefSCC as it has no references.
-  // Validate these properties first.
-  assert(C.size() == 1 && "Dead functions must be in a singular SCC");
-  assert(RC->size() == 1 && "Dead functions must be in a singular RefSCC");
-
-  // Finally clear out all the data structures from the node down through the
-  // components. postorder_ref_scc_iterator will skip empty RefSCCs, so no need
-  // to adjust LazyCallGraph data structures.
-  N.clear();
-  N.G = nullptr;
-  N.F = nullptr;
-  C.clear();
-  RC->clear();
-  RC->G = nullptr;
-
-  // Nothing to delete as all the objects are allocated in stable bump pointer
-  // allocators.
+    N.clear();
+    N.G = nullptr;
+    N.F = nullptr;
+  }
 }
 
 // Gets the Edge::Kind from one function to another by looking at the function's
diff --git a/llvm/lib/Transforms/IPO/Inliner.cpp b/llvm/lib/Transforms/IPO/Inliner.cpp
@@ -740,6 +740,14 @@ InlinerPass::getAdvisor(const ModuleAnalysisManagerCGSCCProxy::Result &MAM,
   return *IAA->getAdvisor();
 }
 
+void makeFunctionBodyUnreachable(Function &F) {
+  F.dropAllReferences();
+  for (BasicBlock &BB : make_early_inc_range(F))
+    BB.eraseFromParent();
+  BasicBlock *BB = BasicBlock::Create(F.getContext(), "", &F);
+  new UnreachableInst(F.getContext(), BB);
+}
+
 PreservedAnalyses InlinerPass::run(LazyCallGraph::SCC &InitialC,
                                    CGSCCAnalysisManager &AM, LazyCallGraph &CG,
                                    CGSCCUpdateResult &UR) {
@@ -988,11 +996,9 @@ PreservedAnalyses InlinerPass::run(LazyCallGraph::SCC &InitialC,
                              }),
               Calls.end());
 
-          // Clear the body and queue the function itself for deletion when we
-          // finish inlining and call graph updates.
-          // Note that after this point, it is an error to do anything other
-          // than use the callee's address or delete it.
-          Callee.dropAllReferences();
+          // Clear the body and queue the function itself for call graph
+          // updating when we finish inlining.
+          makeFunctionBodyUnreachable(Callee);
           assert(!is_contained(DeadFunctions, &Callee) &&
                  "Cannot put cause a function to become dead twice!");
           DeadFunctions.push_back(&Callee);
@@ -1070,7 +1076,7 @@ PreservedAnalyses InlinerPass::run(LazyCallGraph::SCC &InitialC,
   if (!DeadFunctionsInComdats.empty()) {
     filterDeadComdatFunctions(DeadFunctionsInComdats);
     for (auto *Callee : DeadFunctionsInComdats)
-      Callee->dropAllReferences();
+      makeFunctionBodyUnreachable(*Callee);
     DeadFunctions.append(DeadFunctionsInComdats);
   }
 
@@ -1082,25 +1088,18 @@ PreservedAnalyses InlinerPass::run(LazyCallGraph::SCC &InitialC,
   // that is OK as all we do is delete things and add pointers to unordered
   // sets.
   for (Function *DeadF : DeadFunctions) {
+    CG.markDeadFunction(*DeadF);
     // Get the necessary information out of the call graph and nuke the
     // function there. Also, clear out any cached analyses.
     auto &DeadC = *CG.lookupSCC(*CG.lookup(*DeadF));
     FAM.clear(*DeadF, DeadF->getName());
     AM.clear(DeadC, DeadC.getName());
-    auto &DeadRC = DeadC.getOuterRefSCC();
-    CG.removeDeadFunction(*DeadF);
 
     // Mark the relevant parts of the call graph as invalid so we don't visit
     // them.
     UR.InvalidatedSCCs.insert(&DeadC);
-    UR.InvalidatedRefSCCs.insert(&DeadRC);
-
-    // If the updated SCC was the one containing the deleted function, clear it.
-    if (&DeadC == UR.UpdatedC)
-      UR.UpdatedC = nullptr;
 
-    // And delete the actual function from the module.
-    M.getFunctionList().erase(DeadF);
+    UR.DeadFunctions.push_back(DeadF);
 
     ++NumDeleted;
   }
diff --git a/llvm/lib/Transforms/Utils/CallGraphUpdater.cpp b/llvm/lib/Transforms/Utils/CallGraphUpdater.cpp
@@ -67,16 +67,20 @@ bool CallGraphUpdater::finalize() {
 
         FAM.clear(*DeadFn, DeadFn->getName());
         AM->clear(*DeadSCC, DeadSCC->getName());
-        LCG->removeDeadFunction(*DeadFn);
+        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->DeadFunctions.push_back(DeadFn);
+      } else {
+        // The CGSCC infrastructure batch deletes functions at the end of the
+        // call graph walk, so only erase the function if we're not using that
+        // infrastructure.
+        // The function is now really dead and de-attached from everything.
+        DeadFn->eraseFromParent();
       }
-
-      // The function is now really dead and de-attached from everything.
-      DeadFn->eraseFromParent();
     }
   }
 
diff --git a/llvm/test/Other/cgscc-refscc-mutation-order.ll b/llvm/test/Other/cgscc-refscc-mutation-order.ll
@@ -15,8 +15,6 @@
 ; CHECK-NOT: InstCombinePass
 ; CHECK: Running pass: InstCombinePass on f4
 ; CHECK-NOT: InstCombinePass
-; CHECK: Running pass: InstCombinePass on f1
-; CHECK-NOT: InstCombinePass
 
 @a1 = alias void (), ptr @f1
 @a2 = alias void (), ptr @f2
diff --git a/llvm/test/Other/devirt-invalidated.ll b/llvm/test/Other/devirt-invalidated.ll
@@ -1,8 +1,6 @@
 ; RUN: opt -passes='devirt<0>(inline)' < %s -S | FileCheck %s
 
-; CHECK-NOT: internal
 ; CHECK: define void @e()
-; CHECK-NOT: internal
 
 define void @e() {
 entry:
diff --git a/llvm/test/Transforms/Inline/cgscc-cycle-debug.ll b/llvm/test/Transforms/Inline/cgscc-cycle-debug.ll
@@ -13,7 +13,6 @@
 ; CHECK: Running an SCC pass across the RefSCC: [(test1_a, test1_b, test1_c)]
 ; CHECK: Enqueuing the existing SCC in the worklist:(test1_b)
 ; CHECK: Enqueuing a newly formed SCC:(test1_c)
-; CHECK: Enqueuing a new RefSCC in the update worklist: [(test1_b)]
 ; CHECK: Switch an internal ref edge to a call edge from 'test1_a' to 'test1_c'
 ; CHECK: Switch an internal ref edge to a call edge from 'test1_a' to 'test1_a'
 ; CHECK: Re-running SCC passes after a refinement of the current SCC: (test1_c, test1_a)
diff --git a/llvm/unittests/Analysis/CGSCCPassManagerTest.cpp b/llvm/unittests/Analysis/CGSCCPassManagerTest.cpp
diff --git a/llvm/unittests/Analysis/LazyCallGraphTest.cpp b/llvm/unittests/Analysis/LazyCallGraphTest.cpp
