[CSSPGO][llvm-profgen] Reimplement CS profile generator using context trie

[wlei]

Summary

Our investigation showed ProfileMap's key is the bottleneck of the memory consumption for CS profile generation on some large services. This patch tries to optimize it by storing the CS function samples using the context trie tree structure instead of the context frame array ref. Parts of code in ContextTrieNode are reused.

Our experiment on one internal service showed that the context key's memory can be reduced from 80GB to 300MB.

To be compatible with non-CS profiles, the profile writer still needs to use ProfileMap as input, so rebuild the ProfileMap using the context trie in postProcessProfiles.

The optimization is not complete yet, next step is to reimplement Pre-inliner or profile trimmer, after that, ProfileMap should be small to be written.

[hoy]

Our experiment on one internal service showed that the context key's memory can be reduced from 80GB to 300MB.

This is awesome! Thanks for the work!

How about running time?

[wlei]

In D125246#3501359, @hoy wrote:

Our experiment on one internal service showed that the context key's memory can be reduced from 80GB to 300MB.

This is awesome! Thanks for the work!

How about running time?

Just tested the trie based is faster than the old one, E2E running time : 103mins vs 122mins.

The reason for the speed-up should be due to the time spent on getting caller's FunctionSamples. For the new version, it's just call getParentContext() which is O(1). The old version it generate the caller's context frames then call getFunctionProfileForContext(...), which is O(n), n is the context length.

[hoy]

In D125246#3502570, @wlei wrote:

In D125246#3501359, @hoy wrote:

Our experiment on one internal service showed that the context key's memory can be reduced from 80GB to 300MB.

This is awesome! Thanks for the work!

How about running time?

Just tested the trie based is faster than the old one, E2E running time : 103mins vs 122mins.

The reason for the speed-up should be due to the time spent on getting caller's FunctionSamples. For the new version, it's just call getParentContext() which is O(1). The old version it generate the caller's context frames then call getFunctionProfileForContext(...), which is O(n), n is the context length.

Very nice. BTW, looks like this is done to line number based profile generation only. How about pseudo probes? Am I missing anything?

[hoy]

In D125246#3504823, @hoy wrote:

In D125246#3502570, @wlei wrote:

In D125246#3501359, @hoy wrote:

Our experiment on one internal service showed that the context key's memory can be reduced from 80GB to 300MB.

This is awesome! Thanks for the work!

How about running time?

Just tested the trie based is faster than the old one, E2E running time : 103mins vs 122mins.

The reason for the speed-up should be due to the time spent on getting caller's FunctionSamples. For the new version, it's just call getParentContext() which is O(1). The old version it generate the caller's context frames then call getFunctionProfileForContext(...), which is O(n), n is the context length.

Very nice. BTW, looks like this is done to line number based profile generation only. How about pseudo probes? Am I missing anything?

NVM, I didn't notice the changes for pseudo probes.

[hoy]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

894

PLease add more comments about why the processing order matters.

917

Update the comment "But we can't modify ProfileMap while iterating it."?

On the other hand, since we are doing bottom-up processing, we should have no such limitation now?

933–944

nit: perhaps just name this function buildProfileMap since we are using trie by default.

934

use reference type instead of pointer type for Node to explicitly indicate that this is a non-null ptr?

945

Can FProfile be updated in place to avoid copy and deletion?

969

I guess eventually this call will be after the preiliner run, but can this be moved after computeSummaryAndThreshold now?

969

nit: seems that Context is not used anywhere. Perhaps make a wrapper function void buildProfileMapFromContextTrie( ContextTrieNode *Node) and declare Context inside?

llvm/tools/llvm-profgen/ProfileGenerator.h

304

nit: getOrCreateContextNodeForContext

[wlei]

Addressing feedback.

llvm/tools/llvm-profgen/ProfileGenerator.cpp

894

Comments updated.

917

Good point! Now we can just create the new function samples for caller, actually that's already there see the later code creating this.

FunctionSamples &CallerProfile = *getOrCreateFunctionSamples(CallerNode);

933–944

Sounds good!

934

Sounds good!

945

the profile in ProfileMap require the full context as the key, so we can't create the value in ProfileMap in the early time. Or perhaps we can use the move semantics, but FunctionSamples currently doesn't have the move constructor, which will make things complicated. I think after the trimming and pre-inliner, the num of profile will be small, copy and deletion should be fine then?

969

Sounds good!

969

