[PM] Port ArgumentPromotion to the new pass manager.

[chandlerc]

Summary

This requires teaching the LazyCallGraph how to replace functions in the
graph in-place, but otherwise is a fairly direct port of the existing
logic.

The one interesting thing is that LCG can handle much more directly
making this update which allows us to just skip the hooks in the core
argument promotion logic used by the old PM. This is largely a function
of not tracking call instructions and instead tracking the actual edges
in the graph at an abstract level.

Note that this isn't quite ready to submit yet. It exposes a crash in the
inliner that I'm still investigating. But the code seems reasonable enough and
I have both unittesting for the LCG change and direct testing of argpromote.

[davide]

I have one high level comment about the tests (it's not your fault but I think it's a good time to recover). I assume you tried to bootstrap and the testsuite with this change and it's clean.

include/llvm/Analysis/LazyCallGraph.h

234

s/tranfsormations/transformations/

lib/Transforms/IPO/ArgumentPromotion.cpp

1081 ↗

(On Diff #86267)

s/updat e/update/

test/Transforms/ArgumentPromotion/aggregate-promote.ll

2–4 ↗

(On Diff #86267)

Ugh for tests still using grep. If you're fine with that, I would rather do the following:

1. Convert them to FileCheck
2. Expand the output so that they don't accidentally test the wrong thing
3. Make sure the output between old and new PM matches.

Maybe I'm being overcautious here, but I've been bitten already by cases where the output didn't quite match (and was hiding bugs).

[chandlerc]

Strip this down to the just the LazyCallGraph enhancement.

The core ArgumentPromotion change needs to wait for more LCG improvements that
I can effectively separate.

I will also deal with the grep-based tests when dealing with the pass itself.

[chandlerc]

Marking that this is dead. I'll be sending out a fresh series of patches that take us in a different (and IMO much better) direction for argpromote...

include/llvm/Analysis/LazyCallGraph.h

public:
   /// callees, de-duplicate and provide fast testing of whether a function is
   /// a callee, and facilitate iteration of child nodes in the graph.
   class Node {
     friend class LazyCallGraph;
     friend class LazyCallGraph::SCC;
     friend class LazyCallGraph::RefSCC;
 
     LazyCallGraph *G;
-    Function &F;
+    Function *F;
 
     // We provide for the DFS numbering and Tarjan walk lowlink numbers to be
     // stored directly within the node. These are both '-1' when nodes are part
     // of an SCC (or RefSCC), or '0' when not yet reached in a DFS walk.
     int DFSNumber;
     int LowLink;
 
     mutable EdgeVectorT Edges;
class Node {
     void insertEdgeInternal(Node &ChildN, Edge::Kind EK);
 
     /// Internal helper to change an edge kind.
     void setEdgeKind(Function &ChildF, Edge::Kind EK);
 
     /// Internal helper to remove the edge to the given function.
     void removeEdgeInternal(Function &ChildF);
 
+    /// Internal helper to directly replace the function with a new one.
+    ///
+    /// This is used to facilitate tranfsormations which need to replace the
+    /// formal Function object but directly move the body and users from one to
+    /// the other.
+    void replaceFunction(Function &NewF);
+
+    /// Internal helper to replace an edge key with a new one.
+    ///
+    /// This should be used when the function for a particular node in the
+    /// graph gets replaced and we are updating all of the edges to that node
+    /// to use the new function as the key.
+    void replaceEdgeKey(Function &OldTarget, Function &NewTarget);
+
     void clear() {
       Edges.clear();
       EdgeIndexMap.clear();
     }
 
     /// Print the name of this node's function.
     friend raw_ostream &operator<<(raw_ostream &OS, const Node &N) {
-      return OS << N.F.getName();
+      return OS << N.F->getName();
     }
 
     /// Dump the name of this node's function to stderr.
     void dump() const;
 
   public:
     LazyCallGraph &getGraph() const { return *G; }
 
-    Function &getFunction() const { return F; }
+    Function &getFunction() const { return *F; }
 
     edge_iterator begin() const {
       return edge_iterator(Edges.begin(), Edges.end());
     }
     edge_iterator end() const { return edge_iterator(Edges.end(), Edges.end()); }
 
     const Edge &operator[](int i) const { return Edges[i]; }
     const Edge &operator[](Function &F) const {
public:
     /// This is trivial whenever the target is in the same RefSCC as the source
     /// or the edge is an outgoing edge to some descendant RefSCC. In these
     /// cases there is no change to the cyclic structure of the RefSCCs.
     ///
     /// To further make calling this convenient, it also handles inserting
     /// already existing edges.
     void insertTrivialRefEdge(Node &SourceN, Node &TargetN);
 
+    /// Directly replace a node's function with a new function.
+    ///
+    /// This should be used when moving the body and users of a function to
+    /// a new formal function object but not otherwise changing the call graph
+    /// structure in any way.
+    ///
+    /// It requires that the old function in the provided node have zero uses
+    /// and the new function must have calls and references to it establishing
+    /// an equivalent graph.
+    void replaceNodeFunction(Node &N, Function &NewF);
+
     ///@}
   };
 
   /// A post-order depth-first RefSCC iterator over the call graph.
   ///
   /// This iterator triggers the Tarjan DFS-based formation of the RefSCC (and
   /// SCC) DAG for the call graph, walking it lazily in depth-first post-order.
   /// That is, it always visits RefSCCs for the target of a reference edge

lib/Analysis/LazyCallGraph.cpp

static void addEdge(SmallVectorImpl<LazyCallGraph::Edge> &Edges,
   if (!EdgeIndexMap.insert({&F, Edges.size()}).second)
     return;
 
   DEBUG(dbgs() << "    Added callable function: " << F.getName() << "\n");
   Edges.emplace_back(LazyCallGraph::Edge(F, EK));
 }
 
 LazyCallGraph::Node::Node(LazyCallGraph &G, Function &F)
-    : G(&G), F(F), DFSNumber(0), LowLink(0) {
+    : G(&G), F(&F), DFSNumber(0), LowLink(0) {
   DEBUG(dbgs() << "  Adding functions called by '" << F.getName()
                << "' to the graph.\n");
 
   SmallVector<Constant *, 16> Worklist;
   SmallPtrSet<Function *, 4> Callees;
   SmallPtrSet<Constant *, 16> Visited;
 
   // Find all the potential call graph edges in this function. We track both
void LazyCallGraph::Node::removeEdgeInternal(Function &Target) {
   auto IndexMapI = EdgeIndexMap.find(&Target);
   assert(IndexMapI != EdgeIndexMap.end() &&
          "Target not in the edge set for this caller?");
 
   Edges[IndexMapI->second] = Edge();
   EdgeIndexMap.erase(IndexMapI);
 }
 
+void LazyCallGraph::Node::replaceFunction(Function &NewF) {
+  assert(F != &NewF && "Must not replace a function with itself!");
+  F = &NewF;
+}
+
+void LazyCallGraph::Node::replaceEdgeKey(Function &OldTarget, Function &NewTarget) {
+  auto IndexMapI = EdgeIndexMap.find(&OldTarget);
+  assert(IndexMapI != EdgeIndexMap.end() && "Old target not in the edge set!");
+  int Index = IndexMapI->second;
+  assert(&Edges[Index].getFunction() == &NewTarget &&
+         "Edge's target hasn't been updated yet!");
+  // Ok, everything is in order, so remove the old key and insert a new one.
+  EdgeIndexMap.erase(IndexMapI);
+  bool Inserted = EdgeIndexMap.insert({&NewTarget, Index}).second;
+  (void)Inserted;
+  assert(Inserted && "Didn't actually insert a new key!");
+}
+
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
 LLVM_DUMP_METHOD void LazyCallGraph::Node::dump() const {
   dbgs() << *this << '\n';
 }
 #endif
 
 LazyCallGraph::LazyCallGraph(Module &M) : NextDFSNumber(0) {
   DEBUG(dbgs() << "Building CG for module: " << M.getModuleIdentifier()
#endif
 
   // Create the new edge.
   SourceN.Edges.emplace_back(TargetN, Edge::Ref);
 
   // Now that we have the edge, handle the graph fallout.
   handleTrivialEdgeInsertion(SourceN, TargetN);
 }
 
+void LazyCallGraph::RefSCC::replaceNodeFunction(Node &N, Function &NewF) {
+#ifndef NDEBUG
+  // Check that the RefSCC is still valid when we finish.
+  auto ExitVerifier = make_scope_exit([this] { verify(); });
+
+  assert(G->lookupRefSCC(N) == this &&
+         "Cannot replace the function of a node outside this RefSCC.");
+
+  assert(G->NodeMap.find(&NewF) == G->NodeMap.end() &&
+         "Must not have already walked the new function!'");
+  assert(G->EntryIndexMap.find(&NewF) == G->EntryIndexMap.end() &&
+         "Must not have already walked the new function!'");
+
+  // It is important that this replacement not introduce graph changes so we
+  // insist that the caller has already removed every use of the original
+  // function and that all uses of the new function correspond to existing
+  // edges in the graph. The common and expected way to use this is when
+  // replacing the function itself in the IR without changing the call graph
+  // shape and just updating the analysis based on that.
+  Function &OldF = N.getFunction();
+  assert(&OldF != &NewF && "Cannot replace a function with itself!");
+  assert(OldF.use_empty() &&
+         "Must have moved all uses from the old function to the new!");
+  assert(llvm::find(G->RefSCCEntryNodes, &OldF) == G->RefSCCEntryNodes.end() &&
+         "Must not replace a function which hasn't yet been processed into the "
+         "SCC graph!");
+#endif
+
+  N.replaceFunction(NewF);
+
+  // Update various call graph maps.
+  G->NodeMap.erase(&OldF);
+  G->NodeMap[&NewF] = &N;
+  auto EntryI = G->EntryIndexMap.find(&OldF);
+  if (EntryI != G->EntryIndexMap.end()) {
+    int Index = EntryI->second;
+    G->EntryIndexMap.erase(EntryI);
+    G->EntryIndexMap[&NewF] = Index;
+  }
+
+  // Update all of the edges from other nodes in the call graph to this
+  // function.
+  SmallPtrSet<Function *, 4> Updated;
+  SmallPtrSet<User *, 16> Visited;
+  SmallVector<User *, 16> Worklist(NewF.user_begin(), NewF.user_end());
+  for (User *U : Worklist)
+    Visited.insert(U);
+  while (!Worklist.empty()) {
+    User *U = Worklist.pop_back_val();
+    // If the user is a constant, we just recurse down through its users.
+    if (auto *C = dyn_cast<Constant>(U)) {
+      for (User *CU : C->users())
+        if (Visited.insert(CU).second)
+          Worklist.push_back(CU);
+      continue;
+    }
+
+    // We don't care what kind of instruction this is, but it must be an
+    // instruction at the bottom of the recursive walk through the users of the
+    // constants.
+    auto *I = cast<Instruction>(U);
+
+    // Find the function and make sure we haven't processed it yet.
+    Function &UserF = *I->getParent()->getParent();
+    if (!Updated.insert(&UserF).second)
+      // Already handled this node.
+      continue;
+
+    Node *UserN = G->lookup(UserF);
+    if (!UserN)
+      // No need to update a function that hasn't been put into a node.
+      continue;
+
+    assert(UserN->lookup(OldF) &&
+           "No edge from a caller to the replacement node!");
+    assert(UserN->lookup(OldF)->getNode() == &N &&
+           "Edge from a caller doesn't point to the node!");
+
+    UserN->replaceEdgeKey(OldF, NewF);
+  }
+}
+
 void LazyCallGraph::insertEdge(Node &SourceN, Function &Target, Edge::Kind EK) {
   assert(SCCMap.empty() && DFSStack.empty() &&
          "This method cannot be called after SCCs have been formed!");
 
   return SourceN.insertEdgeInternal(Target, EK);
 }
 
 void LazyCallGraph::removeEdge(Node &SourceN, Function &Target) {

unittests/Analysis/LazyCallGraphTest.cpp

TEST(LazyCallGraphTest, HandleBlockAddress) {
 
   LazyCallGraph::Node &F = *CG.lookup(lookupFunction(*M, "f"));
   LazyCallGraph::Node &G = *CG.lookup(lookupFunction(*M, "g"));
   EXPECT_EQ(&FRC, CG.lookupRefSCC(F));
   EXPECT_EQ(&GRC, CG.lookupRefSCC(G));
   EXPECT_TRUE(GRC.isParentOf(FRC));
 }
 
+TEST(LazyCallGraphTest, ReplaceNodeFunction) {
+  LLVMContext Context;
+  // A graph with several different kinds of edges pointing at a particular
+  // function.
+  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"
+               "}\n"
+               "define void @b(i8** %ptr) {\n"
+               "entry:\n"
+               "  store i8* bitcast (void(i8**)* @d to i8*), i8** %ptr\n"
+               "  store i8* bitcast (void(i8**)* @d to i8*), i8** %ptr\n"
+               "  call void @d(i8** %ptr)"
+               "  ret void\n"
+               "}\n"
+               "define void @c(i8** %ptr) {\n"
+               "entry:\n"
+               "  call void @d(i8** %ptr)"
+               "  call void @d(i8** %ptr)"
+               "  store i8* bitcast (void(i8**)* @d to i8*), i8** %ptr\n"
+               "  ret void\n"
+               "}\n"
+               "define void @d(i8** %ptr) {\n"
+               "entry:\n"
+               "  store i8* bitcast (void(i8**)* @b to i8*), i8** %ptr\n"
+               "  call void @c(i8** %ptr)"
+               "  call void @d(i8** %ptr)"
+               "  store i8* bitcast (void(i8**)* @d to i8*), i8** %ptr\n"
+               "  ret void\n"
+               "}\n");
+  LazyCallGraph CG(*M);
+
+  // Force the graph to be fully expanded.
+  auto I = CG.postorder_ref_scc_begin();
+  LazyCallGraph::RefSCC &RC1 = *I++;
+  LazyCallGraph::RefSCC &RC2 = *I++;
+  EXPECT_EQ(CG.postorder_ref_scc_end(), I);
+
+  ASSERT_EQ(2u, RC1.size());
+  LazyCallGraph::SCC &C1 = RC1[0];
+  LazyCallGraph::SCC &C2 = RC1[1];
+
+  LazyCallGraph::Node &AN = *CG.lookup(lookupFunction(*M, "a"));
+  LazyCallGraph::Node &BN = *CG.lookup(lookupFunction(*M, "b"));
+  LazyCallGraph::Node &CN = *CG.lookup(lookupFunction(*M, "c"));
+  LazyCallGraph::Node &DN = *CG.lookup(lookupFunction(*M, "d"));
+  EXPECT_EQ(&C1, CG.lookupSCC(DN));
+  EXPECT_EQ(&C1, CG.lookupSCC(CN));
+  EXPECT_EQ(&C2, CG.lookupSCC(BN));
+  EXPECT_EQ(&RC1, CG.lookupRefSCC(DN));
+  EXPECT_EQ(&RC1, CG.lookupRefSCC(CN));
+  EXPECT_EQ(&RC1, CG.lookupRefSCC(BN));
+  EXPECT_EQ(&RC2, CG.lookupRefSCC(AN));
+
+  // Now we need to build a new function 'e' with the same signature as 'd'.
+  Function &D = DN.getFunction();
+  Function &E = *Function::Create(D.getFunctionType(), D.getLinkage(), "e");
+  D.getParent()->getFunctionList().insert(D.getIterator(), &E);
+
+  // Change each use of 'd' to use 'e'. This is particularly easy as they have
+  // the same type.
+  D.replaceAllUsesWith(&E);
+
+  // Splice the body of the old function into the new one.
+  E.getBasicBlockList().splice(E.begin(), D.getBasicBlockList());
+  // And fix up the one argument.
+  D.arg_begin()->replaceAllUsesWith(&*E.arg_begin());
+  E.arg_begin()->takeName(&*D.arg_begin());
+
+  // Now replace the function in the graph.
+  RC1.replaceNodeFunction(DN, E);
+
+  EXPECT_EQ(&E, &DN.getFunction());
+  EXPECT_EQ(&DN, CN[E].getNode());
+  EXPECT_EQ(&DN, BN[E].getNode());
+}
+
 }