[NewPM][CGSCC] Handle newly added functions not directly referenced by existing nodes

[aeubanks]

Summary

As mentioned in https://reviews.llvm.org/D87798, newly created functions
that aren't referenced by an existing node are not properly handled.

This checks all newly added functions for references to other functions
that don't have a populated node.

Coroutines can split a function into multiple functions. For example,
foo can be split into foo, foo.1, and foo.2. foo may reference foo.1,
but foo.2 is only referenced by foo.1.
coro-retcon-resume-values2.ll actually has an instance of this
happening, which would trigger an assert that a node is unpopulated.

[MaskRay]

llvm/lib/Analysis/CGSCCPassManager.cpp

981–992

The whole loop deserves a comment, along the line of: the newly created nodes are referenced by some SCC nodes via call/ref edges. We conservatively assume the newly created nodes are part of the SCC (as if they call back to the SCC).

983

assert(!NewNode->isPopulated())

995

Having the container in the loop can lead to duplicate traversal. Move NewNodesWorklist & NewNodesVisited outside the loop (and decrease 16 -> 4 may be sufficient)

Additionally, should NewNodesVisited affect the G.getLibFunctions() loop below which uses Visited?

[MaskRay]

coro-retcon-resume-values2.ll actually has an instance of this happening.

Might be better to elaborate this a bit.

[aeubanks]

rebase
add comments

[aeubanks]

llvm/lib/Analysis/CGSCCPassManager.cpp

981–992

Added comments, better?

983

There might be multiple references to a new node added into NewNodes before we reach the new node.

995

Each iteration will traverse a different function. The only potential speedup is if multiple functions reference the same Constant.

The G.getLibFunctions() part below is only relevant to the original function in N, doesn't affect the other functions we are traversing here.

[aeubanks]

Actually I think I need to better handle newly created functions. Currently I'm putting them in the same SCC as the current node, but in reality it might need to be in its own SCC (else we'd break some invariants about all nodes being able to reach all other nodes in the same SCC)

[aeubanks]

https://reviews.llvm.org/D93828 is a more principled way of fixing this

llvm/lib/Analysis/CGSCCPassManager.cpp