I guess eventually this call will be after the preiliner run, but can this be moved after computeSummaryAndThreshold now?

Yeah, we need to move it after preinliner and moreover it should be after SampleContextTrimmer.

computeSummaryAndThreshold is still using ProfileMap and I also found that ProfileMap can also be imported from the SampleReader. Both of them will needs to be rewritten. Maybe leave them a separate diff, just to be easy for the code review, what do you think?

llvm/tools/llvm-profgen/ProfileGenerator.h

304

Fixed!

[hoy]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

945

Makes sense. Eventually the trie should be small after preinliner.

You may want to set Node.getFunctionSamples() to be null after deleting it.

969

computeSummaryAndThreshold is still using ProfileMap

I'm wondering if we can make a version of computeSummaryAndThreshold to run on the trie to avoid generate a giant ProfileMap with all contexts in the middle.

SampleReader is a special path and we don't care about the performance there. It can be treated differently.

Yeah, a separate diff is fine.

[wlei]

rebase
using context tracker directly

[wlei]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

945

Set to null, thanks!

[wenlei]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

894

parent(inliner)'s sample depends on child(inlinee)'s sample, so traverse the tree in post-order.

I thought order was not important because we use head samples to update body samples. But it looks like we use entry samples which could be first body samples.

So this was a bug in previous implementation?

946

Maybe it isn't too complicated to add move ctor for FunctionSamples? We're taking care of SampleContext here, and the rest are two std::map. But we can deal with it in separate change.

Relatedly, this feels more like converting trie to map, rather than building a map from trie, because the process also "destroy" the trie. Who owns FProfile of the trie? While delete it on the fly saves a bit of mem, it also feels a bit weird for it to be deleted here as iiuc the trie may not own the underlying function samples?

[wlei]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

718

The FunctionSamples is new here.

894

Right, we use the callee's body sample to update caller's body sample.

It is a bug in previous implementation, but it should not be critical one, as you said, it's only related with the first body samples.

946

Maybe it isn't too complicated to add move ctor for FunctionSamples? We're taking care of SampleContext here, and the rest are two std::map. But we can deal with it in separate change.

OK, I just realized that we can use the default implicit move ctor of FunctionSamples, it can automatically move the two std::map. Then I think it should work here to add the std::move(..) to avoid copy.

Who owns FProfile of the trie?

It's a little bit tricky.

1. It can be from the ProfileMap(using --llvm-sample-profile as input) or sample reader from compiler side. we can't delete the profile owned by ProfileMap.

2. for profile generation, it's created on on-the-fly on the trie.(see ProfileGenerator.cpp:718). That's why I delete this profile while "destroying" trie.

I guess you meant the memory saving of the on-the-fly deletion is not noticeable(I agree that the memory after pre-inliner should be small), if so, I will separate this into two parts: 1) build map and 2) free the memory(in the future patch).

[wlei]

use std::move to avoid memory copy
remove the "memory free" code, will add in a separate patch.

[hoy]

LGTM.

[wlei]

Put all FSamples in a std::list<FunctionSamples> to avoid memory new and delete.

[wenlei]

Thanks for the change -- good work. I can see how it improves performance of llvm-profgen. A few minor comments.

llvm/tools/llvm-profgen/ProfileGenerator.cpp

718

FSamplesList.emplace_back(FunctionSamples());
->
FSamplesList.emplace_back();

emplace_back should be able to take ctor parameters directly, which avoids a copy/move.

916–931

nit: we can simplify this by having a single populateInferredFunctionSamples function with default argument point to root context node.

946

Thanks for changing it to use move. With that the original trie still gets destroyed because it no longer holds valid FunctionSamples. So perhaps we should name this function "convert" instead of "build"?

And since this is stateful, add an assertion/check to verify? like assert that ProfileMap is empty before we start the conversion, and the trie is still valid (holds valid function samples).

llvm/tools/llvm-profgen/ProfileGenerator.h

304

nit: might want to unify the naming convention, getOrCreateFunctionSamples doesn't have ForXYZ. So getOrCreateContextNode?

[wlei]

addressing Wenlei's feedback

[wenlei]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

954

nit: convertProfileMap -> convertToProfileMap

971–972

what makes function samples empty during conversion?

980–981

I was thinking about checking to make sure things are valid before conversion. Checking after conversion doesn't prevent double conversion where the 2nd conversion would be wrong. We could also have a simple state boolean to represent the statefulness and use that for checks.

[wlei]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

971–972

