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

[Teresa Johnson via llvm-commits]

================
@@ -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()) {
----------------
teresajohnson wrote:

Not sure what you mean exactly. The empty case above is a special case that we need to handle (we might have a non-empty ContextIdsToMove set passed in that is all of the context ids on Edge, so they are not identical situations).

https://github.com/llvm/llvm-project/pull/87112