[CSSPGO][llvm-profgen] Reimplement SampleContextTracker using context trie

[wlei]

Summary

This is the followup patch to https://reviews.llvm.org/D125246 for the SampleContextTracker part. Before the promotion and merging of the context is based on the SampleContext(the array of frame), this causes a lot of cost to the memory. This patch detaches the tracker from using the array ref instead to use the context trie itself. This can save a lot of memory usage and benefit both the compiler's CS inliner and llvm-profgen's pre-inliner.

One structure needs to be specially treated is the FuncToCtxtProfiles, this is used to get all the functionSamples for one function to do the merging and promoting. Before it search each functions' context and traverse the trie to get the node of the context. Now we don't have the context inside the profile, instead we directly use an auxiliary map ProfileToNodeMap for profile , it initialize to create the FunctionSamples to TrieNode relations and keep updating it during promoting and merging the node.

Moreover, I was expecting the results before and after remain the same, but I found that the order of FuncToCtxtProfiles matter and affect the results. This can happen on recursive context case, but the difference should be small. Now we don't have the context, so I just used a vector for the order, the result is still deterministic.

Measured on one huge size(12GB) profile from one of our internal service. The profile similarity difference is 99.999%, and the running time is improved by 3X(debug mode) and the memory is reduced from 170GB to 90GB.

[wlei]

Updating D127031: [CSSPGO][llvm-profgen] Reimplement SampleContextTracker using context trie

[hoy]

Would be good to have some numbers for build time and memory usage in the summary.

llvm/include/llvm/ProfileData/SampleProf.h

688

The field is only useful during context manipulation. Is it possible to store the link in an auxiliary map whereever it is used?

llvm/lib/Transforms/IPO/SampleContextTracker.cpp

93

may still be helpful to keep the debug dump?

209

Is this also done in populateFuncToCtxtMap?

433

nit: use SampleContextFrameVector

[wenlei]

llvm/include/llvm/Transforms/IPO/SampleContextTracker.h

106

Is std::set unnecessary or is it wrong?

llvm/lib/Transforms/IPO/SampleContextTracker.cpp

353

Why we need to check MergedContext on OldContext, but InlinedContext on ContextOnNode?

435

return string();

[wlei]

llvm/include/llvm/ProfileData/SampleProf.h

688

I think we use context for two things 1)context manipulation(inlining) 2) reader and writer based on ProfileMap.
My thought is ContextNode is the same thing to FullContext, i,e. we can always start from this node to iterate trie reversely to get the array of frames. I was wondering if we can use ContextNode to replace the FullContext in the future. And for sample reader and writer things, the key of the map can work as the FullContext, we might not need this field either. That's why I'm thinking to keep it here. WDYT?

llvm/include/llvm/Transforms/IPO/SampleContextTracker.h

106

Before it sorted the elements by A->getContext() < B->getContext();, but now getContext is an empty context, so it ends up with non-deterministic issue and losing elements. While we can recover the context by waling through the trie, but thinking using SmallVector is enough here, the order of traversing trie is deterministic(getAllChildContext() is deterministic).

llvm/lib/Transforms/IPO/SampleContextTracker.cpp

93

moveToChildContext is the function of ContextTrieNode, the getContextString is in SampleContextTracker. Then it requires to pass SampleContextTracker to the moveToChildContext function.

I just realize before it dump the promotion result for each node, but now we don't have the concept of promotion. Maybe dumping the the head node(ToNode) is enough?

209

Good catch!

353

Let me demo by an example.

Given two node with same name foo, say N1 and N2. N1 own a FunctionSample S1, and N2 own S2, i, e. S1->N1, S2->N2

Supposing during merge bar, foo is also merge, say N2 --> N1, then S2's context is marked MergedContext and S2's is update to point to N1. S1 still point to N1.

Then when call getBaseSample for foo to merge all nodes, S2 and S1 both point to N1 now, we can just skip S2 which contains the OldContext marked MergedContext .

433

Done, thanks!

435

Done, thanks!

[wlei]

Addressing reviewers' feedback. Also added the build time and memory usage number in the summary.

[hoy]

llvm/include/llvm/ProfileData/SampleProf.h

688

yeah, I had the same thought that the two fields are kinda redundant to each other. I'm a bit worried about the compile-time memory increase. Would it be possible to use a union for the two fields for now?

2. reader and writer based on ProfileMap.

What is this usage?

I was wondering if we can use ContextNode to replace the FullContext in the future.

It depends. For example, the profile merger probably does not want to always build a tri and merge multiple tries. The profile similarity checker may neither.

llvm/lib/Transforms/IPO/SampleContextTracker.cpp

93

Yeah, I don't have a strong opinion, as long as it makes your debug experience better :)

[wlei]

Make a union type for fullcontext(array) and context trie node.

[wlei]

llvm/include/llvm/ProfileData/SampleProf.h

688

Would it be possible to use a union for the two fields for now?

Yeah, that's good to save memory. I had a try this, see the current version, it can pass all the test. However, I feel like it might be error-prone. The two shared fields have different way of initialization, currently I need to make sure SampleContextFrames is initialized then the ContextTrieNode field will be unknown status(not null). We might need to carefully use it in the future.

reader and writer based on ProfileMap.
What is this usage?

For CS profile, SampleContext is the key of ProfileMap, FullContext is always not empty here.
https://github.com/llvm/llvm-project/blob/main/llvm/lib/ProfileData/SampleProfReader.cpp#L279

It depends. For example, the profile merger probably does not want to always build a tri and merge multiple tries. The profile similarity checker may neither.

Yeah, right now ProfileMap always require the FullContext as the key. wondering if we can only use FullContext as the key, FunctionSample's SampleContext doesn't own the fullContext.

[hoy]

llvm/include/llvm/ProfileData/SampleProf.h

688

Using union is error-prone, yes. My original thought is to use a separate map in ContextTracker that maintains the link from FunctionSamples to its ContextTrieNode. So far I don't see a use of getting ContextTrieNode from FunctionSamples out of the ContextTracker context. It seems encapsulating the map inside ContextTracker is enough. Does this make sense to you?

llvm/tools/llvm-profgen/ProfileGenerator.cpp

996

What are the allocations to be undone? I thought in D125246 we use std:move to transmit the FunctionSamples on the tri to profileMap.

[wlei]

Updating D127031: [CSSPGO][llvm-profgen] Reimplement SampleContextTracker using context trie

llvm/include/llvm/ProfileData/SampleProf.h

688

I see, now I understand why to keep in a separate map, it's less intrusive to the FunctionSamples class. Thanks for your clarification.

llvm/tools/llvm-profgen/ProfileGenerator.cpp

996

I was thinking the std::move 's source object should be freed, it's allocated in CSProfileGenerator::getOrCreateFunctionSamples and it's allocated by new, so we need explicitly delete, though the memory usage should be small.

[hoy]

A general question about the design as I read more into the code. I was trying to figure out if a FunctionSample, its SampleContext and ContextTrieNode stick together as a whole, as also commented in another place. It seems to me multiple FunctionSamples can share the ContextTrieNode. Am I correct? What is the reason for this?

llvm/include/llvm/ProfileData/SampleProf.h

688

Yeah, it looks better now. Thanks for addressing this.

llvm/include/llvm/Transforms/IPO/SampleContextTracker.h

214

type: are -> is

llvm/lib/Transforms/IPO/SampleContextTracker.cpp

345

Should we also check InlinedContext here?

350

I'm a bit confused here. Is ContextOnNode same with OldContext? Should a FunctionSample, its SampleContext and ContextTrieNode stick together as a whole? It may be worth a comment here.

[wlei]

llvm/lib/Transforms/IPO/SampleContextTracker.cpp

345

Answering in the comment below.

350

Supposing Node(NA) owned a FunctionSamples(SA) and NB owns SB, NA is merged into NB.

SA's context is marked MergedContext and Now SA's node is pointed to NB(NA is removed)

so here OldContext is the context of NA, ContextOnNode is the context of NB(might be marked InlinedContext in the future).

will comment using this example.

Should a FunctionSample, its SampleContext and ContextTrieNode stick together as a whole

It seems they can't be sticked together if we want to keep the old status(MergedContext) and the node pointer always point to the new on trie node.

