[llvm] [MemProf] Remove empty edges once after cloning (PR #85320)

Teresa Johnson via llvm-commits llvm-commits at lists.llvm.org
Tue Mar 26 19:21:46 PDT 2024


https://github.com/teresajohnson updated https://github.com/llvm/llvm-project/pull/85320

>From fc1ceaca17046174a6c860f28091195190bb8cd2 Mon Sep 17 00:00:00 2001
From: Teresa Johnson <tejohnson at google.com>
Date: Thu, 14 Mar 2024 14:49:07 -0700
Subject: [PATCH 1/3] [MemProf] Remove empty edges once after cloning

Restructure the handling of edges that become empty during the cloning
process. Instead of removing them as they become empty (no context ids
and alloc type), do this once after all cloning is complete.

This has no effect on the cloning result, but prepares for a follow on
change that does improve the cloning. The structural change here reduces
the diffs for the follow on change, which would be much more difficult
with the previous handling.
---
 .../IPO/MemProfContextDisambiguation.cpp      | 77 ++++++++++++-------
 1 file changed, 48 insertions(+), 29 deletions(-)

diff --git a/llvm/lib/Transforms/IPO/MemProfContextDisambiguation.cpp b/llvm/lib/Transforms/IPO/MemProfContextDisambiguation.cpp
index 271d3ed40030b4..742b62a277a4b9 100644
--- a/llvm/lib/Transforms/IPO/MemProfContextDisambiguation.cpp
+++ b/llvm/lib/Transforms/IPO/MemProfContextDisambiguation.cpp
@@ -306,7 +306,7 @@ class CallsiteContextGraph {
     // True if this node was effectively removed from the graph, in which case
     // its context id set, caller edges, and callee edges should all be empty.
     bool isRemoved() const {
-      assert(ContextIds.empty() ==
+      assert(!ContextIds.empty() ||
              (CalleeEdges.empty() && CallerEdges.empty()));
       return ContextIds.empty();
     }
@@ -349,9 +349,12 @@ class CallsiteContextGraph {
     }
   };
 
-  /// Helper to remove callee edges that have allocation type None (due to not
+  /// Helpers to remove callee edges that have allocation type None (due to not
   /// carrying any context ids) after transformations.
   void removeNoneTypeCalleeEdges(ContextNode *Node);
+  void
+  recursivelyRemoveNoneTypeCalleeEdges(ContextNode *Node,
+                                       DenseSet<const ContextNode *> &Visited);
 
 protected:
   /// Get a list of nodes corresponding to the stack ids in the given callsite
@@ -2427,11 +2430,42 @@ void CallsiteContextGraph<DerivedCCG, FuncTy, CallTy>::
   }
 }
 
+template <typename DerivedCCG, typename FuncTy, typename CallTy>
+void CallsiteContextGraph<DerivedCCG, FuncTy, CallTy>::
+    recursivelyRemoveNoneTypeCalleeEdges(
+        ContextNode *Node, DenseSet<const ContextNode *> &Visited) {
+  auto Inserted = Visited.insert(Node);
+  if (!Inserted.second)
+    return;
+
+  removeNoneTypeCalleeEdges(Node);
+
+  for (auto *Clone : Node->Clones)
+    recursivelyRemoveNoneTypeCalleeEdges(Clone, Visited);
+
+  // The recursive call may remove some of this Node's caller edges.
+  // Iterate over a copy and skip any that were removed.
+  auto CallerEdges = Node->CallerEdges;
+  for (auto &Edge : CallerEdges) {
+    // Skip any that have been removed by an earlier recursive call.
+    if (Edge->Callee == nullptr && Edge->Caller == nullptr) {
+      assert(!std::count(Node->CallerEdges.begin(), Node->CallerEdges.end(),
+                         Edge));
+      continue;
+    }
+    recursivelyRemoveNoneTypeCalleeEdges(Edge->Caller, Visited);
+  }
+}
+
 template <typename DerivedCCG, typename FuncTy, typename CallTy>
 void CallsiteContextGraph<DerivedCCG, FuncTy, CallTy>::identifyClones() {
   DenseSet<const ContextNode *> Visited;
   for (auto &Entry : AllocationCallToContextNodeMap)
     identifyClones(Entry.second, Visited);
+  Visited.clear();
+  for (auto &Entry : AllocationCallToContextNodeMap)
+    recursivelyRemoveNoneTypeCalleeEdges(Entry.second, Visited);
+  check();
 }
 
 // helper function to check an AllocType is cold or notcold or both.
