Add LazyCallGraph API to add function to RefSCC

Summary:
Depends on https://reviews.llvm.org/D70927.

`LazyCallGraph::addNewFunctionIntoSCC` allows users to insert a new
function node into a call graph, into a specific, existing SCC.

Extend this interface such that functions can be added even when they do
not belong in any existing SCC, but instead in a new SCC within an
existing RefSCC.

The ability to insert new functions as part of a RefSCC is necessary for
outlined functions that do not form a strongly connected cycle with the
function they are outlined from. An example of such a function would be the
coroutine funclets 'f.resume', etc., which are outlined from a coroutine 'f'.
Coroutine 'f' only references the funclets' addresses, it does not call
them directly.

Reviewers: jdoerfert, chandlerc, wenlei, hfinkel

Reviewed By: jdoerfert

Subscribers: hfinkel, JonChesterfield, mehdi_amini, hiraditya, llvm-commits

Tags: #llvm

Differential Revision: https://reviews.llvm.org/D72226

Author: modocache

llvm/include/llvm/Analysis/LazyCallGraph.h

@@ -1061,6 +1061,10 @@ class LazyCallGraph {
   /// Introduce a node for the function \p NewF in the SCC \p C.
   void addNewFunctionIntoSCC(Function &NewF, SCC &C);
 
+  /// Introduce a node for the function \p NewF, as a single node in a
+  /// new SCC, in the RefSCC \p RC.
+  void addNewFunctionIntoRefSCC(Function &NewF, RefSCC &RC);
+
   ///@}
 
   ///@{
@@ -1167,6 +1171,13 @@ class LazyCallGraph {
   /// Helper to update pointers back to the graph object during moves.
   void updateGraphPtrs();
 
+  /// Helper to insert a new function, add it to the NodeMap, and populate its
+  /// node.
+  Node &createNode(Function &F);
+
+  /// Helper to add the given Node \p N to the SCCMap, mapped to the SCC \p C.
+  void addNodeToSCC(SCC &C, Node &N);
+
   /// Allocates an SCC and constructs it using the graph allocator.
   ///
   /// The arguments are forwarded to the constructor.

llvm/lib/Analysis/LazyCallGraph.cpp

@@ -1567,12 +1567,17 @@ void LazyCallGraph::removeDeadFunction(Function &F) {
 }
 
 void LazyCallGraph::addNewFunctionIntoSCC(Function &NewF, SCC &C) {
-  Node &CGNode = get(NewF);
-  CGNode.DFSNumber = CGNode.LowLink = -1;
-  CGNode.populate();
-  C.Nodes.push_back(&CGNode);
-  SCCMap[&CGNode] = &C;
-  NodeMap[&NewF] = &CGNode;
+  addNodeToSCC(C, createNode(NewF));
+}
+
+void LazyCallGraph::addNewFunctionIntoRefSCC(Function &NewF, RefSCC &RC) {
+  Node &N = createNode(NewF);
+
+  auto *C = createSCC(RC, SmallVector<Node *, 1>());
+  addNodeToSCC(*C, N);
+
+  RC.SCCIndices[C] = RC.SCCIndices.size();
+  RC.SCCs.push_back(C);
 }
 
 LazyCallGraph::Node &LazyCallGraph::insertInto(Function &F, Node *&MappedN) {
@@ -1589,6 +1594,21 @@ void LazyCallGraph::updateGraphPtrs() {
     RC->G = this;
 }
 
+LazyCallGraph::Node &LazyCallGraph::createNode(Function &F) {
+  assert(!lookup(F) && "node already exists");
+
+  Node &N = get(F);
+  NodeMap[&F] = &N;
+  N.DFSNumber = N.LowLink = -1;
+  N.populate();
+  return N;
+}
+
+void LazyCallGraph::addNodeToSCC(LazyCallGraph::SCC &C, Node &N) {
+  C.Nodes.push_back(&N);
+  SCCMap[&N] = &C;
+}
+
 template <typename RootsT, typename GetBeginT, typename GetEndT,
           typename GetNodeT, typename FormSCCCallbackT>
 void LazyCallGraph::buildGenericSCCs(RootsT &&Roots, GetBeginT &&GetBegin,

llvm/unittests/Analysis/CGSCCPassManagerTest.cpp

@@ -1685,5 +1685,54 @@ TEST_F(CGSCCPassManagerTest, TestUpdateCGAndAnalysisManagerForPasses10) {
   MPM.run(*M, MAM);
 }
 
+TEST_F(CGSCCPassManagerTest, TestInsertionOfNewRefSCC) {
+  std::unique_ptr<Module> M = parseIR("define void @f() {\n"
+                                      "entry:\n"
+                                      "  call void @f()\n"
+                                      "  ret void\n"
+                                      "}\n");
+
+  CGSCCPassManager CGPM(/*DebugLogging*/ true);
+  CGPM.addPass(LambdaSCCPassNoPreserve(
+      [&](LazyCallGraph::SCC &C, CGSCCAnalysisManager &AM, LazyCallGraph &CG,
+          CGSCCUpdateResult &UR) {
+        for (auto &N : C) {
+          auto &F = N.getFunction();
+          if (F.getName() != "f")
+            continue;
+          auto *Call = dyn_cast<CallInst>(F.begin()->begin());
+          if (!Call || Call->getCalledFunction()->getName() != "f")
+            continue;
+
+          // Create a new function 'g'.
+          auto *G = Function::Create(F.getFunctionType(), F.getLinkage(),
+                                     F.getAddressSpace(), "g", F.getParent());
+          BasicBlock::Create(F.getParent()->getContext(), "entry", G);
+          // Instruct the LazyCallGraph to create a new node for 'g', as the
+          // single node in a new SCC, into the call graph. As a result
+          // the call graph is composed of a single RefSCC with two SCCs:
+          // [(f), (g)].
+          CG.addNewFunctionIntoRefSCC(*G, C.getOuterRefSCC());
+
+          // "Demote" the 'f -> f' call egde to a ref edge.
+          // 1. Erase the call edge from 'f' to 'f'.
+          Call->eraseFromParent();
+          // 2. Insert a ref edge from 'f' to 'f'.
+          (void)CastInst::CreatePointerCast(&F,
+                                            Type::getInt8PtrTy(F.getContext()),
+                                            "f.ref", &*F.begin()->begin());
+
+          ASSERT_NO_FATAL_FAILURE(
+              updateCGAndAnalysisManagerForCGSCCPass(CG, C, N, AM, UR))
+              << "Updating the call graph with a demoted, self-referential "
+                 "call edge 'f -> f', and a newly inserted ref edge 'f -> g', "
+                 "caused a fatal failure";
+        }
+      }));
+
+  ModulePassManager MPM(/*DebugLogging*/ true);
+  MPM.addPass(createModuleToPostOrderCGSCCPassAdaptor(std::move(CGPM)));
+  MPM.run(*M, MAM);
+}
 #endif
 } // namespace

llvm/unittests/Analysis/LazyCallGraphTest.cpp

@@ -2210,4 +2210,46 @@ TEST(LazyCallGraphTest, RemoveFunctionWithSpurriousRef) {
   EXPECT_EQ(&RC2, &*I++);
   EXPECT_EQ(CG.postorder_ref_scc_end(), I);
 }
+
+TEST(LazyCallGraphTest, AddNewFunctionIntoRefSCC) {
+  LLVMContext Context;
+  // Build and populate a graph composed of a single, self-referential node.
+  std::unique_ptr<Module> M = parseAssembly(Context, "define void @f() {\n"
+                                                     "entry:\n"
+                                                     "  call void @f()\n"
+                                                     "  ret void\n"
+                                                     "}\n");
+  LazyCallGraph CG = buildCG(*M);
+  CG.buildRefSCCs();
+
+  // At this point 'f' is in the call graph.
+  auto &F = lookupFunction(*M, "f");
+  LazyCallGraph::Node *FN = CG.lookup(F);
+  EXPECT_NE(FN, nullptr);
+
+  // And it has an SCC, of course.
+  auto *FSCC = CG.lookupSCC(*FN);
+  EXPECT_NE(FSCC, nullptr);
+
+  // Now, create a new function 'g'.
+  auto *G = Function::Create(F.getFunctionType(), F.getLinkage(),
+                             F.getAddressSpace(), "g", F.getParent());
+  BasicBlock::Create(F.getParent()->getContext(), "entry", G);
+
+  // Instruct the LazyCallGraph to create a new node for 'g', within the same
+  // RefSCC as 'f', but in a separate SCC.
+  CG.addNewFunctionIntoRefSCC(*G, FSCC->getOuterRefSCC());
+
+  // 'g' should now be in the call graph.
+  LazyCallGraph::Node *GN = CG.lookup(*G);
+  EXPECT_NE(GN, nullptr);
+  // 'g' should have an SCC, composed of the singular node 'g'.
+  // ('f' should not be included in the 'g' SCC.)
+  LazyCallGraph::SCC *GSCC = CG.lookupSCC(*GN);
+  EXPECT_NE(GSCC, nullptr);
+  EXPECT_EQ(GSCC->size(), 1);
+  EXPECT_NE(GSCC, FSCC);
+  // 'g' and 'f' should be part of the same RefSCC.
+  EXPECT_EQ(&GSCC->getOuterRefSCC(), &FSCC->getOuterRefSCC());
+}
 }