[hoy]

llvm/lib/Transforms/IPO/SampleContextTracker.cpp

350

SA's context is marked MergedContext and Now SA's node is pointed to NB(NA is removed)

Why SA's should point to NB? Can NA be kept so that SA always point to it?

When a context is merged to another context, the context should never be used anywhere. What's the benefit of making it point to the other tri node?

[wlei]

llvm/lib/Transforms/IPO/SampleContextTracker.cpp

350

We need to update FSamples to always point to the new ContextTrieNode, that's why we use updateContextNode. This is because we iterate the FuncToCtxtProfiles which is a list of FSamples to be merged to get base sample. Since the the node can be merged dynamically, take the previous example. the NB after merging is the sum of (SA + SB), NA is the old SA, it should be removed. but when to get FA's Node, if we get the old SA, it will be merged to NA again becoming (2*SA +SB). so we need to always make point to the node on the trie. And now both SA and SB are pointing to NB, then here use the oldcontext(MergedContext ) to filter out one.

[hoy]

llvm/lib/Transforms/IPO/SampleContextTracker.cpp

350

but when to get FA's Node, if we get the old SA, it will be merged to NA again becoming (2*SA +SB).

At this time. SA's context should be marked "MergedContext". Can we skip it based on the flag?

[wlei]

llvm/lib/Transforms/IPO/SampleContextTracker.cpp

350

Yeah, that's just done here, that's to skip the FSample(SA) whose old context(NA) is marked MergedContext .

[hoy]

llvm/lib/Transforms/IPO/SampleContextTracker.cpp

350

So do we need to follow SA's link to its context tri node? Feel we wouldn't never need to access SA anymore since all of its data are already merged into another node.

[hoy]

llvm/include/llvm/Transforms/IPO/SampleContextTracker.h

223

Actually my previous question could be made more clear as why we need to maintain this map.

[wlei]

llvm/include/llvm/Transforms/IPO/SampleContextTracker.h

223

I see, so you're suggested that
1 FSamples2NodeMap make sure FSamples always pointed to its belonging context node.
2 even the FSamples whose context is mergedcontext point to a stale trie node, we can filter out by the

SampleContext &OldContext = CSamples->getContext();
// Check to avoid re-merge any context.
if (OldContext.hasState(MergedContext))

So we won't have a redundant node to be merged. We also don't need to update FSample to point to the merged node.

That sounds a clear way!

[hoy]

llvm/include/llvm/Transforms/IPO/SampleContextTracker.h

223

Exactly. Hopefully the logic can be simplified a little bit.

[wlei]

addressing feedback from Hongtao

[hoy]

llvm/lib/Transforms/IPO/SampleContextTracker.cpp

350

ContextOnNode and OldContext should be the same thing now?

573

This is no longer needed?

llvm/tools/llvm-profgen/CSPreInliner.cpp

303

Should we just keep the original condition checks?

Also, where do we actually deallocate the function profile?

llvm/tools/llvm-profgen/ProfiledBinary.h

164

nit: use const for the parameter

[wlei]

Updating D127031: [CSSPGO][llvm-profgen] Reimplement SampleContextTracker using context trie

llvm/lib/Transforms/IPO/SampleContextTracker.cpp

350

Yeah, new the context node is always point to the FunctionSamples, and vice versa. Recovered the original code.

573

This was to be consistent with the branch below(line:608), I just realized that actually the FromNode will be removed from the trie, they both can be removed, removing them.

llvm/tools/llvm-profgen/CSPreInliner.cpp

303

Should we just keep the original condition checks?

do you mean the isBaseContext?, it depends on the FullContext, now we don't have it, so I changed to Node->getParentContext() != &ContextTracker.getRootContext()

Also, where do we actually deallocate the function profile?

It's a bit tricky, the FunctionSample can be from two sources

1. the llvm-sample-profile, the FunctionSamples's lifetime will be along with the reader. we can't explicitly call delete on this.

2. ProfileGenerator( by the new), we need to explicitly call delete, I haven't came up with a way to deallocate this yet. The Node is removed from trie dynamically, I was thinking to use FuncToCtxtProfiles to get all FunctionSamples's pointer but this doesn't include the base context.
   
   (I leave a TODO for this)

llvm/tools/llvm-profgen/ProfiledBinary.h

164

Done!

[hoy]

llvm/tools/llvm-profgen/CSPreInliner.cpp

303

it depends on the FullContext, now we don't have it, so I changed to Node->getParentContext() != &ContextTracker.getRootContext()

I see. That makes sense.

It's a bit tricky, the FunctionSample can be from two sources

It is tricky.

Do we have any memory leakage by resetting the pointer to null but not freeing the data? If so, we may fail the sanitizer tests.

[wlei]

llvm/tools/llvm-profgen/CSPreInliner.cpp

303

Okay.. The FuncToCtxtProfiles also hold the pointers but eventually after ContextTracker is freed, it's still memory leakage.. let me try to free them.

[wlei]

free memory allocated by llvm-profgen

[hoy]

llvm/include/llvm/Transforms/IPO/SampleContextTracker.h

146

PreInliner is a llvm-profgen concept, thus it should not be exposed here.

Also the deallocation should be done by the original allocator, i.e, llvm-profgen.

llvm/tools/llvm-profgen/ProfileGenerator.cpp

1006

I'm wondering if these two deallocations can be combined. Maybe in buildProfileMap we could exclude merged contexts there and then we don't need to do Node->setFunctionSamples(nullptr); inside CSPreInliner::run?

[wlei]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

1006

The merged contexts node is removed from the trie after merging, so we can't traverse the whole trie to delete all the FSamples. That's why I was thinking to leverage the FuncToCtxtProfiles which stores all the FSample pointers at the beginning. And if it's not the preinliner pass, FuncToCtxtProfiles is empty, then we need to traverse the trie to delete them.

[wlei]

move memory free code inside llvm-profgen

[hoy]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

1006

The merged contexts node is removed from the trie after merging, so we can't traverse the whole trie to delete all the FSamples.

I see. I guess we have to keep the diversion for now. A shared_ptr or a centralized FunctionSample storage such as what the reader maintains would typically help. That would require a change in CSProfileGenerator::getOrCreateFunctionSamples . Please leave a todo somewhere to be addressed laster.

[wlei]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

1006

Good idea, I feel that's something feasible in this patch. Based on your idea to change getOrCreateFunctionSamples , I'm wondering if we can just use a vector to store all the FunctionSamples, the trie will naturally make sure the node is unique, so we don't need a Map which is our concern before.

[hoy]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

1006

Yeah, a std::list would work. std::vector can be volatile when growing.

[wlei]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

1006

Yeah, thanks for the suggestion! I will implement it in https://reviews.llvm.org/D125246, that's the first diff
introduce the memory leak.

[wlei]

rebase, memory leak issue is addressed in the https://reviews.llvm.org/D125246

[hoy]

Thanks for addressing all the comments! Looks good to me now except for a couple minor things.

llvm/include/llvm/Transforms/IPO/SampleContextTracker.h

91

Wondering if 16 can be larger than needed in general. Since we are aiming for mem savings, maybe a smaller value is better? Or just use std::vector. Shouldn't be a big deal though.

146

Make this debug-only?

[wlei]

Updating D127031: [CSSPGO][llvm-profgen] Reimplement SampleContextTracker using context trie

llvm/include/llvm/Transforms/IPO/SampleContextTracker.h

91

Sounds good, used std::vector

146

Done, thanks!

[hoy]

LGTM.

[wlei]

rebase

[wenlei]

llvm/include/llvm/ProfileData/SampleProf.h

688

I also like this being an auxiliary data rather than extending with a field. Thanks for working on this

llvm/include/llvm/Transforms/IPO/SampleContextTracker.h

91–92

nit: this comment is more like an explanation of the change (why we no longer need container to order elements), based on previous comment as context. but reader will never have previous comment/code as context.

Saying this because when I read this comment, I actually got confused, and I only understand it after reading the original code..

215

nit: to be consistent with FuncToCtxtProfiles, name this ProfileToNodeMap

llvm/lib/Transforms/IPO/SampleContextTracker.cpp

69

DeleteNode is always false and can be simplified?

