From 5057e0ff5698124875c01a078c9c2b7e4dbc4bd2 Mon Sep 17 00:00:00 2001
From: Teresa Johnson <tejohnson@google.com>
Date: Fri, 21 Feb 2025 12:12:29 -0800
Subject: [PATCH] [MemProf] Refactor backedge computation and invoke earlier

Invoke the backedge computation (refactored as a new method) at the end
of the graph construction, instead of at the start of cloning. That
makes more logical sense, and it also makes it easier to look at the
results in the postbuild dot graph with a follow on change to display
those differently.
---
 .../IPO/MemProfContextDisambiguation.cpp      | 44 ++++++++++++-------
 1 file changed, 29 insertions(+), 15 deletions(-)

diff --git a/llvm/lib/Transforms/IPO/MemProfContextDisambiguation.cpp b/llvm/lib/Transforms/IPO/MemProfContextDisambiguation.cpp
index 8bfbb2c348459..705806b98f8da 100644
--- a/llvm/lib/Transforms/IPO/MemProfContextDisambiguation.cpp
+++ b/llvm/lib/Transforms/IPO/MemProfContextDisambiguation.cpp
@@ -530,6 +530,9 @@ class CallsiteContextGraph {
   /// multiple different callee target functions.
   void handleCallsitesWithMultipleTargets();
 
+  /// Mark backedges via the standard DFS based backedge algorithm.
+  void markBackedges();
+
   // Try to partition calls on the given node (already placed into the AllCalls
   // array) by callee function, creating new copies of Node as needed to hold
   // calls with different callees, and moving the callee edges appropriately.
@@ -740,6 +743,7 @@ class CallsiteContextGraph {
   void moveCalleeEdgeToNewCaller(const std::shared_ptr<ContextEdge> &Edge,
                                  ContextNode *NewCaller);
 
+  /// Recursive helper for marking backedges via DFS.
   void markBackedges(ContextNode *Node, DenseSet<const ContextNode *> &Visited,
                      DenseSet<const ContextNode *> &CurrentStack);
 
@@ -2108,6 +2112,8 @@ ModuleCallsiteContextGraph::ModuleCallsiteContextGraph(
 
   handleCallsitesWithMultipleTargets();
 
+  markBackedges();
+
   // Strip off remaining callsite metadata, no longer needed.
   for (auto &FuncEntry : FuncToCallsWithMetadata)
     for (auto &Call : FuncEntry.second)
@@ -2204,6 +2210,8 @@ IndexCallsiteContextGraph::IndexCallsiteContextGraph(
   updateStackNodes();
 
   handleCallsitesWithMultipleTargets();
+
+  markBackedges();
 }
 
 template <typename DerivedCCG, typename FuncTy, typename CallTy>
@@ -3405,6 +3413,27 @@ void CallsiteContextGraph<DerivedCCG, FuncTy, CallTy>::
 
 // This is the standard DFS based backedge discovery algorithm.
 template <typename DerivedCCG, typename FuncTy, typename CallTy>
+void CallsiteContextGraph<DerivedCCG, FuncTy, CallTy>::markBackedges() {
+  // If we are cloning recursive contexts, find and mark backedges from all root
+  // callers, using the typical DFS based backedge analysis.
+  if (!CloneRecursiveContexts)
+    return;
+  DenseSet<const ContextNode *> Visited;
+  DenseSet<const ContextNode *> CurrentStack;
+  for (auto &Entry : NonAllocationCallToContextNodeMap) {
+    auto *Node = Entry.second;
+    if (Node->isRemoved())
+      continue;
+    // It is a root if it doesn't have callers.
+    if (!Node->CallerEdges.empty())
+      continue;
+    markBackedges(Node, Visited, CurrentStack);
+    assert(CurrentStack.empty());
+  }
+}
+
+// Recursive helper for above markBackedges method.
+template <typename DerivedCCG, typename FuncTy, typename CallTy>
 void CallsiteContextGraph<DerivedCCG, FuncTy, CallTy>::markBackedges(
     ContextNode *Node, DenseSet<const ContextNode *> &Visited,
     DenseSet<const ContextNode *> &CurrentStack) {
@@ -3429,22 +3458,7 @@ void CallsiteContextGraph<DerivedCCG, FuncTy, CallTy>::markBackedges(
 
 template <typename DerivedCCG, typename FuncTy, typename CallTy>
 void CallsiteContextGraph<DerivedCCG, FuncTy, CallTy>::identifyClones() {
-  // If we are cloning recursive contexts, find and mark backedges from all root
-  // callers, using the typical DFS based backedge analysis.
   DenseSet<const ContextNode *> Visited;
-  if (CloneRecursiveContexts) {
-    DenseSet<const ContextNode *> CurrentStack;
-    for (auto &Entry : NonAllocationCallToContextNodeMap) {
-      auto *Node = Entry.second;
-      if (Node->isRemoved())
-        continue;
-      // It is a root if it doesn't have callers.
-      if (!Node->CallerEdges.empty())
-        continue;
-      markBackedges(Node, Visited, CurrentStack);
-      assert(CurrentStack.empty());
-    }
-  }
   for (auto &Entry : AllocationCallToContextNodeMap) {
     Visited.clear();
     identifyClones(Entry.second, Visited, Entry.second->getContextIds());