[LazyCallGraph] Handle spurious ref edges when deleting a dead function

[aeubanks]

Summary

Spurious ref edges are ref edges that still exist in the call graph even
though the corresponding IR reference no longer exists. This can cause
issues when deleting a dead function which has a spurious ref edge
pointed at it because currently we expect the dead function's RefSCC to
be trivial.

In the case that the dead function's RefSCC is not trivial, remove all
ref edges from other nodes in the RefSCC to it.

Removing a ref edge can result in splitting RefSCCs. There's actually no
reason to revisit those RefSCCs because currently we only run passes on
SCCs, and we've already added all SCCs in the RefSCC to the worklist.
(as opposed to removing the ref edge in
updateCGAndAnalysisManagerForPass() which can modify the call graph of
SCCs we have not visited yet). We also don't expect that RefSCC
refinement will allow us to glean any more information for optimization
use. Also, doing so would drastically increase the complexity of
LazyCallGraph::removeDeadFunction(), requiring us to return a list of
invalidated RefSCCs and new RefSCCs to add to the worklist.

Fixes #56503

[aeubanks]

add comment

[asbirlea]

Thanks for resolving this!

llvm/include/llvm/Analysis/LazyCallGraph.h

1068

incomplete comment?

[aeubanks]

remove outdated comment

llvm/include/llvm/Analysis/LazyCallGraph.h

#endif
   /// module because they do not reference it in any way. This is also the unit
   /// where we do mutation of the graph in order to restrict mutations to those
   /// which don't violate this independence.
   ///
   /// A RefSCC contains a DAG of actual SCCs. All the nodes within the RefSCC
   /// are necessarily within some actual SCC that nests within it. Since
   /// a direct call *is* a reference, there will always be at least one RefSCC
   /// around any SCC.
+  ///
+  /// Spurious ref edges, meaning ref edges that still exist in the call graph
+  /// even though the corresponding IR reference no longer exists, are allowed.
+  /// This is mostly to support argument promotion, which can modify a caller to
+  /// no longer pass a function. The only place that needs to specially handle
+  /// this is deleting a dead function/node, otherwise the dead ref edges are
+  /// automatically removed when visiting the function/node no longer containing
+  /// the ref edge.
   class RefSCC {
     friend class LazyCallGraph;
     friend class LazyCallGraph::Node;
 
     LazyCallGraph *G;
 
     /// A postorder list of the inner SCCs.
     SmallVector<SCC *, 4> SCCs;
public:
     ///
     /// The runtime complexity of this method is, in the worst case, O(V+E)
     /// where V is the number of nodes in this RefSCC and E is the number of
     /// edges leaving the nodes in this RefSCC. Note that E includes both edges
     /// within this RefSCC and edges from this RefSCC to child RefSCCs. Some
     /// effort has been made to minimize the overhead of common cases such as
     /// self-edges and edge removals which result in a spanning tree with no
     /// more cycles.
-    SmallVector<RefSCC *, 1> removeInternalRefEdge(Node &SourceN,
-                                                   ArrayRef<Node *> TargetNs);
+    [[nodiscard]] SmallVector<RefSCC *, 1>
+    removeInternalRefEdge(Node &SourceN, ArrayRef<Node *> TargetNs);
 
     /// A convenience wrapper around the above to handle trivial cases of
     /// inserting a new call edge.
     ///
     /// This is trivial whenever the target is in the same SCC as the source or
     /// the edge is an outgoing edge to some descendant SCC. In these cases
     /// there is no change to the cyclic structure of SCCs or RefSCCs.
     ///
#endif
   /// entry points into the call graph.
   ///
   /// 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);
 
   /// Add a new function split/outlined from an existing function.
   ///
   /// The new function may only reference other functions that the original
   /// function did.
   ///
   /// The original function must reference (either directly or indirectly) the
   /// new function.
   ///
   /// The new function may also reference the original function.
   /// It may end up in a parent SCC in the case that the original function's

llvm/lib/Analysis/LazyCallGraph.cpp

assert(!isLibFunction(F) &&
          "Must not remove lib functions from the call graph!");
 
   auto NI = NodeMap.find(&F);
   if (NI == NodeMap.end())
     // Not in the graph at all!
     return;
 
   Node &N = *NI->second;
-  NodeMap.erase(NI);
-
-  // Remove this from the entry edges if present.
-  EntryEdges.removeEdgeInternal(N);
 
   // 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);
+    }
+    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();
+      }
+    }
+  }
+
+  NodeMap.erase(NI);
+  EntryEdges.removeEdgeInternal(N);
   SCCMap.erase(CI);
-  RefSCC &RC = C.getOuterRefSCC();
 
   // 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");
+  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;
+  RC->clear();
+  RC->G = nullptr;
 
   // Nothing to delete as all the objects are allocated in stable bump pointer
   // allocators.
 }
 
 // Gets the Edge::Kind from one function to another by looking at the function's
 // instructions. Asserts if there is no edge.
 // Useful for determining what type of edge should exist between functions when

