[CSSPGO] split context string - compiler changes

[hoy]

Summary

This is the compiler part of the context split work. Please refer to D107299 for the original whole patch.

[wmi]

LGTM, thanks.

[wenlei]

Thanks for working on split context representation. This really makes context first class citizen in the framework, and the speed up it brought us is also very nice.

llvm/include/llvm/ProfileData/SampleProf.h

406

nit: having the word Type in type name does not add value, and the class also doesn't have anything specific for sample. Suggest we call it CallSiteLocation.

There're a few other others with Type in its name too.

423

LineOffset 0 can be a valid call site location given this is offset not absolute line number.

443–444

SampleContextRefType makes people think of SampleContext&, how about SampleContextStorageType -> SampleContextFrameVector, SampleContextRefType -> SampleContextFrames?

Accordingly SampleContext::getFullContext->SampleContext::getContextFrames.

449–458

This needs update.

One thing I was looking for a clarification from comments somewhere but didn't find it: whether context-less profile would have a non-empty FullContext which only contains leaf? From the code it seems answer is no? But then this is not consistent since FullContext of CS profile contains leaf frame.

Or would it be possible to remove leaf frame in FullContext for CS profile too? leaf doesn't need call site location either.

551–552

Can we remove this and replace all its call with getName?

575

Use std::ostringstream to help repeated string append?

581

Seems std::string toString() const is the conventional name for string converter.

592

Sanity check/assert to make sure the new name does not conflict leaf frame in FullContext? Or is this only allowed for non-CS case?

598

Should we assert that CState is not unknown? i.e. even if Context only has one frame, this is still considered context profile (something like [foo] before, as opposed to foo).

648

Clarify in comment whether this is empty for non-CS profile.

648–655

This sequence of copy, edit and compare could be simplified?

if (ThisContext.back().CallerName != ThatContext.back().CallerName)
  return false;

Even better, would it be possible to canonicalize context, so location of leaf frame is always (0, 0)?

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

93–107

This is a comparer only as it doesn't provide sorting directly. So ProfileComparer?

llvm/lib/Transforms/IPO/SampleContextTracker.cpp

67

One step further, ContextFramesToRemove?

307–309

Curious what triggered this order change?

[hoy]

llvm/include/llvm/ProfileData/SampleProf.h

406

CallSiteLocation sounds good.

423

Good catch. Switched to using INT32_MAX as an invalid line offset.

443–444

Sounds good.

449–458

The new context representation should be consistent with the old one for CS profile, i.e, it can contain a leaf frame only. The sample reader creates such leaf-only context. For non-CS profile, the context string will be empty.

Excluding leaf frame from FullContext should also work and it is cleaner by design. But the current way is more natural to implement, for example the IsPrefixOf operation, and some operations in llvm-profgen that appends the leaf probe into an existing calling context. Also the constructor takes in an arrayRef of the underlying vector, which, if separated, will need some special handling in the reader.

551–552

Yes, done.

575

Done.

581

Changed to toString

592

Assert added. This supposed to use by non-CS case.

598

Sounds good, assert added.

648

Done.

648–655

The API can be used to compare A:1 @ B with A:1 @ B:2 @ C, where we have to drop the line number of B:2.

The location of leaf frame was (0,0) previously, now I'm changing it to (-1,0). But the that context can still have a valid line number like B:2 where we have to drop the line number for comparision.

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

93–107

Sounds good.

llvm/lib/Transforms/IPO/SampleContextTracker.cpp

67

Renamed to ContextFramesToRemove.

307–309

The context of an inlined or merged node may be out-of-sync with its tree path so getContextFor will AV in that case. This can happen when functions on a tree path are not processed in top-down order, due to recursions.

[hoy]

Addressing Wenlei's comment.

[hoy]

Updating D108433: [CSSPGO] split context string - compiler changes

[hoy]

Updating D108433: [CSSPGO] split context string - compiler changes

[wenlei]

llvm/include/llvm/ProfileData/SampleProf.h

449–458

makes sense, thanks for clarification. though it is indeed inconsistent with non-CS profile in the sense that FullContext doesn't contain leaf for non-CS, but contain leaf for CS profile.

648–655

Ok, if we have to ignore the location part of B, can we simplified the copy, edit and compare sequence as suggested above still?