IMO, The std::move will move the BodySamples from the trie node's FSamples to the ProfileMap's FSamples. This is to check after the conversion no FSamples left not moved.

980–981

By a valid state, do you mean to check that the trie node's getAllChildren is not empty?

[wlei]

Updating D125246: [CSSPGO][llvm-profgen] Reimplement CS profile generator using context trie

[wenlei]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

971–972

The move-from object are in unspecified state according to standard, and we'd want to avoid accessing move-from objects.

980–981

I was thinking of valid state as an abstract term -- since the conversion destroy the original profile on trie through move etc, after conversion the function samples on the trie is no longer valid. We could simply have a boolean isProfileValidOnTrie as a state, set it to false inside convert function, and also check it is true at the beginning of conversion..

[wlei]

add a valid state check for the trie conversion.

llvm/tools/llvm-profgen/ProfileGenerator.cpp

980–981

Ah, I see, thanks for the clarification!

[wenlei]

lgtm, thanks.

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

bool operator()(FunctionSamples *A, FunctionSamples *B) const {
       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>;
 
+  SampleContextTracker() = default;
   SampleContextTracker(SampleProfileMap &Profiles,
                        const DenseMap<uint64_t, StringRef> *GUIDToFuncNameMap);
   // 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
   // is identified by input DILocation.
   FunctionSamples *getContextSamplesFor(const DILocation *DIL);
   // Query context profile for a given sample contxt of a function.
   FunctionSamples *getContextSamplesFor(const SampleContext &Context);
   // Get all context profile for given function.
   ContextSamplesTy &getAllContextSamplesFor(const Function &Func);
   ContextSamplesTy &getAllContextSamplesFor(StringRef Name);
+  ContextTrieNode *getOrCreateContextPath(const SampleContext &Context,
+                                          bool AllowCreate);
   // Query base profile for a given function. A base profile is a merged view
   // of all context profiles for contexts that are not inlined.
   FunctionSamples *getBaseSamplesFor(const Function &Func,
                                      bool MergeContext = true);
   // Query base profile for a given function by name.
   FunctionSamples *getBaseSamplesFor(StringRef Name, bool MergeContext = true);
   // Retrieve the context trie node for given profile context
   ContextTrieNode *getContextFor(const SampleContext &Context);
   // Get real function name for a given trie node.
   StringRef getFuncNameFor(ContextTrieNode *Node) const;
   // Mark a context profile as inlined when function is inlined.
   // This makes sure that inlined context profile will be excluded in
   // function's base profile.
   void markContextSamplesInlined(const FunctionSamples *InlinedSamples);
   ContextTrieNode &getRootContext();
   void promoteMergeContextSamplesTree(const Instruction &Inst,
                                       StringRef CalleeName);
   // Dump the internal context profile trie.
   void dump();
 
 private:
   ContextTrieNode *getContextFor(const DILocation *DIL);
   ContextTrieNode *getCalleeContextFor(const DILocation *DIL,
                                        StringRef CalleeName);
-  ContextTrieNode *getOrCreateContextPath(const SampleContext &Context,
-                                          bool AllowCreate);
   ContextTrieNode *getTopLevelContextNode(StringRef FName);
   ContextTrieNode &addTopLevelContextNode(StringRef FName);
   ContextTrieNode &promoteMergeContextSamplesTree(ContextTrieNode &NodeToPromo);
   void mergeContextNode(ContextTrieNode &FromNode, ContextTrieNode &ToNode,
                         uint32_t ContextFramesToRemove);
   ContextTrieNode &
   promoteMergeContextSamplesTree(ContextTrieNode &FromNode,
                                  ContextTrieNode &ToNodeParent,

llvm/tools/llvm-profgen/ProfileGenerator.h

while (I <= MaxDedupSize) {
       I++;
       Context.resize(End);
       MaxDedupSize = std::min(static_cast<uint32_t>(End / 2), MaxDedupSize);
     }
   }
 
 private:
   void generateLineNumBasedProfile();
-  // Lookup or create FunctionSamples for the context
-  FunctionSamples &
-  getFunctionProfileForContext(const SampleContextFrameVector &Context,
+
+  FunctionSamples *getOrCreateFunctionSamples(ContextTrieNode *ContextNode,
                                               bool WasLeafInlined = false);
+
+  // Lookup or create ContextTrieNode for the context, FunctionSamples is
+  // created inside this function.
+  ContextTrieNode *getOrCreateContextNode(const SampleContextFrames Context,
+                                          bool WasLeafInlined = false);
+
   // For profiled only functions, on-demand compute their inline context
   // function byte size which is used by the pre-inliner.
   void computeSizeForProfiledFunctions();
   // Post processing for profiles before writing out, such as mermining
   // and trimming cold profiles, running preinliner on profiles.
   void postProcessProfiles();
 
   void populateBodySamplesForFunction(FunctionSamples &FunctionProfile,
                                       const RangeSample &RangeCounters);
-  void populateBoundarySamplesForFunction(SampleContextFrames ContextId,
-                                          FunctionSamples *CallerProfile,
+
+  void populateBoundarySamplesForFunction(ContextTrieNode *CallerNode,
                                           const BranchSample &BranchCounters);
-  void populateInferredFunctionSamples();
+
+  void populateInferredFunctionSamples(ContextTrieNode &Node);
+
+  void updateFunctionSamples();
 
   void generateProbeBasedProfile();
 
   // Fill in function body samples from probes
   void populateBodySamplesWithProbes(const RangeSample &RangeCounter,
                                      SampleContextFrames ContextStack);
   // Fill in boundary samples for a call probe
   void populateBoundarySamplesWithProbes(const BranchSample &BranchCounter,
                                          SampleContextFrames ContextStack);
+
+  ContextTrieNode *
+  getContextNodeForLeafProbe(SampleContextFrames ContextStack,
+                             const MCDecodedPseudoProbe *LeafProbe);
+
   // Helper function to get FunctionSamples for the leaf probe
   FunctionSamples &
   getFunctionProfileForLeafProbe(SampleContextFrames ContextStack,
                                  const MCDecodedPseudoProbe *LeafProbe);
 
+  void convertToProfileMap(ContextTrieNode &Node,
+                           SampleContextFrameVector &Context);
+
+  void convertToProfileMap();
+
+  ContextTrieNode &getRootContext() { return ContextTracker.getRootContext(); };
+
+  // The container for holding the FunctionSamples used by context trie.
+  std::list<FunctionSamples> FSamplesList;
+
   // Underlying context table serves for sample profile writer.
   std::unordered_set<SampleContextFrameVector, SampleContextFrameHash> Contexts;
 
+  SampleContextTracker ContextTracker;
+
+  bool IsProfileValidOnTrie = true;
+
 public:
   // Deduplicate adjacent repeated context sequences up to a given sequence
   // length. -1 means no size limit.
   static int32_t MaxCompressionSize;
   static int MaxContextDepth;
 };
 
 } // end namespace sampleprof
 } // end namespace llvm
 
 #endif