87

any reason to make this a member of SampleContextTracker rather than member of ContextTrieNode? asking because the naming moveToChildContext is assuming this is the parent to move to.

217

I'm seeing many repetition of tree BFS and the boiler plate code. Perhaps it makes sense to create iterator API for SampleContextTracker to go over all nodes.

llvm/tools/llvm-profgen/CSPreInliner.cpp

65

populateFuncToCtxtMap isn't really the responsibility of preinliner. can we move this earlier to right after profile generation on trie?

preinliner should be one client of context tracker and it can assume complete, fully functioning context tracker, including auxiliary maps.

[wenlei]

llvm/lib/Transforms/IPO/SampleContextTracker.cpp

87

Ok, I guess it's because setContextNode needs to update the map for context tracker. In that case, we probably want to make context tracker own the move completely, because ContextTrieNode::moveToChildContext now isn't complete without updating the map. What I'm suggesting is to move ContextTrieNode::moveToChildContext into this function completely.

Also rename SampleContextTracker::moveToChildContext to something like ContextTrieNode::moveContextSamples since this is no longer the parent to move to.

[wlei]

Addressing Wenlei's feeback: add an iterator API for SampleContextTracker and other refactors.

[wenlei]

Thanks for the change, and the iterator cleans up quite a bit of duplication. There's one more case left, see inline comment.

llvm/lib/Transforms/IPO/SampleContextTracker.cpp

86

We can use the iterator for this as well?

[wlei]

llvm/lib/Transforms/IPO/SampleContextTracker.cpp

86

We can use iterator for this. But we need to call setParentContext to set Parent for all childNode, there is still the below for loop inside, it's not like others very straightforward good for the readability. WDYT?

for(.....) {
   ....
  for (auto &It : Node->getAllChildContext()) {
    ContextTrieNode *ChildNode = &It.second;
    ChildNode->setParentContext(Node);
 }

[wenlei]

lgtm, thanks.

llvm/lib/Transforms/IPO/SampleContextTracker.cpp

86

Ok, then it works either way. Keep current version is also good because this loop involves modification.

llvm/include/llvm/ProfileData/SampleProf.h

while (!ContextRemain.empty()) {
       ChildContext = ContextSplit.first;
       ContextRemain = ContextSplit.second;
       LineLocation CallSiteLoc(0, 0);
       decodeContextString(ChildContext, CalleeName, CallSiteLoc);
       Context.emplace_back(CalleeName, CallSiteLoc);
     }
   }
 
-  // Promote context by removing top frames with the length of
-  // `ContextFramesToRemove`. Note that with array representation of context,
-  // the promotion is effectively a slice operation with first
-  // `ContextFramesToRemove` elements removed from left.
-  void promoteOnPath(uint32_t ContextFramesToRemove) {
-    assert(ContextFramesToRemove <= FullContext.size() &&
-           "Cannot remove more than the whole context");
-    FullContext = FullContext.drop_front(ContextFramesToRemove);
-  }
-
   // Decode context string for a frame to get function name and location.
   // `ContextStr` is in the form of `FuncName:StartLine.Discriminator`.
   static void decodeContextString(StringRef ContextStr, StringRef &FName,
                                   LineLocation &LineLoc) {
     // Get function name
     auto EntrySplit = ContextStr.split(':');
     FName = EntrySplit.first;
 

 private:
   /// Mangled name of the function.
   StringRef Name;
   // Full context including calling context and leaf function name
   SampleContextFrames FullContext;
   // State of the associated sample profile
   uint32_t State;
   // Attribute of the associated sample profile
   uint32_t Attributes;
 };
 
 static inline hash_code hash_value(const SampleContext &arg) {
   return arg.hasContext() ? hash_value(arg.getContextFrames())
                           : hash_value(arg.getName());
 }
 
 class FunctionSamples;

llvm/include/llvm/Transforms/IPO/SampleContextTracker.h

 
 #ifndef LLVM_TRANSFORMS_IPO_SAMPLECONTEXTTRACKER_H
 #define LLVM_TRANSFORMS_IPO_SAMPLECONTEXTTRACKER_H
 
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ProfileData/SampleProf.h"
 #include <map>
+#include <queue>
 #include <vector>
 
 namespace llvm {
 class CallBase;
 class DILocation;
 class Function;
 class Instruction;
 
ContextTrieNode(ContextTrieNode *Parent = nullptr,
       : ParentContext(Parent), FuncName(FName), FuncSamples(FSamples),
         CallSiteLoc(CallLoc){};
   ContextTrieNode *getChildContext(const LineLocation &CallSite,
                                    StringRef ChildName);
   ContextTrieNode *getHottestChildContext(const LineLocation &CallSite);
   ContextTrieNode *getOrCreateChildContext(const LineLocation &CallSite,
                                            StringRef ChildName,
                                            bool AllowCreate = true);
-
-  ContextTrieNode &moveToChildContext(const LineLocation &CallSite,
-                                      ContextTrieNode &&NodeToMove,
-                                      uint32_t ContextFramesToRemove,
-                                      bool DeleteNode = true);
   void removeChildContext(const LineLocation &CallSite, StringRef ChildName);
   std::map<uint64_t, ContextTrieNode> &getAllChildContext();
   StringRef getFuncName() const;
   FunctionSamples *getFunctionSamples() const;
   void setFunctionSamples(FunctionSamples *FSamples);
   Optional<uint32_t> getFunctionSize() const;
   void addFunctionSize(uint32_t FSize);
   LineLocation getCallSiteLoc() const;
   ContextTrieNode *getParentContext() const;
   void setParentContext(ContextTrieNode *Parent);
+  void setCallSiteLoc(const LineLocation &Loc);
   void dumpNode();
   void dumpTree();
 
 private:
   // Map line+discriminator location to child context
   std::map<uint64_t, ContextTrieNode> AllChildContext;
 
   // Link to parent context node

 // provides interfaces used by sample profile loader to query context profile or
 // base profile for given function or location; it also manages context tree
 // manipulation that is needed to accommodate inline decisions so we have
 // accurate post-inline profile for functions. Internally context profiles
 // are organized in a trie, with each node representing profile for specific
 // calling context and the context is identified by path from root to the node.
 class SampleContextTracker {
 public:
-  struct ProfileComparer {
-    bool operator()(FunctionSamples *A, FunctionSamples *B) const {
-      // Sort function profiles by the number of total samples and their
-      // contexts.
-      if (A->getTotalSamples() == B->getTotalSamples())
-        return A->getContext() < B->getContext();
-      return A->getTotalSamples() > B->getTotalSamples();
-    }
-  };
-
-  // Keep profiles of a function sorted so that they will be processed/promoted
-  // deterministically.
-  using ContextSamplesTy = std::set<FunctionSamples *, ProfileComparer>;
+  using ContextSamplesTy = std::vector<FunctionSamples *>;
 
   SampleContextTracker() = default;
   SampleContextTracker(SampleProfileMap &Profiles,
                        const DenseMap<uint64_t, StringRef> *GUIDToFuncNameMap);
+  // Populate the FuncToCtxtProfiles map after the trie is built.
+  void populateFuncToCtxtMap();
   // Query context profile for a specific callee with given name at a given
   // call-site. The full context is identified by location of call instruction.
   FunctionSamples *getCalleeContextSamplesFor(const CallBase &Inst,
                                               StringRef CalleeName);
   // Get samples for indirect call targets for call site at given location.
   std::vector<const FunctionSamples *>
   getIndirectCalleeContextSamplesFor(const DILocation *DIL);
   // Query context profile for a given location. The full context
public:
   // function's base profile.
   void markContextSamplesInlined(const FunctionSamples *InlinedSamples);
   ContextTrieNode &getRootContext();
   void promoteMergeContextSamplesTree(const Instruction &Inst,
                                       StringRef CalleeName);
 
   // Create a merged conext-less profile map.
   void createContextLessProfileMap(SampleProfileMap &ContextLessProfiles);
+  ContextTrieNode *
+  getContextNodeForProfile(const FunctionSamples *FSamples) const {
+    auto I = ProfileToNodeMap.find(FSamples);
+    if (I == ProfileToNodeMap.end())
+      return nullptr;
+    return I->second;
+  }
+  StringMap<ContextSamplesTy> &getFuncToCtxtProfiles() {
+    return FuncToCtxtProfiles;
+  }
+
+  class Iterator : public std::iterator<std::forward_iterator_tag,
+                                        const ContextTrieNode *> {
+    std::queue<ContextTrieNode *> NodeQueue;
+
+  public:
+    explicit Iterator() = default;
+    explicit Iterator(ContextTrieNode *Node) { NodeQueue.push(Node); }
+    Iterator &operator++() {
+      assert(!NodeQueue.empty() && "Iterator already at the end");
+      ContextTrieNode *Node = NodeQueue.front();
+      NodeQueue.pop();
+      for (auto &It : Node->getAllChildContext())
+        NodeQueue.push(&It.second);
+      return *this;
+    }
+
+    Iterator operator++(int) {
+      assert(!NodeQueue.empty() && "Iterator already at the end");
+      Iterator Ret = *this;
+      ++(*this);
+      return Ret;
+    }
+    bool operator==(const Iterator &Other) const {
+      if (NodeQueue.empty() && Other.NodeQueue.empty())
+        return true;
+      if (NodeQueue.empty() || Other.NodeQueue.empty())
+        return false;
+      return NodeQueue.front() == Other.NodeQueue.front();
+    }
+    bool operator!=(const Iterator &Other) const { return !(*this == Other); }
+    ContextTrieNode *operator*() const {
+      assert(!NodeQueue.empty() && "Invalid access to end iterator");
+      return NodeQueue.front();
+    }
+  };
+
+  Iterator begin() { return Iterator(&RootContext); }
+  Iterator end() { return Iterator(); }
+
+#ifndef NDEBUG
+  // Get a context string from root to current node.
+  std::string getContextString(const FunctionSamples &FSamples) const;
+  std::string getContextString(ContextTrieNode *Node) const;
+#endif
   // Dump the internal context profile trie.
   void dump();
 
 private:
   ContextTrieNode *getContextFor(const DILocation *DIL);
   ContextTrieNode *getCalleeContextFor(const DILocation *DIL,
                                        StringRef CalleeName);
   ContextTrieNode *getTopLevelContextNode(StringRef FName);
   ContextTrieNode &addTopLevelContextNode(StringRef FName);
   ContextTrieNode &promoteMergeContextSamplesTree(ContextTrieNode &NodeToPromo);
-  void mergeContextNode(ContextTrieNode &FromNode, ContextTrieNode &ToNode,
-                        uint32_t ContextFramesToRemove);
+  void mergeContextNode(ContextTrieNode &FromNode, ContextTrieNode &ToNode);
   ContextTrieNode &
   promoteMergeContextSamplesTree(ContextTrieNode &FromNode,
-                                 ContextTrieNode &ToNodeParent,
-                                 uint32_t ContextFramesToRemove);
+                                 ContextTrieNode &ToNodeParent);
+  ContextTrieNode &moveContextSamples(ContextTrieNode &ToNodeParent,
+                                      const LineLocation &CallSite,
+                                      ContextTrieNode &&NodeToMove);
+  void setContextNode(const FunctionSamples *FSample, ContextTrieNode *Node) {
+    ProfileToNodeMap[FSample] = Node;
+  }
   // Map from function name to context profiles (excluding base profile)
   StringMap<ContextSamplesTy> FuncToCtxtProfiles;
 
+  // Map from current FunctionSample to the belonged context trie.
+  std::unordered_map<const FunctionSamples *, ContextTrieNode *>
+      ProfileToNodeMap;
+
   // Map from function guid to real function names. Only used in md5 mode.
   const DenseMap<uint64_t, StringRef> *GUIDToFuncNameMap;
 
   // Root node for context trie tree
   ContextTrieNode RootContext;
 };
 
 } // end namespace llvm
 #endif // LLVM_TRANSFORMS_IPO_SAMPLECONTEXTTRACKER_H