@@ -2446,7 +2480,7 @@ template <typename DerivedCCG, typename FuncTy, typename CallTy>
 void CallsiteContextGraph<DerivedCCG, FuncTy, CallTy>::identifyClones(
     ContextNode *Node, DenseSet<const ContextNode *> &Visited) {
   if (VerifyNodes)
-    checkNode<DerivedCCG, FuncTy, CallTy>(Node);
+    checkNode<DerivedCCG, FuncTy, CallTy>(Node, /*CheckEdges=*/false);
   assert(!Node->CloneOf);
 
   // If Node as a null call, then either it wasn't found in the module (regular
@@ -2507,8 +2541,16 @@ void CallsiteContextGraph<DerivedCCG, FuncTy, CallTy>::identifyClones(
   std::stable_sort(Node->CallerEdges.begin(), Node->CallerEdges.end(),
                    [&](const std::shared_ptr<ContextEdge> &A,
                        const std::shared_ptr<ContextEdge> &B) {
-                     assert(checkColdOrNotCold(A->AllocTypes) &&
-                            checkColdOrNotCold(B->AllocTypes));
+                     // Nodes with non-empty context ids should be sorted before
+                     // those with empty context ids.
+                     if (A->ContextIds.empty())
+                       // Either B ContextIds are non-empty (in which case we
+                       // should return false because B < A), or B ContextIds
+                       // are empty, in which case they are equal, and we should
+                       // maintain the original relative ordering.
+                       return false;
+                     if (B->ContextIds.empty())
+                       return true;
 
                      if (A->AllocTypes == B->AllocTypes)
                        // Use the first context id for each edge as a
@@ -2587,39 +2629,16 @@ void CallsiteContextGraph<DerivedCCG, FuncTy, CallTy>::identifyClones(
            Node->AllocTypes != (uint8_t)AllocationType::None);
     // Sanity check that no alloc types on clone or its edges are None.
     assert(Clone->AllocTypes != (uint8_t)AllocationType::None);
-    assert(llvm::none_of(
-        Clone->CallerEdges, [&](const std::shared_ptr<ContextEdge> &E) {
-          return E->AllocTypes == (uint8_t)AllocationType::None;
-        }));
   }
 
-  // Cloning may have resulted in some cloned callee edges with type None,
-  // because they aren't carrying any contexts. Remove those edges.
-  for (auto *Clone : Node->Clones) {
-    removeNoneTypeCalleeEdges(Clone);
-    if (VerifyNodes)
-      checkNode<DerivedCCG, FuncTy, CallTy>(Clone);
-  }
   // We should still have some context ids on the original Node.
   assert(!Node->ContextIds.empty());
 
-  // Remove any callee edges that ended up with alloc type None after creating
-  // clones and updating callee edges.
-  removeNoneTypeCalleeEdges(Node);
-
   // Sanity check that no alloc types on node or edges are None.
   assert(Node->AllocTypes != (uint8_t)AllocationType::None);
-  assert(llvm::none_of(Node->CalleeEdges,
-                       [&](const std::shared_ptr<ContextEdge> &E) {
-                         return E->AllocTypes == (uint8_t)AllocationType::None;
-                       }));
-  assert(llvm::none_of(Node->CallerEdges,
-                       [&](const std::shared_ptr<ContextEdge> &E) {
-                         return E->AllocTypes == (uint8_t)AllocationType::None;
-                       }));
 
   if (VerifyNodes)
-    checkNode<DerivedCCG, FuncTy, CallTy>(Node);
+    checkNode<DerivedCCG, FuncTy, CallTy>(Node, /*CheckEdges=*/false);
 }
 
 void ModuleCallsiteContextGraph::updateAllocationCall(

>From d420ac5a8c5881a5f3f763ceb6498b70d4d24c66 Mon Sep 17 00:00:00 2001
From: Teresa Johnson <tejohnson at google.com>
Date: Tue, 26 Mar 2024 11:09:26 -0700
Subject: [PATCH 2/3] Address comments

---
 llvm/lib/Transforms/IPO/MemProfContextDisambiguation.cpp | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/Transforms/IPO/MemProfContextDisambiguation.cpp b/llvm/lib/Transforms/IPO/MemProfContextDisambiguation.cpp
index 742b62a277a4b9..888377055d586d 100644
--- a/llvm/lib/Transforms/IPO/MemProfContextDisambiguation.cpp
+++ b/llvm/lib/Transforms/IPO/MemProfContextDisambiguation.cpp
@@ -306,7 +306,11 @@ class CallsiteContextGraph {
     // True if this node was effectively removed from the graph, in which case
     // its context id set, caller edges, and callee edges should all be empty.
     bool isRemoved() const {
+      // Note that we can have non-empty context ids with empty caller and
+      // callee edges if the graph ends up with a single node.
       assert(!ContextIds.empty() ||
+             // Since the context ids were empty we should have empty callee and
+             // caller edges.
              (CalleeEdges.empty() && CallerEdges.empty()));
       return ContextIds.empty();
     }
@@ -2465,7 +2469,8 @@ void CallsiteContextGraph<DerivedCCG, FuncTy, CallTy>::identifyClones() {
   Visited.clear();
   for (auto &Entry : AllocationCallToContextNodeMap)
     recursivelyRemoveNoneTypeCalleeEdges(Entry.second, Visited);
-  check();
+  if (VerifyCCG)
+    check();
 }
 
 // helper function to check an AllocType is cold or notcold or both.

>From 43db3755f84a33e3cc1b6ad6bd6ddf5cc74dd010 Mon Sep 17 00:00:00 2001
From: Teresa Johnson <tejohnson at google.com>
Date: Tue, 26 Mar 2024 19:21:21 -0700
Subject: [PATCH 3/3] Restructure assert

---
 llvm/lib/Transforms/IPO/MemProfContextDisambiguation.cpp | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/llvm/lib/Transforms/IPO/MemProfContextDisambiguation.cpp b/llvm/lib/Transforms/IPO/MemProfContextDisambiguation.cpp
index 888377055d586d..43ae74938a24f0 100644
--- a/llvm/lib/Transforms/IPO/MemProfContextDisambiguation.cpp
+++ b/llvm/lib/Transforms/IPO/MemProfContextDisambiguation.cpp
@@ -308,10 +308,10 @@ class CallsiteContextGraph {
     bool isRemoved() const {
       // Note that we can have non-empty context ids with empty caller and
       // callee edges if the graph ends up with a single node.
-      assert(!ContextIds.empty() ||
-             // Since the context ids were empty we should have empty callee and
-             // caller edges.
-             (CalleeEdges.empty() && CallerEdges.empty()));
+      if (ContextIds.empty())
+        assert(CalleeEdges.empty() && CallerEdges.empty() &&
+               "Context ids empty but at least one of callee and caller edges "
+               "were not!");
       return ContextIds.empty();
     }
 



More information about the llvm-commits mailing list