(maybe i'm missing something, but it looks to me that you answered the 2nd question only?)

[hoy]

llvm/include/llvm/ProfileData/SampleProf.h

648–655

Do you mean something like below?

if (ThisContext.back().CallerName != ThatContext.back().CallerName)
  return false;

if (ThisContext.back().hasLineLocation()) 
   return ThisContext.back().linelocation == ThatContext.back().linelocation 
   
return true;

[hoy]

llvm/include/llvm/ProfileData/SampleProf.h

648–655

return true --> return ThisContext.drop_back() == ThatContext.drop_back()

[wenlei]

llvm/include/llvm/ProfileData/SampleProf.h

648–655

Yes if actually expect ThisContext (the prefix) to have location for its leaf (i thought the prefix's leaf won't have location as shown in your example), otherwise we could omit the location compare too.

[hoy]

llvm/include/llvm/ProfileData/SampleProf.h

648–655

That's true with the current uses of this function. I'm adding an assert for that and removing the line location check. Supporting the location compare is slow, especially this function is frequently used in profile reading.

[modimo]

llvm/include/llvm/ProfileData/SampleProf.h

543

The implicit operator here seems unnecessary given you have getContextFrames below providing the same thing. Is this needed?

612

The == and != operators are not always opposites. If another instance of this object has the same State/Name/FullContext == and != will both evaluate as true. Should != check on all of these fields/be a direct negation of ==?

[hoy]

Updating D108433: [CSSPGO] split context string - compiler changes

[hoy]

llvm/include/llvm/ProfileData/SampleProf.h

543

It's needed for an implicit type cast where a SampleContext object is used as a SampleContextFrames object.

612

Can you elaborate more? Why would != return true for two objects having the same State/Name/FullContext?

[wenlei]

llvm/include/llvm/ProfileData/SampleProf.h

426

This one can now be removed as well.

[hoy]

Removing hasLineLocation.

[wlei]

llvm/include/llvm/ProfileData/SampleProf.h

695

is this more heavy operator than direct implementation of >(like <).

say if two contexts are very long and only the last frame is different.

it will take 2 * n for (*this != That) plus !(*this < That), rather than n.

[hoy]

llvm/include/llvm/ProfileData/SampleProf.h

423

On the second thought, sticking with LineOffset 0 for leaf frame is more natural. We have a bench places in the compiler/reader/writer/llvm-profgen that uses 0 as leaf line location.

By default, SampleContext should represent a valid call context with the leaf frame, so most of the time we should just automatically strip off the leaf line location., such as when doing prefix check, getting the full context string. However there is a case in llvm-profgen where SampleContext is used to represent intermediate context for pseudo probes when the leaf probe is not ready. This causes diversion, especially in the full context string printing. Given that's the only place we diverge, I'm adding a flag to getContextString to tolerate that, instead of introducing the hasLineLocation API.

So with the latest change, by default, we should assume the last element of the context vector is the leaf frame, except in the special llvm-profgen case. Does that make sense?

[wenlei]

llvm/include/llvm/ProfileData/SampleProf.h

26

I don't see DJB hash used anywhere, remove as well?

llvm/lib/Transforms/IPO/SampleContextTracker.cpp

307–309

Could you share a concrete example for such case? I was expecting context to be in-sync with its position in trie, because when we promote a node in trie, we also adjust its context accordingly. Does it also happen before this change?

[wenlei]

llvm/include/llvm/ProfileData/SampleProf.h

423

I think as long as we don't use 0 line offset to indicate no line info, it's fine.

Actually if we use CallSiteLocation to represent leaf, this really isn't a call site (call site always has a line location), but rather a frame. The comment is also inaccurate: "Represent a callsite with caller function name and line location".

We have typedef std::pair<std::string, LineLocation> FrameLocation; in llvm-profgen, once we have string buffers owning the underlying strings, that FrameLocation can be merged too.

Btw, I hope we can unify string converter to be toString. Right now we have CallSiteLocation::str(), SampleContext::getContextString() and SampleContext::toString.

SampleContext::toString properly dispatches for CS and non-CS case, but SampleContext::getContextString() doesn't. How about merge the two and let toString take a boolean?

[wlei]

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

98

I remember you mentioned this is the bottleneck, so how about adding another check to filter more cases, specifically use the frame size(or context string size). like the above code.
I'm also curious why here it changed to use std::set instead of SmallVector, I guess it will change to semantic of the original code if the context has two same frame function name.

[hoy]

llvm/include/llvm/ProfileData/SampleProf.h

26

Removed.

423

Right CallSiteLocation is really for a callsite, we just take advantage of it to represent a leaf frame in SampleContext, and we intentionally ignore the line location for leaf frame. This isn't ideal, but easier to implement. Alternatively, the LineLocation field can be made an Optional field, but that adds the complexity most of the time we don't need.

The FrameLocation type is being replaced by D108437.

I'm renaming CallSiteLocation::str() as toString. SampleContext::getContextString() is a static function that works on an ArrayRef object. This is mostly for the sake of intermediate contexts (without leaf probe) in llvm-profgen. SampleContext::toString is not supposed to work with intermediate contexts. I can move getContextString out of SampleContext if it is confusing here.

695

Good point. I just removed the > operator which is only used by sorting profiles for dumping. The sorting will use < instead.

llvm/lib/Transforms/IPO/SampleContextTracker.cpp

307–309

Unfortunately I don't seem to find that extreme case after adjusting the context promotion order, i.e, with a comparer. On the other hand, if we are not going to do anything for an inlined or merged context, we can bail out earlier without taking time to find its context node?

[hoy]

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

98

This is not the bottleneck. The bottleneck is in sample reader when trying to load contexts out of the current module. There we sort all contexts in the profile first and selectively load the contexts which appear to be a subtree of a given context.

Using ordered set here is to enforce an order of promoting the contexts of a function later. We have seen cases that different promoting order led to different codegen, especially when recursion is involved.

[hoy]

Addressing comments.

[wenlei]

llvm/include/llvm/ProfileData/SampleProf.h

423

Ok, keeping static function getContextString is fine.

Given we have SampleContextFrames, SampleContextFrameVector, perhaps this CallSiteLocation should be renamed to ContextFrame to reflect what it really is?

llvm/lib/Transforms/IPO/SampleContextTracker.cpp

307–309

Avoid context trie look up sounds good. Just wanted to make sure we don't have lurking issues as I don't see how they can be out of sync - it could be some other issues.

[hoy]

llvm/include/llvm/ProfileData/SampleProf.h

423

ContextFrame sounds good. How about SampleContextFrame?

[wenlei]

llvm/include/llvm/ProfileData/SampleProf.h

423

Sounds good too, though it may look too similar to SampleContextFrames, which can cause confusion. But either way is fine.

[hoy]

Renaming CallsiteLocation to SampleContextFrame.

[wenlei]

lgtm, thanks

[wlei]

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

98

I see, thanks for the clarification.

[modimo]

llvm/include/llvm/ProfileData/SampleProf.h

612

Ah I missed that it was comparing the dereferenced this which is invoking the == operator directly. Thought it was comparing the pointers.

[hoy]

Updating D108433: [CSSPGO] split context string - compiler changes

llvm/include/llvm/ProfileData/SampleProf.h

 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/StringSet.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/GlobalValue.h"
 #include "llvm/IR/Module.h"
 #include "llvm/Support/Allocator.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/ErrorOr.h"
 #include "llvm/Support/MathExtras.h"
 #include "llvm/Support/raw_ostream.h"
 #include <algorithm>
 #include <cstdint>
+#include <list>
 #include <map>
 #include <set>
+#include <sstream>
 #include <string>
 #include <system_error>
+#include <unordered_map>
 #include <utility>
 
 namespace llvm {
 
 const std::error_category &sampleprof_category();
 
 enum class sampleprof_error {
   success = 0,

 // value of enum. Only append new ones.
 enum SecType {
   SecInValid = 0,
   SecProfSummary = 1,
   SecNameTable = 2,
   SecProfileSymbolList = 3,
   SecFuncOffsetTable = 4,
   SecFuncMetadata = 5,
+  SecCSNameTable = 6,
   // marker for the first type of profile.
   SecFuncProfileFirst = 32,
   SecLBRProfile = SecFuncProfileFirst
 };
 
 static inline std::string getSecName(SecType Type) {
   switch (Type) {
   case SecInValid:
     return "InvalidSection";
   case SecProfSummary:
     return "ProfileSummarySection";
   case SecNameTable:
     return "NameTableSection";
   case SecProfileSymbolList:
     return "ProfileSymbolListSection";
   case SecFuncOffsetTable:
     return "FuncOffsetTableSection";
   case SecFuncMetadata:
     return "FunctionMetadata";
+  case SecCSNameTable:
+    return "CSNameTableSection";
   case SecLBRProfile:
     return "LBRProfileSection";
   }
   llvm_unreachable("A SecType has no name for output");
 }
 
 // Entry type of section header table used by SampleProfileExtBinaryBaseReader
 // and SampleProfileExtBinaryBaseWriter.

 
 // Attribute of context associated with FunctionSamples
 enum ContextAttributeMask {
   ContextNone = 0x0,
   ContextWasInlined = 0x1,      // Leaf of context was inlined in previous build
   ContextShouldBeInlined = 0x2, // Leaf of context should be inlined
 };
 
+// Represent a callsite with caller function name and line location
+struct SampleContextFrame {
+  StringRef CallerName;
+  LineLocation Callsite;
+
+  SampleContextFrame() : Callsite(0, 0) {}
+
+  SampleContextFrame(StringRef CallerName, LineLocation Callsite)
+      : CallerName(CallerName), Callsite(Callsite) {}
+
+  bool operator==(const SampleContextFrame &That) const {
+    return Callsite == That.Callsite && CallerName == That.CallerName;
+  }
+
+  bool operator!=(const SampleContextFrame &That) const {
+    return !(*this == That);
+  }
+
+  std::string toString(bool OutputLineLocation) const {
+    std::ostringstream OContextStr;
+    OContextStr << CallerName.str();
+    if (OutputLineLocation) {
+      OContextStr << ":" << Callsite.LineOffset;
+      if (Callsite.Discriminator)
+        OContextStr << "." << Callsite.Discriminator;
+    }
+    return OContextStr.str();
+  }
+};
+
+static inline hash_code hash_value(const SampleContextFrame &arg) {
+  return hash_combine(arg.CallerName, arg.Callsite.LineOffset,
+                      arg.Callsite.Discriminator);
+}
+
+using SampleContextFrameVector = SmallVector<SampleContextFrame, 10>;
+using SampleContextFrames = ArrayRef<SampleContextFrame>;
+
+struct SampleContextFrameHash {
+  uint64_t operator()(const SampleContextFrameVector &S) const {
+    return hash_combine_range(S.begin(), S.end());
+  }
+};
+
 // Sample context for FunctionSamples. It consists of the calling context,
 // the function name and context state. Internally sample context is represented
-// using StringRef, which is also the input for constructing a `SampleContext`.
+// using ArrayRef, which is also the input for constructing a `SampleContext`.
 // It can accept and represent both full context string as well as context-less
 // function name.
-// Example of full context string (note the wrapping `[]`):
-//    `[main:3 @ _Z5funcAi:1 @ _Z8funcLeafi]`
-// Example of context-less function name (same as AutoFDO):
-//    `_Z8funcLeafi`
+// For a CS profile, a full context vector can look like:
+//    `main:3 _Z5funcAi:1 _Z8funcLeafi`
+// For a base CS profile without calling context, the context vector should only
+// contain the leaf frame name.
+// For a non-CS profile, the context vector should be empty.
 class SampleContext {
 public:
   SampleContext() : State(UnknownContext), Attributes(ContextNone) {}
-  SampleContext(StringRef ContextStr, ContextStateMask CState = UnknownContext)
-      : Attributes(ContextNone) {
-    setContext(ContextStr, CState);
-  }
 
-  // Promote context by removing top frames (represented by `ContextStrToRemove`).
-  // Note that with string representation of context, the promotion is effectively
-  // a substr operation with `ContextStrToRemove` removed from left.
-  void promoteOnPath(StringRef ContextStrToRemove) {
-    assert(FullContext.startswith(ContextStrToRemove));
+  SampleContext(StringRef Name)
+      : Name(Name), State(UnknownContext), Attributes(ContextNone) {}
 
-    // Remove leading context and frame separator " @ ".
-    FullContext = FullContext.substr(ContextStrToRemove.size() + 3);
-    CallingContext = CallingContext.substr(ContextStrToRemove.size() + 3);
+  SampleContext(SampleContextFrames Context,
+                ContextStateMask CState = RawContext)
+      : Attributes(ContextNone) {
+    assert(!Context.empty() && "Context is empty");
+    setContext(Context, CState);
   }
 
-  // Split the top context frame (left-most substr) from context.
-  static std::pair<StringRef, StringRef>
-  splitContextString(StringRef ContextStr) {
-    return ContextStr.split(" @ ");
+  // Give a context string, decode and populate internal states like
+  // Function name, Calling context and context state. Example of input
+  // `ContextStr`: `[main:3 @ _Z5funcAi:1 @ _Z8funcLeafi]`
+  SampleContext(StringRef ContextStr,
+                std::list<SampleContextFrameVector> &CSNameTable,
+                ContextStateMask CState = RawContext)
+      : Attributes(ContextNone) {
+    assert(!ContextStr.empty());
+    // Note that `[]` wrapped input indicates a full context string, otherwise
+    // it's treated as context-less function name only.
+    bool HasContext = ContextStr.startswith("[");
+    if (!HasContext) {
+      State = UnknownContext;
+      Name = ContextStr;
+    } else {
+      // Remove encapsulating '[' and ']' if any
+      ContextStr = ContextStr.substr(1, ContextStr.size() - 2);
+      CSNameTable.emplace_back();
+      SampleContextFrameVector &Context = CSNameTable.back();
+      /// Create a context vector from a given context string and save it in
+      /// `Context`.
+      StringRef ContextRemain = ContextStr;
+      StringRef ChildContext;
+      StringRef CalleeName;
+      while (!ContextRemain.empty()) {
+        auto ContextSplit = ContextRemain.split(" @ ");
+        ChildContext = ContextSplit.first;
+        ContextRemain = ContextSplit.second;
+        LineLocation CallSiteLoc(0, 0);
+        decodeContextString(ChildContext, CalleeName, CallSiteLoc);
+        Context.emplace_back(CalleeName, CallSiteLoc);
+      }
+
+      setContext(Context, CState);
     }
-
-  // Reconstruct a new context with the last k frames, return the context-less
-  // name if K = 1
-  StringRef getContextWithLastKFrames(uint32_t K) {
-    if (K == 1)
-      return getNameWithoutContext();
-
-    size_t I = FullContext.size();
-    while (K--) {
-      I = FullContext.find_last_of(" @ ", I);
-      if (I == StringRef::npos)
-        return FullContext;
-      I -= 2;
-    }
-    return FullContext.slice(I + 3, StringRef::npos);
+
+  // 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(':');
if (!EntrySplit.second.empty()) {
       LineLoc.LineOffset = LineOffset;
 
       // Get discriminator
       if (!LocSplit.second.empty())
         LocSplit.second.getAsInteger(10, LineLoc.Discriminator);
     }
   }
 
-  operator StringRef() const { return FullContext; }
+  operator SampleContextFrames() const { return FullContext; }
   bool hasAttribute(ContextAttributeMask A) { return Attributes & (uint32_t)A; }
   void setAttribute(ContextAttributeMask A) { Attributes |= (uint32_t)A; }
   uint32_t getAllAttributes() { return Attributes; }
   void setAllAttributes(uint32_t A) { Attributes = A; }
   bool hasState(ContextStateMask S) { return State & (uint32_t)S; }
   void setState(ContextStateMask S) { State |= (uint32_t)S; }
   void clearState(ContextStateMask S) { State &= (uint32_t)~S; }
   bool hasContext() const { return State != UnknownContext; }
-  bool isBaseContext() const { return CallingContext.empty(); }
-  StringRef getNameWithoutContext() const { return Name; }
-  StringRef getCallingContext() const { return CallingContext; }
-  StringRef getNameWithContext() const { return FullContext; }
+  bool isBaseContext() const { return FullContext.size() == 1; }
+  StringRef getName() const { return Name; }
+  SampleContextFrames getContextFrames() const { return FullContext; }
 
-private:
-  // Give a context string, decode and populate internal states like
-  // Function name, Calling context and context state. Example of input
-  // `ContextStr`: `[main:3 @ _Z5funcAi:1 @ _Z8funcLeafi]`
-  void setContext(StringRef ContextStr, ContextStateMask CState) {
-    assert(!ContextStr.empty());
-    // Note that `[]` wrapped input indicates a full context string, otherwise
-    // it's treated as context-less function name only.
-    bool HasContext = ContextStr.startswith("[");
-    if (!HasContext && CState == UnknownContext) {
-      State = UnknownContext;
-      Name = FullContext = ContextStr;
-    } else {
-      // Assume raw context profile if unspecified
-      if (CState == UnknownContext)
-        State = RawContext;
-      else
-        State = CState;
+  static std::string getContextString(SampleContextFrames Context,
+                                      bool IncludeLeafLineLocation = false) {
+    std::ostringstream OContextStr;
+    for (uint32_t I = 0; I < Context.size(); I++) {
+      if (OContextStr.str().size()) {
+        OContextStr << " @ ";
+      }
+      OContextStr << Context[I].toString(I != Context.size() - 1 ||
+                                         IncludeLeafLineLocation);
+    }
+    return OContextStr.str();
+  }
+
+  std::string toString() const {
+    if (!hasContext())
+      return Name.str();
+    return getContextString(FullContext, false);
+  }
+
+  uint64_t getHashCode() const {
+    return hasContext() ? hash_value(getContextFrames())
+                        : hash_value(getName());
+  }
+
+  /// Set the name of the function.
+  void setName(StringRef FunctionName) {
+    assert(FullContext.empty() &&
+           "setName should only be called for non-CS profile");
+    Name = FunctionName;
+  }
+
+  void setContext(SampleContextFrames Context,
+                  ContextStateMask CState = RawContext) {
+    assert(CState != UnknownContext);
+    FullContext = Context;
+    Name = Context.back().CallerName;
+  }
 
-      // Remove encapsulating '[' and ']' if any
-      if (HasContext)
-        FullContext = ContextStr.substr(1, ContextStr.size() - 2);
-      else
-        FullContext = ContextStr;
-
-      // Caller is to the left of callee in context string
-      auto NameContext = FullContext.rsplit(" @ ");
-      if (NameContext.second.empty()) {
-        Name = NameContext.first;
-        CallingContext = NameContext.second;
-      } else {
-        Name = NameContext.second;
-        CallingContext = NameContext.first;
-      }
+  bool operator==(const SampleContext &That) const {
+    return State == That.State && Name == That.Name &&
+           FullContext == That.FullContext;
+  }
+
+  bool operator!=(const SampleContext &That) const { return !(*this == That); }
+
+  bool operator<(const SampleContext &That) const {
+    if (State != That.State)
+      return State < That.State;
+
+    if (!hasContext()) {
+      return (Name.compare(That.Name)) == -1;
+    }
+
+    uint64_t I = 0;
+    while (I < std::min(FullContext.size(), That.FullContext.size())) {
+      auto &Context1 = FullContext[I];
+      auto &Context2 = That.FullContext[I];
+      auto V = Context1.CallerName.compare(Context2.CallerName);
+      if (V)
+        return V == -1;
+      if (Context1.Callsite != Context2.Callsite)
+        return Context1.Callsite < Context2.Callsite;
+      I++;
+    }
+
+    return FullContext.size() < That.FullContext.size();
+
+  struct Hash {
+    uint64_t operator()(const SampleContext &Context) const {
+      return Context.getHashCode();
     }
+  };
+
+  bool IsPrefixOf(const SampleContext &That) {
+    auto ThisContext = FullContext;
+    auto ThatContext = That.FullContext;
+    if (ThatContext.size() < ThisContext.size())
+      return false;
+    ThatContext = ThatContext.take_front(ThisContext.size());
+    // Compare Leaf frame first
+    if (ThisContext.back().CallerName != ThatContext.back().CallerName)
+      return false;
+    // Compare leading context
+    return ThisContext.drop_back() == ThatContext.drop_back();
   }
 
-  // Full context string including calling context and leaf function name
-  StringRef FullContext;
-  // Function name for the associated sample profile
+private:
+  /// Mangled name of the function.
   StringRef Name;
-  // Calling context (leaf function excluded) for the associated sample profile
-  StringRef CallingContext;
+  // 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;
 class SampleProfileReaderItaniumRemapper;
 
 using BodySampleMap = std::map<LineLocation, SampleRecord>;
 // NOTE: Using a StringMap here makes parsed profiles consume around 17% more
 // memory, which is *very* significant for large profiles.
 using FunctionSamplesMap = std::map<std::string, FunctionSamples, std::less<>>;
 using CallsiteSampleMap = std::map<LineLocation, FunctionSamplesMap>;
 
 /// Representation of the samples collected for a function.
 ///
 /// This data structure contains all the collected samples for the body
 /// of a function. Each sample corresponds to a LineLocation instance
 /// within the body of the function.
 class FunctionSamples {
 public:
   FunctionSamples() = default;
 
   void print(raw_ostream &OS = dbgs(), unsigned Indent = 0) const;
   void dump() const;
 
   sampleprof_error addTotalSamples(uint64_t Num, uint64_t Weight = 1) {
     bool Overflowed;
     TotalSamples =
         SaturatingMultiplyAdd(Num, Weight, TotalSamples, &Overflowed);
     return Overflowed ? sampleprof_error::counter_overflow
                       : sampleprof_error::success;
   }
 
   void setTotalSamples(uint64_t Num) { TotalSamples = Num; }
 
   sampleprof_error addHeadSamples(uint64_t Num, uint64_t Weight = 1) {
     bool Overflowed;
     TotalHeadSamples =
         SaturatingMultiplyAdd(Num, Weight, TotalHeadSamples, &Overflowed);
     return Overflowed ? sampleprof_error::counter_overflow
                       : sampleprof_error::success;
   }
 
   sampleprof_error addBodySamples(uint32_t LineOffset, uint32_t Discriminator,
                                   uint64_t Num, uint64_t Weight = 1) {
     return BodySamples[LineLocation(LineOffset, Discriminator)].addSamples(
         Num, Weight);
   }
for (const auto &C : getCallsiteSamples())
         MaxCount = std::max(MaxCount, F.second.getMaxCountInside());
     return MaxCount;
   }
 
   /// Merge the samples in \p Other into this one.
   /// Optionally scale samples by \p Weight.
   sampleprof_error merge(const FunctionSamples &Other, uint64_t Weight = 1) {
     sampleprof_error Result = sampleprof_error::success;
-    Name = Other.getName();
     if (!GUIDToFuncNameMap)
       GUIDToFuncNameMap = Other.GUIDToFuncNameMap;
-    if (Context.getNameWithContext().empty())
+    if (Context.getName().empty())
       Context = Other.getContext();
     if (FunctionHash == 0) {
       // Set the function hash code for the target profile.
       FunctionHash = Other.getFunctionHash();
     } else if (FunctionHash != Other.getFunctionHash()) {
       // The two profiles coming with different valid hash codes indicates
       // either:
       // 1. They are same-named static functions from different compilation
void findInlinedFunctions(DenseSet<GlobalValue::GUID> &S,
                             uint64_t Threshold) const {
     if (TotalSamples <= Threshold)
       return;
     auto isDeclaration = [](const Function *F) {
       return !F || F->isDeclaration();
     };
     if (isDeclaration(SymbolMap.lookup(getFuncName()))) {
       // Add to the import list only when it's defined out of module.
-      S.insert(getGUID(Name));
+      S.insert(getGUID(getName()));
     }
     // Import hot CallTargets, which may not be available in IR because full
     // profile annotation cannot be done until backend compilation in ThinLTO.
     for (const auto &BS : BodySamples)
       for (const auto &TS : BS.second.getCallTargets())
         if (TS.getValue() > Threshold) {
           const Function *Callee = SymbolMap.lookup(getFuncName(TS.getKey()));
           if (isDeclaration(Callee))
             S.insert(getGUID(TS.getKey()));
         }
     for (const auto &CS : CallsiteSamples)
       for (const auto &NameFS : CS.second)
         NameFS.second.findInlinedFunctions(S, SymbolMap, Threshold);
   }
 
   /// Set the name of the function.
-  void setName(StringRef FunctionName) { Name = FunctionName; }
+  void setName(StringRef FunctionName) { Context.setName(FunctionName); }
 
   /// Return the function name.
-  StringRef getName() const { return Name; }
-
-  /// Return function name with context.
-  StringRef getNameWithContext() const {
-    return FunctionSamples::ProfileIsCS ? Context.getNameWithContext() : Name;
-  }
+  StringRef getName() const { return Context.getName(); }
 
   /// Return the original function name.
-  StringRef getFuncName() const { return getFuncName(Name); }
+  StringRef getFuncName() const { return getFuncName(getName()); }
 
   void setFunctionHash(uint64_t Hash) { FunctionHash = Hash; }
 
   uint64_t getFunctionHash() const { return FunctionHash; }
 
   /// Return the canonical name for a function, taking into account
   /// suffix elision policy attributes.
   static StringRef getCanonicalFnName(const Function &F) {
static uint64_t getGUID(StringRef Name) {
     return UseMD5 ? std::stoull(Name.data()) : Function::getGUID(Name);
   }
 
   // Find all the names in the current FunctionSamples including names in
   // all the inline instances and names of call targets.
   void findAllNames(DenseSet<StringRef> &NameSet) const;
 
 private:
-  /// Mangled name of the function.
-  StringRef Name;
-
   /// CFG hash value for the function.
   uint64_t FunctionHash = 0;
 
   /// Calling context for function profile
   mutable SampleContext Context;
 
   /// Total number of samples collected inside this function.
   ///
private:
   /// map will contain two entries.  One for all the samples collected
   /// in the call to bar() at line offset 1, the other for all the samples
   /// collected in the call to baz() at line offset 8.
   CallsiteSampleMap CallsiteSamples;
 };
 
 raw_ostream &operator<<(raw_ostream &OS, const FunctionSamples &FS);
 
-using NameFunctionSamples = std::pair<StringRef, const FunctionSamples *>;
+using SampleProfileMap =
+    std::unordered_map<SampleContext, FunctionSamples, SampleContext::Hash>;
 
-void sortFuncProfiles(const StringMap<FunctionSamples> &ProfileMap,
+using NameFunctionSamples = std::pair<SampleContext, const FunctionSamples *>;
+
+void sortFuncProfiles(const SampleProfileMap &ProfileMap,
                       std::vector<NameFunctionSamples> &SortedProfiles);
 
 /// Sort a LocationT->SampleT map by LocationT.
 ///
 /// It produces a sorted list of <LocationT, SampleT> records by ascending
 /// order of LocationT.
 template <class LocationT, class SampleT> class SampleSorter {
 public:
private:
   SamplesWithLocList V;
 };
 
 /// SampleContextTrimmer impelements helper functions to trim, merge cold
 /// context profiles. It also supports context profile canonicalization to make
 /// sure ProfileMap's key is consistent with FunctionSample's name/context.
 class SampleContextTrimmer {
 public:
-  SampleContextTrimmer(StringMap<FunctionSamples> &Profiles)
-      : ProfileMap(Profiles){};
+  SampleContextTrimmer(SampleProfileMap &Profiles) : ProfileMap(Profiles){};
   // Trim and merge cold context profile when requested.
   void trimAndMergeColdContextProfiles(uint64_t ColdCountThreshold,
                                        bool TrimColdContext,
                                        bool MergeColdContext,
                                        uint32_t ColdContextFrameLength);
   // Canonicalize context profile name and attributes.
   void canonicalizeContextProfiles();
 
 private:
-  StringMap<FunctionSamples> &ProfileMap;
+  SampleProfileMap &ProfileMap;
 };
 
 /// ProfileSymbolList records the list of function symbols shown up
 /// in the binary used to generate the profile. It is useful to
 /// to discriminate a function being so cold as not to shown up
 /// in the profile and a function newly added.
 class ProfileSymbolList {
 public:
private:
   // writing it into profile. The variable is unused when the symbol
   // list is read from an existing profile.
   bool ToCompress = false;
   DenseSet<StringRef> Syms;
   BumpPtrAllocator Allocator;
 };
 
 } // end namespace sampleprof
+
+using namespace sampleprof;
+// Provide DenseMapInfo for SampleContext.
+template <> struct DenseMapInfo<SampleContext> {
+  static inline SampleContext getEmptyKey() { return SampleContext(); }
+
+  static inline SampleContext getTombstoneKey() { return SampleContext("@"); }
+
+  static unsigned getHashValue(const SampleContext &Val) {
+    return Val.getHashCode();
+  }
+
+  static bool isEqual(const SampleContext &LHS, const SampleContext &RHS) {
+    return LHS == RHS;
+  }
+};
 } // end namespace llvm
 
 #endif // LLVM_PROFILEDATA_SAMPLEPROF_H

llvm/include/llvm/Transforms/IPO/ProfiledCallGraph.h

struct ProfiledCallGraphNode {
   std::set<ProfiledCallGraphNode *, ProfiledCallGraphNodeComparer> Callees;
 };
 
 class ProfiledCallGraph {
 public:
   using iterator = std::set<ProfiledCallGraphNode *>::iterator;
 
   // Constructor for non-CS profile.
-  ProfiledCallGraph(StringMap<FunctionSamples> &ProfileMap) {
+  ProfiledCallGraph(SampleProfileMap &ProfileMap) {
     assert(!FunctionSamples::ProfileIsCS && "CS profile is not handled here");
     for (const auto &Samples : ProfileMap) {
       addProfiledCalls(Samples.second);
     }
   }
 
   // Constructor for CS profile.
   ProfiledCallGraph(SampleContextTracker &ContextTracker) {

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

ContextTrieNode *getChildContext(const LineLocation &CallSite,
                                    StringRef CalleeName);
   ContextTrieNode *getHottestChildContext(const LineLocation &CallSite);
   ContextTrieNode *getOrCreateChildContext(const LineLocation &CallSite,
                                            StringRef CalleeName,
                                            bool AllowCreate = true);
 
   ContextTrieNode &moveToChildContext(const LineLocation &CallSite,
                                       ContextTrieNode &&NodeToMove,
-                                      StringRef ContextStrToRemove,
+                                      uint32_t ContextFramesToRemove,
                                       bool DeleteNode = true);
   void removeChildContext(const LineLocation &CallSite, StringRef CalleeName);
   std::map<uint32_t, ContextTrieNode> &getAllChildContext();
   StringRef getFuncName() const;
   FunctionSamples *getFunctionSamples() const;
   void setFunctionSamples(FunctionSamples *FSamples);
   LineLocation getCallSiteLoc() const;
   ContextTrieNode *getParentContext() const;

 // 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:
-  using ContextSamplesTy = SmallVector<FunctionSamples *, 16>;
+  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>;
 
-  SampleContextTracker(StringMap<FunctionSamples> &Profiles);
+  SampleContextTracker(SampleProfileMap &Profiles);
   // 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
private:
   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,
-                        StringRef ContextStrToRemove);
-  ContextTrieNode &promoteMergeContextSamplesTree(ContextTrieNode &FromNode,
-                                                  ContextTrieNode &ToNodeParent,
-                                                  StringRef ContextStrToRemove);
+                        uint32_t ContextFramesToRemove);
+  ContextTrieNode &
+  promoteMergeContextSamplesTree(ContextTrieNode &FromNode,
+                                 uint32_t ContextFramesToRemove);
 
   // Map from function name to context profiles (excluding base profile)
   StringMap<ContextSamplesTy> FuncToCtxtProfiles;
 
   // Root node for context trie tree
   ContextTrieNode RootContext;
 };
 
 } // end namespace llvm
 #endif // LLVM_TRANSFORMS_IPO_SAMPLECONTEXTTRACKER_H

llvm/lib/ProfileData/ProfileSummaryBuilder.cpp

 std::unique_ptr<ProfileSummary> SampleProfileSummaryBuilder::getSummary() {
   computeDetailedSummary();
   return std::make_unique<ProfileSummary>(
       ProfileSummary::PSK_Sample, DetailedSummary, TotalCount, MaxCount, 0,
       MaxFunctionCount, NumCounts, NumFunctions);
 }
 
 std::unique_ptr<ProfileSummary>
 SampleProfileSummaryBuilder::computeSummaryForProfiles(
-    const StringMap<sampleprof::FunctionSamples> &Profiles) {
+    const SampleProfileMap &Profiles) {
   assert(NumFunctions == 0 &&
          "This can only be called on an empty summary builder");
-  StringMap<sampleprof::FunctionSamples> ContextLessProfiles;
-  const StringMap<sampleprof::FunctionSamples> *ProfilesToUse = &Profiles;
+  sampleprof::SampleProfileMap ContextLessProfiles;
+  const sampleprof::SampleProfileMap *ProfilesToUse = &Profiles;
   // For CSSPGO, context-sensitive profile effectively split a function profile
   // into many copies each representing the CFG profile of a particular calling
   // context. That makes the count distribution looks more flat as we now have
   // more function profiles each with lower counts, which in turn leads to lower
   // hot thresholds. To compensate for that, by defauly we merge context
-  // profiles before coumputing profile summary.
+  // profiles before computing profile summary.
   if (UseContextLessSummary || (sampleprof::FunctionSamples::ProfileIsCS &&
                                 !UseContextLessSummary.getNumOccurrences())) {
     for (const auto &I : Profiles) {
       ContextLessProfiles[I.second.getName()].merge(I.second);
     }
     ProfilesToUse = &ContextLessProfiles;
   }
 

llvm/lib/ProfileData/SampleProf.cpp

 }
 
 raw_ostream &llvm::sampleprof::operator<<(raw_ostream &OS,
                                           const FunctionSamples &FS) {
   FS.print(OS);
   return OS;
 }
 
 void sampleprof::sortFuncProfiles(
     const StringMap<FunctionSamples> &ProfileMap,
     std::vector<NameFunctionSamples> &SortedProfiles) {
   for (const auto &I : ProfileMap) {
     assert(I.getKey() == I.second.getNameWithContext() &&
            "Inconsistent profile map");
     SortedProfiles.push_back(
         std::make_pair(I.second.getNameWithContext(), &I.second));
   }
   llvm::stable_sort(SortedProfiles, [](const NameFunctionSamples &A,
                                        const NameFunctionSamples &B) {
     if (A.second->getTotalSamples() == B.second->getTotalSamples())
-      return A.first > B.first;
+      return A.first < B.first;
     return A.second->getTotalSamples() > B.second->getTotalSamples();
   });
 }
 
 unsigned FunctionSamples::getOffset(const DILocation *DIL) {
   return (DIL->getLine() - DIL->getScope()->getSubprogram()->getLine()) &
       0xffff;
 }
const FunctionSamples *FunctionSamples::findFunctionSamples(
   const FunctionSamples *FS = this;
   for (int i = S.size() - 1; i >= 0 && FS != nullptr; i--) {
     FS = FS->findFunctionSamplesAt(S[i].first, S[i].second, Remapper);
   }
   return FS;
 }
 
 void FunctionSamples::findAllNames(DenseSet<StringRef> &NameSet) const {
-  NameSet.insert(Name);
+  NameSet.insert(getName());
   for (const auto &BS : BodySamples)
     for (const auto &TS : BS.second.getCallTargets())
       NameSet.insert(TS.getKey());
 
   for (const auto &CS : CallsiteSamples) {
     for (const auto &NameFS : CS.second) {
       NameSet.insert(NameFS.first);
       NameFS.second.findAllNames(NameSet);
void SampleContextTrimmer::trimAndMergeColdContextProfiles(
 
   // Filter the cold profiles from ProfileMap and move them into a tmp
   // container
   std::vector<std::pair<StringRef, const FunctionSamples *>> ColdProfiles;
   for (const auto &I : ProfileMap) {
     const FunctionSamples &FunctionProfile = I.second;
     if (FunctionProfile.getTotalSamples() >= ColdCountThreshold)
       continue;
     ColdProfiles.emplace_back(I.getKey(), &I.second);
   }
 
   // Remove the cold profile from ProfileMap and merge them into
   // MergedProfileMap by the last K frames of context
   StringMap<FunctionSamples> MergedProfileMap;
   for (const auto &I : ColdProfiles) {
     if (MergeColdContext) {
       auto Ret = MergedProfileMap.try_emplace(
           I.second->getContext().getContextWithLastKFrames(
               ColdContextFrameLength),
           FunctionSamples());
       FunctionSamples &MergedProfile = Ret.first->second;
       MergedProfile.merge(*I.second);
     }
     ProfileMap.erase(I.first);
   }
 
   // Move the merged profiles into ProfileMap;
   for (const auto &I : MergedProfileMap) {
     // Filter the cold merged profile
     if (TrimColdContext && I.second.getTotalSamples() < ColdCountThreshold &&
         ProfileMap.find(I.getKey()) == ProfileMap.end())
       continue;
     // Merge the profile if the original profile exists, otherwise just insert
     // as a new profile
     auto Ret = ProfileMap.try_emplace(I.getKey(), FunctionSamples());
     if (Ret.second) {
       SampleContext FContext(Ret.first->first(), RawContext);
       FunctionSamples &FProfile = Ret.first->second;
       FProfile.setContext(FContext);
       FProfile.setName(FContext.getNameWithoutContext());
     }
     FunctionSamples &OrigProfile = Ret.first->second;
     OrigProfile.merge(I.second);
   }
 }
 
 void SampleContextTrimmer::canonicalizeContextProfiles() {
   std::vector<StringRef> ProfilesToBeRemoved;
   StringMap<FunctionSamples> ProfilesToBeAdded;
   for (auto &I : ProfileMap) {
     FunctionSamples &FProfile = I.second;
     StringRef ContextStr = FProfile.getNameWithContext();
     if (I.first() == ContextStr)
       continue;
 
     // Use the context string from FunctionSamples to update the keys of
     // ProfileMap. They can get out of sync after context profile promotion
     // through pre-inliner.
     // Duplicate the function profile for later insertion to avoid a conflict
     // caused by a context both to be add and to be removed. This could happen
     // when a context is promoted to another context which is also promoted to
     // the third context. For example, given an original context A @ B @ C that
     // is promoted to B @ C and the original context B @ C which is promoted to
     // just C, adding B @ C to the profile map while removing same context (but
     // with different profiles) from the map can cause a conflict if they are
     // not handled in a right order. This can be solved by just caching the
     // profiles to be added.
     auto Ret = ProfilesToBeAdded.try_emplace(ContextStr, FProfile);
     (void)Ret;
     assert(Ret.second && "Context conflict during canonicalization");
     ProfilesToBeRemoved.push_back(I.first());
   }
 
   for (auto &I : ProfilesToBeRemoved) {
     ProfileMap.erase(I);
   }
 
   for (auto &I : ProfilesToBeAdded) {
     ProfileMap.try_emplace(I.first(), I.second);
   }
 }
 
 std::error_code ProfileSymbolList::write(raw_ostream &OS) {
   // Sort the symbols before output. If doing compression.
   // It will make the compression much more effective.
   std::vector<StringRef> SortedList(Syms.begin(), Syms.end());
   llvm::sort(SortedList);

llvm/lib/Transforms/IPO/SampleContextTracker.cpp

for (auto &It : AllChildContext) {
     }
   }
 
   return ChildNodeRet;
 }
 
 ContextTrieNode &ContextTrieNode::moveToChildContext(
     const LineLocation &CallSite, ContextTrieNode &&NodeToMove,
-    StringRef ContextStrToRemove, bool DeleteNode) {
+    uint32_t ContextFramesToRemove, bool DeleteNode) {
   uint32_t Hash = nodeHash(NodeToMove.getFuncName(), CallSite);
   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;
 
   // 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);
   NodeToUpdate.push(&NewNode);
 
   while (!NodeToUpdate.empty()) {
     ContextTrieNode *Node = NodeToUpdate.front();
     NodeToUpdate.pop();
     FunctionSamples *FSamples = Node->getFunctionSamples();
 
     if (FSamples) {
-      FSamples->getContext().promoteOnPath(ContextStrToRemove);
+      FSamples->getContext().promoteOnPath(ContextFramesToRemove);
       FSamples->getContext().setState(SyntheticContext);
-      LLVM_DEBUG(dbgs() << "  Context promoted to: " << FSamples->getContext()
-                        << "\n");
+      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);
     }
   }
ContextTrieNode *ContextTrieNode::getOrCreateChildContext(
   if (!AllowCreate)
     return nullptr;
 
   AllChildContext[Hash] = ContextTrieNode(this, CalleeName, nullptr, CallSite);
   return &AllChildContext[Hash];
 }
 
 // Profiler tracker than manages profiles and its associated context
-SampleContextTracker::SampleContextTracker(
-    StringMap<FunctionSamples> &Profiles) {
+SampleContextTracker::SampleContextTracker(SampleProfileMap &Profiles) {
   for (auto &FuncSample : Profiles) {
     FunctionSamples *FSamples = &FuncSample.second;
-    SampleContext Context(FuncSample.first(), RawContext);
-    LLVM_DEBUG(dbgs() << "Tracking Context for function: " << Context << "\n");
+    SampleContext Context = FuncSample.first;
+    LLVM_DEBUG(dbgs() << "Tracking Context for function: " << Context.toString()
+                      << "\n");
     if (!Context.isBaseContext())
-      FuncToCtxtProfiles[Context.getNameWithoutContext()].push_back(FSamples);
+      FuncToCtxtProfiles[Context.getName()].insert(FSamples);
     ContextTrieNode *NewNode = getOrCreateContextPath(Context, true);
     assert(!NewNode->getFunctionSamples() &&
            "New node can't have sample profile");
     NewNode->setFunctionSamples(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);
 
   // 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() << "\n";
+      dbgs() << "  Callee context found: " << FSamples->getContext().toString()
+             << "\n";
     });
     return FSamples;
   }
 
   return nullptr;
 }
 
 std::vector<const FunctionSamples *>
if (MergeContext) {
     LLVM_DEBUG(dbgs() << "  Merging context profile into base profile: " << Name
                       << "\n");
 
     // 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();
-      ContextTrieNode *FromNode = getContextFor(Context);
-      if (FromNode == Node)
-        continue;
-
       // Skip inlined context profile and also don't re-merge any context
       if (Context.hasState(InlinedContext) || Context.hasState(MergedContext))
         continue;
 
+      ContextTrieNode *FromNode = getContextFor(Context);
+      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() << "\n");
+                    << InlinedSamples->getContext().toString() << "\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() << "\n");
+                    << FromSamples->getContext().toString() << "\n");
 
   assert(!FromSamples->getContext().hasState(InlinedContext) &&
          "Shouldn't promote inlined context profile");
-  StringRef ContextStrToRemove = FromSamples->getContext().getCallingContext();
+  uint32_t ContextFramesToRemove =
+      FromSamples->getContext().getContextFrames().size() - 1;
   return promoteMergeContextSamplesTree(NodeToPromo, RootContext,
-                                        ContextStrToRemove);
+                                        ContextFramesToRemove);
 }
 
 void SampleContextTracker::dump() {
   dbgs() << "Context Profile Tree:\n";
   std::queue<ContextTrieNode *> NodeQueue;
   NodeQueue.push(&RootContext);
 
   while (!NodeQueue.empty()) {
ContextTrieNode *SampleContextTracker::getContextFor(const DILocation *DIL) {
 
   return nullptr;
 }
 
 ContextTrieNode *
 SampleContextTracker::getOrCreateContextPath(const SampleContext &Context,
                                              bool AllowCreate) {
   ContextTrieNode *ContextNode = &RootContext;
-  StringRef ContextRemain = Context;
-  StringRef ChildContext;
-  StringRef CalleeName;
   LineLocation CallSiteLoc(0, 0);
 
-  while (ContextNode && !ContextRemain.empty()) {
-    auto ContextSplit = SampleContext::splitContextString(ContextRemain);
-    ChildContext = ContextSplit.first;
-    ContextRemain = ContextSplit.second;
-    LineLocation NextCallSiteLoc(0, 0);
-    SampleContext::decodeContextString(ChildContext, CalleeName,
-                                       NextCallSiteLoc);
-
+  for (auto &Callsite : Context.getContextFrames()) {
     // Create child node at parent line/disc location
     if (AllowCreate) {
-      ContextNode =
-          ContextNode->getOrCreateChildContext(CallSiteLoc, CalleeName);
+      ContextNode = ContextNode->getOrCreateChildContext(CallSiteLoc,
+                                                         Callsite.CallerName);
     } else {
-      ContextNode = ContextNode->getChildContext(CallSiteLoc, CalleeName);
+      ContextNode =
+          ContextNode->getChildContext(CallSiteLoc, Callsite.CallerName);
     }
-    CallSiteLoc = NextCallSiteLoc;
+    CallSiteLoc = Callsite.Callsite;
   }
 
   assert((!AllowCreate || ContextNode) &&
          "Node must exist if creation is allowed");
   return ContextNode;
 }
 
 ContextTrieNode *SampleContextTracker::getTopLevelContextNode(StringRef FName) {
   assert(!FName.empty() && "Top level node query must provide valid name");
   return RootContext.getChildContext(LineLocation(0, 0), FName);
 }
 
 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,
-                                            StringRef ContextStrToRemove) {
+                                            uint32_t ContextFramesToRemove) {
   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);
   } else if (FromSamples) {
     // Transfer FromSamples from FromNode to ToNode
     ToNode.setFunctionSamples(FromSamples);
     FromSamples->getContext().setState(SyntheticContext);
-    FromSamples->getContext().promoteOnPath(ContextStrToRemove);
+    FromSamples->getContext().promoteOnPath(ContextFramesToRemove);
     FromNode.setFunctionSamples(nullptr);
   }
 }
 
 ContextTrieNode &SampleContextTracker::promoteMergeContextSamplesTree(
     ContextTrieNode &FromNode, ContextTrieNode &ToNodeParent,
-    StringRef ContextStrToRemove) {
-  assert(!ContextStrToRemove.empty() && "Context to remove can't be empty");
+    uint32_t ContextFramesToRemove) {
+  assert(ContextFramesToRemove && "Context to remove can't be empty");
 
   // 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), ContextStrToRemove, false);
+        NewCallSiteLoc, std::move(FromNode), ContextFramesToRemove, false);
   } else {
     // Destination node exists, merge samples for the context tree
-    mergeContextNode(FromNode, *ToNode, ContextStrToRemove);
+    mergeContextNode(FromNode, *ToNode, ContextFramesToRemove);
     LLVM_DEBUG({
       if (ToNode->getFunctionSamples())
         dbgs() << "  Context promoted and merged to: "
-               << ToNode->getFunctionSamples()->getContext() << "\n";
+               << ToNode->getFunctionSamples()->getContext().toString() << "\n";
     });
 
     // Recursively promote and merge children
     for (auto &It : FromNode.getAllChildContext()) {
       ContextTrieNode &FromChildNode = It.second;
       promoteMergeContextSamplesTree(FromChildNode, *ToNode,
-                                     ContextStrToRemove);
+                                     ContextFramesToRemove);
     }
 
     // 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;
 }
 } // namespace llvm

llvm/test/Transforms/SampleProfile/csspgo-import-list.ll

 ; Make sure Import GUID list for ThinLTO properly set for CSSPGO
 ; RUN: opt < %s -passes='thinlto-pre-link<O2>' -pgo-kind=pgo-sample-use-pipeline -sample-profile-file=%S/Inputs/csspgo-import-list.prof -S | FileCheck %s
-; RUN: opt < %s -passes='thinlto-pre-link<O2>' -pgo-kind=pgo-sample-use-pipeline -sample-profile-file=%S/Inputs/csspgo-import-list.prof.extbin -S | FileCheck %s
+; RUN: llvm-profdata merge --sample --extbinary %S/Inputs/csspgo-import-list.prof -o %t.prof
+; RUN: opt < %s -passes='thinlto-pre-link<O2>' -pgo-kind=pgo-sample-use-pipeline -sample-profile-file=%t.prof -S | FileCheck %s
 
 declare i32 @_Z5funcBi(i32 %x)
 declare i32 @_Z5funcAi(i32 %x)
 
 define dso_local i32 @main() local_unnamed_addr #0 !dbg !18 {
 entry:
   br label %for.body, !dbg !25
 

llvm/test/Transforms/SampleProfile/profile-context-tracker-debug.ll

 ; INLINE-ALL-NEXT: Getting callee context for instr:   %call.i1 = tail call i32 @_Z3fibi
 ; INLINE-ALL-NEXT: Getting callee context for instr:   %call5.i = tail call i32 @_Z3fibi
 ; INLINE-ALL-DAG:  Getting base profile for function: _Z5funcAi
 ; INLINE-ALL-DAG-NEXT:   Merging context profile into base profile: _Z5funcAi
 ; INLINE-ALL-DAG:  Getting base profile for function: _Z5funcBi
 ; INLINE-ALL-DAG-NEXT:   Merging context profile into base profile: _Z5funcBi
 ; INLINE-ALL-DAG-NEXT:   Found context tree root to promote: external:10 @ _Z5funcBi
 ; INLINE-ALL-DAG-NEXT:   Context promoted to: _Z5funcBi
+; INLINE-ALL-DAG-NEXT:   Found context tree root to promote: externalA:17 @ _Z5funcBi
+; INLINE-ALL-DAG-NEXT:   Context promoted and merged to: _Z5funcBi
 ; INLINE-ALL-DAG-NEXT:   Found context tree root to promote: main:3.1 @ _Z5funcBi
 ; INLINE-ALL-DAG-NEXT:   Context promoted and merged to: _Z5funcBi
 ; INLINE-ALL-DAG-NEXT:   Context promoted to: _Z5funcBi:1 @ _Z8funcLeafi
-; INLINE-ALL-DAG-NEXT:   Found context tree root to promote: externalA:17 @ _Z5funcBi
-; INLINE-ALL-DAG-NEXT:   Context promoted and merged to: _Z5funcBi
 ; INLINE-ALL-DAG-NEXT:   Getting callee context for instr:   %call = tail call i32 @_Z8funcLeafi
 ; INLINE-ALL-DAG-NEXT:   Callee context found: _Z5funcBi:1 @ _Z8funcLeafi
 ; INLINE-ALL-DAG-NEXT:   Marking context profile as inlined: _Z5funcBi:1 @ _Z8funcLeafi
 ; INLINE-ALL-DAG-NEXT:   Getting callee context for instr:   %call.i = tail call i32 @_Z3fibi
 ; INLINE-ALL-DAG-NEXT:   Getting callee context for instr:   %call5.i = tail call i32 @_Z3fibi
 ; INLINE-ALL:     Getting base profile for function: _Z8funcLeafi
 ; INLINE-ALL-NEXT:   Merging context profile into base profile: _Z8funcLeafi
 

 ; INLINE-HOT-DAG-NEXT:   Callee context found: _Z5funcAi:1 @ _Z8funcLeafi
 ; INLINE-HOT-DAG-NEXT:   Marking context profile as inlined: _Z5funcAi:1 @ _Z8funcLeafi
 ; INLINE-HOT-DAG-NEXT:   Getting callee context for instr:   %call.i = tail call i32 @_Z3fibi(i32 %tmp.i) #2, !dbg !62
 ; INLINE-HOT-DAG-NEXT:   Getting callee context for instr:   %call5.i = tail call i32 @_Z3fibi(i32 %tmp1.i) #2, !dbg !69
 ; INLINE-HOT-DAG: Getting base profile for function: _Z5funcBi
 ; INLINE-HOT-DAG-NEXT:   Merging context profile into base profile: _Z5funcBi
 ; INLINE-HOT-DAG-NEXT:   Found context tree root to promote: external:10 @ _Z5funcBi
 ; INLINE-HOT-DAG-NEXT:   Context promoted to: _Z5funcBi
+; INLINE-HOT-DAG-NEXT:   Found context tree root to promote: externalA:17 @ _Z5funcBi
+; INLINE-HOT-DAG-NEXT:   Context promoted and merged to: _Z5funcBi
 ; INLINE-HOT-DAG-NEXT:   Found context tree root to promote: main:3.1 @ _Z5funcBi
 ; INLINE-HOT-DAG-NEXT:   Context promoted and merged to: _Z5funcBi
 ; INLINE-HOT-DAG-NEXT:   Context promoted to: _Z5funcBi:1 @ _Z8funcLeafi
-; INLINE-HOT-DAG-NEXT:   Found context tree root to promote: externalA:17 @ _Z5funcBi
-; INLINE-HOT-DAG-NEXT:   Context promoted and merged to: _Z5funcBi
 ; INLINE-HOT-DAG-NEXT:   Getting callee context for instr:   %call = tail call i32 @_Z8funcLeafi
 ; INLINE-HOT-DAG-NEXT:   Callee context found: _Z5funcBi:1 @ _Z8funcLeafi
 ; INLINE-HOT-DAG-NEXT:   Marking context profile as inlined: _Z5funcBi:1 @ _Z8funcLeafi
 ; INLINE-HOT-DAG-NEXT:   Getting callee context for instr:   %call.i = tail call i32 @_Z3fibi
 ; INLINE-HOT-DAG-NEXT:   Getting callee context for instr:   %call5.i = tail call i32 @_Z3fibi
 ; INLINE-HOT:      Getting base profile for function: _Z8funcLeafi
 ; INLINE-HOT-NEXT:   Merging context profile into base profile: _Z8funcLeafi