static LazyCallGraph::SCC &updateCGAndAnalysisManagerForPass(
 
   // First walk the function and handle all called functions. We do this first
   // because if there is a single call edge, whether there are ref edges is
   // irrelevant.
   for (Instruction &I : instructions(F)) {
     if (auto *CB = dyn_cast<CallBase>(&I)) {
       if (Function *Callee = CB->getCalledFunction()) {
         if (Visited.insert(Callee).second && !Callee->isDeclaration()) {
-          Node *CalleeN = G.lookup(*Callee);
-          if (!CalleeN) {
-            CalleeN = &G.get(*Callee);
+          Node *CalleeN = &G.get(*Callee);
+          if (!CalleeN->isPopulated()) {
             NewNodes.insert(CalleeN);
           }
           Edge *E = N->lookup(*CalleeN);
           assert((E || !FunctionPass) &&
                  "No function transformations should introduce *new* "
                  "call edges! Any new calls should be modeled as "
                  "promoted existing ref edges!");
           bool Inserted = RetainedEdges.insert(CalleeN).second;
static LazyCallGraph::SCC &updateCGAndAnalysisManagerForPass(
   // Now walk all references.
   for (Instruction &I : instructions(F))
     for (Value *Op : I.operand_values())
       if (auto *OpC = dyn_cast<Constant>(Op))
         if (Visited.insert(OpC).second)
           Worklist.push_back(OpC);
 
   auto VisitRef = [&](Function &Referee) {
-    Node *RefereeN = G.lookup(Referee);
-    if (!RefereeN) {
-      RefereeN = &G.get(Referee);
+    Node *RefereeN = &G.get(Referee);
+    if (!RefereeN->isPopulated()) {
       NewNodes.insert(RefereeN);
     }
     Edge *E = N->lookup(*RefereeN);
     assert((E || !FunctionPass) &&
            "No function transformations should introduce *new* ref "
            "edges! Any new ref edges would require IPO which "
            "function passes aren't allowed to do!");
     bool Inserted = RetainedEdges.insert(RefereeN).second;
     (void)Inserted;
     assert(Inserted && "We should never visit a function twice.");
     if (!E)
       NewRefEdges.insert(RefereeN);
     else if (E->isCall())
       DemotedCallTargets.insert(RefereeN);
   };
   LazyCallGraph::visitReferences(Worklist, Visited, VisitRef);
 
-  for (Node *NewNode : NewNodes)
+  // Find potential new nodes referenced by previous new nodes that wouldn't be
+  // found by only scanning the original function.
+  while (!NewNodes.empty()) {
+    Node *NewNode = NewNodes.pop_back_val();
+    if (NewNode->isPopulated())
+      continue;
+    // We treat newly added functions as initially being in the same SCC as the
+    // current function. Newly added functions shouldn't reference functions
+    // that the original function did not.
+    // FIXME: We currently skip running the CGSCC pipeline on newly added
+    // functions. Currently coro-split works around this by manually adding the
+    // current SCC into UR.CWorklist.
     G.initNode(*NewNode, *C);
 
+    SmallVector<Constant *, 4> NewNodeWorklist;
+    SmallPtrSet<Constant *, 4> NewNodeVisited;
+
+    for (Instruction &I : instructions(NewNode->getFunction()))
+      for (Value *Op : I.operand_values())
+        if (auto *C = dyn_cast<Constant>(Op))
+          if (NewNodeVisited.insert(C).second)
+            NewNodeWorklist.push_back(C);
+
+    auto CheckNewNode = [&](Function &Referee) {
+      auto &N = G.get(Referee);
+      if (!N.isPopulated())
+        NewNodes.insert(&N);
+    };
+
+    LazyCallGraph::visitReferences(NewNodeWorklist, NewNodeVisited,
+                                   CheckNewNode);
+  }
+
   // Handle new ref edges.
   for (Node *RefTarget : NewRefEdges) {
     SCC &TargetC = *G.lookupSCC(*RefTarget);
     RefSCC &TargetRC = TargetC.getOuterRefSCC();
     (void)TargetRC;
     // TODO: This only allows trivial edges to be added for now.
     assert((RC == &TargetRC ||
            RC->isAncestorOf(TargetRC)) && "New ref edge is not trivial!");

llvm/test/Transforms/Coroutines/coro-retcon-resume-values2.ll

 ; RUN: opt < %s -coro-split -coro-cleanup -S | FileCheck %s
+; RUN: opt < %s -passes=coro-split,coro-cleanup -S | FileCheck %s
 
 define i8* @f(i8* %buffer, i32 %n) "coroutine.presplit"="1" {
 entry:
   %id = call token @llvm.coro.id.retcon(i32 8, i32 4, i8* %buffer, i8* bitcast (i8* (i8*, i32)* @prototype to i8*), i8* bitcast (i8* (i32)* @allocate to i8*), i8* bitcast (void (i8*)* @deallocate to i8*))
   %hdl = call i8* @llvm.coro.begin(token %id, i8* null)
   %value0 = call i32 (...) @llvm.coro.suspend.retcon.i32()
   %sum0 = call i32 @add(i32 %n, i32 %value0)
   %value1 = call i32 (...) @llvm.coro.suspend.retcon.i32()

llvm/unittests/Analysis/CGSCCPassManagerTest.cpp

CGPM.addPass(LambdaSCCPassNoPreserve(
         }
       }));
 
   ModulePassManager MPM(/*DebugLogging*/ true);
   MPM.addPass(createModuleToPostOrderCGSCCPassAdaptor(std::move(CGPM)));
   MPM.run(*M, MAM);
 }
 
+TEST_F(CGSCCPassManagerTest, TestInsertionOfNewRefSCCIndirect) {
+  std::unique_ptr<Module> M = parseIR("define void @f() {\n"
+                                      "entry:\n"
+                                      "  ret void\n"
+                                      "}\n");
+
+  CGSCCPassManager CGPM(/*DebugLogging*/ true);
+  CGPM.addPass(LambdaSCCPassNoPreserve([&](LazyCallGraph::SCC &C,
+                                           CGSCCAnalysisManager &AM,
+                                           LazyCallGraph &CG,
+                                           CGSCCUpdateResult &UR) {
+    auto &FAM =
+        AM.getResult<FunctionAnalysisManagerCGSCCProxy>(C, CG).getManager();
+
+    for (auto &N : C) {
+      auto &F = N.getFunction();
+      if (F.getName() != "f")
+        continue;
+
+      // Create new functions g and h.
+      auto *G = Function::Create(F.getFunctionType(), F.getLinkage(),
+                                 F.getAddressSpace(), "g", F.getParent());
+      auto *H = Function::Create(F.getFunctionType(), F.getLinkage(),
+                                 F.getAddressSpace(), "h", F.getParent());
+      // Create f -> g -> h call edges.
+      BasicBlock *GBB =
+          BasicBlock::Create(F.getParent()->getContext(), "entry", G);
+      BasicBlock *HBB =
+          BasicBlock::Create(F.getParent()->getContext(), "entry", H);
+      (void)CallInst::Create(G, {}, "", &F.getEntryBlock().front());
+      (void)CallInst::Create(H, {}, "", GBB);
+      (void)ReturnInst::Create(G->getContext(), GBB);
+      (void)ReturnInst::Create(H->getContext(), HBB);
+
+      ASSERT_NO_FATAL_FAILURE(
+          updateCGAndAnalysisManagerForCGSCCPass(CG, C, N, AM, UR, FAM))
+          << "Updating the call graph with f -> g -> h caused a fatal failure";
+    }
+  }));
+
+  ModulePassManager MPM(/*DebugLogging*/ true);
+  MPM.addPass(createModuleToPostOrderCGSCCPassAdaptor(std::move(CGPM)));
+  MPM.run(*M, MAM);
+}
+
 TEST_F(CGSCCPassManagerTest, TestInsertionOfNewRefSCCMutuallyRecursive) {
   std::unique_ptr<Module> M = parseIR("define void @f() {\n"
                                       "entry:\n"
                                       "  ret void\n"
                                       "}\n");
 
   CGSCCPassManager CGPM(/*DebugLogging*/ true);
   CGPM.addPass(LambdaSCCPassNoPreserve([&](LazyCallGraph::SCC &C,