[MemProf] Context disambiguation cloning pass [patch 3/4]

[tejohnson]

Summary

Applies cloning decisions to the IR, cloning functions and updating
calls. For Regular LTO, the IR is updated directly during function
assignment, whereas for ThinLTO it is recorded in the summary index
(a subsequent patch will apply to the IR via the index during the
ThinLTO backend.

The function assignment and cloning proceeds greedily, and we create new
clones as needed when we find an incompatible assignment of function
clones to callsite clones (i.e. when different callers need to invoke
different combinations of callsite clones).

Depends on D140949.

[tejohnson]

Rebase on top of patch 1 and 2 changes.

[tejohnson]

Misc minor fixes from building with a large application

[tejohnson]

Rebase on top of committed patches 1a/1b, and rebased patch 2. Reduce new test cases.

[snehasish]

Still need to look into assignFunctions but before I do I'm wondering if this patch can be split into two with applyImport and assignFunctions in separate patches. What do you think?

llvm/lib/Transforms/IPO/MemProfContextDisambiguation.cpp

109

Add a function comment?

2843

I think this will silently skip the if condition on error since the bool operation for the condition will set the error as checked [1]. Wondering if this is intended since the previous line exits on error.

[1] https://github.com/llvm/llvm-project/blob/main/llvm/include/llvm/Support/Error.h#L234

2925

auto& Call to avoid a copy?

3353

Is IsMemProfClone more accurate?

3354

Can we move this to a common const string somewhere? It would be good to use the same constant when assigning the name for the clone in getMemProfFuncName.

3377

Can we move this functor to a static (or private) helper method?

I think we can return the ValueToValueMap vector as a return value and the number of clones created is implicit from the size of the vector.

3397

Can we use make_unique instead of new here? Looking at new made me wonder if there was a leak.

3444

Can we move L2855-L2904 to a helper which returns a function summary (or nullptr)?

3526

Can we skip the rest of the loop body here since MemProfMD will be false below? I think we just need to duplicate the metadata stripping code.

3538

Just increment MIBIter here and below we can use MIBIter->StackIdIndices (since I don't see other uses of MIB). We already use this spelling in the assert below.

3562

Add a comment for the param /*NumClones=*/

3605

IIUC we index with J-1 since the loop which populates the VMap starts from I=1 on L2807. This seems a little brittle. Perhaps add a comment?

[tejohnson]

In D141077#4282228, @snehasish wrote:

Still need to look into assignFunctions but before I do I'm wondering if this patch can be split into two with applyImport and assignFunctions in separate patches. What do you think?

Yes, looks like it will be straightforward to remove that part (which also means removing a few tests and a few changes to other files). What I will do is update this patch with the rebased version (on top of the committed patch 2). Then I will address the 1 comment that I see that is in the ThinLTO backend specific handling, so we don't lose that. Then I will strip out the ThinLTO backend code from here, in preparation for moving it to another patch, and finally, I will address the rest of your comments. I'll let you know when it is ready to re-review.

[tejohnson]

Rebase

[tejohnson]

llvm/lib/Transforms/IPO/MemProfContextDisambiguation.cpp

2843

Yes this is a good point, we should also exit if it isn't parseable as a ModuleSummaryIndex. Modified this to ExitOnErr. Will update the patch with this change now, but then will be moving this code out as it is ThinLTO backend specific.

[tejohnson]

Address comment in ThinLTO backend handling.

[tejohnson]

I just realized that there are in fact comments in applyImports. Will address those as well before I move that code out...

[tejohnson]

Went ahead and addressed the rest of the comments.

llvm/lib/Transforms/IPO/MemProfContextDisambiguation.cpp

3377

I moved the large chunk of this that does the actual cloning into a static helper that returns the VMaps array. But I kept the outer functor since it does some checking to see if we even need to do any cloning (since we call it multiple times per function, called from a couple different places) - this avoids the needs for duplicating the initial checking at each site.

[tejohnson]

Address remaining comments

[tejohnson]

Remove code that applies the cloning to the ThinLTO backends, it will be added in a new patch.

[tejohnson]

I have updated this to remove applyImports and associated code, along with the parts of the ThinLTO tests that were testing the IR transformations. This one is ready for review again.

[snehasish]

Looks good overall, still thinking if we can make assignFunctions less monolithic and easier to follow.

llvm/lib/Transforms/IPO/MemProfContextDisambiguation.cpp

1364

Add a comment why? I believe we use CloneNo == 0 (as convention) to identify the original version thus returning the base name instead of adding the suffix.

2805

I had trouble finding what NV is. I think it would be clearer if we used ore::NV instead of a using namespace ore; directive above. Many uses within LLVM seem to retain the ore qualifier too [1].

[1] https://github.com/search?q=repo%3Allvm%2Fllvm-project%20ore%3A%3ANV&type=code

2835

Maybe inline the comment into the assert and add the reason it can't be an allocation? E.g. assert(CI && "Caller cannot be an allocation because ... ");

2902

Thanks for the detailed inline comments. Can we summarize the heuristic as pseudo-code in a function comment above the implementation too?

2907

Since we don't need the ordering properties here, can we switch to llvm::DenseMap? It would probably be faster and have lower memory overhead than std::map.

2919

for (auto &[Func, CallsWithMetadata] : FuncToCallsWithMetadata) to make it clearer in the usage below.

2933

nit: inline the comment into the assert with &&

2957

[optional] This sounds like we are implementing a worklist. The loop below uses indices and conditions based on them, I wonder if it would be clearer to restructure as a traditional worklist e.g. while(!ClonesWorklist.empty()) { ...}. I think some of the conditions inside can be rewritten and if the resulting code would be easier to follow. Not a major concern but wanted to get your opinion.

2978

Can we hoist the contents of this if condition outside the loop? Then we can start the iteration at index 1 and remove this check.

3025

Inline the comment into the assert with &&?

3085

I think this condition is simple enough to be merged with the condition above with if (Callee == Clone || !Callee->hasCall()) continue.

3264

Should we flip this default? It seems like the callsites with true outnumber the ones without.

llvm/test/ThinLTO/X86/memprof-basic.ll

112

I don't the addition of these attributes is necessary for this patch. Is this left over from the split or just added to make the test more resilient? Perhaps split them out and commit as NFC for the latter.

[tejohnson]

llvm/lib/Transforms/IPO/MemProfContextDisambiguation.cpp

2902

I wasn't sure how much detail you wanted here. PTAL and let me know what I added is helpful.

2957

Good idea, switched to a worklist approach. The uses of "I" below are switched to uses of a new variable NodeCloneCount that effectively counts from 1, so the uses are changed slightly.

2978

No, because we don't always get here for I==0. We would only get here for the I==0 case if the FuncClonesToCallMap is empty, i.e. this callsite's function has not been cloned previously. The function might have already been cloned if we previously processed a different callsite from the same function, in which case we will handle this clone like any other clone in the later code that assigns it to an existing function version (which will be the original copy in practice, but we utilize the same code to do this as for all other clones).

llvm/test/ThinLTO/X86/memprof-basic.ll

112

Ah, this is leftover from the split. It is needed for the IR checking. I will remove this change from here (and from other tests) and eventually they will show up in patch 4 when I rebase.

[tejohnson]

Address comments

[snehasish]

lgtm, remaining comments are (mostly) cosmetic.

llvm/lib/Transforms/IPO/MemProfContextDisambiguation.cpp

2299

Could this and the one below just be a static cast since we've already checked the type with "is"?

2310

I thought that the usage of dyn_cast required the addition of enum kind to the struct [1]?

https://github.com/llvm/llvm-project/blob/main/llvm/include/llvm/IR/ModuleSummaryIndex.h#L380

2345

Missing a word here?

2416

Prefer std::deque which has the same interface with better performance due to spatial locality of elements (I believe it uses small fixed size preallocated arrays internally).

2500

nit: auto* here and a few uses below where the type is a pointer.

2767

Should we enforce a tail call to avoid accidental memory usage increase? Is this a concern for the size of inputs we are working with?

2771

Add a comment for the work done in this DFS?

2902

I wasn't sure how much detail you wanted here. PTAL and let me know what I added is helpful.

Looks great!

2978

No, because we don't always get here for I==0. We would only get here for the I==0 case if the FuncClonesToCallMap is empty, i.e. this callsite's function has not been cloned previously. The function might have already been cloned if we previously processed a different callsite from the same function, in which case we will handle this clone like any other clone in the later code that assigns it to an existing function version (which will be the original copy in practice, but we utilize the same code to do this as for all other clones).

I assumed we could add && FuncClonesToCallMap.empty() to the (new) condition outside the loop. If I'm still missing something then I think we don't have enough test coverage since I was able to move the contents outside (in the prior version of the patch) and all tests pass.

Not a strong opinion since the worklist changes make it a easier to follow.

[tejohnson]

llvm/lib/Transforms/IPO/MemProfContextDisambiguation.cpp

2299

CallTy is an IndexCall struct which derived from PointerUnion, and PointerUnion supports a dyn_cast() method but not a static_cast one.

2310

This is the dyn_cast() implemented by PointerUnion (see above comment about CallTy here).

2500

These are actually shared_ptr not raw pointer. Some can be auto & however, I will add that.

2767

I don't see any uses of the musttail attribute currently in LLVM, so I'm not sure if that is a good idea. I looked at the optimized code, and this is not being marked as a tail call (however interestingly the earlier call to updateAllocationCall is). But this isn't a recursive call, so I'm not sure it matters too much.

2978

I assumed we could add && FuncClonesToCallMap.empty() to the (new) condition outside the loop. If I'm still missing something then I think we don't have enough test coverage since I was able to move the contents outside (in the prior version of the patch) and all tests pass.

Not a strong opinion since the worklist changes make it a easier to follow.

Oh I see what you meant. I think since the worklist approach simplified the logic, I'd prefer to keep all of the handling inside the worklist loop.

[tejohnson]

Address comments

[tejohnson]

Reverted in 6fbf022908c104a380fd1854fb96eafc64509366 due to bot failure: https://lab.llvm.org/buildbot/#/builders/245/builds/7888

Looks like the remarks output is correct but in a different order. Must have some nondeterminism somewhere, will need to investigate.

[tejohnson]

llvm/lib/Transforms/IPO/MemProfContextDisambiguation.cpp

2771

I believe it is this loop that is causing the unstable remarks ordering I'm seeing, which is emitted by the IR updating code invoked during the search. The key into this map is a pointer so we don't have a guaranteed stable ordering across runs. I could use a MapVector instead of a std::map, but rather than carry that extra space overhead throughout, probably it makes the most sense to build and then iterate a vector of all the allocation nodes (the map values) and sort by their OrigStackOrAllocId field, which is assigned in increasing order as they are inserted in the map and will be stable.

[tejohnson]

Fix for nondeterminism.

[snehasish]

Fix lgtm.

[tejohnson]

Reopening so I can eventually upload fix for instability surfaced by expensive checks bot.

[tejohnson]

llvm/lib/Transforms/IPO/MemProfContextDisambiguation.cpp

2771

I had to revert due to a different bot failure, in new test funcassigncloning.ll. This was in the expensive checks bot (which among other things also randomly shuffles elements before sorting). I built with expensive checks on but cannot reproduce the failure, even running thousands of times.

However, looking at the excerpt of the failure the bot provided, I realized that the difference was a different ordering of some newly cloned nodes coming out of node cloning, which was not part of this patch. I realized it is almost certainly due to another iteration over AllocationCallToContextNodeMap, this time in identifyClones(), which would affect the ordering of the cloning. I decided it is probably safest to switch this map (as well as NonAllocationCallToContextNodeMap to avoid any problems there in the future) to MapVector.

Since this affects the prior patch, I will send a patch for just that change along with the part of funcassigncloning.ll that doesn't rely on this patch (including the checking which triggered the failure). I can then remove the change I made here to sort the allocation nodes.

[tejohnson]

llvm/lib/Transforms/IPO/MemProfContextDisambiguation.cpp

2771

That fix was made in D149924 which went in this morning, and the expensive checks bot succeeded with it. I'm going to rebase this one on top of that, remove the change made here to use a vector, incorporate some another bot fix I made earlier on top of this patch (to require asserts in the tests that use -stats), and will submit after retesting.

[tejohnson]

Rebase on top of the committed D149924.

llvm/include/llvm/Transforms/IPO/MemProfContextDisambiguation.h

 #include "llvm/IR/GlobalValue.h"
 #include "llvm/IR/PassManager.h"
 #include <functional>
 
 namespace llvm {
 class GlobalValueSummary;
 class Module;
 class ModuleSummaryIndex;
+class OptimizationRemarkEmitter;
 
 class MemProfContextDisambiguation
     : public PassInfoMixin<MemProfContextDisambiguation> {
   /// Run the context disambiguator on \p M, returns true if any changes made.
-  bool processModule(Module &M);
+  bool processModule(
+      Module &M,
+      function_ref<OptimizationRemarkEmitter &(Function *)> OREGetter);
 
 public:
   MemProfContextDisambiguation() {}
 
   PreservedAnalyses run(Module &M, ModuleAnalysisManager &AM);
 
   void run(ModuleSummaryIndex &Index,
            function_ref<bool(GlobalValue::GUID, const GlobalValueSummary *)>
                isPrevailing);
 };
 } // end namespace llvm
 
 #endif // LLVM_TRANSFORMS_IPO_MEMPROF_CONTEXT_DISAMBIGUATION_H

llvm/lib/Transforms/IPO/MemProfContextDisambiguation.cpp

 #include "llvm/Transforms/IPO/MemProfContextDisambiguation.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/SetOperations.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/Statistic.h"
 #include "llvm/Analysis/MemoryProfileInfo.h"
 #include "llvm/Analysis/ModuleSummaryAnalysis.h"
+#include "llvm/Analysis/OptimizationRemarkEmitter.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/ModuleSummaryIndex.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/GraphWriter.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Transforms/IPO.h"
+#include "llvm/Transforms/Utils/Cloning.h"
 #include <sstream>
 #include <vector>
 using namespace llvm;
 using namespace llvm::memprof;
 
 #define DEBUG_TYPE "memprof-context-disambiguation"
 
+STATISTIC(FunctionClonesAnalysis,
+          "Number of function clones created during whole program analysis");
+STATISTIC(AllocTypeNotCold, "Number of not cold static allocations (possibly "
+                            "cloned) during whole program analysis");
+STATISTIC(AllocTypeCold, "Number of cold static allocations (possibly cloned) "
+                         "during whole program analysis");
+
 static cl::opt<std::string> DotFilePathPrefix(
     "memprof-dot-file-path-prefix", cl::init(""), cl::Hidden,
     cl::value_desc("filename"),
     cl::desc("Specify the path prefix of the MemProf dot files."));
 
 static cl::opt<bool> ExportToDot("memprof-export-to-dot", cl::init(false),
                                  cl::Hidden,
                                  cl::desc("Export graph to dot files."));
public:
 
   /// Main entry point to perform analysis and transformations on graph.
   bool process();
 
   /// Perform cloning on the graph necessary to uniquely identify the allocation
   /// behavior of an allocation based on its context.
   void identifyClones();
 
+  /// Assign callsite clones to functions, cloning functions as needed to
+  /// accommodate the combinations of their callsite clones reached by callers.
+  /// For regular LTO this clones functions and callsites in the IR, but for
+  /// ThinLTO the cloning decisions are noted in the summaries and applied
+  /// later.
+  bool assignFunctions();
+
   void dump() const;
   void print(raw_ostream &OS) const;
 
   friend raw_ostream &operator<<(raw_ostream &OS,
                                  const CallsiteContextGraph &CCG) {
     CCG.print(OS);
     return OS;
   }
return static_cast<DerivedCCG *>(this)->getStackIdsWithContextNodesForCall(
         Call);
   }
 
   /// Get the last stack id in the context for callsite.
   uint64_t getLastStackId(CallTy Call) {
     return static_cast<DerivedCCG *>(this)->getLastStackId(Call);
   }
 
+  /// Update the allocation call to record type of allocated memory.
+  void updateAllocationCall(CallInfo &Call, AllocationType AllocType) {
+    AllocType == AllocationType::Cold ? AllocTypeCold++ : AllocTypeNotCold++;
+    static_cast<DerivedCCG *>(this)->updateAllocationCall(Call, AllocType);
+  }
+
+  /// Update non-allocation call to invoke (possibly cloned) function
+  /// CalleeFunc.
+  void updateCall(CallInfo &CallerCall, FuncInfo CalleeFunc) {
+    static_cast<DerivedCCG *>(this)->updateCall(CallerCall, CalleeFunc);
+  }
+
+  /// Clone the given function for the given callsite, recording mapping of all
+  /// of the functions tracked calls to their new versions in the CallMap.
+  /// Assigns new clones to clone number CloneNo.
+  FuncInfo cloneFunctionForCallsite(
+      FuncInfo &Func, CallInfo &Call, std::map<CallInfo, CallInfo> &CallMap,
+      std::vector<CallInfo> &CallsWithMetadataInFunc, unsigned CloneNo) {
+    return static_cast<DerivedCCG *>(this)->cloneFunctionForCallsite(
+        Func, Call, CallMap, CallsWithMetadataInFunc, CloneNo);
+  }
+
   /// Gets a label to use in the dot graph for the given call clone in the given
   /// function.
   std::string getLabel(const FuncTy *Func, const CallTy Call,
                        unsigned CloneNo) const {
     return static_cast<const DerivedCCG *>(this)->getLabel(Func, Call, CloneNo);
   }
 
   /// Helpers to find the node corresponding to the given call or stackid.

 using CallInfo =
     typename CallsiteContextGraph<DerivedCCG, FuncTy, CallTy>::CallInfo;
 
 /// CRTP derived class for graphs built from IR (regular LTO).
 class ModuleCallsiteContextGraph
     : public CallsiteContextGraph<ModuleCallsiteContextGraph, Function,
                                   Instruction *> {
 public:
-  ModuleCallsiteContextGraph(Module &M);
+  ModuleCallsiteContextGraph(
+      Module &M,
+      function_ref<OptimizationRemarkEmitter &(Function *)> OREGetter);
 
 private:
   friend CallsiteContextGraph<ModuleCallsiteContextGraph, Function,
                               Instruction *>;
 
   uint64_t getStackId(uint64_t IdOrIndex) const;
   bool calleeMatchesFunc(Instruction *Call, const Function *Func);
   uint64_t getLastStackId(Instruction *Call);
   std::vector<uint64_t> getStackIdsWithContextNodesForCall(Instruction *Call);
+  void updateAllocationCall(CallInfo &Call, AllocationType AllocType);
+  void updateCall(CallInfo &CallerCall, FuncInfo CalleeFunc);
+  CallsiteContextGraph<ModuleCallsiteContextGraph, Function,
+                       Instruction *>::FuncInfo
+  cloneFunctionForCallsite(FuncInfo &Func, CallInfo &Call,
+                           std::map<CallInfo, CallInfo> &CallMap,
+                           std::vector<CallInfo> &CallsWithMetadataInFunc,
+                           unsigned CloneNo);
   std::string getLabel(const Function *Func, const Instruction *Call,
                        unsigned CloneNo) const;
 
   const Module &Mod;
+  function_ref<OptimizationRemarkEmitter &(Function *)> OREGetter;
 };
 
 /// Represents a call in the summary index graph, which can either be an
 /// allocation or an interior callsite node in an allocation's context.
 /// Holds a pointer to the corresponding data structure in the index.
 struct IndexCall : public PointerUnion<CallsiteInfo *, AllocInfo *> {
   IndexCall() : PointerUnion() {}
   IndexCall(std::nullptr_t) : IndexCall() {}

 private:
   friend CallsiteContextGraph<IndexCallsiteContextGraph, FunctionSummary,
                               IndexCall>;
 
   uint64_t getStackId(uint64_t IdOrIndex) const;
   bool calleeMatchesFunc(IndexCall &Call, const FunctionSummary *Func);
   uint64_t getLastStackId(IndexCall &Call);
   std::vector<uint64_t> getStackIdsWithContextNodesForCall(IndexCall &Call);
+  void updateAllocationCall(CallInfo &Call, AllocationType AllocType);
+  void updateCall(CallInfo &CallerCall, FuncInfo CalleeFunc);
+  CallsiteContextGraph<IndexCallsiteContextGraph, FunctionSummary,
+                       IndexCall>::FuncInfo
+  cloneFunctionForCallsite(FuncInfo &Func, CallInfo &Call,
+                           std::map<CallInfo, CallInfo> &CallMap,
+                           std::vector<CallInfo> &CallsWithMetadataInFunc,
+                           unsigned CloneNo);
   std::string getLabel(const FunctionSummary *Func, const IndexCall &Call,
                        unsigned CloneNo) const;
 
   // Saves mapping from function summaries containing memprof records back to
   // its VI, for use in checking and debugging.
   std::map<const FunctionSummary *, ValueInfo> FSToVIMap;
 
   const ModuleSummaryIndex &Index;

 uint64_t IndexCallsiteContextGraph::getLastStackId(IndexCall &Call) {
   assert(isa<CallsiteInfo *>(Call.getBase()));
   CallStack<CallsiteInfo, SmallVector<unsigned>::const_iterator>
       CallsiteContext(dyn_cast_if_present<CallsiteInfo *>(Call.getBase()));
   // Need to convert index into stack id.
   return Index.getStackIdAtIndex(CallsiteContext.back());
 }
 
+static const std::string MemProfCloneSuffix = ".memprof.";
+
 static std::string getMemProfFuncName(Twine Base, unsigned CloneNo) {
+  // We use CloneNo == 0 to refer to the original version, which doesn't get
+  // renamed with a suffix.
   if (!CloneNo)
     return Base.str();
-  return (Base + ".memprof." + Twine(CloneNo)).str();
+  return (Base + MemProfCloneSuffix + Twine(CloneNo)).str();
 }
 
 std::string ModuleCallsiteContextGraph::getLabel(const Function *Func,
                                                  const Instruction *Call,
                                                  unsigned CloneNo) const {
   return (Twine(Call->getFunction()->getName()) + " -> " +
           cast<CallBase>(Call)->getCalledFunction()->getName())
       .str();
for (auto IdOrIndex : CallsiteContext) {
     ContextNode *Node = getNodeForStackId(StackId);
     if (!Node)
       break;
     StackIds.push_back(StackId);
   }
   return StackIds;
 }
 
-ModuleCallsiteContextGraph::ModuleCallsiteContextGraph(Module &M) : Mod(M) {
+ModuleCallsiteContextGraph::ModuleCallsiteContextGraph(
+    Module &M, function_ref<OptimizationRemarkEmitter &(Function *)> OREGetter)
+    : Mod(M), OREGetter(OREGetter) {
   for (auto &F : M) {
     std::vector<CallInfo> CallsWithMetadata;
     for (auto &BB : F) {
       for (auto &I : BB) {
         if (!isa<CallBase>(I))
           continue;
         if (auto *MemProfMD = I.getMetadata(LLVMContext::MD_memprof)) {
           CallsWithMetadata.push_back(&I);
static void checkEdge(
   // Confirm that alloc type is not None and that we have at least one context
   // id.
   assert(Edge->AllocTypes != (uint8_t)AllocationType::None);
   assert(!Edge->ContextIds.empty());
 }
 
 template <typename DerivedCCG, typename FuncTy, typename CallTy>
 static void checkNode(const ContextNode<DerivedCCG, FuncTy, CallTy> *Node,
-                      bool CheckEdges = false) {
+                      bool CheckEdges = true) {
   if (Node->isRemoved())
     return;
   // Node's context ids should be the union of both its callee and caller edge
   // context ids.
   if (Node->CallerEdges.size()) {
     auto EI = Node->CallerEdges.begin();
     auto &FirstEdge = *EI;
     EI++;
if (Node->CalleeEdges.size()) {
     assert(Node->ContextIds == CalleeEdgeContextIds);
   }
 }
 
 template <typename DerivedCCG, typename FuncTy, typename CallTy>
 void CallsiteContextGraph<DerivedCCG, FuncTy, CallTy>::check() const {
   using GraphType = const CallsiteContextGraph<DerivedCCG, FuncTy, CallTy> *;
   for (const auto Node : nodes<GraphType>(this)) {
-    checkNode<DerivedCCG, FuncTy, CallTy>(Node);
+    checkNode<DerivedCCG, FuncTy, CallTy>(Node, /*CheckEdges=*/false);
     for (auto &Edge : Node->CallerEdges)
       checkEdge<DerivedCCG, FuncTy, CallTy>(Edge);
   }
 }
 
 template <typename DerivedCCG, typename FuncTy, typename CallTy>
 struct GraphTraits<const CallsiteContextGraph<DerivedCCG, FuncTy, CallTy> *> {
   using GraphType = const CallsiteContextGraph<DerivedCCG, FuncTy, CallTy> *;
for (auto &OldCalleeEdge : OldCallee->CalleeEdges) {
     }
     auto NewEdge = std::make_shared<ContextEdge>(
         OldCalleeEdge->Callee, NewCallee,
         computeAllocType(EdgeContextIdsToMove), EdgeContextIdsToMove);
     NewCallee->CalleeEdges.push_back(NewEdge);
     NewEdge->Callee->CallerEdges.push_back(NewEdge);
   }
   if (VerifyCCG) {
-    checkNode<DerivedCCG, FuncTy, CallTy>(OldCallee);
-    checkNode<DerivedCCG, FuncTy, CallTy>(NewCallee);
+    checkNode<DerivedCCG, FuncTy, CallTy>(OldCallee, /*CheckEdges=*/false);
+    checkNode<DerivedCCG, FuncTy, CallTy>(NewCallee, /*CheckEdges=*/false);
     for (const auto &OldCalleeEdge : OldCallee->CalleeEdges)
-      checkNode<DerivedCCG, FuncTy, CallTy>(OldCalleeEdge->Callee);
+      checkNode<DerivedCCG, FuncTy, CallTy>(OldCalleeEdge->Callee,
+                                            /*CheckEdges=*/false);
     for (const auto &NewCalleeEdge : NewCallee->CalleeEdges)
-      checkNode<DerivedCCG, FuncTy, CallTy>(NewCalleeEdge->Callee);
+      checkNode<DerivedCCG, FuncTy, CallTy>(NewCalleeEdge->Callee,
+                                            /*CheckEdges=*/false);
   }
 }
 
 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);
 }
 
 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, /*CheckEdges=*/true);
+    checkNode<DerivedCCG, FuncTy, CallTy>(Node);
   assert(!Node->CloneOf);
 
   // If Node as a null call, then either it wasn't found in the module (regular
   // LTO) or summary index (ThinLTO), or there were other conditions blocking
   // cloning (e.g. recursion, calls multiple targets, etc).
   // Do this here so that we don't try to recursively clone callers below, which
   // isn't useful at least for this node.
   if (!Node->hasCall())
assert(llvm::none_of(
         }));
   }
 
   // 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, /*CheckEdges=*/true);
+      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, /*CheckEdges=*/true);
+    checkNode<DerivedCCG, FuncTy, CallTy>(Node);
+}
+
+static std::string getAllocTypeAttributeString(AllocationType Type) {
+  switch (Type) {
+  case AllocationType::NotCold:
+    return "notcold";
+    break;
+  case AllocationType::Cold:
+    return "cold";
+    break;
+  default:
+    dbgs() << "Unexpected alloc type " << (uint8_t)Type;
+    assert(false);
+  }
+  llvm_unreachable("invalid alloc type");
+}
+
+void ModuleCallsiteContextGraph::updateAllocationCall(
+    CallInfo &Call, AllocationType AllocType) {
+  std::string AllocTypeString = getAllocTypeAttributeString(AllocType);
+  auto A = llvm::Attribute::get(Call.call()->getFunction()->getContext(),
+                                "memprof", AllocTypeString);
+  cast<CallBase>(Call.call())->addFnAttr(A);
+  OREGetter(Call.call()->getFunction())
+      .emit(OptimizationRemark(DEBUG_TYPE, "MemprofAttribute", Call.call())
+            << ore::NV("AllocationCall", Call.call()) << " in clone "
+            << ore::NV("Caller", Call.call()->getFunction())
+            << " marked with memprof allocation attribute "
+            << ore::NV("Attribute", AllocTypeString));
+}
+
+void IndexCallsiteContextGraph::updateAllocationCall(CallInfo &Call,
+                                                     AllocationType AllocType) {
+  auto *AI = Call.call().dyn_cast<AllocInfo *>();
+  assert(AI);
+  assert(AI->Versions.size() > Call.cloneNo());
+  AI->Versions[Call.cloneNo()] = (uint8_t)AllocType;
+}
+
+void ModuleCallsiteContextGraph::updateCall(CallInfo &CallerCall,
+                                            FuncInfo CalleeFunc) {
+  if (CalleeFunc.cloneNo() > 0)
+    cast<CallBase>(CallerCall.call())->setCalledFunction(CalleeFunc.func());
+  OREGetter(CallerCall.call()->getFunction())
+      .emit(OptimizationRemark(DEBUG_TYPE, "MemprofCall", CallerCall.call())
+            << ore::NV("Call", CallerCall.call()) << " in clone "
+            << ore::NV("Caller", CallerCall.call()->getFunction())
+            << " assigned to call function clone "
+            << ore::NV("Callee", CalleeFunc.func()));
+}
+
+void IndexCallsiteContextGraph::updateCall(CallInfo &CallerCall,
+                                           FuncInfo CalleeFunc) {
+  auto *CI = CallerCall.call().dyn_cast<CallsiteInfo *>();
+  assert(CI &&
+         "Caller cannot be an allocation which should not have profiled calls");
+  assert(CI->Clones.size() > CallerCall.cloneNo());
+  CI->Clones[CallerCall.cloneNo()] = CalleeFunc.cloneNo();
+}
+
+CallsiteContextGraph<ModuleCallsiteContextGraph, Function,
+                     Instruction *>::FuncInfo
+ModuleCallsiteContextGraph::cloneFunctionForCallsite(
+    FuncInfo &Func, CallInfo &Call, std::map<CallInfo, CallInfo> &CallMap,
+    std::vector<CallInfo> &CallsWithMetadataInFunc, unsigned CloneNo) {
+  // Use existing LLVM facilities for cloning and obtaining Call in clone
+  ValueToValueMapTy VMap;
+  auto *NewFunc = CloneFunction(Func.func(), VMap);
+  std::string Name = getMemProfFuncName(Func.func()->getName(), CloneNo);
+  assert(!Func.func()->getParent()->getFunction(Name));
+  NewFunc->setName(Name);
+  for (auto &Inst : CallsWithMetadataInFunc) {
+    // This map always has the initial version in it.
+    assert(Inst.cloneNo() == 0);
+    CallMap[Inst] = {cast<Instruction>(VMap[Inst.call()]), CloneNo};
+  }
+  OREGetter(Func.func())
+      .emit(OptimizationRemark(DEBUG_TYPE, "MemprofClone", Func.func())
+            << "created clone " << ore::NV("NewFunction", NewFunc));
+  return {NewFunc, CloneNo};
+}
+
+CallsiteContextGraph<IndexCallsiteContextGraph, FunctionSummary,
+                     IndexCall>::FuncInfo
+IndexCallsiteContextGraph::cloneFunctionForCallsite(
+    FuncInfo &Func, CallInfo &Call, std::map<CallInfo, CallInfo> &CallMap,
+    std::vector<CallInfo> &CallsWithMetadataInFunc, unsigned CloneNo) {
+  // Check how many clones we have of Call (and therefore function).
+  // The next clone number is the current size of versions array.
+  // Confirm this matches the CloneNo provided by the caller, which is based on
+  // the number of function clones we have.
+  assert(CloneNo ==
+         (Call.call().is<AllocInfo *>()
+              ? Call.call().dyn_cast<AllocInfo *>()->Versions.size()
+              : Call.call().dyn_cast<CallsiteInfo *>()->Clones.size()));
+  // Walk all the instructions in this function. Create a new version for
+  // each (by adding an entry to the Versions/Clones summary array), and copy
+  // over the version being called for the function clone being cloned here.
+  // Additionally, add an entry to the CallMap for the new function clone,
+  // mapping the original call (clone 0, what is in CallsWithMetadataInFunc)
+  // to the new call clone.
+  for (auto &Inst : CallsWithMetadataInFunc) {
+    // This map always has the initial version in it.
+    assert(Inst.cloneNo() == 0);
+    if (auto *AI = Inst.call().dyn_cast<AllocInfo *>()) {
+      assert(AI->Versions.size() == CloneNo);
+      // We assign the allocation type later (in updateAllocationCall), just add
+      // an entry for it here.
+      AI->Versions.push_back(0);
+    } else {
+      auto *CI = Inst.call().dyn_cast<CallsiteInfo *>();
+      assert(CI && CI->Clones.size() == CloneNo);
+      // We assign the clone number later (in updateCall), just add an entry for
+      // it here.
+      CI->Clones.push_back(0);
+    }
+    CallMap[Inst] = {Inst.call(), CloneNo};
+  }
+  return {Func.func(), CloneNo};
+}
+
+// This method assigns cloned callsites to functions, cloning the functions as
+// needed. The assignment is greedy and proceeds roughly as follows:
+//
+// For each function Func:
+//   For each call with graph Node having clones:
+//     Initialize ClonesWorklist to Node and its clones
+//     Initialize NodeCloneCount to 0
+//     While ClonesWorklist is not empty:
+//        Clone = pop front ClonesWorklist
+//        NodeCloneCount++
+//        If Func has been cloned less than NodeCloneCount times:
+//           If NodeCloneCount is 1:
+//             Assign Clone to original Func
+//             Continue
+//           Create a new function clone
+//           If other callers not assigned to call a function clone yet:
+//              Assign them to call new function clone
+//              Continue
+//           Assign any other caller calling the cloned version to new clone
+//
+//        For each caller of Clone:
+//           If caller is assigned to call a specific function clone:
+//             If we cannot assign Clone to that function clone:
+//               Create new callsite Clone NewClone
+//               Add NewClone to ClonesWorklist
+//               Continue
+//             Assign Clone to existing caller's called function clone
+//           Else:
+//             If Clone not already assigned to a function clone:
+//                Assign to first function clone without assignment
+//             Assign caller to selected function clone
+template <typename DerivedCCG, typename FuncTy, typename CallTy>
+bool CallsiteContextGraph<DerivedCCG, FuncTy, CallTy>::assignFunctions() {
+  bool Changed = false;
+
+  // Keep track of the assignment of nodes (callsites) to function clones they
+  // call.
+  DenseMap<ContextNode *, FuncInfo> CallsiteToCalleeFuncCloneMap;
+
+  // Update caller node to call function version CalleeFunc, by recording the
+  // assignment in CallsiteToCalleeFuncCloneMap.
+  auto RecordCalleeFuncOfCallsite = [&](ContextNode *Caller,
+                                        const FuncInfo &CalleeFunc) {
+    assert(Caller->hasCall());
+    CallsiteToCalleeFuncCloneMap[Caller] = CalleeFunc;
+  };
+
+  // Walk all functions for which we saw calls with memprof metadata, and handle
+  // cloning for each of its calls.
+  for (auto &[Func, CallsWithMetadata] : FuncToCallsWithMetadata) {
+    FuncInfo OrigFunc(Func);
+    // Map from each clone of OrigFunc to a map of remappings of each call of
+    // interest (from original uncloned call to the corresponding cloned call in
+    // that function clone).
+    std::map<FuncInfo, std::map<CallInfo, CallInfo>> FuncClonesToCallMap;
+    for (auto &Call : CallsWithMetadata) {
+      ContextNode *Node = getNodeForInst(Call);
+      // Skip call if we do not have a node for it (all uses of its stack ids
+      // were either on inlined chains or pruned from the MIBs), or if we did
+      // not create any clones for it.
+      if (!Node || Node->Clones.empty())
+        continue;
+      assert(Node->hasCall() &&
+             "Not having a call should have prevented cloning");
+
+      // Track the assignment of function clones to clones of the current
+      // callsite Node being handled.
+      std::map<FuncInfo, ContextNode *> FuncCloneToCurNodeCloneMap;
+
+      // Assign callsite version CallsiteClone to function version FuncClone,
+      // and also assign (possibly cloned) Call to CallsiteClone.
+      auto AssignCallsiteCloneToFuncClone = [&](const FuncInfo &FuncClone,
+                                                CallInfo &Call,
+                                                ContextNode *CallsiteClone,
+                                                bool IsAlloc) {
+        // Record the clone of callsite node assigned to this function clone.
+        FuncCloneToCurNodeCloneMap[FuncClone] = CallsiteClone;
+
+        assert(FuncClonesToCallMap.count(FuncClone));
+        std::map<CallInfo, CallInfo> &CallMap = FuncClonesToCallMap[FuncClone];
+        CallInfo CallClone(Call);
+        if (CallMap.count(Call))
+          CallClone = CallMap[Call];
+        CallsiteClone->setCall(CallClone);
+      };
+
+      // Keep track of the clones of callsite Node that need to be assigned to
+      // function clones. This list may be expanded in the loop body below if we
+      // find additional cloning is required.
+      std::deque<ContextNode *> ClonesWorklist;
+      // Ignore original Node if we moved all of its contexts to clones.
+      if (!Node->ContextIds.empty())
+        ClonesWorklist.push_back(Node);
+      ClonesWorklist.insert(ClonesWorklist.end(), Node->Clones.begin(),
+                            Node->Clones.end());
+
+      // Now walk through all of the clones of this callsite Node that we need,
+      // and determine the assignment to a corresponding clone of the current
+      // function (creating new function clones as needed).
+      unsigned NodeCloneCount = 0;
+      while (!ClonesWorklist.empty()) {
+        ContextNode *Clone = ClonesWorklist.front();
+        ClonesWorklist.pop_front();
+        NodeCloneCount++;
+        if (VerifyNodes)
+          checkNode<DerivedCCG, FuncTy, CallTy>(Clone);
+
+        // Need to create a new function clone if we have more callsite clones
+        // than existing function clones, which would have been assigned to an
+        // earlier clone in the list (we assign callsite clones to function
+        // clones greedily).
+        if (FuncClonesToCallMap.size() < NodeCloneCount) {
+          // If this is the first callsite copy, assign to original function.
+          if (NodeCloneCount == 1) {
+            // Since FuncClonesToCallMap is empty in this case, no clones have
+            // been created for this function yet, and no callers should have
+            // been assigned a function clone for this callee node yet.
+            assert(llvm::none_of(
+                Clone->CallerEdges, [&](const std::shared_ptr<ContextEdge> &E) {
+                  return CallsiteToCalleeFuncCloneMap.count(E->Caller);
+                }));
+            // Initialize with empty call map, assign Clone to original function
+            // and its callers, and skip to the next clone.
+            FuncClonesToCallMap[OrigFunc] = {};
+            AssignCallsiteCloneToFuncClone(
+                OrigFunc, Call, Clone,
+                AllocationCallToContextNodeMap.count(Call));
+            for (auto &CE : Clone->CallerEdges) {
+              // Ignore any caller that does not have a recorded callsite Call.
+              if (!CE->Caller->hasCall())
+                continue;
+              RecordCalleeFuncOfCallsite(CE->Caller, OrigFunc);
+            }
+            continue;
+          }
+
+          // First locate which copy of OrigFunc to clone again. If a caller
+          // of this callsite clone was already assigned to call a particular
+          // function clone, we need to redirect all of those callers to the
+          // new function clone, and update their other callees within this
+          // function.
+          FuncInfo PreviousAssignedFuncClone;
+          auto EI = llvm::find_if(
+              Clone->CallerEdges, [&](const std::shared_ptr<ContextEdge> &E) {
+                return CallsiteToCalleeFuncCloneMap.count(E->Caller);
+              });
+          bool CallerAssignedToCloneOfFunc = false;
+          if (EI != Clone->CallerEdges.end()) {
+            const std::shared_ptr<ContextEdge> &Edge = *EI;
+            PreviousAssignedFuncClone =
+                CallsiteToCalleeFuncCloneMap[Edge->Caller];
+            CallerAssignedToCloneOfFunc = true;
+          }
+
+          // Clone function and save it along with the CallInfo map created
+          // during cloning in the FuncClonesToCallMap.
+          std::map<CallInfo, CallInfo> NewCallMap;
+          unsigned CloneNo = FuncClonesToCallMap.size();
+          assert(CloneNo > 0 && "Clone 0 is the original function, which "
+                                "should already exist in the map");
+          FuncInfo NewFuncClone = cloneFunctionForCallsite(
+              OrigFunc, Call, NewCallMap, CallsWithMetadata, CloneNo);
+          FuncClonesToCallMap.emplace(NewFuncClone, std::move(NewCallMap));
+          FunctionClonesAnalysis++;
+          Changed = true;
+
+          // If no caller callsites were already assigned to a clone of this
+          // function, we can simply assign this clone to the new func clone
+          // and update all callers to it, then skip to the next clone.
+          if (!CallerAssignedToCloneOfFunc) {
+            AssignCallsiteCloneToFuncClone(
+                NewFuncClone, Call, Clone,
+                AllocationCallToContextNodeMap.count(Call));
+            for (auto &CE : Clone->CallerEdges) {
+              // Ignore any caller that does not have a recorded callsite Call.
+              if (!CE->Caller->hasCall())
+                continue;
+              RecordCalleeFuncOfCallsite(CE->Caller, NewFuncClone);
+            }
+            continue;
+          }
+
+          // We may need to do additional node cloning in this case.
+          // Reset the CallsiteToCalleeFuncCloneMap entry for any callers
+          // that were previously assigned to call PreviousAssignedFuncClone,
+          // to record that they now call NewFuncClone.
+          for (auto CE : Clone->CallerEdges) {
+            // Ignore any caller that does not have a recorded callsite Call.
+            if (!CE->Caller->hasCall())
+              continue;
+
+            if (!CallsiteToCalleeFuncCloneMap.count(CE->Caller) ||
+                // We subsequently fall through to later handling that
+                // will perform any additional cloning required for
+                // callers that were calling other function clones.
+                CallsiteToCalleeFuncCloneMap[CE->Caller] !=
+                    PreviousAssignedFuncClone)
+              continue;
+
+            RecordCalleeFuncOfCallsite(CE->Caller, NewFuncClone);
+
+            // If we are cloning a function that was already assigned to some
+            // callers, then essentially we are creating new callsite clones
+            // of the other callsites in that function that are reached by those
+            // callers. Clone the other callees of the current callsite's caller
+            // that were already assigned to PreviousAssignedFuncClone
+            // accordingly. This is important since we subsequently update the
+            // calls from the nodes in the graph and their assignments to callee
+            // functions recorded in CallsiteToCalleeFuncCloneMap.
+            for (auto CalleeEdge : CE->Caller->CalleeEdges) {
+              // Skip any that have been removed on an earlier iteration when
+              // cleaning up newly None type callee edges.
+              if (!CalleeEdge)
+                continue;
+              ContextNode *Callee = CalleeEdge->Callee;
+              // Skip the current callsite, we are looking for other
+              // callsites Caller calls, as well as any that does not have a
+              // recorded callsite Call.
+              if (Callee == Clone || !Callee->hasCall())
+                continue;
+              ContextNode *NewClone = moveEdgeToNewCalleeClone(CalleeEdge);
+              removeNoneTypeCalleeEdges(NewClone);
+              // Moving the edge may have resulted in some none type
+              // callee edges on the original Callee.
+              removeNoneTypeCalleeEdges(Callee);
+              assert(NewClone->AllocTypes != (uint8_t)AllocationType::None);
+              // If the Callee node was already assigned to call a specific
+              // function version, make sure its new clone is assigned to call
+              // that same function clone.
+              if (CallsiteToCalleeFuncCloneMap.count(Callee))
+                RecordCalleeFuncOfCallsite(
+                    NewClone, CallsiteToCalleeFuncCloneMap[Callee]);
+              // Update NewClone with the new Call clone of this callsite's Call
+              // created for the new function clone created earlier.
+              // Recall that we have already ensured when building the graph
+              // that each caller can only call callsites within the same
+              // function, so we are guaranteed that Callee Call is in the
+              // current OrigFunc.
+              // CallMap is set up as indexed by original Call at clone 0.
+              CallInfo OrigCall(Callee->getOrigNode()->Call);
+              OrigCall.setCloneNo(0);
+              std::map<CallInfo, CallInfo> &CallMap =
+                  FuncClonesToCallMap[NewFuncClone];
+              assert(CallMap.count(OrigCall));
+              CallInfo NewCall(CallMap[OrigCall]);
+              assert(NewCall);
+              NewClone->setCall(NewCall);
+            }
+          }
+          // Fall through to handling below to perform the recording of the
+          // function for this callsite clone. This enables handling of cases
+          // where the callers were assigned to different clones of a function.
+        }
+
+        // See if we can use existing function clone. Walk through
+        // all caller edges to see if any have already been assigned to
+        // a clone of this callsite's function. If we can use it, do so. If not,
+        // because that function clone is already assigned to a different clone
+        // of this callsite, then we need to clone again.
+        // Basically, this checking is needed to handle the case where different
+        // caller functions/callsites may need versions of this function
+        // containing different mixes of callsite clones across the different
+        // callsites within the function. If that happens, we need to create
+        // additional function clones to handle the various combinations.
+        //
+        // Keep track of any new clones of this callsite created by the
+        // following loop, as well as any existing clone that we decided to
+        // assign this clone to.
+        std::map<FuncInfo, ContextNode *> FuncCloneToNewCallsiteCloneMap;
+        FuncInfo FuncCloneAssignedToCurCallsiteClone;
+        // We need to be able to remove Edge from CallerEdges, so need to adjust
+        // iterator in the loop.
+        for (auto EI = Clone->CallerEdges.begin();
+             EI != Clone->CallerEdges.end();) {
+          auto Edge = *EI;
+          // Ignore any caller that does not have a recorded callsite Call.
+          if (!Edge->Caller->hasCall()) {
+            EI++;
+            continue;
+          }
+          // If this caller already assigned to call a version of OrigFunc, need
+          // to ensure we can assign this callsite clone to that function clone.
+          if (CallsiteToCalleeFuncCloneMap.count(Edge->Caller)) {
+            FuncInfo FuncCloneCalledByCaller =
+                CallsiteToCalleeFuncCloneMap[Edge->Caller];
+            // First we need to confirm that this function clone is available
+            // for use by this callsite node clone.
+            //
+            // While FuncCloneToCurNodeCloneMap is built only for this Node and
+            // its callsite clones, one of those callsite clones X could have
+            // been assigned to the same function clone called by Edge's caller
+            // - if Edge's caller calls another callsite within Node's original
+            // function, and that callsite has another caller reaching clone X.
+            // We need to clone Node again in this case.
+            if ((FuncCloneToCurNodeCloneMap.count(FuncCloneCalledByCaller) &&
+                 FuncCloneToCurNodeCloneMap[FuncCloneCalledByCaller] !=
+                     Clone) ||
+                // Detect when we have multiple callers of this callsite that
+                // have already been assigned to specific, and different, clones
+                // of OrigFunc (due to other unrelated callsites in Func they
+                // reach via call contexts). Is this Clone of callsite Node
+                // assigned to a different clone of OrigFunc? If so, clone Node
+                // again.
+                (FuncCloneAssignedToCurCallsiteClone &&
+                 FuncCloneAssignedToCurCallsiteClone !=
+                     FuncCloneCalledByCaller)) {
+              // We need to use a different newly created callsite clone, in
+              // order to assign it to another new function clone on a
+              // subsequent iteration over the Clones array (adjusted below).
+              // Note we specifically do not reset the
+              // CallsiteToCalleeFuncCloneMap entry for this caller, so that
+              // when this new clone is processed later we know which version of
+              // the function to copy (so that other callsite clones we have
+              // assigned to that function clone are properly cloned over). See
+              // comments in the function cloning handling earlier.
+
+              // Check if we already have cloned this callsite again while
+              // walking through caller edges, for a caller calling the same
+              // function clone. If so, we can move this edge to that new clone
+              // rather than creating yet another new clone.
+              if (FuncCloneToNewCallsiteCloneMap.count(
+                      FuncCloneCalledByCaller)) {
+                ContextNode *NewClone =
+                    FuncCloneToNewCallsiteCloneMap[FuncCloneCalledByCaller];
+                moveEdgeToExistingCalleeClone(Edge, NewClone, &EI);
+                // Cleanup any none type edges cloned over.
+                removeNoneTypeCalleeEdges(NewClone);
+              } else {
+                // Create a new callsite clone.
+                ContextNode *NewClone = moveEdgeToNewCalleeClone(Edge, &EI);
+                removeNoneTypeCalleeEdges(NewClone);
+                FuncCloneToNewCallsiteCloneMap[FuncCloneCalledByCaller] =
+                    NewClone;
+                // Add to list of clones and process later.
+                ClonesWorklist.push_back(NewClone);
+                assert(EI == Clone->CallerEdges.end() ||
+                       Clone->AllocTypes != (uint8_t)AllocationType::None);
+                assert(NewClone->AllocTypes != (uint8_t)AllocationType::None);
+              }
+              // Moving the caller edge may have resulted in some none type
+              // callee edges.
+              removeNoneTypeCalleeEdges(Clone);
+              // We will handle the newly created callsite clone in a subsequent
+              // iteration over this Node's Clones. Continue here since we
+              // already adjusted iterator EI while moving the edge.
+              continue;
+            }
+
+            // Otherwise, we can use the function clone already assigned to this
+            // caller.
+            if (!FuncCloneAssignedToCurCallsiteClone) {
+              FuncCloneAssignedToCurCallsiteClone = FuncCloneCalledByCaller;
+              // Assign Clone to FuncCloneCalledByCaller
+              AssignCallsiteCloneToFuncClone(
+                  FuncCloneCalledByCaller, Call, Clone,
+                  AllocationCallToContextNodeMap.count(Call));
+            } else
+              // Don't need to do anything - callsite is already calling this
+              // function clone.
+              assert(FuncCloneAssignedToCurCallsiteClone ==
+                     FuncCloneCalledByCaller);
+
+          } else {
+            // We have not already assigned this caller to a version of
+            // OrigFunc. Do the assignment now.
+
+            // First check if we have already assigned this callsite clone to a
+            // clone of OrigFunc for another caller during this iteration over
+            // its caller edges.
+            if (!FuncCloneAssignedToCurCallsiteClone) {
+              // Find first function in FuncClonesToCallMap without an assigned
+              // clone of this callsite Node. We should always have one
+              // available at this point due to the earlier cloning when the
+              // FuncClonesToCallMap size was smaller than the clone number.
+              for (auto &CF : FuncClonesToCallMap) {
+                if (!FuncCloneToCurNodeCloneMap.count(CF.first)) {
+                  FuncCloneAssignedToCurCallsiteClone = CF.first;
+                  break;
+                }
+              }
+              assert(FuncCloneAssignedToCurCallsiteClone);
+              // Assign Clone to FuncCloneAssignedToCurCallsiteClone
+              AssignCallsiteCloneToFuncClone(
+                  FuncCloneAssignedToCurCallsiteClone, Call, Clone,
+                  AllocationCallToContextNodeMap.count(Call));
+            } else
+              assert(FuncCloneToCurNodeCloneMap
+                         [FuncCloneAssignedToCurCallsiteClone] == Clone);
+            // Update callers to record function version called.
+            RecordCalleeFuncOfCallsite(Edge->Caller,
+                                       FuncCloneAssignedToCurCallsiteClone);
+          }
+
+          EI++;
+        }
+      }
+      if (VerifyCCG) {
+        checkNode<DerivedCCG, FuncTy, CallTy>(Node);
+        for (const auto &PE : Node->CalleeEdges)
+          checkNode<DerivedCCG, FuncTy, CallTy>(PE->Callee);
+        for (const auto &CE : Node->CallerEdges)
+          checkNode<DerivedCCG, FuncTy, CallTy>(CE->Caller);
+        for (auto *Clone : Node->Clones) {
+          checkNode<DerivedCCG, FuncTy, CallTy>(Clone);
+          for (const auto &PE : Clone->CalleeEdges)
+            checkNode<DerivedCCG, FuncTy, CallTy>(PE->Callee);
+          for (const auto &CE : Clone->CallerEdges)
+            checkNode<DerivedCCG, FuncTy, CallTy>(CE->Caller);
+        }
+      }
+    }
+  }
+
+  auto UpdateCalls = [&](ContextNode *Node,
+                         DenseSet<const ContextNode *> &Visited,
+                         auto &&UpdateCalls) {
+    auto Inserted = Visited.insert(Node);
+    if (!Inserted.second)
+      return;
+
+    for (auto *Clone : Node->Clones)
+      UpdateCalls(Clone, Visited, UpdateCalls);
+
+    for (auto &Edge : Node->CallerEdges)
+      UpdateCalls(Edge->Caller, Visited, UpdateCalls);
+
+    // Skip if either no call to update, or if we ended up with no context ids
+    // (we moved all edges onto other clones).
+    if (!Node->hasCall() || Node->ContextIds.empty())
+      return;
+
+    if (Node->IsAllocation) {
+      updateAllocationCall(Node->Call, allocTypeToUse(Node->AllocTypes));
+      return;
+    }
+
+    if (!CallsiteToCalleeFuncCloneMap.count(Node))
+      return;
+
+    auto CalleeFunc = CallsiteToCalleeFuncCloneMap[Node];
+    updateCall(Node->Call, CalleeFunc);
+  };
+
+  // Performs DFS traversal starting from allocation nodes to update calls to
+  // reflect cloning decisions recorded earlier. For regular LTO this will
+  // update the actual calls in the IR to call the appropriate function clone
+  // (and add attributes to allocation calls), whereas for ThinLTO the decisions
+  // are recorded in the summary entries.
+  DenseSet<const ContextNode *> Visited;
+  for (auto &Entry : AllocationCallToContextNodeMap)
+    UpdateCalls(Entry.second, Visited, UpdateCalls);
+
+  return Changed;
 }
 
 template <typename DerivedCCG, typename FuncTy, typename CallTy>
 bool CallsiteContextGraph<DerivedCCG, FuncTy, CallTy>::process() {
   if (DumpCCG) {
     dbgs() << "CCG before cloning:\n";
     dbgs() << *this;
   }
bool CallsiteContextGraph<DerivedCCG, FuncTy, CallTy>::process() {
   if (VerifyCCG) {
     check();
   }
 
   if (DumpCCG) {
     dbgs() << "CCG after cloning:\n";
     dbgs() << *this;
   }
   if (ExportToDot)
     exportToDot("cloned");
 
-  return false;
+  bool Changed = assignFunctions();
+
+  if (DumpCCG) {
+    dbgs() << "CCG after assigning function clones:\n";
+    dbgs() << *this;
+  }
+  if (ExportToDot)
+    exportToDot("clonefuncassign");
+
+  return Changed;
 }
 
-bool MemProfContextDisambiguation::processModule(Module &M) {
+bool MemProfContextDisambiguation::processModule(
+    Module &M,
+    function_ref<OptimizationRemarkEmitter &(Function *)> OREGetter) {
   bool Changed = false;
 
-  ModuleCallsiteContextGraph CCG(M);
+  ModuleCallsiteContextGraph CCG(M, OREGetter);
   Changed = CCG.process();
 
   return Changed;
 }
 
 PreservedAnalyses MemProfContextDisambiguation::run(Module &M,
                                                     ModuleAnalysisManager &AM) {
-  if (!processModule(M))
+  auto &FAM = AM.getResult<FunctionAnalysisManagerModuleProxy>(M).getManager();
+  auto OREGetter = [&](Function *F) -> OptimizationRemarkEmitter & {
+    return FAM.getResult<OptimizationRemarkEmitterAnalysis>(*F);
+  };
+  if (!processModule(M, OREGetter))
     return PreservedAnalyses::all();
   return PreservedAnalyses::none();
 }
 
 void MemProfContextDisambiguation::run(
     ModuleSummaryIndex &Index,
     function_ref<bool(GlobalValue::GUID, const GlobalValueSummary *)>
         isPrevailing) {
   IndexCallsiteContextGraph CCG(Index, isPrevailing);
   CCG.process();
 }

llvm/test/ThinLTO/X86/memprof-basic.ll

 ; RUN: opt -thinlto-bc %s >%t.o
 ; RUN: llvm-lto2 run %t.o -enable-memprof-context-disambiguation \
 ; RUN:  -r=%t.o,main,plx \
 ; RUN:  -r=%t.o,_ZdaPv, \
 ; RUN:  -r=%t.o,sleep, \
 ; RUN:  -r=%t.o,_Znam, \
 ; RUN:  -memprof-verify-ccg -memprof-verify-nodes -memprof-dump-ccg \
 ; RUN:  -memprof-export-to-dot -memprof-dot-file-path-prefix=%t. \
-; RUN:  -o %t.out 2>&1 | FileCheck %s --check-prefix=DUMP
+; RUN:  -stats -pass-remarks=memprof-context-disambiguation -save-temps \
+; RUN:  -o %t.out 2>&1 | FileCheck %s --check-prefix=DUMP \
+; RUN:  --check-prefix=STATS
 
 ; RUN:  cat %t.ccg.postbuild.dot | FileCheck %s --check-prefix=DOT
 ;; We should have cloned bar, baz, and foo, for the cold memory allocation.
 ; RUN:  cat %t.ccg.cloned.dot | FileCheck %s --check-prefix=DOTCLONED
 
 
+;; Try again but with distributed ThinLTO
+; RUN: llvm-lto2 run %t.o -enable-memprof-context-disambiguation \
+; RUN:  -thinlto-distributed-indexes \
+; RUN:  -r=%t.o,main,plx \
+; RUN:  -r=%t.o,_ZdaPv, \
+; RUN:  -r=%t.o,sleep, \
+; RUN:  -r=%t.o,_Znam, \
+; RUN:  -memprof-verify-ccg -memprof-verify-nodes -memprof-dump-ccg \
+; RUN:  -memprof-export-to-dot -memprof-dot-file-path-prefix=%t2. \
+; RUN:  -stats -pass-remarks=memprof-context-disambiguation \
+; RUN:  -o %t2.out 2>&1 | FileCheck %s --check-prefix=DUMP \
+; RUN:  --check-prefix=STATS
+
+; RUN:  cat %t2.ccg.postbuild.dot | FileCheck %s --check-prefix=DOT
+;; We should have cloned bar, baz, and foo, for the cold memory allocation.
+; RUN:  cat %t2.ccg.cloned.dot | FileCheck %s --check-prefix=DOTCLONED
+
+;; Check distributed index
+; RUN: llvm-dis %t.o.thinlto.bc -o - | FileCheck %s --check-prefix=DISTRIB
+
 source_filename = "memprof-basic.ll"
 target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
 target triple = "x86_64-unknown-linux-gnu"
 
 define i32 @main() {
 entry:
   %call = call ptr @_Z3foov(), !callsite !0
   %call1 = call ptr @_Z3foov(), !callsite !1

 entry:
   %call = call ptr @_Z3bazv(), !callsite !9
   ret ptr null
 }
 
 ; uselistorder directives
 uselistorder ptr @_Z3foov, { 1, 0 }
 
 !0 = !{i64 8632435727821051414}
 !1 = !{i64 -3421689549917153178}
 !2 = !{!3, !5}
 !3 = !{!4, !"notcold"}
 !4 = !{i64 9086428284934609951, i64 -5964873800580613432, i64 2732490490862098848, i64 8632435727821051414}
 !5 = !{!6, !"cold"}
 !6 = !{i64 9086428284934609951, i64 -5964873800580613432, i64 2732490490862098848, i64 -3421689549917153178}
 !7 = !{i64 9086428284934609951}
 !8 = !{i64 -5964873800580613432}

 ; DUMP:   AllocTypes: Cold
 ; DUMP:   ContextIds: 2
 ; DUMP:   CalleeEdges:
 ; DUMP:   CallerEdges:
 ; DUMP:     Edge from Callee [[BAR2]] to Caller: [[BAZ2]] AllocTypes: Cold ContextIds: 2
 ; DUMP:   Clone of [[BAR]]
 
 
+; STATS: 1 memprof-context-disambiguation - Number of cold static allocations (possibly cloned)
+; STATS: 1 memprof-context-disambiguation - Number of not cold static allocations (possibly cloned)
+; STATS: 3 memprof-context-disambiguation - Number of function clones created during whole program analysis
+
+
 ; DOT: digraph "postbuild" {
 ; DOT:  label="postbuild";
 ; DOT:   Node[[BAR:0x[a-z0-9]+]] [shape=record,tooltip="N[[BAR]] ContextIds: 1 2",fillcolor="mediumorchid1",style="filled",style="filled",label="{OrigId: Alloc0\n_Z3barv -\> alloc}"];
 ; DOT:   Node[[BAZ:0x[a-z0-9]+]] [shape=record,tooltip="N[[BAZ]] ContextIds: 1 2",fillcolor="mediumorchid1",style="filled",style="filled",label="{OrigId: 12481870273128938184\n_Z3bazv -\> _Z3barv}"];
 ; DOT:  Node[[BAZ]] -> Node[[BAR]][tooltip="ContextIds: 1 2",fillcolor="mediumorchid1"];
 ; DOT:   Node[[FOO:0x[a-z0-9]+]] [shape=record,tooltip="N[[FOO]] ContextIds: 1 2",fillcolor="mediumorchid1",style="filled",style="filled",label="{OrigId: 2732490490862098848\n_Z3foov -\> _Z3bazv}"];
 ; DOT:  Node[[FOO]] -> Node[[BAZ]][tooltip="ContextIds: 1 2",fillcolor="mediumorchid1"];
 ; DOT:   Node[[MAIN1:0x[a-z0-9]+]] [shape=record,tooltip="N[[MAIN1]] ContextIds: 1",fillcolor="brown1",style="filled",style="filled",label="{OrigId: 8632435727821051414\nmain -\> _Z3foov}"];

 ; DOTCLONED:   Node[[MAIN2:0x[a-z0-9]+]] [shape=record,tooltip="N[[MAIN2]] ContextIds: 2",fillcolor="cyan",style="filled",style="filled",label="{OrigId: 15025054523792398438\nmain -\> _Z3foov}"];
 ; DOTCLONED:  Node[[MAIN2]] -> Node[[FOO2:0x[a-z0-9]+]][tooltip="ContextIds: 2",fillcolor="cyan"];
 ; DOTCLONED:   Node[[FOO2]] [shape=record,tooltip="N[[FOO2]] ContextIds: 2",fillcolor="cyan",style="filled",color="blue",style="filled,bold,dashed",label="{OrigId: 0\n_Z3foov -\> _Z3bazv}"];
 ; DOTCLONED:  Node[[FOO2]] -> Node[[BAZ2:0x[a-z0-9]+]][tooltip="ContextIds: 2",fillcolor="cyan"];
 ; DOTCLONED:   Node[[BAZ2]] [shape=record,tooltip="N[[BAZ2]] ContextIds: 2",fillcolor="cyan",style="filled",color="blue",style="filled,bold,dashed",label="{OrigId: 0\n_Z3bazv -\> _Z3barv}"];
 ; DOTCLONED:  Node[[BAZ2]] -> Node[[BAR2:0x[a-z0-9]+]][tooltip="ContextIds: 2",fillcolor="cyan"];
 ; DOTCLONED:   Node[[BAR2]] [shape=record,tooltip="N[[BAR2]] ContextIds: 2",fillcolor="cyan",style="filled",color="blue",style="filled,bold,dashed",label="{OrigId: Alloc0\n_Z3barv -\> alloc}"];
 ; DOTCLONED: }
+
+
+; DISTRIB: ^[[BAZ:[0-9]+]] = gv: (guid: 5878270615442837395, {{.*}} callsites: ((callee: ^[[BAR:[0-9]+]], clones: (0, 1)
+; DISTRIB: ^[[FOO:[0-9]+]] = gv: (guid: 6731117468105397038, {{.*}} callsites: ((callee: ^[[BAZ]], clones: (0, 1)
+; DISTRIB: ^[[BAR]] = gv: (guid: 9832687305761716512, {{.*}} allocs: ((versions: (notcold, cold)
+; DISTRIB: ^[[MAIN:[0-9]+]] = gv: (guid: 15822663052811949562, {{.*}} callsites: ((callee: ^[[FOO]], clones: (0), {{.*}} (callee: ^[[FOO]], clones: (1)

llvm/test/ThinLTO/X86/memprof-duplicate-context-ids.ll

 ;; Test callsite context graph generation for call graph with with MIBs
 ;; that have pruned contexts that partially match multiple inlined
 ;; callsite contexts, requiring duplication of context ids and nodes
-;; while matching callsite nodes onto the graph.
+;; while matching callsite nodes onto the graph. Also tests graph and IR
+;; cloning.
 ;;
 ;; Original code looks like:
 ;;
 ;; char *D() {
 ;;   return new char[10];
 ;; }
 ;;
 ;; char *F() {

 ; RUN: opt -thinlto-bc %s >%t.o
 ; RUN: llvm-lto2 run %t.o -enable-memprof-context-disambiguation \
 ; RUN:  -r=%t.o,main,plx \
 ; RUN:  -r=%t.o,_ZdaPv, \
 ; RUN:  -r=%t.o,sleep, \
 ; RUN:  -r=%t.o,_Znam, \
 ; RUN:  -memprof-verify-ccg -memprof-verify-nodes -memprof-dump-ccg \
 ; RUN:  -memprof-export-to-dot -memprof-dot-file-path-prefix=%t. \
-; RUN:  -o %t.out 2>&1 | FileCheck %s --check-prefix=DUMP
+; RUN:  -stats -pass-remarks=memprof-context-disambiguation -save-temps \
+; RUN:  -o %t.out 2>&1 | FileCheck %s --check-prefix=DUMP \
+; RUN:  --check-prefix=STATS
 
 ; RUN:  cat %t.ccg.prestackupdate.dot | FileCheck %s --check-prefix=DOTPRE
 ; RUN:  cat %t.ccg.postbuild.dot | FileCheck %s --check-prefix=DOTPOST
 ;; We should clone D once for the cold allocations via C.
 ; RUN:  cat %t.ccg.cloned.dot | FileCheck %s --check-prefix=DOTCLONED
 
 
+;; Try again but with distributed ThinLTO
+; RUN: llvm-lto2 run %t.o -enable-memprof-context-disambiguation \
+; RUN:  -thinlto-distributed-indexes \
+; RUN:  -r=%t.o,main,plx \
+; RUN:  -r=%t.o,_ZdaPv, \
+; RUN:  -r=%t.o,sleep, \
+; RUN:  -r=%t.o,_Znam, \
+; RUN:  -memprof-verify-ccg -memprof-verify-nodes -memprof-dump-ccg \
+; RUN:  -memprof-export-to-dot -memprof-dot-file-path-prefix=%t2. \
+; RUN:  -stats -pass-remarks=memprof-context-disambiguation \
+; RUN:  -o %t2.out 2>&1 | FileCheck %s --check-prefix=DUMP \
+; RUN:  --check-prefix=STATS
+
+; RUN:  cat %t.ccg.prestackupdate.dot | FileCheck %s --check-prefix=DOTPRE
+; RUN:  cat %t.ccg.postbuild.dot | FileCheck %s --check-prefix=DOTPOST
+;; We should clone D once for the cold allocations via C.
+; RUN:  cat %t.ccg.cloned.dot | FileCheck %s --check-prefix=DOTCLONED
+
+;; Check distributed index
+; RUN: llvm-dis %t.o.thinlto.bc -o - | FileCheck %s --check-prefix=DISTRIB
+
 source_filename = "duplicate-context-ids.ll"
 target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
 target triple = "x86_64-unknown-linux-gnu"
 
 define internal ptr @_Z1Dv() {
 entry:
   %call = call ptr @_Znam(i64 0), !memprof !0, !callsite !5
   ret ptr null

 }
 
 define internal ptr @_Z1Ev() {
 entry:
   %call.i = call ptr @_Z1Dv(), !callsite !9
   ret ptr null
 }
 
-declare i32 @main()
+define i32 @main() {
+entry:
+  call ptr @_Z1Bv()
+  call ptr @_Z1Ev()
+  call ptr @_Z1Fv()
+  ret i32 0
+}
 
 declare void @_ZdaPv()
 
 declare i32 @sleep()
 
 !0 = !{!1, !3}
 !1 = !{!2, !"cold"}
 !2 = !{i64 6541423618768552252, i64 -6270142974039008131}

 ; DUMP:   CalleeEdges:
 ; DUMP:   CallerEdges:
 ; DUMP:     Edge from Callee [[D2]] to Caller: [[E:0x[a-z0-9]+]] AllocTypes: Cold ContextIds: 1
 ; DUMP:     Edge from Callee [[D2]] to Caller: [[C2:0x[a-z0-9]+]] AllocTypes: Cold ContextIds: 3
 ; DUMP:     Edge from Callee [[D2]] to Caller: [[B:0x[a-z0-9]+]] AllocTypes: Cold ContextIds: 4
 ; DUMP:         Clone of [[D]]
 
 
+; STATS: 1 memprof-context-disambiguation - Number of cold static allocations (possibly cloned)
+; STATS: 1 memprof-context-disambiguation - Number of not cold static allocations (possibly cloned)
+; STATS: 1 memprof-context-disambiguation - Number of function clones created during whole program analysis
+
+
 ; DOTPRE: digraph "prestackupdate" {
 ; DOTPRE:   label="prestackupdate";
 ; DOTPRE:   Node[[D:0x[a-z0-9]+]] [shape=record,tooltip="N[[D]] ContextIds: 1 2",fillcolor="mediumorchid1",style="filled",style="filled",label="{OrigId: Alloc0\n_Z1Dv -\> alloc}"];
 ; DOTPRE:   Node[[C:0x[a-z0-9]+]] [shape=record,tooltip="N[[C]] ContextIds: 1",fillcolor="cyan",style="filled",style="filled",label="{OrigId: 12176601099670543485\nnull call (external)}"];
 ; DOTPRE:   Node[[C]] -> Node[[D]][tooltip="ContextIds: 1",fillcolor="cyan"];
 ; DOTPRE:   Node[[F:0x[a-z0-9]+]] [shape=record,tooltip="N[[F]] ContextIds: 2",fillcolor="brown1",style="filled",style="filled",label="{OrigId: 13543580133643026784\nnull call (external)}"];
 ; DOTPRE:   Node[[F]] -> Node[[D]][tooltip="ContextIds: 2",fillcolor="brown1"];
 ; DOTPRE: }

 ; DOTCLONED:  Node[[C:0x[a-z0-9]+]] [shape=record,tooltip="N[[C]] ContextIds: 3",fillcolor="cyan",style="filled",style="filled",label="{OrigId: 0\n_Z1Cv -\> _Z1Dv}"];
 ; DOTCLONED:  Node[[C]] -> Node[[D2:0x[a-z0-9]+]][tooltip="ContextIds: 3",fillcolor="cyan"];
 ; DOTCLONED:  Node[[B:0x[a-z0-9]+]] [shape=record,tooltip="N[[B]] ContextIds: 4",fillcolor="cyan",style="filled",style="filled",label="{OrigId: 0\n_Z1Bv -\> _Z1Dv}"];
 ; DOTCLONED:  Node[[B]] -> Node[[D2]][tooltip="ContextIds: 4",fillcolor="cyan"];
 ; DOTCLONED:  Node[[E:0x[a-z0-9]+]] [shape=record,tooltip="N[[E]] ContextIds: 1",fillcolor="cyan",style="filled",style="filled",label="{OrigId: 0\n_Z1Ev -\> _Z1Dv}"];
 ; DOTCLONED:  Node[[E]] -> Node[[D2]][tooltip="ContextIds: 1",fillcolor="cyan"];
 ; DOTCLONED:   Node[[D2]] [shape=record,tooltip="N[[D2]] ContextIds: 1 3 4",fillcolor="cyan",style="filled",color="blue",style="filled,bold,dashed",label="{OrigId: Alloc0\n_Z1Dv -\> alloc}"];
 ; DOTCLONED: }
+
+; DISTRIB: ^[[C:[0-9]+]] = gv: (guid: 1643923691937891493, {{.*}} callsites: ((callee: ^[[D:[0-9]+]], clones: (1)
+; DISTRIB: ^[[D]] = gv: (guid: 4881081444663423788, {{.*}} allocs: ((versions: (notcold, cold)
+; DISTRIB: ^[[B:[0-9]+]] = gv: (guid: 14590037969532473829, {{.*}} callsites: ((callee: ^[[D]], clones: (1)
+; DISTRIB: ^[[F:[0-9]+]] = gv: (guid: 17035303613541779335, {{.*}} callsites: ((callee: ^[[D]], clones: (0)
+; DISTRIB: ^[[E:[0-9]+]] = gv: (guid: 17820708772846654376, {{.*}} callsites: ((callee: ^[[D]], clones: (1)

llvm/test/ThinLTO/X86/memprof-funcassigncloning.ll

 ;;   return 0;
 ;; }
 ;;
 ;; Code compiled with -mllvm -memprof-min-lifetime-cold-threshold=5 so that the
 ;; memory freed after sleep(10) results in cold lifetimes.
 ;;
 ;; The IR was then reduced using llvm-reduce with the expected FileCheck input.
 
+;; -stats requires asserts
+; REQUIRES: asserts
+
 
 ; RUN: opt -thinlto-bc %s >%t.o
 ; RUN: llvm-lto2 run %t.o -enable-memprof-context-disambiguation \
 ; RUN:  -r=%t.o,main,plx \
 ; RUN:  -r=%t.o,_ZdaPv, \
 ; RUN:  -r=%t.o,sleep, \
 ; RUN:  -r=%t.o,_Znam, \
 ; RUN:  -memprof-verify-ccg -memprof-verify-nodes -memprof-dump-ccg \
-; RUN:  -o %t.out 2>&1 | FileCheck %s --check-prefix=DUMP
+; RUN:  -stats -pass-remarks=memprof-context-disambiguation -save-temps \
+; RUN:  -o %t.out 2>&1 | FileCheck %s --check-prefix=DUMP \
+; RUN:  --check-prefix=STATS
 
 
 ;; Try again but with distributed ThinLTO
 ; RUN: llvm-lto2 run %t.o -enable-memprof-context-disambiguation \
 ; RUN:  -thinlto-distributed-indexes \
 ; RUN:  -r=%t.o,main,plx \
 ; RUN:  -r=%t.o,_ZdaPv, \
 ; RUN:  -r=%t.o,sleep, \
 ; RUN:  -r=%t.o,_Znam, \
 ; RUN:  -memprof-verify-ccg -memprof-verify-nodes -memprof-dump-ccg \
-; RUN:  -o %t2.out 2>&1 | FileCheck %s --check-prefix=DUMP
+; RUN:  -stats -pass-remarks=memprof-context-disambiguation \
+; RUN:  -o %t2.out 2>&1 | FileCheck %s --check-prefix=DUMP \
+; RUN:  --check-prefix=STATS
 
 
 source_filename = "funcassigncloning.ll"
 target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
 target triple = "x86_64-unknown-linux-gnu"
 
 ; Function Attrs: noinline optnone
 define internal void @_Z1EPPcS0_(ptr %buf1, ptr %buf2) {

 ; DUMP:                 AllocType 1 StackIds: 0
 ; DUMP:         (clone 0)
 ; DUMP:   AllocTypes: Cold
 ; DUMP:   ContextIds: 5
 ; DUMP:   CalleeEdges:
 ; DUMP:   CallerEdges:
 ; DUMP:     Edge from Callee [[ENEW2CLONE]] to Caller: [[C]] AllocTypes: Cold ContextIds: 5
 ; DUMP:   Clone of [[ENEW2ORIG]]
+
+
+; STATS: 2 memprof-context-disambiguation - Number of cold static allocations (possibly cloned)
+; STATS: 4 memprof-context-disambiguation - Number of not cold static allocations (possibly cloned)
+; STATS: 3 memprof-context-disambiguation - Number of function clones created during whole program analysis

llvm/test/ThinLTO/X86/memprof-indirectcall.ll

 ;; Tests callsite context graph generation for call graph containing indirect
 ;; calls. Currently this should result in conservative behavior, such that the
 ;; indirect call receives a null call in its graph node, to prevent subsequent
-;; cloning.
+;; cloning. Also tests graph and IR cloning.
 ;;
 ;; Original code looks like:
 ;;
 ;; char *foo() {
 ;;   return new char[10];
 ;; }
 ;; class A {
 ;; public:

 ; RUN:  -r=%t.o,main,plx \
 ; RUN:  -r=%t.o,sleep, \
 ; RUN:  -r=%t.o,_Znam, \
 ; RUN:  -r=%t.o,_ZdaPv, \
 ; RUN:  -r=%t.o,_ZTVN10__cxxabiv120__si_class_type_infoE, \
 ; RUN:  -r=%t.o,_ZTVN10__cxxabiv117__class_type_infoE, \
 ; RUN:  -memprof-verify-ccg -memprof-verify-nodes -memprof-dump-ccg \
 ; RUN:  -memprof-export-to-dot -memprof-dot-file-path-prefix=%t. \
-; RUN:  -o %t.out 2>&1 | FileCheck %s --check-prefix=DUMP
+; RUN:  -stats -pass-remarks=memprof-context-disambiguation -save-temps \
+; RUN:  -o %t.out 2>&1 | FileCheck %s --check-prefix=DUMP \
+; RUN:  --check-prefix=STATS
 
 ; RUN:  cat %t.ccg.postbuild.dot | FileCheck %s --check-prefix=DOT
 ;; We should only create a single clone of foo, for the direct call
 ;; from main allocating cold memory.
 ; RUN:  cat %t.ccg.cloned.dot | FileCheck %s --check-prefix=DOTCLONED
 
 
+;; Try again but with distributed ThinLTO
+; RUN: llvm-lto2 run %t.o -enable-memprof-context-disambiguation \
+; RUN:  -thinlto-distributed-indexes \
+; RUN:  -r=%t.o,main,plx \
+; RUN:  -r=%t.o,_ZdaPv, \
+; RUN:  -r=%t.o,sleep, \
+; RUN:  -r=%t.o,_Znam, \
+; RUN:  -r=%t.o,_ZTVN10__cxxabiv120__si_class_type_infoE, \
+; RUN:  -r=%t.o,_ZTVN10__cxxabiv117__class_type_infoE, \
+; RUN:  -memprof-verify-ccg -memprof-verify-nodes -memprof-dump-ccg \
+; RUN:  -memprof-export-to-dot -memprof-dot-file-path-prefix=%t2. \
+; RUN:  -stats -pass-remarks=memprof-context-disambiguation \
+; RUN:  -o %t2.out 2>&1 | FileCheck %s --check-prefix=DUMP \
+; RUN:  --check-prefix=STATS
+
+; RUN:  cat %t.ccg.postbuild.dot | FileCheck %s --check-prefix=DOT
+;; We should only create a single clone of foo, for the direct call
+;; from main allocating cold memory.
+; RUN:  cat %t.ccg.cloned.dot | FileCheck %s --check-prefix=DOTCLONED
+
 source_filename = "indirectcall.ll"
 target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
 target triple = "x86_64-unknown-linux-gnu"
 
 @_ZTVN10__cxxabiv120__si_class_type_infoE = external global ptr
 @_ZTVN10__cxxabiv117__class_type_infoE = external global ptr
 
 define internal ptr @_Z3barP1A(ptr %a) {

 ; DUMP:   AllocTypes: Cold
 ; DUMP:   ContextIds: 6
 ; DUMP:   CalleeEdges:
 ; DUMP:   CallerEdges:
 ; DUMP:     Edge from Callee [[FOO2]] to Caller: [[MAIN2]] AllocTypes: Cold ContextIds: 6
 ; DUMP:   Clone of [[FOO]]
 
 
+; STATS: 1 memprof-context-disambiguation - Number of cold static allocations (possibly cloned)
+; STATS: 1 memprof-context-disambiguation - Number of not cold static allocations (possibly cloned)
+; STATS: 1 memprof-context-disambiguation - Number of function clones created during whole program analysis
+
+
 ; DOT: digraph "postbuild" {
 ; DOT:  label="postbuild";
 ; DOT:   Node[[FOO:0x[a-z0-9]+]] [shape=record,tooltip="N[[FOO]] ContextIds: 1 2 3 4 5 6",fillcolor="mediumorchid1",style="filled",style="filled",label="{OrigId: Alloc0\n_Z3foov -\> alloc}"];
 ; DOT:   Node[[AX:0x[a-z0-9]+]] [shape=record,tooltip="N[[AX]] ContextIds: 1 2",fillcolor="mediumorchid1",style="filled",style="filled",label="{OrigId: 8256774051149711748\n_ZN1A1xEv -\> _Z3foov}"];
 ; DOT:  Node[[AX]] -> Node[[FOO]][tooltip="ContextIds: 1 2",fillcolor="mediumorchid1"];
 ; DOT:   Node[[BAR:0x[a-z0-9]+]] [shape=record,tooltip="N[[BAR]] ContextIds: 1 2 4 5",fillcolor="mediumorchid1",style="filled",style="filled",label="{OrigId: 13626499562959447861\nnull call (external)}"];
 ; DOT:  Node[[BAR]] -> Node[[AX]][tooltip="ContextIds: 1 2",fillcolor="mediumorchid1"];
 ; DOT:  Node[[BAR]] -> Node[[BX:0x[a-z0-9]+]][tooltip="ContextIds: 4 5",fillcolor="mediumorchid1"];

llvm/test/ThinLTO/X86/memprof-inlined.ll

 ;; Test callsite context graph generation for call graph with two memprof
 ;; contexts and partial inlining, requiring generation of a new fused node to
 ;; represent the inlined sequence while matching callsite nodes onto the graph.
+;; Also tests graph and IR cloning.
 ;;
 ;; Original code looks like:
 ;;
 ;; char *bar() {
 ;;   return new char[10];
 ;; }
 ;;
 ;; char *baz() {

 ; RUN: opt -thinlto-bc %s >%t.o
 ; RUN: llvm-lto2 run %t.o -enable-memprof-context-disambiguation \
 ; RUN:  -r=%t.o,main,plx \
 ; RUN:  -r=%t.o,_ZdaPv, \
 ; RUN:  -r=%t.o,sleep, \
 ; RUN:  -r=%t.o,_Znam, \
 ; RUN:  -memprof-verify-ccg -memprof-verify-nodes -memprof-dump-ccg \
 ; RUN:  -memprof-export-to-dot -memprof-dot-file-path-prefix=%t. \
-; RUN:  -o %t.out 2>&1 | FileCheck %s --check-prefix=DUMP
+; RUN:  -stats -pass-remarks=memprof-context-disambiguation -save-temps \
+; RUN:  -o %t.out 2>&1 | FileCheck %s --check-prefix=DUMP \
+; RUN:  --check-prefix=STATS
 
 ; RUN:  cat %t.ccg.postbuild.dot | FileCheck %s --check-prefix=DOT
 ;; We should create clones for foo and bar for the call from main to allocate
 ;; cold memory.
 ; RUN:  cat %t.ccg.cloned.dot | FileCheck %s --check-prefix=DOTCLONED
 
 
+;; Try again but with distributed ThinLTO
+; RUN: llvm-lto2 run %t.o -enable-memprof-context-disambiguation \
+; RUN:  -thinlto-distributed-indexes \
+; RUN:  -r=%t.o,main,plx \
+; RUN:  -r=%t.o,_ZdaPv, \
+; RUN:  -r=%t.o,sleep, \
+; RUN:  -r=%t.o,_Znam, \
+; RUN:  -memprof-verify-ccg -memprof-verify-nodes -memprof-dump-ccg \
+; RUN:  -memprof-export-to-dot -memprof-dot-file-path-prefix=%t2. \
+; RUN:  -stats -pass-remarks=memprof-context-disambiguation \
+; RUN:  -o %t2.out 2>&1 | FileCheck %s --check-prefix=DUMP \
+; RUN:  --check-prefix=STATS
+
+; RUN:  cat %t.ccg.postbuild.dot | FileCheck %s --check-prefix=DOT
+;; We should create clones for foo and bar for the call from main to allocate
+;; cold memory.
+; RUN:  cat %t.ccg.cloned.dot | FileCheck %s --check-prefix=DOTCLONED
+
 source_filename = "inlined.ll"
 target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
 target triple = "x86_64-unknown-linux-gnu"
 
 define internal ptr @_Z3barv() {
 entry:
   %call = call ptr @_Znam(i64 0), !memprof !0, !callsite !5
   ret ptr null

 ; DUMP:   AllocTypes: Cold
 ; DUMP:   ContextIds: 4
 ; DUMP:   CalleeEdges:
 ; DUMP:   CallerEdges:
 ; DUMP:     Edge from Callee [[BAR2]] to Caller: [[FOO3]] AllocTypes: Cold ContextIds: 4
 ; DUMP:         Clone of [[BAR]]
 
 
+; STATS: 1 memprof-context-disambiguation - Number of cold static allocations (possibly cloned)
+; STATS: 2 memprof-context-disambiguation - Number of not cold static allocations (possibly cloned)
+; STATS: 2 memprof-context-disambiguation - Number of function clones created during whole program analysis
+
+
 ; DOT: digraph "postbuild" {
 ; DOT:  label="postbuild";
 ; DOT:   Node[[BAZ:0x[a-z0-9]+]] [shape=record,tooltip="N[[BAZ]] ContextIds: 1 2",fillcolor="mediumorchid1",style="filled",style="filled",label="{OrigId: Alloc0\n_Z3bazv -\> alloc}"];
 ; DOT:   Node[[FOO:0x[a-z0-9]+]] [shape=record,tooltip="N[[FOO]] ContextIds: 1 2",fillcolor="mediumorchid1",style="filled",style="filled",label="{OrigId: 2732490490862098848\nnull call (external)}"];
 ; DOT:  Node[[FOO]] -> Node[[BAZ]][tooltip="ContextIds: 1 2",fillcolor="mediumorchid1"];
 ; DOT:   Node[[MAIN1:0x[a-z0-9]+]] [shape=record,tooltip="N[[MAIN1]] ContextIds: 1 3",fillcolor="brown1",style="filled",style="filled",label="{OrigId: 8632435727821051414\nmain -\> _Z3foov}"];
 ; DOT:  Node[[MAIN1]] -> Node[[FOO]][tooltip="ContextIds: 1",fillcolor="brown1"];
 ; DOT:  Node[[MAIN1]] -> Node[[FOO2:0x[a-z0-9]+]][tooltip="ContextIds: 3",fillcolor="brown1"];

llvm/test/Transforms/MemProfContextDisambiguation/basic.ll

 ;; Test callsite context graph generation for simple call graph with
-;; two memprof contexts and no inlining.
+;; two memprof contexts and no inlining, as well as graph and IR cloning.
 ;;
 ;; Original code looks like:
 ;;
 ;; char *bar() {
 ;;   return new char[10];
 ;; }
 ;;
 ;; char *baz() {

 ;; The IR was then reduced using llvm-reduce with the expected FileCheck input.
 
 ;; -stats requires asserts
 ; REQUIRES: asserts
 
 ; RUN: opt -passes=memprof-context-disambiguation \
 ; RUN:  -memprof-verify-ccg -memprof-verify-nodes -memprof-dump-ccg \
 ; RUN:  -memprof-export-to-dot -memprof-dot-file-path-prefix=%t. \
-; RUN:  %s -S 2>&1 | FileCheck %s --check-prefix=DUMP
+; RUN:  -stats -pass-remarks=memprof-context-disambiguation \
+; RUN:  %s -S 2>&1 | FileCheck %s --check-prefix=DUMP --check-prefix=IR \
+; RUN:  --check-prefix=STATS --check-prefix=REMARKS
 
 ; RUN:  cat %t.ccg.postbuild.dot | FileCheck %s --check-prefix=DOT
 ;; We should have cloned bar, baz, and foo, for the cold memory allocation.
 ; RUN:  cat %t.ccg.cloned.dot | FileCheck %s --check-prefix=DOTCLONED
 
 target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
 target triple = "x86_64-unknown-linux-gnu"
 

 ; DUMP:   AllocTypes: Cold
 ; DUMP:   ContextIds: 2
 ; DUMP:   CalleeEdges:
 ; DUMP:   CallerEdges:
 ; DUMP:     Edge from Callee [[BAR2]] to Caller: [[BAZ2]] AllocTypes: Cold ContextIds: 2
 ; DUMP:   Clone of [[BAR]]
 
 
+; REMARKS: created clone _Z3barv.memprof.1
+; REMARKS: created clone _Z3bazv.memprof.1
+; REMARKS: created clone _Z3foov.memprof.1
+; REMARKS: call in clone main assigned to call function clone _Z3foov.memprof.1
+; REMARKS: call in clone _Z3foov.memprof.1 assigned to call function clone _Z3bazv.memprof.1
+; REMARKS: call in clone _Z3bazv.memprof.1 assigned to call function clone _Z3barv.memprof.1
+; REMARKS: call in clone _Z3barv.memprof.1 marked with memprof allocation attribute cold
+; REMARKS: call in clone main assigned to call function clone _Z3foov
+; REMARKS: call in clone _Z3foov assigned to call function clone _Z3bazv
+; REMARKS: call in clone _Z3bazv assigned to call function clone _Z3barv
+; REMARKS: call in clone _Z3barv marked with memprof allocation attribute notcold
+
+
+; IR: define {{.*}} @main
+;; The first call to foo does not allocate cold memory. It should call the
+;; original functions, which ultimately call the original allocation decorated
+;; with a "notcold" attribute.
+; IR:   call {{.*}} @_Z3foov()
+;; The second call to foo allocates cold memory. It should call cloned functions
+;; which ultimately call a cloned allocation decorated with a "cold" attribute.
+; IR:   call {{.*}} @_Z3foov.memprof.1()
+; IR: define internal {{.*}} @_Z3barv()
+; IR:   call {{.*}} @_Znam(i64 noundef 10) #[[NOTCOLD:[0-9]+]]
+; IR: define internal {{.*}} @_Z3bazv()
+; IR:   call {{.*}} @_Z3barv()
+; IR: define internal {{.*}} @_Z3foov()
+; IR:   call {{.*}} @_Z3bazv()
+; IR: define internal {{.*}} @_Z3barv.memprof.1()
+; IR:   call {{.*}} @_Znam(i64 noundef 10) #[[COLD:[0-9]+]]
+; IR: define internal {{.*}} @_Z3bazv.memprof.1()
+; IR:   call {{.*}} @_Z3barv.memprof.1()
+; IR: define internal {{.*}} @_Z3foov.memprof.1()
+; IR:   call {{.*}} @_Z3bazv.memprof.1()
+; IR: attributes #[[NOTCOLD]] = { builtin "memprof"="notcold" }
+; IR: attributes #[[COLD]] = { builtin "memprof"="cold" }
+
+
+; STATS: 1 memprof-context-disambiguation - Number of cold static allocations (possibly cloned)
+; STATS: 1 memprof-context-disambiguation - Number of not cold static allocations (possibly cloned)
+; STATS: 3 memprof-context-disambiguation - Number of function clones created during whole program analysis
+
+
 ; DOT: digraph "postbuild" {
 ; DOT:  label="postbuild";
 ; DOT:   Node[[BAR:0x[a-z0-9]+]] [shape=record,tooltip="N[[BAR]] ContextIds: 1 2",fillcolor="mediumorchid1",style="filled",style="filled",label="{OrigId: Alloc0\n_Z3barv -\> _Znam}"];
 ; DOT:   Node[[BAZ:0x[a-z0-9]+]] [shape=record,tooltip="N[[BAZ]] ContextIds: 1 2",fillcolor="mediumorchid1",style="filled",style="filled",label="{OrigId: 12481870273128938184\n_Z3bazv -\> _Z3barv}"];
 ; DOT:  Node[[BAZ]] -> Node[[BAR]][tooltip="ContextIds: 1 2",fillcolor="mediumorchid1"];
 ; DOT:   Node[[FOO:0x[a-z0-9]+]] [shape=record,tooltip="N[[FOO]] ContextIds: 1 2",fillcolor="mediumorchid1",style="filled",style="filled",label="{OrigId: 2732490490862098848\n_Z3foov -\> _Z3bazv}"];
 ; DOT:  Node[[FOO]] -> Node[[BAZ]][tooltip="ContextIds: 1 2",fillcolor="mediumorchid1"];
 ; DOT:   Node[[MAIN1:0x[a-z0-9]+]] [shape=record,tooltip="N[[MAIN1]] ContextIds: 1",fillcolor="brown1",style="filled",style="filled",label="{OrigId: 8632435727821051414\nmain -\> _Z3foov}"];

llvm/test/Transforms/MemProfContextDisambiguation/duplicate-context-ids.ll

 ;; Test callsite context graph generation for call graph with with MIBs
 ;; that have pruned contexts that partially match multiple inlined
 ;; callsite contexts, requiring duplication of context ids and nodes
-;; while matching callsite nodes onto the graph.
+;; while matching callsite nodes onto the graph. Also tests graph and IR
+;; cloning.
 ;;
 ;; Original code looks like:
 ;;
 ;; char *D() {
 ;;   return new char[10];
 ;; }
 ;;
 ;; char *F() {

 ;; The IR was then reduced using llvm-reduce with the expected FileCheck input.
 
 ;; -stats requires asserts
 ; REQUIRES: asserts
 
 ; RUN: opt -passes=memprof-context-disambiguation \
 ; RUN:  -memprof-verify-ccg -memprof-verify-nodes -memprof-dump-ccg \
 ; RUN:  -memprof-export-to-dot -memprof-dot-file-path-prefix=%t. \
-; RUN:  %s -S 2>&1 | FileCheck %s --check-prefix=DUMP
+; RUN:  -stats -pass-remarks=memprof-context-disambiguation \
+; RUN:  %s -S 2>&1 | FileCheck %s --check-prefix=DUMP --check-prefix=IR \
+; RUN:  --check-prefix=STATS --check-prefix=REMARKS
 
 ; RUN:  cat %t.ccg.prestackupdate.dot | FileCheck %s --check-prefix=DOTPRE
 ; RUN:  cat %t.ccg.postbuild.dot | FileCheck %s --check-prefix=DOTPOST
 ;; We should clone D once for the cold allocations via C.
 ; RUN:  cat %t.ccg.cloned.dot | FileCheck %s --check-prefix=DOTCLONED
 
 target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
 target triple = "x86_64-unknown-linux-gnu"

 ; DUMP:   ContextIds: 1 3 4
 ; DUMP:   CalleeEdges:
 ; DUMP:   CallerEdges:
 ; DUMP:     Edge from Callee [[D2]] to Caller: [[E:0x[a-z0-9]+]] AllocTypes: Cold ContextIds: 1
 ; DUMP:     Edge from Callee [[D2]] to Caller: [[C2:0x[a-z0-9]+]] AllocTypes: Cold ContextIds: 3
 ; DUMP:     Edge from Callee [[D2]] to Caller: [[B:0x[a-z0-9]+]] AllocTypes: Cold ContextIds: 4
 ; DUMP:         Clone of [[D]]
 
+; REMARKS: created clone _Z1Dv.memprof.1
+; REMARKS: call in clone _Z1Ev assigned to call function clone _Z1Dv.memprof.1
+; REMARKS: call in clone _Z1Cv assigned to call function clone _Z1Dv.memprof.1
+; REMARKS: call in clone _Z1Bv assigned to call function clone _Z1Dv.memprof.1
+; REMARKS: call in clone _Z1Dv.memprof.1 marked with memprof allocation attribute cold
+; REMARKS: call in clone _Z1Fv assigned to call function clone _Z1Dv
+; REMARKS: call in clone _Z1Dv marked with memprof allocation attribute notcold
+
+
+;; The allocation via F does not allocate cold memory. It should call the
+;; original D, which ultimately call the original allocation decorated
+;; with a "notcold" attribute.
+; IR: define internal {{.*}} @_Z1Dv()
+; IR:   call {{.*}} @_Znam(i64 noundef 10) #[[NOTCOLD:[0-9]+]]
+; IR: define internal {{.*}} @_Z1Fv()
+; IR:   call {{.*}} @_Z1Dv()
+;; The allocations via B and E allocate cold memory. They should call the
+;; cloned D, which ultimately call the cloned allocation decorated with a
+;; "cold" attribute.
+; IR: define internal {{.*}} @_Z1Bv()
+; IR:   call {{.*}} @_Z1Dv.memprof.1()
+; IR: define internal {{.*}} @_Z1Ev()
+; IR:   call {{.*}} @_Z1Dv.memprof.1()
+; IR: define internal {{.*}} @_Z1Dv.memprof.1()
+; IR:   call {{.*}} @_Znam(i64 noundef 10) #[[COLD:[0-9]+]]
+; IR: attributes #[[NOTCOLD]] = { builtin "memprof"="notcold" }
+; IR: attributes #[[COLD]] = { builtin "memprof"="cold" }
+
+
+; STATS: 1 memprof-context-disambiguation - Number of cold static allocations (possibly cloned)
+; STATS: 1 memprof-context-disambiguation - Number of not cold static allocations (possibly cloned)
+; STATS: 1 memprof-context-disambiguation - Number of function clones created during whole program analysis
+
 
 ; DOTPRE: digraph "prestackupdate" {
 ; DOTPRE:   label="prestackupdate";
 ; DOTPRE:   Node[[D:0x[a-z0-9]+]] [shape=record,tooltip="N[[D]] ContextIds: 1 2",fillcolor="mediumorchid1",style="filled",style="filled",label="{OrigId: Alloc0\n_Z1Dv -\> _Znam}"];
 ; DOTPRE:   Node[[C:0x[a-z0-9]+]] [shape=record,tooltip="N[[C]] ContextIds: 1",fillcolor="cyan",style="filled",style="filled",label="{OrigId: 12176601099670543485\nnull call (external)}"];
 ; DOTPRE:   Node[[C]] -> Node[[D]][tooltip="ContextIds: 1",fillcolor="cyan"];
 ; DOTPRE:   Node[[F:0x[a-z0-9]+]] [shape=record,tooltip="N[[F]] ContextIds: 2",fillcolor="brown1",style="filled",style="filled",label="{OrigId: 13543580133643026784\nnull call (external)}"];
 ; DOTPRE:   Node[[F]] -> Node[[D]][tooltip="ContextIds: 2",fillcolor="brown1"];

llvm/test/Transforms/MemProfContextDisambiguation/funcassigncloning.ll

 ;;   return 0;
 ;; }
 ;;
 ;; Code compiled with -mllvm -memprof-min-lifetime-cold-threshold=5 so that the
 ;; memory freed after sleep(10) results in cold lifetimes.
 ;;
 ;; The IR was then reduced using llvm-reduce with the expected FileCheck input.
 
+;; -stats requires asserts
+; REQUIRES: asserts
+
 ; RUN: opt -passes=memprof-context-disambiguation \
 ; RUN:  -memprof-verify-ccg -memprof-verify-nodes -memprof-dump-ccg \
-; RUN:  %s -S 2>&1 | FileCheck %s --check-prefix=DUMP
+; RUN:  -stats -pass-remarks=memprof-context-disambiguation \
+; RUN:  %s -S 2>&1 | FileCheck %s --check-prefix=DUMP --check-prefix=IR \
+; RUN:  --check-prefix=STATS --check-prefix=REMARKS
 
 
 target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
 target triple = "x86_64-unknown-linux-gnu"
 
 define internal void @_Z1EPPcS0_(ptr %buf1, ptr %buf2) #0 {
 entry:
   %call = call noalias noundef nonnull ptr @_Znam(i64 noundef 10) #6, !memprof !0, !callsite !7

 ; DUMP: Node [[ENEW2CLONE]]
 ; DUMP:     %call1 = call noalias noundef nonnull ptr @_Znam(i64 noundef 10) #6 (clone 0)
 ; DUMP:   AllocTypes: Cold
 ; DUMP:   ContextIds: 5
 ; DUMP:   CalleeEdges:
 ; DUMP:   CallerEdges:
 ; DUMP:     Edge from Callee [[ENEW2CLONE]] to Caller: [[C]] AllocTypes: Cold ContextIds: 5
 ; DUMP:   Clone of [[ENEW2ORIG]]
+
+
+;; We greedily create a clone of E that is initially used by the clones of the
+;; first call to new. However, we end up with an incompatible set of callers
+;; given the second call to new which has clones with a different combination of
+;; callers. Eventually, we create 2 more clones, and the first clone becomes dead.
+; REMARKS: created clone _Z1EPPcS0_.memprof.1
+; REMARKS: created clone _Z1EPPcS0_.memprof.2
+; REMARKS: created clone _Z1EPPcS0_.memprof.3
+; REMARKS: call in clone _Z1DPPcS0_ assigned to call function clone _Z1EPPcS0_.memprof.2
+; REMARKS: call in clone _Z1EPPcS0_.memprof.2 marked with memprof allocation attribute cold
+; REMARKS: call in clone _Z1CPPcS0_ assigned to call function clone _Z1EPPcS0_.memprof.3
+; REMARKS: call in clone _Z1EPPcS0_.memprof.3 marked with memprof allocation attribute notcold
+; REMARKS: call in clone _Z1BPPcS0_ assigned to call function clone _Z1EPPcS0_
+; REMARKS: call in clone _Z1EPPcS0_ marked with memprof allocation attribute notcold
+; REMARKS: call in clone _Z1EPPcS0_.memprof.2 marked with memprof allocation attribute notcold
+; REMARKS: call in clone _Z1EPPcS0_.memprof.3 marked with memprof allocation attribute cold
+; REMARKS: call in clone _Z1EPPcS0_ marked with memprof allocation attribute notcold
+
+
+;; Original version of E is used for the non-cold allocations, both from B.
+; IR: define internal {{.*}} @_Z1EPPcS0_(
+; IR:   call {{.*}} @_Znam(i64 noundef 10) #[[NOTCOLD:[0-9]+]]
+; IR:   call {{.*}} @_Znam(i64 noundef 10) #[[NOTCOLD]]
+; IR: define internal {{.*}} @_Z1BPPcS0_(
+; IR:   call {{.*}} @_Z1EPPcS0_(
+;; C calls a clone of E with the first new allocating cold memory and the
+;; second allocating non-cold memory.
+; IR: define internal {{.*}} @_Z1CPPcS0_(
+; IR:   call {{.*}} @_Z1EPPcS0_.memprof.3(
+;; D calls a clone of E with the first new allocating non-cold memory and the
+;; second allocating cold memory.
+; IR: define internal {{.*}} @_Z1DPPcS0_(
+; IR:   call {{.*}} @_Z1EPPcS0_.memprof.2(
+;; Transient clone that will get removed as it ends up with no callers.
+;; Its calls to new never get updated with a memprof attribute as a result.
+; IR: define internal {{.*}} @_Z1EPPcS0_.memprof.1(
+; IR:   call {{.*}} @_Znam(i64 noundef 10) #[[DEFAULT:[0-9]+]]
+; IR:   call {{.*}} @_Znam(i64 noundef 10) #[[DEFAULT]]
+; IR: define internal {{.*}} @_Z1EPPcS0_.memprof.2(
+; IR:   call {{.*}} @_Znam(i64 noundef 10) #[[COLD:[0-9]+]]
+; IR:   call {{.*}} @_Znam(i64 noundef 10) #[[NOTCOLD]]
+; IR: define internal {{.*}} @_Z1EPPcS0_.memprof.3(
+; IR:   call {{.*}} @_Znam(i64 noundef 10) #[[NOTCOLD]]
+; IR:   call {{.*}} @_Znam(i64 noundef 10) #[[COLD]]
+; IR: attributes #[[NOTCOLD]] = { builtin "memprof"="notcold" }
+; IR: attributes #[[DEFAULT]] = { builtin }
+; IR: attributes #[[COLD]] = { builtin "memprof"="cold" }
+
+
+; STATS: 2 memprof-context-disambiguation - Number of cold static allocations (possibly cloned)
+; STATS: 4 memprof-context-disambiguation - Number of not cold static allocations (possibly cloned)
+; STATS: 3 memprof-context-disambiguation - Number of function clones created during whole program analysis

llvm/test/Transforms/MemProfContextDisambiguation/indirectcall.ll

 ;; Tests callsite context graph generation for call graph containing indirect
 ;; calls. Currently this should result in conservative behavior, such that the
 ;; indirect call receives a null call in its graph node, to prevent subsequent
-;; cloning.
+;; cloning. Also tests graph and IR cloning.
 ;;
 ;; Original code looks like:
 ;;
 ;; char *foo() {
 ;;   return new char[10];
 ;; }
 ;; class A {
 ;; public:

 ;; The IR was then reduced using llvm-reduce with the expected FileCheck input.
 
 ;; -stats requires asserts
 ; REQUIRES: asserts
 
 ; RUN: opt -passes=memprof-context-disambiguation \
 ; RUN:  -memprof-verify-ccg -memprof-verify-nodes -memprof-dump-ccg \
 ; RUN:  -memprof-export-to-dot -memprof-dot-file-path-prefix=%t. \
-; RUN:  %s -S 2>&1 | FileCheck %s --check-prefix=DUMP
+; RUN:  -stats -pass-remarks=memprof-context-disambiguation \
+; RUN:  %s -S 2>&1 | FileCheck %s --check-prefix=DUMP --check-prefix=IR \
+; RUN:  --check-prefix=STATS --check-prefix=REMARKS
 
 ; RUN:  cat %t.ccg.postbuild.dot | FileCheck %s --check-prefix=DOT
 ;; We should only create a single clone of foo, for the direct call
 ;; from main allocating cold memory.
 ; RUN:  cat %t.ccg.cloned.dot | FileCheck %s --check-prefix=DOTCLONED
 
 
 target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"

 ; DUMP:   AllocTypes: Cold
 ; DUMP:   ContextIds: 6
 ; DUMP:   CalleeEdges:
 ; DUMP:   CallerEdges:
 ; DUMP:     Edge from Callee [[FOO2]] to Caller: [[MAIN2]] AllocTypes: Cold ContextIds: 6
 ; DUMP:   Clone of [[FOO]]
 
 
+; REMARKS: created clone _Z3foov.memprof.1
+; REMARKS: call in clone main assigned to call function clone _Z3foov.memprof.1
+; REMARKS: call in clone _Z3foov.memprof.1 marked with memprof allocation attribute cold
+; REMARKS: call in clone _ZN1A1xEv assigned to call function clone _Z3foov
+; REMARKS: call in clone _ZN1B1xEv assigned to call function clone _Z3foov
+; REMARKS: call in clone main assigned to call function clone _Z3foov
+; REMARKS: call in clone _Z3foov marked with memprof allocation attribute notcold
+
+
+; IR: define {{.*}} @main(
+; IR:   call {{.*}} @_Z3foov()
+;; Only the second call to foo, which allocates cold memory via direct calls,
+;; is replaced with a call to a clone that calls a cold allocation.
+; IR:   call {{.*}} @_Z3foov.memprof.1()
+; IR:   call {{.*}} @_Z3barP1A(
+; IR:   call {{.*}} @_Z3barP1A(
+; IR:   call {{.*}} @_Z3barP1A(
+; IR:   call {{.*}} @_Z3barP1A(
+; IR: define internal {{.*}} @_ZN1A1xEv(
+; IR:   call {{.*}} @_Z3foov()
+; IR: define internal {{.*}} @_ZN1B1xEv(
+; IR:   call {{.*}} @_Z3foov()
+; IR: define internal {{.*}} @_Z3foov()
+; IR:   call {{.*}} @_Znam(i64 noundef 10) #[[NOTCOLD:[0-9]+]]
+; IR: define internal {{.*}} @_Z3foov.memprof.1()
+; IR:   call {{.*}} @_Znam(i64 noundef 10) #[[COLD:[0-9]+]]
+; IR: attributes #[[NOTCOLD]] = { builtin "memprof"="notcold" }
+; IR: attributes #[[COLD]] = { builtin "memprof"="cold" }
+
+
+; STATS: 1 memprof-context-disambiguation - Number of cold static allocations (possibly cloned)
+; STATS: 1 memprof-context-disambiguation - Number of not cold static allocations (possibly cloned)
+; STATS: 1 memprof-context-disambiguation - Number of function clones created during whole program analysis
+
+
 ; DOT: digraph "postbuild" {
 ; DOT:  label="postbuild";
 ; DOT:   Node[[FOO:0x[a-z0-9]+]] [shape=record,tooltip="N[[FOO]] ContextIds: 1 2 3 4 5 6",fillcolor="mediumorchid1",style="filled",style="filled",label="{OrigId: Alloc0\n_Z3foov -\> _Znam}"];
 ; DOT:   Node[[AX:0x[a-z0-9]+]] [shape=record,tooltip="N[[AX]] ContextIds: 1 2",fillcolor="mediumorchid1",style="filled",style="filled",label="{OrigId: 8256774051149711748\n_ZN1A1xEv -\> _Z3foov}"];
 ; DOT:  Node[[AX]] -> Node[[FOO]][tooltip="ContextIds: 1 2",fillcolor="mediumorchid1"];
 ; DOT:   Node[[BAR:0x[a-z0-9]+]] [shape=record,tooltip="N[[BAR]] ContextIds: 1 2 4 5",fillcolor="mediumorchid1",style="filled",style="filled",label="{OrigId: 13626499562959447861\nnull call (external)}"];
 ; DOT:  Node[[BAR]] -> Node[[AX]][tooltip="ContextIds: 1 2",fillcolor="mediumorchid1"];
 ; DOT:  Node[[BAR]] -> Node[[BX:0x[a-z0-9]+]][tooltip="ContextIds: 4 5",fillcolor="mediumorchid1"];

llvm/test/Transforms/MemProfContextDisambiguation/inlined.ll

 ;; Test callsite context graph generation for call graph with two memprof
 ;; contexts and partial inlining, requiring generation of a new fused node to
 ;; represent the inlined sequence while matching callsite nodes onto the graph.
+;; Also tests graph and IR cloning.
 ;;
 ;; Original code looks like:
 ;;
 ;; char *bar() {
 ;;   return new char[10];
 ;; }
 ;;
 ;; char *baz() {

 ;; The IR was then reduced using llvm-reduce with the expected FileCheck input.
 
 ;; -stats requires asserts
 ; REQUIRES: asserts
 
 ; RUN: opt -passes=memprof-context-disambiguation \
 ; RUN:  -memprof-verify-ccg -memprof-verify-nodes -memprof-dump-ccg \
 ; RUN:  -memprof-export-to-dot -memprof-dot-file-path-prefix=%t. \
-; RUN:  %s -S 2>&1 | FileCheck %s --check-prefix=DUMP
+; RUN:  -stats -pass-remarks=memprof-context-disambiguation \
+; RUN:  %s -S 2>&1 | FileCheck %s --check-prefix=DUMP --check-prefix=IR \
+; RUN:  --check-prefix=STATS --check-prefix=REMARKS
 
 ; RUN:  cat %t.ccg.postbuild.dot | FileCheck %s --check-prefix=DOT
 ;; We should create clones for foo and bar for the call from main to allocate
 ;; cold memory.
 ; RUN:  cat %t.ccg.cloned.dot | FileCheck %s --check-prefix=DOTCLONED
 
 
 target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"

 ; DUMP:   AllocTypes: Cold
 ; DUMP:   ContextIds: 2
 ; DUMP:   CalleeEdges:
 ; DUMP:   CallerEdges:
 ; DUMP:     Edge from Callee [[BAR2]] to Caller: [[FOO3]] AllocTypes: Cold ContextIds: 2
 ; DUMP:         Clone of [[BAR]]
 
 
+; REMARKS: created clone _Z3barv.memprof.1
+; REMARKS: created clone _Z3foov.memprof.1
+; REMARKS: call in clone main assigned to call function clone _Z3foov.memprof.1
+; REMARKS: call in clone _Z3foov.memprof.1 assigned to call function clone _Z3barv.memprof.1
+; REMARKS: call in clone _Z3barv.memprof.1 marked with memprof allocation attribute cold
+; REMARKS: call in clone main assigned to call function clone _Z3foov
+; REMARKS: call in clone _Z3foov assigned to call function clone _Z3barv
+; REMARKS: call in clone _Z3barv marked with memprof allocation attribute notcold
+; REMARKS: call in clone _Z3bazv marked with memprof allocation attribute notcold
+
+
+; IR: define internal {{.*}} @_Z3barv()
+; IR:   call {{.*}} @_Znam(i64 noundef 10) #[[NOTCOLD:[0-9]+]]
+; IR: define internal {{.*}} @_Z3foov()
+; IR:   call {{.*}} @_Z3barv()
+; IR: define {{.*}} @main()
+;; The first call to foo does not allocate cold memory. It should call the
+;; original functions, which ultimately call the original allocation decorated
+;; with a "notcold" attribute.
+; IR:   call {{.*}} @_Z3foov()
+;; The second call to foo allocates cold memory. It should call cloned functions
+;; which ultimately call a cloned allocation decorated with a "cold" attribute.
+; IR:   call {{.*}} @_Z3foov.memprof.1()
+; IR: define internal {{.*}} @_Z3barv.memprof.1()
+; IR:   call {{.*}} @_Znam(i64 noundef 10) #[[COLD:[0-9]+]]
+; IR: define internal {{.*}} @_Z3foov.memprof.1()
+; IR:   call {{.*}} @_Z3barv.memprof.1()
+; IR: attributes #[[NOTCOLD]] = { builtin "memprof"="notcold" }
+; IR: attributes #[[COLD]] = { builtin "memprof"="cold" }
+
+
+; STATS: 1 memprof-context-disambiguation - Number of cold static allocations (possibly cloned)
+; STATS: 2 memprof-context-disambiguation - Number of not cold static allocations (possibly cloned)
+; STATS: 2 memprof-context-disambiguation - Number of function clones created during whole program analysis
+
+
 ; DOT: digraph "postbuild" {
 ; DOT:  label="postbuild";
 ; DOT:   Node[[BAR:0x[a-z0-9]+]] [shape=record,tooltip="N[[BAR]] ContextIds: 1 2",fillcolor="mediumorchid1",style="filled",style="filled",label="{OrigId: Alloc0\n_Z3barv -\> _Znam}"];
 ; DOT:   Node[[FOO:0x[a-z0-9]+]] [shape=record,tooltip="N[[FOO]] ContextIds: 3 4",fillcolor="mediumorchid1",style="filled",style="filled",label="{OrigId: 2732490490862098848\nnull call (external)}"];
 ; DOT:  Node[[FOO]] -> Node[[BAZ:0x[a-z0-9]+]][tooltip="ContextIds: 3 4",fillcolor="mediumorchid1"];
 ; DOT:   Node[[MAIN1:0x[a-z0-9]+]] [shape=record,tooltip="N[[MAIN1]] ContextIds: 1 3",fillcolor="brown1",style="filled",style="filled",label="{OrigId: 8632435727821051414\nmain -\> _Z3foov}"];
 ; DOT:  Node[[MAIN1]] -> Node[[FOO]][tooltip="ContextIds: 3",fillcolor="brown1"];
 ; DOT:  Node[[MAIN1]] -> Node[[FOO2:0x[a-z0-9]+]][tooltip="ContextIds: 1",fillcolor="brown1"];