[MemProf] Perform cloning for each allocation separately (#87112)

Restructures the cloning slightly to perform all cloning for each
allocation separately. The prior algorithm would sometimes miss cloning
opportunities in cases where trimmed cold contexts partially overlapped
with longer contexts for different allocations.

Most of the change is isolated to the helpers that move edges to new or
existing clones, which now support moving a subset of context ids.

Author: teresajohnson

llvm/lib/Transforms/IPO/MemProfContextDisambiguation.cpp

@@ -526,25 +526,30 @@ class CallsiteContextGraph {
   /// Create a clone of Edge's callee and move Edge to that new callee node,
   /// performing the necessary context id and allocation type updates.
   /// If callee's caller edge iterator is supplied, it is updated when removing
-  /// the edge from that list.
+  /// the edge from that list. If ContextIdsToMove is non-empty, only that
+  /// subset of Edge's ids are moved to an edge to the new callee.
   ContextNode *
   moveEdgeToNewCalleeClone(const std::shared_ptr<ContextEdge> &Edge,
-                           EdgeIter *CallerEdgeI = nullptr);
+                           EdgeIter *CallerEdgeI = nullptr,
+                           DenseSet<uint32_t> ContextIdsToMove = {});
 
   /// Change the callee of Edge to existing callee clone NewCallee, performing
   /// the necessary context id and allocation type updates.
   /// If callee's caller edge iterator is supplied, it is updated when removing
-  /// the edge from that list.
+  /// the edge from that list. If ContextIdsToMove is non-empty, only that
+  /// subset of Edge's ids are moved to an edge to the new callee.
   void moveEdgeToExistingCalleeClone(const std::shared_ptr<ContextEdge> &Edge,
                                      ContextNode *NewCallee,
                                      EdgeIter *CallerEdgeI = nullptr,
-                                     bool NewClone = false);
+                                     bool NewClone = false,
+                                     DenseSet<uint32_t> ContextIdsToMove = {});
 
   /// Recursively perform cloning on the graph for the given Node and its
   /// callers, in order to uniquely identify the allocation behavior of an
-  /// allocation given its context.
-  void identifyClones(ContextNode *Node,
-                      DenseSet<const ContextNode *> &Visited);
+  /// allocation given its context. The context ids of the allocation being
+  /// processed are given in AllocContextIds.
+  void identifyClones(ContextNode *Node, DenseSet<const ContextNode *> &Visited,
+                      const DenseSet<uint32_t> &AllocContextIds);
 
   /// Map from each context ID to the AllocationType assigned to that context.
   std::map<uint32_t, AllocationType> ContextIdToAllocationType;