llvm/lib/Transforms/IPO/SampleContextTracker.cpp

if (Samples->getTotalSamples() > MaxCalleeSamples) {
       ChildNodeRet = &ChildNode;
       MaxCalleeSamples = Samples->getTotalSamples();
     }
   }
 
   return ChildNodeRet;
 }
 
-ContextTrieNode &ContextTrieNode::moveToChildContext(
-    const LineLocation &CallSite, ContextTrieNode &&NodeToMove,
-    uint32_t ContextFramesToRemove, bool DeleteNode) {
+ContextTrieNode &
+SampleContextTracker::moveContextSamples(ContextTrieNode &ToNodeParent,
+                                         const LineLocation &CallSite,
+                                         ContextTrieNode &&NodeToMove) {
   uint64_t Hash =
       FunctionSamples::getCallSiteHash(NodeToMove.getFuncName(), CallSite);
+  std::map<uint64_t, ContextTrieNode> &AllChildContext =
+      ToNodeParent.getAllChildContext();
   assert(!AllChildContext.count(Hash) && "Node to remove must exist");
-  LineLocation OldCallSite = NodeToMove.CallSiteLoc;
-  ContextTrieNode &OldParentContext = *NodeToMove.getParentContext();
   AllChildContext[Hash] = NodeToMove;
   ContextTrieNode &NewNode = AllChildContext[Hash];
-  NewNode.CallSiteLoc = CallSite;
+  NewNode.setCallSiteLoc(CallSite);
 
   // Walk through nodes in the moved the subtree, and update
   // FunctionSamples' context as for the context promotion.
   // We also need to set new parant link for all children.
   std::queue<ContextTrieNode *> NodeToUpdate;
-  NewNode.setParentContext(this);
+  NewNode.setParentContext(&ToNodeParent);
   NodeToUpdate.push(&NewNode);
 
   while (!NodeToUpdate.empty()) {
     ContextTrieNode *Node = NodeToUpdate.front();
     NodeToUpdate.pop();
     FunctionSamples *FSamples = Node->getFunctionSamples();
 
     if (FSamples) {
-      FSamples->getContext().promoteOnPath(ContextFramesToRemove);
+      setContextNode(FSamples, Node);
       FSamples->getContext().setState(SyntheticContext);
-      LLVM_DEBUG(dbgs() << "  Context promoted to: "
-                        << FSamples->getContext().toString() << "\n");
     }
 
     for (auto &It : Node->getAllChildContext()) {
       ContextTrieNode *ChildNode = &It.second;
       ChildNode->setParentContext(Node);
       NodeToUpdate.push(ChildNode);
     }
   }
 
-  // Original context no longer needed, destroy if requested.
-  if (DeleteNode)
-    OldParentContext.removeChildContext(OldCallSite, NewNode.getFuncName());
-
   return NewNode;
 }
 
 void ContextTrieNode::removeChildContext(const LineLocation &CallSite,
                                          StringRef CalleeName) {
   uint64_t Hash = FunctionSamples::getCallSiteHash(CalleeName, CallSite);
   // Note this essentially calls dtor and destroys that child context
   AllChildContext.erase(Hash);

 ContextTrieNode *ContextTrieNode::getParentContext() const {
   return ParentContext;
 }
 
 void ContextTrieNode::setParentContext(ContextTrieNode *Parent) {
   ParentContext = Parent;
 }
 
+void ContextTrieNode::setCallSiteLoc(const LineLocation &Loc) {
+  CallSiteLoc = Loc;
+}
+
 void ContextTrieNode::dumpNode() {
   dbgs() << "Node: " << FuncName << "\n"
          << "  Callsite: " << CallSiteLoc << "\n"
          << "  Size: " << FuncSize << "\n"
          << "  Children:\n";
 
   for (auto &It : AllChildContext) {
     dbgs() << "    Node: " << It.second.getFuncName() << "\n";
SampleContextTracker::SampleContextTracker(
     SampleProfileMap &Profiles,
     const DenseMap<uint64_t, StringRef> *GUIDToFuncNameMap)
     : GUIDToFuncNameMap(GUIDToFuncNameMap) {
   for (auto &FuncSample : Profiles) {
     FunctionSamples *FSamples = &FuncSample.second;
     SampleContext Context = FuncSample.first;
     LLVM_DEBUG(dbgs() << "Tracking Context for function: " << Context.toString()
                       << "\n");
-    if (!Context.isBaseContext())
-      FuncToCtxtProfiles[Context.getName()].insert(FSamples);
     ContextTrieNode *NewNode = getOrCreateContextPath(Context, true);
     assert(!NewNode->getFunctionSamples() &&
            "New node can't have sample profile");
     NewNode->setFunctionSamples(FSamples);
   }
+  populateFuncToCtxtMap();
+}
+
+void SampleContextTracker::populateFuncToCtxtMap() {
+  for (auto *Node : *this) {
+    FunctionSamples *FSamples = Node->getFunctionSamples();
+    if (FSamples) {
+      FSamples->getContext().setState(RawContext);
+      setContextNode(FSamples, Node);
+      FuncToCtxtProfiles[Node->getFuncName()].push_back(FSamples);
+    }
+  }
 }
 
 FunctionSamples *
 SampleContextTracker::getCalleeContextSamplesFor(const CallBase &Inst,
                                                  StringRef CalleeName) {
   LLVM_DEBUG(dbgs() << "Getting callee context for instr: " << Inst << "\n");
   DILocation *DIL = Inst.getDebugLoc();
   if (!DIL)
     return nullptr;
 
   CalleeName = FunctionSamples::getCanonicalFnName(CalleeName);
   // Convert real function names to MD5 names, if the input profile is
   // MD5-based.
   std::string FGUID;
   CalleeName = getRepInFormat(CalleeName, FunctionSamples::UseMD5, FGUID);
 
   // For indirect call, CalleeName will be empty, in which case the context
   // profile for callee with largest total samples will be returned.
   ContextTrieNode *CalleeContext = getCalleeContextFor(DIL, CalleeName);
   if (CalleeContext) {
     FunctionSamples *FSamples = CalleeContext->getFunctionSamples();
     LLVM_DEBUG(if (FSamples) {
-      dbgs() << "  Callee context found: " << FSamples->getContext().toString()
+      dbgs() << "  Callee context found: " << getContextString(CalleeContext)
              << "\n";
     });
     return FSamples;
   }
 
   return nullptr;
 }
 
if (MergeContext) {
     // We have profile for function under different contexts,
     // create synthetic base profile and merge context profiles
     // into base profile.
     for (auto *CSamples : FuncToCtxtProfiles[Name]) {
       SampleContext &Context = CSamples->getContext();
       // Skip inlined context profile and also don't re-merge any context
       if (Context.hasState(InlinedContext) || Context.hasState(MergedContext))
         continue;
 
-      ContextTrieNode *FromNode = getContextFor(Context);
+      ContextTrieNode *FromNode = getContextNodeForProfile(CSamples);
       if (FromNode == Node)
         continue;
 
       ContextTrieNode &ToNode = promoteMergeContextSamplesTree(*FromNode);
       assert((!Node || Node == &ToNode) && "Expect only one base profile");
       Node = &ToNode;
     }
   }
 
   // Still no profile even after merge/promotion (if allowed)
   if (!Node)
     return nullptr;
 
   return Node->getFunctionSamples();
 }
 
 void SampleContextTracker::markContextSamplesInlined(
     const FunctionSamples *InlinedSamples) {
   assert(InlinedSamples && "Expect non-null inlined samples");
   LLVM_DEBUG(dbgs() << "Marking context profile as inlined: "
-                    << InlinedSamples->getContext().toString() << "\n");
+                    << getContextString(*InlinedSamples) << "\n");
   InlinedSamples->getContext().setState(InlinedContext);
 }
 
 ContextTrieNode &SampleContextTracker::getRootContext() { return RootContext; }
 
 void SampleContextTracker::promoteMergeContextSamplesTree(
     const Instruction &Inst, StringRef CalleeName) {
   LLVM_DEBUG(dbgs() << "Promoting and merging context tree for instr: \n"
ContextTrieNode &SampleContextTracker::promoteMergeContextSamplesTree(
     ContextTrieNode &NodeToPromo) {
   // Promote the input node to be directly under root. This can happen
   // when we decided to not inline a function under context represented
   // by the input node. The promote and merge is then needed to reflect
   // the context profile in the base (context-less) profile.
   FunctionSamples *FromSamples = NodeToPromo.getFunctionSamples();
   assert(FromSamples && "Shouldn't promote a context without profile");
   LLVM_DEBUG(dbgs() << "  Found context tree root to promote: "
-                    << FromSamples->getContext().toString() << "\n");
+                    << getContextString(&NodeToPromo) << "\n");
 
   assert(!FromSamples->getContext().hasState(InlinedContext) &&
          "Shouldn't promote inlined context profile");
-  uint32_t ContextFramesToRemove =
-      FromSamples->getContext().getContextFrames().size() - 1;
-  return promoteMergeContextSamplesTree(NodeToPromo, RootContext,
-                                        ContextFramesToRemove);
+  return promoteMergeContextSamplesTree(NodeToPromo, RootContext);
+}
+
+#ifndef NDEBUG
+std::string
+SampleContextTracker::getContextString(const FunctionSamples &FSamples) const {
+  return getContextString(getContextNodeForProfile(&FSamples));
+}
+
+std::string
+SampleContextTracker::getContextString(ContextTrieNode *Node) const {
+  SampleContextFrameVector Res;
+  if (Node == &RootContext)
+    return std::string();
+  Res.emplace_back(Node->getFuncName(), LineLocation(0, 0));
+
+  ContextTrieNode *PreNode = Node;
+  Node = Node->getParentContext();
+  while (Node && Node != &RootContext) {
+    Res.emplace_back(Node->getFuncName(), PreNode->getCallSiteLoc());
+    PreNode = Node;
+    Node = Node->getParentContext();
   }
 
+  std::reverse(Res.begin(), Res.end());
+
+  return SampleContext::getContextString(Res);
+}
+#endif
+
 void SampleContextTracker::dump() { RootContext.dumpTree(); }
 
 StringRef SampleContextTracker::getFuncNameFor(ContextTrieNode *Node) const {
   if (!FunctionSamples::UseMD5)
     return Node->getFuncName();
   assert(GUIDToFuncNameMap && "GUIDToFuncNameMap needs to be populated first");
   return GUIDToFuncNameMap->lookup(std::stoull(Node->getFuncName().data()));
 }

 }
 
 ContextTrieNode &SampleContextTracker::addTopLevelContextNode(StringRef FName) {
   assert(!getTopLevelContextNode(FName) && "Node to add must not exist");
   return *RootContext.getOrCreateChildContext(LineLocation(0, 0), FName);
 }
 
 void SampleContextTracker::mergeContextNode(ContextTrieNode &FromNode,
-                                            ContextTrieNode &ToNode,
-                                            uint32_t ContextFramesToRemove) {
+                                            ContextTrieNode &ToNode) {
   FunctionSamples *FromSamples = FromNode.getFunctionSamples();
   FunctionSamples *ToSamples = ToNode.getFunctionSamples();
   if (FromSamples && ToSamples) {
     // Merge/duplicate FromSamples into ToSamples
     ToSamples->merge(*FromSamples);
     ToSamples->getContext().setState(SyntheticContext);
     FromSamples->getContext().setState(MergedContext);
     if (FromSamples->getContext().hasAttribute(ContextShouldBeInlined))
       ToSamples->getContext().setAttribute(ContextShouldBeInlined);
   } else if (FromSamples) {
     // Transfer FromSamples from FromNode to ToNode
     ToNode.setFunctionSamples(FromSamples);
+    setContextNode(FromSamples, &ToNode);
     FromSamples->getContext().setState(SyntheticContext);
-    FromSamples->getContext().promoteOnPath(ContextFramesToRemove);
-    FromNode.setFunctionSamples(nullptr);
   }
 }
 
 ContextTrieNode &SampleContextTracker::promoteMergeContextSamplesTree(
-    ContextTrieNode &FromNode, ContextTrieNode &ToNodeParent,
-    uint32_t ContextFramesToRemove) {
-  assert(ContextFramesToRemove && "Context to remove can't be empty");
+    ContextTrieNode &FromNode, ContextTrieNode &ToNodeParent) {
 
   // Ignore call site location if destination is top level under root
   LineLocation NewCallSiteLoc = LineLocation(0, 0);
   LineLocation OldCallSiteLoc = FromNode.getCallSiteLoc();
   ContextTrieNode &FromNodeParent = *FromNode.getParentContext();
   ContextTrieNode *ToNode = nullptr;
   bool MoveToRoot = (&ToNodeParent == &RootContext);
   if (!MoveToRoot) {
     NewCallSiteLoc = OldCallSiteLoc;
   }
 
   // Locate destination node, create/move if not existing
   ToNode = ToNodeParent.getChildContext(NewCallSiteLoc, FromNode.getFuncName());
   if (!ToNode) {
     // Do not delete node to move from its parent here because
     // caller is iterating over children of that parent node.
-    ToNode = &ToNodeParent.moveToChildContext(
-        NewCallSiteLoc, std::move(FromNode), ContextFramesToRemove, false);
+    ToNode =
+        &moveContextSamples(ToNodeParent, NewCallSiteLoc, std::move(FromNode));
+    LLVM_DEBUG({
+      dbgs() << "  Context promoted and merged to: " << getContextString(ToNode)
+             << "\n";
+    });
   } else {
     // Destination node exists, merge samples for the context tree
-    mergeContextNode(FromNode, *ToNode, ContextFramesToRemove);
+    mergeContextNode(FromNode, *ToNode);
     LLVM_DEBUG({
       if (ToNode->getFunctionSamples())
         dbgs() << "  Context promoted and merged to: "
-               << ToNode->getFunctionSamples()->getContext().toString() << "\n";
+               << getContextString(ToNode) << "\n";
     });
 
     // Recursively promote and merge children
     for (auto &It : FromNode.getAllChildContext()) {
       ContextTrieNode &FromChildNode = It.second;
-      promoteMergeContextSamplesTree(FromChildNode, *ToNode,
-                                     ContextFramesToRemove);
+      promoteMergeContextSamplesTree(FromChildNode, *ToNode);
     }
 
     // Remove children once they're all merged
     FromNode.getAllChildContext().clear();
   }
 
   // For root of subtree, remove itself from old parent too
   if (MoveToRoot)
     FromNodeParent.removeChildContext(OldCallSiteLoc, ToNode->getFuncName());
 
   return *ToNode;
 }
 
 void SampleContextTracker::createContextLessProfileMap(
     SampleProfileMap &ContextLessProfiles) {
-  std::queue<ContextTrieNode *> NodeQueue;
-  NodeQueue.push(&RootContext);
-
-  while (!NodeQueue.empty()) {
-    ContextTrieNode *Node = NodeQueue.front();
+  for (auto *Node : *this) {
     FunctionSamples *FProfile = Node->getFunctionSamples();
-    NodeQueue.pop();
-
-    if (FProfile) {
-      // Profile's context can be empty, use ContextNode's func name.
+    if (FProfile)
       ContextLessProfiles[Node->getFuncName()].merge(*FProfile);
-    }
-
-    for (auto &It : Node->getAllChildContext())
-      NodeQueue.push(&It.second);
-  }
 }
 } // namespace llvm

llvm/lib/Transforms/IPO/SampleProfile.cpp

void SampleProfileLoader::findExternalInlineCandidate(
 
   // For AutoFDO profile, retrieve candidate profiles by walking over
   // the nested inlinee profiles.
   if (!FunctionSamples::ProfileIsCS) {
     Samples->findInlinedFunctions(InlinedGUIDs, SymbolMap, Threshold);
     return;
   }
 
-  ContextTrieNode *Caller =
-      ContextTracker->getContextFor(Samples->getContext());
+  ContextTrieNode *Caller = ContextTracker->getContextNodeForProfile(Samples);
   std::queue<ContextTrieNode *> CalleeList;
   CalleeList.push(Caller);
   while (!CalleeList.empty()) {
     ContextTrieNode *Node = CalleeList.front();
     CalleeList.pop();
     FunctionSamples *CalleeSample = Node->getFunctionSamples();
     // For CSSPGO profile, retrieve candidate profile by walking over the
     // trie built for context profile. Note that also take call targets

llvm/tools/llvm-profgen/CSPreInliner.h

 // Pre-compilation inliner based on context-sensitive profile.
 // The PreInliner estimates inline decision using hotness from profile
 // and cost estimation from machine code size. It helps merges context
 // profile globally and achieves better post-inine profile quality, which
 // otherwise won't be possible for ThinLTO. It also reduce context profile
 // size by only keep context that is estimated to be inlined.
 class CSPreInliner {
 public:
-  CSPreInliner(SampleProfileMap &Profiles, ProfiledBinary &Binary,
+  CSPreInliner(SampleContextTracker &Tracker, ProfiledBinary &Binary,
                ProfileSummary *Summary);
   void run();
 
 private:
   bool getInlineCandidates(ProfiledCandidateQueue &CQueue,
                            const FunctionSamples *FCallerContextSamples);
   std::vector<StringRef> buildTopDownOrder();
   void processFunction(StringRef Name);
   bool shouldInline(ProfiledInlineCandidate &Candidate);
-  uint32_t getFuncSize(const FunctionSamples &FSamples);
+  uint32_t getFuncSize(const ContextTrieNode *ContextNode);
   bool UseContextCost;
-  SampleContextTracker ContextTracker;
-  SampleProfileMap &ProfileMap;
+  SampleContextTracker &ContextTracker;
   ProfiledBinary &Binary;
   ProfileSummary *Summary;
 };
 
 } // end namespace sampleprof
 } // end namespace llvm
 
 #endif

llvm/tools/llvm-profgen/CSPreInliner.cpp

cl::opt<bool> UseContextCostForPreInliner(
     "use-context-cost-for-preinliner", cl::Hidden, cl::init(true),
     cl::desc("Use context-sensitive byte size cost for preinliner decisions"));
 
 static cl::opt<bool> SamplePreInlineReplay(
     "csspgo-replay-preinline", cl::Hidden, cl::init(false),
     cl::desc(
         "Replay previous inlining and adjust context profile accordingly"));
 
-CSPreInliner::CSPreInliner(SampleProfileMap &Profiles, ProfiledBinary &Binary,
-                           ProfileSummary *Summary)
+CSPreInliner::CSPreInliner(SampleContextTracker &Tracker,
+                           ProfiledBinary &Binary, ProfileSummary *Summary)
     : UseContextCost(UseContextCostForPreInliner),
       // TODO: Pass in a guid-to-name map in order for
       // ContextTracker.getFuncNameFor to work, if `Profiles` can have md5 codes
       // as their profile context.
-      ContextTracker(Profiles, nullptr), ProfileMap(Profiles), Binary(Binary),
-      Summary(Summary) {
+      ContextTracker(Tracker), Binary(Binary), Summary(Summary) {
   // Set default preinliner hot/cold call site threshold tuned with CSSPGO.
   // for good performance with reasonable profile size.
   if (!SampleHotCallSiteThreshold.getNumOccurrences())
     SampleHotCallSiteThreshold = 1500;
   if (!SampleColdCallSiteThreshold.getNumOccurrences())
     SampleColdCallSiteThreshold = 0;
   if (!ProfileInlineLimitMax.getNumOccurrences())
     ProfileInlineLimitMax = 3000;
 }
bool CSPreInliner::getInlineCandidates(ProfiledCandidateQueue &CQueue,
                                        const FunctionSamples *CallerSamples) {
   assert(CallerSamples && "Expect non-null caller samples");
 
   // Ideally we want to consider everything a function calls, but as far as
   // context profile is concerned, only those frames that are children of
   // current one in the trie is relavent. So we walk the trie instead of call
   // targets from function profile.
   ContextTrieNode *CallerNode =
-      ContextTracker.getContextFor(CallerSamples->getContext());
+      ContextTracker.getContextNodeForProfile(CallerSamples);
 
   bool HasNewCandidate = false;
   for (auto &Child : CallerNode->getAllChildContext()) {
     ContextTrieNode *CalleeNode = &Child.second;
     FunctionSamples *CalleeSamples = CalleeNode->getFunctionSamples();
     if (!CalleeSamples)
       continue;
 
     // Call site count is more reliable, so we look up the corresponding call
     // target profile in caller's context profile to retrieve call site count.
     uint64_t CalleeEntryCount = CalleeSamples->getEntrySamples();
     uint64_t CallsiteCount = 0;
     LineLocation Callsite = CalleeNode->getCallSiteLoc();
     if (auto CallTargets = CallerSamples->findCallTargetMapAt(Callsite)) {
       SampleRecord::CallTargetMap &TargetCounts = CallTargets.get();
       auto It = TargetCounts.find(CalleeSamples->getName());
       if (It != TargetCounts.end())
         CallsiteCount = It->second;
     }
 
     // TODO: call site and callee entry count should be mostly consistent, add
     // check for that.
     HasNewCandidate = true;
-    uint32_t CalleeSize = getFuncSize(*CalleeSamples);
+    uint32_t CalleeSize = getFuncSize(CalleeNode);
     CQueue.emplace(CalleeSamples, std::max(CallsiteCount, CalleeEntryCount),
                    CalleeSize);
   }
 
   return HasNewCandidate;
 }
 
-uint32_t CSPreInliner::getFuncSize(const FunctionSamples &FSamples) {
-  if (UseContextCost) {
-    return Binary.getFuncSizeForContext(FSamples.getContext());
-  }
+uint32_t CSPreInliner::getFuncSize(const ContextTrieNode *ContextNode) {
+  if (UseContextCost)
+    return Binary.getFuncSizeForContext(ContextNode);
 
-  return FSamples.getBodySamples().size();
+  return ContextNode->getFunctionSamples()->getBodySamples().size();
 }
 
 bool CSPreInliner::shouldInline(ProfiledInlineCandidate &Candidate) {
   // If replay inline is requested, simply follow the inline decision of the
   // profiled binary.
   if (SamplePreInlineReplay)
     return Candidate.CalleeSamples->getContext().hasAttribute(
         ContextWasInlined);
bool CSPreInliner::shouldInline(ProfiledInlineCandidate &Candidate) {
   return (Candidate.SizeCost < SampleThreshold);
 }
 
 void CSPreInliner::processFunction(const StringRef Name) {
   FunctionSamples *FSamples = ContextTracker.getBaseSamplesFor(Name);
   if (!FSamples)
     return;
 
-  unsigned FuncSize = getFuncSize(*FSamples);
+  unsigned FuncSize =
+      getFuncSize(ContextTracker.getContextNodeForProfile(FSamples));
   unsigned FuncFinalSize = FuncSize;
   unsigned SizeLimit = FuncSize * ProfileInlineGrowthLimit;
   SizeLimit = std::min(SizeLimit, (unsigned)ProfileInlineLimitMax);
   SizeLimit = std::max(SizeLimit, (unsigned)ProfileInlineLimitMin);
 
   LLVM_DEBUG(dbgs() << "Process " << Name
                     << " for context-sensitive pre-inlining (pre-inline size: "
                     << FuncSize << ", size limit: " << SizeLimit << ")\n");
if ((ShouldInline = shouldInline(Candidate))) {
       ContextTracker.markContextSamplesInlined(Candidate.CalleeSamples);
       Candidate.CalleeSamples->getContext().setAttribute(
           ContextShouldBeInlined);
       FuncFinalSize += Candidate.SizeCost;
       getInlineCandidates(CQueue, Candidate.CalleeSamples);
     } else {
       ++PreInlNumCSNotInlined;
     }
-    LLVM_DEBUG(dbgs() << (ShouldInline ? "  Inlined" : "  Outlined")
-                      << " context profile for: "
-                      << Candidate.CalleeSamples->getContext().toString()
-                      << " (callee size: " << Candidate.SizeCost
-                      << ", call count:" << Candidate.CallsiteCount << ")\n");
+    LLVM_DEBUG(
+        dbgs() << (ShouldInline ? "  Inlined" : "  Outlined")
+               << ContextTracker.getContextString(*Candidate.CalleeSamples)
   }
 
   if (!CQueue.empty()) {
     if (SizeLimit == (unsigned)ProfileInlineLimitMax)
       ++PreInlNumCSInlinedHitMaxLimit;
     else if (SizeLimit == (unsigned)ProfileInlineLimitMin)
       ++PreInlNumCSInlinedHitMinLimit;
     else
       ++PreInlNumCSInlinedHitGrowthLimit;
   }
 
   LLVM_DEBUG({
     if (!CQueue.empty())
       dbgs() << "  Inline candidates ignored due to size limit (inliner "
                 "original size: "
              << FuncSize << ", inliner final size: " << FuncFinalSize
              << ", size limit: " << SizeLimit << ")\n";
 
     while (!CQueue.empty()) {
       ProfiledInlineCandidate Candidate = CQueue.top();
       CQueue.pop();
       bool WasInlined =
           Candidate.CalleeSamples->getContext().hasAttribute(ContextWasInlined);
-      dbgs() << "    " << Candidate.CalleeSamples->getContext().toString()
+      dbgs() << "    "
+             << ContextTracker.getContextString(*Candidate.CalleeSamples)
              << " (candidate size:" << Candidate.SizeCost
              << ", call count: " << Candidate.CallsiteCount << ", previously "
              << (WasInlined ? "inlined)\n" : "not inlined)\n");
     }
   });
 }
 
 void CSPreInliner::run() {
 #ifndef NDEBUG
-  auto printProfileNames = [](SampleProfileMap &Profiles, bool IsInput) {
-    dbgs() << (IsInput ? "Input" : "Output") << " context-sensitive profiles ("
-           << Profiles.size() << " total):\n";
-    for (auto &It : Profiles) {
-      const FunctionSamples &Samples = It.second;
-      dbgs() << "  [" << Samples.getContext().toString() << "] "
-             << Samples.getTotalSamples() << ":" << Samples.getHeadSamples()
-             << "\n";
+  auto printProfileNames = [](SampleContextTracker &ContextTracker,
+                              bool IsInput) {
+    uint32_t Size = 0;
+    for (auto *Node : ContextTracker) {
+      FunctionSamples *FSamples = Node->getFunctionSamples();
+      if (FSamples) {
+        Size++;
+        dbgs() << "  [" << ContextTracker.getContextString(Node) << "] "
+               << FSamples->getTotalSamples() << ":"
+               << FSamples->getHeadSamples() << "\n";
       }
+    }
+    dbgs() << (IsInput ? "Input" : "Output") << " context-sensitive profiles ("
+           << Size << " total):\n";
   };
 #endif
 
-  LLVM_DEBUG(printProfileNames(ProfileMap, true));
+  LLVM_DEBUG(printProfileNames(ContextTracker, true));
 
   // Execute global pre-inliner to estimate a global top-down inline
   // decision and merge profiles accordingly. This helps with profile
   // merge for ThinLTO otherwise we won't be able to merge profiles back
   // to base profile across module/thin-backend boundaries.
   // It also helps better compress context profile to control profile
   // size, as we now only need context profile for functions going to
   // be inlined.
   for (StringRef FuncName : buildTopDownOrder()) {
     processFunction(FuncName);
   }
 
   // Not inlined context profiles are merged into its base, so we can
   // trim out such profiles from the output.
-  std::vector<SampleContext> ProfilesToBeRemoved;
-  for (auto &It : ProfileMap) {
-    SampleContext &Context = It.second.getContext();
-    if (!Context.isBaseContext() && !Context.hasState(InlinedContext)) {
-      assert(Context.hasState(MergedContext) &&
-             "Not inlined context profile should be merged already");
-      ProfilesToBeRemoved.push_back(It.first);
+  for (auto *Node : ContextTracker) {
+    FunctionSamples *FProfile = Node->getFunctionSamples();
+    if (FProfile &&
+        (Node->getParentContext() != &ContextTracker.getRootContext() &&
+         !FProfile->getContext().hasState(InlinedContext))) {
+      Node->setFunctionSamples(nullptr);
     }
   }
-
-  for (auto &ContextName : ProfilesToBeRemoved) {
-    ProfileMap.erase(ContextName);
-  }
-
-  // Make sure ProfileMap's key is consistent with FunctionSamples' name.
-  SampleContextTrimmer(ProfileMap).canonicalizeContextProfiles();
-
   FunctionSamples::ProfileIsPreInlined = true;
 
-  LLVM_DEBUG(printProfileNames(ProfileMap, false));
+  LLVM_DEBUG(printProfileNames(ContextTracker, false));
 }

llvm/tools/llvm-profgen/ProfileGenerator.cpp

for (const auto &FS : ProfileMap) {
     if (auto *Func = Binary->getBinaryFunction(FS.first.getName()))
       ProfiledFunctions.insert(Func);
   }
   return true;
 }
 
 bool CSProfileGenerator::collectFunctionsFromLLVMProfile(
     std::unordered_set<const BinaryFunction *> &ProfiledFunctions) {
-  std::queue<ContextTrieNode *> NodeQueue;
-  NodeQueue.push(&getRootContext());
-  while (!NodeQueue.empty()) {
-    ContextTrieNode *Node = NodeQueue.front();
-    NodeQueue.pop();
-
+  for (auto *Node : ContextTracker) {
     if (!Node->getFuncName().empty())
       if (auto *Func = Binary->getBinaryFunction(Node->getFuncName()))
         ProfiledFunctions.insert(Func);
-
-    for (auto &It : Node->getAllChildContext())
-      NodeQueue.push(&It.second);
   }
   return true;
 }
 
 FunctionSamples &
 ProfileGenerator::getTopLevelFunctionProfile(StringRef FuncName) {
   SampleContext Context(FuncName);
   auto Ret = ProfileMap.emplace(Context, FunctionSamples());
void CSProfileGenerator::computeSizeForProfiledFunctions() {
   for (auto *Func : Binary->getProfiledFunctions())
     Binary->computeInlinedContextSizeForFunc(Func);
 
   // Flush the symbolizer to save memory.
   Binary->flushSymbolizer();
 }
 
 void CSProfileGenerator::updateFunctionSamples() {
-  std::queue<ContextTrieNode *> NodeQueue;
-  NodeQueue.push(&getRootContext());
-
-  while (!NodeQueue.empty()) {
-    ContextTrieNode *Node = NodeQueue.front();
-    NodeQueue.pop();
-
+  for (auto *Node : ContextTracker) {
     FunctionSamples *FSamples = Node->getFunctionSamples();
     if (FSamples) {
       if (UpdateTotalSamples)
         FSamples->updateTotalSamples();
       FSamples->updateCallsiteSamples();
     }
-
-    for (auto &It : Node->getAllChildContext())
-      NodeQueue.push(&It.second);
   }
 }
 
 void CSProfileGenerator::generateLineNumBasedProfile() {
   for (const auto &CI : *SampleCounters) {
     const auto *CtxKey = cast<StringBasedCtxKey>(CI.first.getPtr());
 
     ContextTrieNode *ContextNode = &getRootContext();
void CSProfileGenerator::convertToProfileMap() {
   IsProfileValidOnTrie = false;
 }
 
 void CSProfileGenerator::postProcessProfiles() {
   // Compute hot/cold threshold based on profile. This will be used for cold
   // context profile merging/trimming.
   computeSummaryAndThreshold();
 
-  convertToProfileMap();
-
   // Run global pre-inliner to adjust/merge context profile based on estimated
   // inline decisions.
   if (EnableCSPreInliner) {
-    CSPreInliner(ProfileMap, *Binary, Summary.get()).run();
+    ContextTracker.populateFuncToCtxtMap();
+    CSPreInliner(ContextTracker, *Binary, Summary.get()).run();
     // Turn off the profile merger by default unless it is explicitly enabled.
     if (!CSProfMergeColdContext.getNumOccurrences())
       CSProfMergeColdContext = false;
   }
 
+  convertToProfileMap();
+
   // Trim and merge cold context profile using cold threshold above.
   if (TrimColdProfile || CSProfMergeColdContext) {
     SampleContextTrimmer(ProfileMap)
         .trimAndMergeColdContextProfiles(
             HotCountThreshold, TrimColdProfile, CSProfMergeColdContext,
             CSProfMaxColdContextDepth, EnableCSPreInliner);
   }
 
   // Merge function samples of CS profile to calculate profile density.
   sampleprof::SampleProfileMap ContextLessProfiles;
   for (const auto &I : ProfileMap) {
     ContextLessProfiles[I.second.getName()].merge(I.second);
   }
 
   calculateAndShowDensity(ContextLessProfiles);
   if (GenCSNestedProfile) {
     CSProfileConverter CSConverter(ProfileMap);
     CSConverter.convertProfiles();

llvm/tools/llvm-profgen/ProfiledBinary.h

 public:
   // Add instruction with given size to a context
   void addInstructionForContext(const SampleContextFrameVector &Context,
                                 uint32_t InstrSize);
 
   // Get function size with a specific context. When there's no exact match
   // for the given context, try to retrieve the size of that function from
   // closest matching context.
-  uint32_t getFuncSizeForContext(const SampleContext &Context);
+  uint32_t getFuncSizeForContext(const ContextTrieNode *Context);
 
   // For inlinees that are full optimized away, we can establish zero size using
   // their remaining probes.
   void trackInlineesOptimizedAway(MCPseudoProbeDecoder &ProbeDecoder);
 
   using ProbeFrameStack = SmallVector<std::pair<StringRef, uint32_t>>;
   void trackInlineesOptimizedAway(MCPseudoProbeDecoder &ProbeDecoder,
                               MCDecodedPseudoProbeInlineTree &ProbeNode,
public:
 
   BinaryFunction *getBinaryFunction(StringRef FName) {
     auto I = BinaryFunctions.find(FName.str());
     if (I == BinaryFunctions.end())
       return nullptr;
     return &I->second;
   }
 
-  uint32_t getFuncSizeForContext(SampleContext &Context) {
-    return FuncSizeTracker.getFuncSizeForContext(Context);
+  uint32_t getFuncSizeForContext(const ContextTrieNode *ContextNode) {
+    return FuncSizeTracker.getFuncSizeForContext(ContextNode);
   }
 
   // Load the symbols from debug table and populate into symbol list.
   void populateSymbolListFromDWARF(ProfileSymbolList &SymbolList);
 
   const SampleContextFrameVector &
   getFrameLocationStack(uint64_t Offset, bool UseProbeDiscriminator = false) {
     auto I = Offset2LocStackMap.emplace(Offset, SampleContextFrameVector());

llvm/tools/llvm-profgen/ProfiledBinary.cpp

for (const auto &Callsite : reverse(Context)) {
     CurNode = CurNode->getOrCreateChildContext(CallsiteLoc, CallerName);
     IsLeaf = false;
   }
 
   CurNode->addFunctionSize(InstrSize);
 }
 
 uint32_t
-BinarySizeContextTracker::getFuncSizeForContext(const SampleContext &Context) {
+BinarySizeContextTracker::getFuncSizeForContext(const ContextTrieNode *Node) {
   ContextTrieNode *CurrNode = &RootContext;
   ContextTrieNode *PrevNode = nullptr;
-  SampleContextFrames Frames = Context.getContextFrames();
-  int32_t I = Frames.size() - 1;
+
   Optional<uint32_t> Size;
 
   // Start from top-level context-less function, traverse down the reverse
   // context trie to find the best/longest match for given context, then
   // retrieve the size.
-
-  while (CurrNode && I >= 0) {
-    // Process from leaf function to callers (added to context).
-    const auto &ChildFrame = Frames[I--];
+  LineLocation CallSiteLoc(0, 0);
+  while (CurrNode && Node->getParentContext() != nullptr) {
     PrevNode = CurrNode;
-    CurrNode =
-        CurrNode->getChildContext(ChildFrame.Location, ChildFrame.FuncName);
+    CurrNode = CurrNode->getChildContext(CallSiteLoc, Node->getFuncName());
     if (CurrNode && CurrNode->getFunctionSize())
       Size = CurrNode->getFunctionSize().getValue();
+    CallSiteLoc = Node->getCallSiteLoc();
+    Node = Node->getParentContext();
   }
 
   // If we traversed all nodes along the path of the context and haven't
   // found a size yet, pivot to look for size from sibling nodes, i.e size
   // of inlinee under different context.
   if (!Size) {
     if (!CurrNode)
       CurrNode = PrevNode;