llvm/test/Analysis/LazyCallGraph/remove-dead-function-spurious-ref-edge.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -passes=argpromotion,inline -S %s | FileCheck %s
+
+; argpromo removes @b's parameter (removing @c's reference to @a without updating the ref edge in the call graph), then the inliner inlines @a into @d and attempts to remove @a.
+
+define internal void @a() alwaysinline {
+  call void @e(ptr @c)
+  ret void
+}
+
+define internal void @b(ptr) noinline {
+; CHECK-LABEL: @b(
+; CHECK-NEXT:    ret void
+;
+  ret void
+}
+
+define internal void @c() noinline {
+; CHECK-LABEL: @c(
+; CHECK-NEXT:    call void @b()
+; CHECK-NEXT:    ret void
+;
+  call void @b(ptr @a)
+  ret void
+}
+
+define void @d() {
+; CHECK-LABEL: @d(
+; CHECK-NEXT:    call void @e(ptr @c)
+; CHECK-NEXT:    ret void
+;
+  call void @a()
+  ret void
+}
+
+declare void @e(ptr);

llvm/unittests/Analysis/LazyCallGraphTest.cpp

TEST(LazyCallGraphTest, ReplaceNodeFunction) {
   // Now replace the function in the graph.
   RC1.replaceNodeFunction(DN, E);
 
   EXPECT_EQ(&E, &DN.getFunction());
   EXPECT_EQ(&DN, &(*CN)[DN].getNode());
   EXPECT_EQ(&DN, &(*BN)[DN].getNode());
 }
 
-TEST(LazyCallGraphTest, RemoveFunctionWithSpurriousRef) {
+TEST(LazyCallGraphTest, RemoveFunctionWithSpuriousRef) {
   LLVMContext Context;
   // A graph with a couple of RefSCCs.
   std::unique_ptr<Module> M =
       parseAssembly(Context,
                     "define void @a(i8** %ptr) {\n"
                     "entry:\n"
                     "  store i8* bitcast (void(i8**)* @d to i8*), i8** %ptr\n"
                     "  ret void\n"
TEST(LazyCallGraphTest, RemoveFunctionWithSpuriousRef) {
   // The only observable change should be that the RefSCC is gone from the
   // postorder sequence.
   I = CG.postorder_ref_scc_begin();
   EXPECT_EQ(&RC1, &*I++);
   EXPECT_EQ(&RC2, &*I++);
   EXPECT_EQ(CG.postorder_ref_scc_end(), I);
 }
 
+TEST(LazyCallGraphTest, RemoveFunctionWithSpuriousRefRecursive) {
+  LLVMContext Context;
+  std::unique_ptr<Module> M =
+      parseAssembly(Context, "define void @a(ptr %p) {\n"
+                             "  store ptr @b, ptr %p\n"
+                             "  ret void\n"
+                             "}\n"
+                             "define void @b(ptr %p) {\n"
+                             "  store ptr @c, ptr %p\n"
+                             "  ret void\n"
+                             "}\n"
+                             "define void @c(ptr %p) {\n"
+                             "  ret void\n"
+                             "}\n");
+  LazyCallGraph CG = buildCG(*M);
+
+  LazyCallGraph::Node &AN = CG.get(lookupFunction(*M, "a"));
+  LazyCallGraph::Node &BN = CG.get(lookupFunction(*M, "b"));
+  LazyCallGraph::Node &CN = CG.get(lookupFunction(*M, "c"));
+  AN.populate();
+  BN.populate();
+  CN.populate();
+  // Insert spurious ref edge.
+  CG.insertEdge(CN, AN, LazyCallGraph::Edge::Ref);
+
+  // Force the graph to be fully expanded.
+  CG.buildRefSCCs();
+  auto I = CG.postorder_ref_scc_begin();
+  LazyCallGraph::RefSCC &RC = *I++;
+  EXPECT_EQ(CG.postorder_ref_scc_end(), I);
+
+  ASSERT_EQ(RC.size(), 3);
+
+  EXPECT_EQ(&RC, CG.lookupRefSCC(AN));
+  EXPECT_EQ(&RC, CG.lookupRefSCC(BN));
+  EXPECT_EQ(&RC, CG.lookupRefSCC(CN));
+
+  // Now delete 'a'. There are no uses of this function but there are
+  // spurious references.
+  CG.removeDeadFunction(AN.getFunction());
+
+  // The only observable change should be that the RefSCC is gone from the
+  // postorder sequence.
+  I = CG.postorder_ref_scc_begin();
+  EXPECT_EQ(CG.lookupRefSCC(CN), &*I++);
+  EXPECT_EQ(CG.lookupRefSCC(BN), &*I++);
+  EXPECT_EQ(CG.postorder_ref_scc_end(), I);
+}
+
+TEST(LazyCallGraphTest, RemoveFunctionWithSpuriousRefRecursive2) {
+  LLVMContext Context;
+  std::unique_ptr<Module> M =
+      parseAssembly(Context, "define void @a(ptr %p) {\n"
+                             "  store ptr @b, ptr %p\n"
+                             "  ret void\n"
+                             "}\n"
+                             "define void @b(ptr %p) {\n"
+                             "  store ptr @c, ptr %p\n"
+                             "  ret void\n"
+                             "}\n"
+                             "define void @c(ptr %p) {\n"
+                             "  store ptr @b, ptr %p\n"
+                             "  store ptr @d, ptr %p\n"
+                             "  ret void\n"
+                             "}\n"
+                             "define void @d(ptr %p) {\n"
+                             "  ret void\n"
+                             "}\n");
+  LazyCallGraph CG = buildCG(*M);
+
+  LazyCallGraph::Node &AN = CG.get(lookupFunction(*M, "a"));
+  LazyCallGraph::Node &BN = CG.get(lookupFunction(*M, "b"));
+  LazyCallGraph::Node &CN = CG.get(lookupFunction(*M, "c"));
+  LazyCallGraph::Node &DN = CG.get(lookupFunction(*M, "d"));
+  AN.populate();
+  BN.populate();
+  CN.populate();
+  DN.populate();
+  // Insert spurious ref edge.
+  CG.insertEdge(DN, AN, LazyCallGraph::Edge::Ref);
+
+  // Force the graph to be fully expanded.
+  CG.buildRefSCCs();
+  auto I = CG.postorder_ref_scc_begin();
+  LazyCallGraph::RefSCC &RC = *I++;
+  EXPECT_EQ(CG.postorder_ref_scc_end(), I);
+
+  ASSERT_EQ(4, RC.size());
+
+  EXPECT_EQ(&RC, CG.lookupRefSCC(AN));
+  EXPECT_EQ(&RC, CG.lookupRefSCC(BN));
+  EXPECT_EQ(&RC, CG.lookupRefSCC(CN));
+  EXPECT_EQ(&RC, CG.lookupRefSCC(DN));
+
+  // Now delete 'a'. There are no uses of this function but there are
+  // spurious references.
+  CG.removeDeadFunction(AN.getFunction());
+
+  // The only observable change should be that the RefSCC is gone from the
+  // postorder sequence.
+  I = CG.postorder_ref_scc_begin();
+  EXPECT_EQ(CG.lookupRefSCC(DN), &*I++);
+  EXPECT_EQ(CG.lookupRefSCC(CN), &*I);
+  EXPECT_EQ(CG.lookupRefSCC(BN), &*I++);
+  EXPECT_EQ(CG.postorder_ref_scc_end(), I);
+}
+
+TEST(LazyCallGraphTest, RemoveFunctionWithSpuriousRefRecursive3) {
+  LLVMContext Context;
+  std::unique_ptr<Module> M =
+      parseAssembly(Context, "define void @a(ptr %p) {\n"
+                             "  store ptr @b, ptr %p\n"
+                             "  ret void\n"
+                             "}\n"
+                             "define void @b(ptr %p) {\n"
+                             "  store ptr @c, ptr %p\n"
+                             "  ret void\n"
+                             "}\n"
+                             "define void @c(ptr %p) {\n"
+                             "  ret void\n"
+                             "}\n");
+  LazyCallGraph CG = buildCG(*M);
+
+  LazyCallGraph::Node &AN = CG.get(lookupFunction(*M, "a"));
+  LazyCallGraph::Node &BN = CG.get(lookupFunction(*M, "b"));
+  LazyCallGraph::Node &CN = CG.get(lookupFunction(*M, "c"));
+  AN.populate();
+  BN.populate();
+  CN.populate();
+  // Insert spurious ref edges.
+  CG.insertEdge(CN, AN, LazyCallGraph::Edge::Ref);
+  CG.insertEdge(BN, AN, LazyCallGraph::Edge::Ref);
+
+  // Force the graph to be fully expanded.
+  CG.buildRefSCCs();
+  auto I = CG.postorder_ref_scc_begin();
+  LazyCallGraph::RefSCC &RC = *I++;
+  EXPECT_EQ(CG.postorder_ref_scc_end(), I);
+
+  ASSERT_EQ(RC.size(), 3);
+
+  EXPECT_EQ(&RC, CG.lookupRefSCC(AN));
+  EXPECT_EQ(&RC, CG.lookupRefSCC(BN));
+  EXPECT_EQ(&RC, CG.lookupRefSCC(CN));
+
+  // Now delete 'a'. There are no uses of this function but there are
+  // spurious references.
+  CG.removeDeadFunction(AN.getFunction());
+
+  // The only observable change should be that the RefSCC is gone from the
+  // postorder sequence.
+  I = CG.postorder_ref_scc_begin();
+  EXPECT_EQ(CG.lookupRefSCC(CN), &*I++);
+  EXPECT_EQ(CG.lookupRefSCC(BN), &*I++);
+  EXPECT_EQ(CG.postorder_ref_scc_end(), I);
+}
+
 TEST(LazyCallGraphTest, AddSplitFunction1) {
   LLVMContext Context;
   std::unique_ptr<Module> M = parseAssembly(Context, "define void @f() {\n"
                                                      "  ret void\n"
                                                      "}\n");
   LazyCallGraph CG = buildCG(*M);
 
   Function &F = lookupFunction(*M, "f");