@@ -2358,39 +2363,99 @@ void CallsiteContextGraph<DerivedCCG, FuncTy, CallTy>::exportToDot(
 template <typename DerivedCCG, typename FuncTy, typename CallTy>
 typename CallsiteContextGraph<DerivedCCG, FuncTy, CallTy>::ContextNode *
 CallsiteContextGraph<DerivedCCG, FuncTy, CallTy>::moveEdgeToNewCalleeClone(
-    const std::shared_ptr<ContextEdge> &Edge, EdgeIter *CallerEdgeI) {
+    const std::shared_ptr<ContextEdge> &Edge, EdgeIter *CallerEdgeI,
+    DenseSet<uint32_t> ContextIdsToMove) {
   ContextNode *Node = Edge->Callee;
   NodeOwner.push_back(
       std::make_unique<ContextNode>(Node->IsAllocation, Node->Call));
   ContextNode *Clone = NodeOwner.back().get();
   Node->addClone(Clone);
   assert(NodeToCallingFunc.count(Node));
   NodeToCallingFunc[Clone] = NodeToCallingFunc[Node];
-  moveEdgeToExistingCalleeClone(Edge, Clone, CallerEdgeI, /*NewClone=*/true);
+  moveEdgeToExistingCalleeClone(Edge, Clone, CallerEdgeI, /*NewClone=*/true,
+                                ContextIdsToMove);
   return Clone;
 }
 
 template <typename DerivedCCG, typename FuncTy, typename CallTy>
 void CallsiteContextGraph<DerivedCCG, FuncTy, CallTy>::
     moveEdgeToExistingCalleeClone(const std::shared_ptr<ContextEdge> &Edge,
                                   ContextNode *NewCallee, EdgeIter *CallerEdgeI,
-                                  bool NewClone) {
+                                  bool NewClone,
+                                  DenseSet<uint32_t> ContextIdsToMove) {
   // NewCallee and Edge's current callee must be clones of the same original
   // node (Edge's current callee may be the original node too).
   assert(NewCallee->getOrigNode() == Edge->Callee->getOrigNode());
-  auto &EdgeContextIds = Edge->getContextIds();
+
   ContextNode *OldCallee = Edge->Callee;
-  if (CallerEdgeI)
-    *CallerEdgeI = OldCallee->CallerEdges.erase(*CallerEdgeI);
-  else
-    OldCallee->eraseCallerEdge(Edge.get());
-  Edge->Callee = NewCallee;
-  NewCallee->CallerEdges.push_back(Edge);
-  // Don't need to update Edge's context ids since we are simply reconnecting
-  // it.
-  set_subtract(OldCallee->ContextIds, EdgeContextIds);
-  NewCallee->ContextIds.insert(EdgeContextIds.begin(), EdgeContextIds.end());
-  NewCallee->AllocTypes |= Edge->AllocTypes;
+
+  // We might already have an edge to the new callee from earlier cloning for a
+  // different allocation. If one exists we will reuse it.
+  auto ExistingEdgeToNewCallee = NewCallee->findEdgeFromCaller(Edge->Caller);
+
+  // Callers will pass an empty ContextIdsToMove set when they want to move the
+  // edge. Copy in Edge's ids for simplicity.
+  if (ContextIdsToMove.empty())
+    ContextIdsToMove = Edge->getContextIds();
+
+  // If we are moving all of Edge's ids, then just move the whole Edge.
+  // Otherwise only move the specified subset, to a new edge if needed.
+  if (Edge->getContextIds().size() == ContextIdsToMove.size()) {
+    // Moving the whole Edge.
+    if (CallerEdgeI)
+      *CallerEdgeI = OldCallee->CallerEdges.erase(*CallerEdgeI);
+    else
+      OldCallee->eraseCallerEdge(Edge.get());
+    if (ExistingEdgeToNewCallee) {
+      // Since we already have an edge to NewCallee, simply move the ids
+      // onto it, and remove the existing Edge.
+      ExistingEdgeToNewCallee->getContextIds().insert(ContextIdsToMove.begin(),
+                                                      ContextIdsToMove.end());
+      ExistingEdgeToNewCallee->AllocTypes |= Edge->AllocTypes;
+      assert(Edge->ContextIds == ContextIdsToMove);
+      Edge->ContextIds.clear();
+      Edge->AllocTypes = (uint8_t)AllocationType::None;
+      Edge->Caller->eraseCalleeEdge(Edge.get());
+    } else {
+      // Otherwise just reconnect Edge to NewCallee.
+      Edge->Callee = NewCallee;
+      NewCallee->CallerEdges.push_back(Edge);
+      // Don't need to update Edge's context ids since we are simply
+      // reconnecting it.
+    }
+    // In either case, need to update the alloc types on New Callee.
+    NewCallee->AllocTypes |= Edge->AllocTypes;
+  } else {
+    // Only moving a subset of Edge's ids.
+    if (CallerEdgeI)
+      ++CallerEdgeI;
+    // Compute the alloc type of the subset of ids being moved.
+    auto CallerEdgeAllocType = computeAllocType(ContextIdsToMove);
+    if (ExistingEdgeToNewCallee) {
+      // Since we already have an edge to NewCallee, simply move the ids
+      // onto it.
+      ExistingEdgeToNewCallee->getContextIds().insert(ContextIdsToMove.begin(),
+                                                      ContextIdsToMove.end());
+      ExistingEdgeToNewCallee->AllocTypes |= CallerEdgeAllocType;
+    } else {
+      // Otherwise, create a new edge to NewCallee for the ids being moved.
+      auto NewEdge = std::make_shared<ContextEdge>(
+          NewCallee, Edge->Caller, CallerEdgeAllocType, ContextIdsToMove);
+      Edge->Caller->CalleeEdges.push_back(NewEdge);
+      NewCallee->CallerEdges.push_back(NewEdge);
+    }
+    // In either case, need to update the alloc types on NewCallee, and remove
+    // those ids and update the alloc type on the original Edge.
+    NewCallee->AllocTypes |= CallerEdgeAllocType;
+    set_subtract(Edge->ContextIds, ContextIdsToMove);
+    Edge->AllocTypes = computeAllocType(Edge->ContextIds);
+  }
+  // Now perform some updates that are common to all cases: the NewCallee gets
+  // the moved ids added, and we need to remove those ids from OldCallee and
+  // update its alloc type (NewCallee alloc type updates handled above).
+  NewCallee->ContextIds.insert(ContextIdsToMove.begin(),
+                               ContextIdsToMove.end());
+  set_subtract(OldCallee->ContextIds, ContextIdsToMove);
   OldCallee->AllocTypes = computeAllocType(OldCallee->ContextIds);
   // OldCallee alloc type should be None iff its context id set is now empty.
   assert((OldCallee->AllocTypes == (uint8_t)AllocationType::None) ==
@@ -2402,7 +2467,7 @@ void CallsiteContextGraph<DerivedCCG, FuncTy, CallTy>::
     // The context ids moving to the new callee are the subset of this edge's
     // context ids and the context ids on the caller edge being moved.
     DenseSet<uint32_t> EdgeContextIdsToMove =
-        set_intersection(OldCalleeEdge->getContextIds(), EdgeContextIds);
+        set_intersection(OldCalleeEdge->getContextIds(), ContextIdsToMove);
     set_subtract(OldCalleeEdge->getContextIds(), EdgeContextIdsToMove);
     OldCalleeEdge->AllocTypes =
         computeAllocType(OldCalleeEdge->getContextIds());
@@ -2468,8 +2533,10 @@ void CallsiteContextGraph<DerivedCCG, FuncTy, CallTy>::
 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);
+  for (auto &Entry : AllocationCallToContextNodeMap) {
+    Visited.clear();
+    identifyClones(Entry.second, Visited, Entry.second->ContextIds);
+  }
   Visited.clear();
   for (auto &Entry : AllocationCallToContextNodeMap)
     recursivelyRemoveNoneTypeCalleeEdges(Entry.second, Visited);
@@ -2487,7 +2554,8 @@ bool checkColdOrNotCold(uint8_t AllocType) {
 
 template <typename DerivedCCG, typename FuncTy, typename CallTy>
 void CallsiteContextGraph<DerivedCCG, FuncTy, CallTy>::identifyClones(
-    ContextNode *Node, DenseSet<const ContextNode *> &Visited) {
+    ContextNode *Node, DenseSet<const ContextNode *> &Visited,
+    const DenseSet<uint32_t> &AllocContextIds) {
   if (VerifyNodes)
     checkNode<DerivedCCG, FuncTy, CallTy>(Node, /*CheckEdges=*/false);
   assert(!Node->CloneOf);
@@ -2521,7 +2589,7 @@ void CallsiteContextGraph<DerivedCCG, FuncTy, CallTy>::identifyClones(
       }
       // Ignore any caller we previously visited via another edge.
       if (!Visited.count(Edge->Caller) && !Edge->Caller->CloneOf) {
-        identifyClones(Edge->Caller, Visited);
+        identifyClones(Edge->Caller, Visited, AllocContextIds);
       }
     }
   }
@@ -2584,13 +2652,23 @@ void CallsiteContextGraph<DerivedCCG, FuncTy, CallTy>::identifyClones(
     if (hasSingleAllocType(Node->AllocTypes) || Node->CallerEdges.size() <= 1)
       break;
 
+    // Only need to process the ids along this edge pertaining to the given
+    // allocation.
+    auto CallerEdgeContextsForAlloc =
+        set_intersection(CallerEdge->getContextIds(), AllocContextIds);
+    if (CallerEdgeContextsForAlloc.empty()) {
+      ++EI;
+      continue;
+    }
+    auto CallerAllocTypeForAlloc = computeAllocType(CallerEdgeContextsForAlloc);
+
     // Compute the node callee edge alloc types corresponding to the context ids
     // for this caller edge.
     std::vector<uint8_t> CalleeEdgeAllocTypesForCallerEdge;
     CalleeEdgeAllocTypesForCallerEdge.reserve(Node->CalleeEdges.size());
     for (auto &CalleeEdge : Node->CalleeEdges)
       CalleeEdgeAllocTypesForCallerEdge.push_back(intersectAllocTypes(
-          CalleeEdge->getContextIds(), CallerEdge->getContextIds()));
+          CalleeEdge->getContextIds(), CallerEdgeContextsForAlloc));
 
     // Don't clone if doing so will not disambiguate any alloc types amongst
     // caller edges (including the callee edges that would be cloned).
@@ -2605,7 +2683,7 @@ void CallsiteContextGraph<DerivedCCG, FuncTy, CallTy>::identifyClones(
     // disambiguated by splitting out different context ids.
     assert(CallerEdge->AllocTypes != (uint8_t)AllocationType::None);
     assert(Node->AllocTypes != (uint8_t)AllocationType::None);
-    if (allocTypeToUse(CallerEdge->AllocTypes) ==
+    if (allocTypeToUse(CallerAllocTypeForAlloc) ==
             allocTypeToUse(Node->AllocTypes) &&
         allocTypesMatch<DerivedCCG, FuncTy, CallTy>(
             CalleeEdgeAllocTypesForCallerEdge, Node->CalleeEdges)) {
@@ -2618,7 +2696,7 @@ void CallsiteContextGraph<DerivedCCG, FuncTy, CallTy>::identifyClones(
     ContextNode *Clone = nullptr;
     for (auto *CurClone : Node->Clones) {
       if (allocTypeToUse(CurClone->AllocTypes) !=
-          allocTypeToUse(CallerEdge->AllocTypes))
+          allocTypeToUse(CallerAllocTypeForAlloc))
         continue;
 
       if (!allocTypesMatch<DerivedCCG, FuncTy, CallTy>(
@@ -2630,9 +2708,11 @@ void CallsiteContextGraph<DerivedCCG, FuncTy, CallTy>::identifyClones(
 
     // The edge iterator is adjusted when we move the CallerEdge to the clone.
     if (Clone)
-      moveEdgeToExistingCalleeClone(CallerEdge, Clone, &EI);
+      moveEdgeToExistingCalleeClone(CallerEdge, Clone, &EI, /*NewClone=*/false,
+                                    CallerEdgeContextsForAlloc);
     else
-      Clone = moveEdgeToNewCalleeClone(CallerEdge, &EI);
+      Clone =
+          moveEdgeToNewCalleeClone(CallerEdge, &EI, CallerEdgeContextsForAlloc);
 
     assert(EI == Node->CallerEdges.end() ||
            Node->AllocTypes != (uint8_t)AllocationType::None);