llvm/tools/llvm-profgen/ProfileGenerator.cpp

 }
 
 void ProfileGeneratorBase::calculateAndShowDensity(
     const SampleProfileMap &Profiles) {
   double Density = calculateDensity(Profiles, HotCountThreshold);
   showDensitySuggestion(Density);
 }
 
-FunctionSamples &CSProfileGenerator::getFunctionProfileForContext(
-    const SampleContextFrameVector &Context, bool WasLeafInlined) {
-  auto I = ProfileMap.find(SampleContext(Context));
-  if (I == ProfileMap.end()) {
-    // Save the new context for future references.
-    SampleContextFrames NewContext = *Contexts.insert(Context).first;
-    SampleContext FContext(NewContext, RawContext);
-    auto Ret = ProfileMap.emplace(FContext, FunctionSamples());
-    if (WasLeafInlined)
-      FContext.setAttribute(ContextWasInlined);
-    FunctionSamples &FProfile = Ret.first->second;
-    FProfile.setContext(FContext);
-    return Ret.first->second;
-  } else {
-    // Update ContextWasInlined attribute for existing contexts.
-    // The current function can be called in two ways:
-    //  - when processing a probe of the current frame
-    //  - when processing the entry probe of an inlinee's frame, which
-    //    is then used to update the callsite count of the current frame.
-    // The two can happen in any order, hence here we are making sure
-    // `ContextWasInlined` is always set as expected.
-    // TODO: Note that the former does not always happen if no probes of the
-    // current frame has samples, and if the latter happens, we could lose the
-    // attribute. This should be fixed.
-    if (WasLeafInlined)
-      I->second.getContext().setAttribute(ContextWasInlined);
-  }
+FunctionSamples *
+CSProfileGenerator::getOrCreateFunctionSamples(ContextTrieNode *ContextNode,
+                                               bool WasLeafInlined) {
+  FunctionSamples *FProfile = ContextNode->getFunctionSamples();
+  if (!FProfile) {
+    FSamplesList.emplace_back();
+    FProfile = &FSamplesList.back();
+    FProfile->setName(ContextNode->getFuncName());
+    ContextNode->setFunctionSamples(FProfile);
+  }
+    FProfile->getContext().setAttribute(ContextWasInlined);
+  return FProfile;
 
-  return I->second;
+ContextTrieNode *
+CSProfileGenerator::getOrCreateContextNode(const SampleContextFrames Context,
+                                           bool WasLeafInlined) {
+  ContextTrieNode *ContextNode =
+      ContextTracker.getOrCreateContextPath(Context, true);
+  getOrCreateFunctionSamples(ContextNode, WasLeafInlined);
+  return ContextNode;
 }
 
 void CSProfileGenerator::generateProfile() {
   FunctionSamples::ProfileIsCS = true;
 
   collectProfiledFunctions();
 
   if (Binary->usePseudoProbes())
void CSProfileGenerator::computeSizeForProfiledFunctions() {
   std::unordered_set<const BinaryFunction *> ProfiledFunctions;
   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();
+
+    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());
 
-    FunctionSamples *FunctionProfile = nullptr;
+    ContextTrieNode *ContextNode = &getRootContext();
     // Sample context will be empty if the jump is an external-to-internal call
     // pattern, the head samples should be added for the internal function.
     if (!CtxKey->Context.empty()) {
       // Get or create function profile for the range
-      FunctionProfile = &getFunctionProfileForContext(CtxKey->Context,
-                                                      CtxKey->WasLeafInlined);
+      ContextNode =
+          getOrCreateContextNode(CtxKey->Context, CtxKey->WasLeafInlined);
       // Fill in function body samples
-      populateBodySamplesForFunction(*FunctionProfile, CI.second.RangeCounter);
+      populateBodySamplesForFunction(*ContextNode->getFunctionSamples(),
+                                     CI.second.RangeCounter);
     }
     // Fill in boundary sample counts as well as call site samples for calls
-    populateBoundarySamplesForFunction(CtxKey->Context, FunctionProfile,
-                                       CI.second.BranchCounter);
+    populateBoundarySamplesForFunction(ContextNode, CI.second.BranchCounter);
   }
   // Fill in call site value sample for inlined calls and also use context to
   // infer missing samples. Since we don't have call count for inlined
   // functions, we estimate it from inlinee's profile using the entry of the
   // body sample.
-  populateInferredFunctionSamples();
+  populateInferredFunctionSamples(getRootContext());
 
   updateFunctionSamples();
 }
 
 void CSProfileGenerator::populateBodySamplesForFunction(
     FunctionSamples &FunctionProfile, const RangeSample &RangeCounter) {
   // Compute disjoint ranges first, so we can use MAX
   // for calculating count for each location.
do {
         updateBodySamplesforFunctionProfile(FunctionProfile, *LeafLoc, Count);
         FunctionProfile.addTotalSamples(Count);
       }
     } while (IP.advance() && IP.Address <= RangeEnd);
   }
 }
 
 void CSProfileGenerator::populateBoundarySamplesForFunction(
-    SampleContextFrames ContextId, FunctionSamples *CallerProfile,
-    const BranchSample &BranchCounters) {
+    ContextTrieNode *Node, const BranchSample &BranchCounters) {
 
   for (const auto &Entry : BranchCounters) {
     uint64_t SourceOffset = Entry.first.first;
     uint64_t TargetOffset = Entry.first.second;
     uint64_t Count = Entry.second;
     assert(Count != 0 && "Unexpected zero weight branch");
 
     StringRef CalleeName = getCalleeNameForOffset(TargetOffset);
     if (CalleeName.size() == 0)
       continue;
 
-    SampleContextFrameVector CalleeCtx;
-    if (CallerProfile) {
-      assert(!ContextId.empty() &&
-             "CallerProfile is null only if ContextId is empty");
+    ContextTrieNode *CallerNode = Node;
+    LineLocation CalleeCallSite(0, 0);
+    if (CallerNode != &getRootContext()) {
       // Record called target sample and its count
       auto LeafLoc = Binary->getInlineLeafFrameLoc(SourceOffset);
       if (LeafLoc) {
-        CallerProfile->addCalledTargetSamples(
+        CallerNode->getFunctionSamples()->addCalledTargetSamples(
             LeafLoc->Location.LineOffset,
             getBaseDiscriminator(LeafLoc->Location.Discriminator), CalleeName,
             Count);
-
         // Record head sample for called target(callee)
-        CalleeCtx.append(ContextId.begin(), ContextId.end());
-        assert(CalleeCtx.back().FuncName == LeafLoc->FuncName &&
-               "Leaf function name doesn't match");
-        CalleeCtx.back() = *LeafLoc;
-      }
-    }
-    CalleeCtx.emplace_back(CalleeName, LineLocation(0, 0));
-    FunctionSamples &CalleeProfile = getFunctionProfileForContext(CalleeCtx);
-    CalleeProfile.addHeadSamples(Count);
+        CalleeCallSite = LeafLoc->Location;
       }
     }
 
-static SampleContextFrame
-getCallerContext(SampleContextFrames CalleeContext,
-                 SampleContextFrameVector &CallerContext) {
-  assert(CalleeContext.size() > 1 && "Unexpected empty context");
-  CalleeContext = CalleeContext.drop_back();
-  CallerContext.assign(CalleeContext.begin(), CalleeContext.end());
-  SampleContextFrame CallerFrame = CallerContext.back();
-  CallerContext.back().Location = LineLocation(0, 0);
-  return CallerFrame;
+    ContextTrieNode *CalleeNode =
+        CallerNode->getOrCreateChildContext(CalleeCallSite, CalleeName);
+    FunctionSamples *CalleeProfile = getOrCreateFunctionSamples(CalleeNode);
+    CalleeProfile->addHeadSamples(Count);
+  }
 }
 
-void CSProfileGenerator::populateInferredFunctionSamples() {
-  for (const auto &Item : ProfileMap) {
-    const auto &CalleeContext = Item.first;
-    const FunctionSamples &CalleeProfile = Item.second;
+void CSProfileGenerator::populateInferredFunctionSamples(
+    ContextTrieNode &Node) {
+  // There is no call jmp sample between the inliner and inlinee, we need to use
+  // the inlinee's context to infer inliner's context, i.e. parent(inliner)'s
+  // sample depends on child(inlinee)'s sample, so traverse the tree in
+  // post-order.
+  for (auto &It : Node.getAllChildContext())
+    populateInferredFunctionSamples(It.second);
 
-    // If we already have head sample counts, we must have value profile
-    // for call sites added already. Skip to avoid double counting.
-    if (CalleeProfile.getHeadSamples())
-      continue;
-    // If we don't have context, nothing to do for caller's call site.
-    // This could happen for entry point function.
-    if (CalleeContext.isBaseContext())
-      continue;
-
-    // Infer Caller's frame loc and context ID through string splitting
-    SampleContextFrameVector CallerContextId;
-    SampleContextFrame &&CallerLeafFrameLoc =
-        getCallerContext(CalleeContext.getContextFrames(), CallerContextId);
-    SampleContextFrames CallerContext(CallerContextId);
-
-    // It's possible that we haven't seen any sample directly in the caller,
-    // in which case CallerProfile will not exist. But we can't modify
-    // ProfileMap while iterating it.
-    // TODO: created function profile for those callers too
-    if (ProfileMap.find(CallerContext) == ProfileMap.end())
-      continue;
-    FunctionSamples &CallerProfile = ProfileMap[CallerContext];
+  FunctionSamples *CalleeProfile = Node.getFunctionSamples();
+  if (!CalleeProfile)
+    return;
+  if (CalleeProfile->getHeadSamples())
+    return;
+  ContextTrieNode *CallerNode = Node.getParentContext();
+  if (CallerNode == &getRootContext())
+    return;
 
-    // Since we don't have call count for inlined functions, we
-    // estimate it from inlinee's profile using entry body sample.
-    uint64_t EstimatedCallCount = CalleeProfile.getEntrySamples();
-    // If we don't have samples with location, use 1 to indicate live.
-    if (!EstimatedCallCount && !CalleeProfile.getBodySamples().size())
-      EstimatedCallCount = 1;
-    CallerProfile.addCalledTargetSamples(
-        CallerLeafFrameLoc.Location.LineOffset,
-        CallerLeafFrameLoc.Location.Discriminator,
-        CalleeProfile.getContext().getName(), EstimatedCallCount);
-    CallerProfile.addBodySamples(CallerLeafFrameLoc.Location.LineOffset,
-                                 CallerLeafFrameLoc.Location.Discriminator,
-                                 EstimatedCallCount);
-    CallerProfile.addTotalSamples(EstimatedCallCount);
-  }
+  LineLocation CallerLeafFrameLoc = Node.getCallSiteLoc();
+  FunctionSamples &CallerProfile = *getOrCreateFunctionSamples(CallerNode);
+  uint64_t EstimatedCallCount = CalleeProfile->getEntrySamples();
+  if (!EstimatedCallCount && !CalleeProfile->getBodySamples().size())
+  CallerProfile.addCalledTargetSamples(CallerLeafFrameLoc.LineOffset,
+                                       CallerLeafFrameLoc.Discriminator,
+                                       Node.getFuncName(), EstimatedCallCount);
+  CallerProfile.addBodySamples(CallerLeafFrameLoc.LineOffset,
+                               CallerLeafFrameLoc.Discriminator,
+
+void CSProfileGenerator::convertToProfileMap(
+    ContextTrieNode &Node, SampleContextFrameVector &Context) {
+  FunctionSamples *FProfile = Node.getFunctionSamples();
+  if (FProfile) {
+    Context.emplace_back(Node.getFuncName(), LineLocation(0, 0));
+    // Save the new context for future references.
+    SampleContextFrames NewContext = *Contexts.insert(Context).first;
+    auto Ret = ProfileMap.emplace(NewContext, std::move(*FProfile));
+    FunctionSamples &NewProfile = Ret.first->second;
+    NewProfile.getContext().setContext(NewContext);
+    Context.pop_back();
+  }
+
+  for (auto &It : Node.getAllChildContext()) {
+    ContextTrieNode &ChildNode = It.second;
+    Context.emplace_back(Node.getFuncName(), ChildNode.getCallSiteLoc());
+    convertToProfileMap(ChildNode, Context);
+    Context.pop_back();
+  }
+}
+
+void CSProfileGenerator::convertToProfileMap() {
+  assert(ProfileMap.empty() &&
+         "ProfileMap should be empty before converting from the trie");
+  assert(IsProfileValidOnTrie &&
+         "Do not convert the trie twice, it's already destroyed");
+
+  SampleContextFrameVector Context;
+  for (auto &It : getRootContext().getAllChildContext())
+    convertToProfileMap(It.second, Context);
+
+  IsProfileValidOnTrie = false;
 }
 
 void CSProfileGenerator::postProcessProfiles() {
+  if (SampleCounters)
+    convertToProfileMap();
+
   // Compute hot/cold threshold based on profile. This will be used for cold
   // context profile merging/trimming.
   computeSummaryAndThreshold();
 
   // Run global pre-inliner to adjust/merge context profile based on estimated
   // inline decisions.
   if (EnableCSPreInliner) {
     CSPreInliner(ProfileMap, *Binary, Summary.get()).run();
     // Turn off the profile merger by default unless it is explicitly enabled.
     if (!CSProfMergeColdContext.getNumOccurrences())
       CSProfMergeColdContext = false;
   }
 
   // Trim and merge cold context profile using cold threshold above.
   if (TrimColdProfile || CSProfMergeColdContext) {
     SampleContextTrimmer(ProfileMap)
         .trimAndMergeColdContextProfiles(
             HotCountThreshold, TrimColdProfile, CSProfMergeColdContext,
             CSProfMaxColdContextDepth, EnableCSPreInliner);
std::unordered_map<MCDecodedPseudoProbeInlineTree *,
       FrameSamples;
   for (const auto &PI : ProbeCounter) {
     const MCDecodedPseudoProbe *Probe = PI.first;
     uint64_t Count = PI.second;
     // Disjoint ranges have introduce zero-filled gap that
     // doesn't belong to current context, filter them out.
     if (!Probe->isBlock() || Count == 0)
       continue;
-    FunctionSamples &FunctionProfile =
-        getFunctionProfileForLeafProbe(ContextStack, Probe);
+
+    ContextTrieNode *ContextNode =
+        getContextNodeForLeafProbe(ContextStack, Probe);
+    FunctionSamples &FunctionProfile = *ContextNode->getFunctionSamples();
     // Record the current frame and FunctionProfile whenever samples are
     // collected for non-danglie probes. This is for reporting all of the
     // zero count probes of the frame later.
     FrameSamples[Probe->getInlineTreeNode()].insert(&FunctionProfile);
     FunctionProfile.addBodySamplesForProbe(Probe->getIndex(), Count);
     FunctionProfile.addTotalSamples(Count);
     if (Probe->isEntry()) {
       FunctionProfile.addHeadSamples(Count);
       // Look up for the caller's function profile
       const auto *InlinerDesc = Binary->getInlinerDescForProbe(Probe);
-      SampleContextFrames CalleeContextId =
-          FunctionProfile.getContext().getContextFrames();
-      if (InlinerDesc != nullptr && CalleeContextId.size() > 1) {
+      ContextTrieNode *CallerNode = ContextNode->getParentContext();
+      if (InlinerDesc != nullptr && CallerNode != &getRootContext()) {
         // Since the context id will be compressed, we have to use callee's
         // context id to infer caller's context id to ensure they share the
         // same context prefix.
-        SampleContextFrameVector CallerContextId;
-        SampleContextFrame &&CallerLeafFrameLoc =
-            getCallerContext(CalleeContextId, CallerContextId);
-        uint64_t CallerIndex = CallerLeafFrameLoc.Location.LineOffset;
+        uint64_t CallerIndex = ContextNode->getCallSiteLoc().LineOffset;
         assert(CallerIndex &&
                "Inferred caller's location index shouldn't be zero!");
         FunctionSamples &CallerProfile =
-            getFunctionProfileForContext(CallerContextId);
+            *getOrCreateFunctionSamples(CallerNode);
         CallerProfile.setFunctionHash(InlinerDesc->FuncHash);
         CallerProfile.addBodySamples(CallerIndex, 0, Count);
         CallerProfile.addTotalSamples(Count);
-        CallerProfile.addCalledTargetSamples(
-            CallerIndex, 0, FunctionProfile.getContext().getName(), Count);
+        CallerProfile.addCalledTargetSamples(CallerIndex, 0,
+                                             ContextNode->getFuncName(), Count);
       }
     }
   }
 
   // Assign zero count for remaining probes without sample hits to
   // differentiate from probes optimized away, of which the counts are unknown
   // and will be inferred by the compiler.
   for (auto &I : FrameSamples) {
for (const auto &BI : BranchCounter) {
     StringRef CalleeName = getCalleeNameForOffset(TargetOffset);
     if (CalleeName.size() == 0)
       continue;
     FunctionProfile.addCalledTargetSamples(CallProbe->getIndex(), 0, CalleeName,
                                            Count);
   }
 }
 
-FunctionSamples &CSProfileGenerator::getFunctionProfileForLeafProbe(
+ContextTrieNode *CSProfileGenerator::getContextNodeForLeafProbe(
     SampleContextFrames ContextStack, const MCDecodedPseudoProbe *LeafProbe) {
 
   // Explicitly copy the context for appending the leaf context
   SampleContextFrameVector NewContextStack(ContextStack.begin(),
                                            ContextStack.end());
   Binary->getInlineContextForProbe(LeafProbe, NewContextStack, true);
   // For leaf inlined context with the top frame, we should strip off the top
   // frame's probe id, like:
   // Inlined stack: [foo:1, bar:2], the ContextId will be "foo:1 @ bar"
   auto LeafFrame = NewContextStack.back();
   LeafFrame.Location = LineLocation(0, 0);
   NewContextStack.pop_back();
   // Compress the context string except for the leaf frame
   CSProfileGenerator::compressRecursionContext(NewContextStack);
   CSProfileGenerator::trimContext(NewContextStack);
   NewContextStack.push_back(LeafFrame);
 
   const auto *FuncDesc = Binary->getFuncDescForGUID(LeafProbe->getGuid());
   bool WasLeafInlined = LeafProbe->getInlineTreeNode()->hasInlineSite();
-  FunctionSamples &FunctionProile =
-      getFunctionProfileForContext(NewContextStack, WasLeafInlined);
-  FunctionProile.setFunctionHash(FuncDesc->FuncHash);
-  return FunctionProile;
+  ContextTrieNode *ContextNode =
+      getOrCreateContextNode(NewContextStack, WasLeafInlined);
+  ContextNode->getFunctionSamples()->setFunctionHash(FuncDesc->FuncHash);
+  return ContextNode;
+}
+
+FunctionSamples &CSProfileGenerator::getFunctionProfileForLeafProbe(
+    SampleContextFrames ContextStack, const MCDecodedPseudoProbe *LeafProbe) {
+  return *getContextNodeForLeafProbe(ContextStack, LeafProbe)
+              ->getFunctionSamples();
 }
 
 } // end namespace sampleprof
 } // end namespace llvm