[CSSPGO] split context string III - llvm-profgen changes

[hoy]

Summary

This is the profile llvm-profgen part of the context split work. Please refer to D107299 for the original whole patch.

[wmi]

llvm/tools/llvm-profgen/ProfiledBinary.h

124

I guess Offset2LocStackMap will serve as the immutable storage of all SampleContextStorageType. Is there way to enforce its immutability?

[wmi]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

335

Add a message.

347

Add a message.

[hoy]

llvm/tools/llvm-profgen/ProfiledBinary.h

124

Yes, it should be immutable once after symbolization is done. The checking of that may not be as easy as the CSName table in the reader. On the other hand, a std::unordered_map operation should not invalidate the values unless it removes or replaces the value. It should just work as other fields like the FuncStartAddrMap above. What do you think?

[hoy]

Addressing Wei's comment.

[wmi]

llvm/tools/llvm-profgen/ProfiledBinary.h

124

Ah, ok, just find that iterators will be invalidated for unordered_map rehashing but pointers and references will still be valid. So it is fine.

[hoy]

Updating D108437: [CSSPGO] split context string III - llvm-profgen changes

[wmi]

LGTM.

llvm/include/llvm/MC/MCPseudoProbe.h

171

Add a comment to explain the uint32_t contains the pseudo index.

[wenlei]

llvm/tools/llvm-profgen/CallContext.h

21

Does this using statement do anything?

llvm/tools/llvm-profgen/PerfReader.cpp

423

nit: NYI is confusing. This is simply not expected, instead of "Not yet implemented", right? we don't have a 3rd ContextKey type.

llvm/tools/llvm-profgen/ProfileGenerator.cpp

373

The name CallerContextRef probably is left over from before we renamed into SampleContextFrames. Rename CallerContextRef accordingly, same for other places. And all other places where Context is called Name better be renamed accordingly too.

llvm/tools/llvm-profgen/ProfiledBinary.cpp

425

emplace_back with variadic args and avoid temp Callsite

llvm/tools/llvm-profgen/ProfiledBinary.h

283

nit: emplace_back with variadic args

284

The actual string for Callsite.first is owned by MCPseudoProbeDecoder, right? Specifically, GUID2FuncDescMap.

[hoy]

Address Wenlei's comment.

[wenlei]

lgtm, thanks.

llvm/test/tools/llvm-profgen/recursion-compression-pseudoprobe.test

12 ↗

(On Diff #369361)

This order change is due to sortFuncProfiles change in D108435? Does the test change belong to this MD5 patch or D108435?

llvm/tools/llvm-profgen/ProfileGenerator.cpp

373

getFunctionProfileForLeafProbe has two instance of ContextRef still.

[hoy]

llvm/test/tools/llvm-profgen/recursion-compression-pseudoprobe.test

12 ↗

(On Diff #369361)

The order change is due to the flip of using > to < for sortFuncProfiles in D108433, so that we don't have to provide a > operator.

llvm/tools/llvm-profgen/CallContext.h

21

oops, looks like it's from a string replacement. Removed.

llvm/tools/llvm-profgen/PerfReader.cpp

423

Makes sense. Message changed.

llvm/tools/llvm-profgen/ProfileGenerator.cpp

373

Good catch. Renamed to CallerFrames.

373

Renamed to NewContext.

llvm/tools/llvm-profgen/ProfiledBinary.cpp

425

fixed

llvm/tools/llvm-profgen/ProfiledBinary.h

284

Yes, the string is owned by the probe decoder, unlike the strings returned from the symbolizer which have to be saved somewhere in the profiled binary.

[wenlei]

llvm/test/tools/llvm-profgen/recursion-compression-pseudoprobe.test

12 ↗

(On Diff #369361)

I was asking why this test change is here since the sortFuncProfiles is in the 2nd patch. But if you're going to commit these three patches together, it's not going to make a difference anyways.

[hoy]

llvm/test/tools/llvm-profgen/recursion-compression-pseudoprobe.test

12 ↗

(On Diff #369361)

Oh yeah, I'll commit all patches as a whole. Will update the original patch before checking them in.

llvm/include/llvm/MC/MCPseudoProbe.h

MCPseudoProbe(MCSymbol *Label, uint64_t Guid, uint64_t Index, uint64_t Type,
     assert(Attributes <= 0xFF &&
            "Probe attributes too big to encode, exceeding 2^16");
   }
 
   MCSymbol *getLabel() const { return Label; }
   void emit(MCObjectStreamer *MCOS, const MCPseudoProbe *LastProbe) const;
 };
 
+using MCPseduoProbeFrameLocation = std::pair<StringRef, uint32_t>;
+
 class MCDecodedPseudoProbe : public MCPseudoProbeBase {
   uint64_t Address;
   MCDecodedPseudoProbeInlineTree *InlineTree;
 
 public:
   MCDecodedPseudoProbe(uint64_t Ad, uint64_t G, uint32_t I, PseudoProbeType K,
                        uint8_t At, MCDecodedPseudoProbeInlineTree *Tree)
       : MCPseudoProbeBase(G, I, At, static_cast<uint8_t>(K)), Address(Ad),
         InlineTree(Tree){};
 
   uint64_t getAddress() const { return Address; }
 
   void setAddress(uint64_t Addr) { Address = Addr; }
 
   MCDecodedPseudoProbeInlineTree *getInlineTreeNode() const {
     return InlineTree;
   }
 
   // Get the inlined context by traversing current inline tree backwards,
   // each tree node has its InlineSite which is taken as the context.
   // \p ContextStack is populated in root to leaf order
-  void getInlineContext(SmallVectorImpl<std::string> &ContextStack,
-                        const GUIDProbeFunctionMap &GUID2FuncMAP,
-                        bool ShowName) const;
+  void
+  getInlineContext(SmallVectorImpl<MCPseduoProbeFrameLocation> &ContextStack,
+                   const GUIDProbeFunctionMap &GUID2FuncMAP) const;
 
   // Helper function to get the string from context stack
-  std::string getInlineContextStr(const GUIDProbeFunctionMap &GUID2FuncMAP,
-                                  bool ShowName) const;
+  std::string
+  getInlineContextStr(const GUIDProbeFunctionMap &GUID2FuncMAP) const;
 
   // Print pseudo probe while disassembling
   void print(raw_ostream &OS, const GUIDProbeFunctionMap &GUID2FuncMAP,
              bool ShowName) const;
 };
 
 template <typename ProbeType, typename DerivedProbeInlineTreeType>
 class MCPseudoProbeInlineTreeBase {
public:
 
   // Helper function to populate one probe's inline stack into
   // \p InlineContextStack.
   // Current leaf location info will be added if IncludeLeaf is true
   // Example:
   //  Current probe(bar:3) inlined at foo:2 then inlined at main:1
   //  IncludeLeaf = true,  Output: [main:1, foo:2, bar:3]
   //  IncludeLeaf = false, Output: [main:1, foo:2]
-  void
-  getInlineContextForProbe(const MCDecodedPseudoProbe *Probe,
-                           SmallVectorImpl<std::string> &InlineContextStack,
-                           bool IncludeLeaf) const;
+  void getInlineContextForProbe(
+      const MCDecodedPseudoProbe *Probe,
+      SmallVectorImpl<MCPseduoProbeFrameLocation> &InlineContextStack,
 
   const AddressProbesMap &getAddress2ProbesMap() const {
     return Address2ProbesMap;
   }
 
   AddressProbesMap &getAddress2ProbesMap() { return Address2ProbesMap; }
 
   const GUIDProbeFunctionMap &getGUID2FuncDescMap() const {

llvm/lib/MC/MCPseudoProbe.cpp

 }
 
 void MCPseudoProbeFuncDesc::print(raw_ostream &OS) {
   OS << "GUID: " << FuncGUID << " Name: " << FuncName << "\n";
   OS << "Hash: " << FuncHash << "\n";
 }
 
 void MCDecodedPseudoProbe::getInlineContext(
-    SmallVectorImpl<std::string> &ContextStack,
-    const GUIDProbeFunctionMap &GUID2FuncMAP, bool ShowName) const {
+    SmallVectorImpl<MCPseduoProbeFrameLocation> &ContextStack,
+    const GUIDProbeFunctionMap &GUID2FuncMAP) const {
   uint32_t Begin = ContextStack.size();
   MCDecodedPseudoProbeInlineTree *Cur = InlineTree;
   // It will add the string of each node's inline site during iteration.
   // Note that it won't include the probe's belonging function(leaf location)
   while (Cur->hasInlineSite()) {
-    std::string ContextStr;
-    if (ShowName) {
-      StringRef FuncName =
-          getProbeFNameForGUID(GUID2FuncMAP, std::get<0>(Cur->ISite));
-      ContextStr += FuncName.str();
-    } else {
-      ContextStr += Twine(std::get<0>(Cur->ISite)).str();
-    }
-    ContextStr += ":";
-    ContextStr += Twine(std::get<1>(Cur->ISite)).str();
-    ContextStack.emplace_back(ContextStr);
+    ContextStack.emplace_back(
+        MCPseduoProbeFrameLocation(FuncName, std::get<1>(Cur->ISite)));
     Cur = static_cast<MCDecodedPseudoProbeInlineTree *>(Cur->Parent);
   }
   // Make the ContextStack in caller-callee order
   std::reverse(ContextStack.begin() + Begin, ContextStack.end());
 }
 
 std::string MCDecodedPseudoProbe::getInlineContextStr(
-    const GUIDProbeFunctionMap &GUID2FuncMAP, bool ShowName) const {
+    const GUIDProbeFunctionMap &GUID2FuncMAP) const {
   std::ostringstream OContextStr;
-  SmallVector<std::string, 16> ContextStack;
-  getInlineContext(ContextStack, GUID2FuncMAP, ShowName);
-  for (auto &CxtStr : ContextStack) {
+  SmallVector<MCPseduoProbeFrameLocation, 16> ContextStack;
+  getInlineContext(ContextStack, GUID2FuncMAP);
+  for (auto &Cxt : ContextStack) {
     if (OContextStr.str().size())
       OContextStr << " @ ";
-    OContextStr << CxtStr;
+    OContextStr << Cxt.first.str() << ":" << Cxt.second;
   }
   return OContextStr.str();
 }
 
 static const char *PseudoProbeTypeStr[3] = {"Block", "IndirectCall",
                                             "DirectCall"};
 
 void MCDecodedPseudoProbe::print(raw_ostream &OS,
                                  const GUIDProbeFunctionMap &GUID2FuncMAP,
                                  bool ShowName) const {
   OS << "FUNC: ";
   if (ShowName) {
     StringRef FuncName = getProbeFNameForGUID(GUID2FuncMAP, Guid);
     OS << FuncName.str() << " ";
   } else {
     OS << Guid << " ";
   }
   OS << "Index: " << Index << "  ";
   OS << "Type: " << PseudoProbeTypeStr[static_cast<uint8_t>(Type)] << "  ";
-  std::string InlineContextStr = getInlineContextStr(GUID2FuncMAP, ShowName);
+  std::string InlineContextStr = getInlineContextStr(GUID2FuncMAP);
   if (InlineContextStr.size()) {
     OS << "Inlined: @ ";
     OS << InlineContextStr;
   }
   OS << "\n";
 }
 
 template <typename T> ErrorOr<T> MCPseudoProbeDecoder::readUnencodedNumber() {

 MCPseudoProbeDecoder::getFuncDescForGUID(uint64_t GUID) const {
   auto It = GUID2FuncDescMap.find(GUID);
   assert(It != GUID2FuncDescMap.end() && "Function descriptor doesn't exist");
   return &It->second;
 }
 
 void MCPseudoProbeDecoder::getInlineContextForProbe(
     const MCDecodedPseudoProbe *Probe,
-    SmallVectorImpl<std::string> &InlineContextStack, bool IncludeLeaf) const {
-  Probe->getInlineContext(InlineContextStack, GUID2FuncDescMap, true);
+    SmallVectorImpl<MCPseduoProbeFrameLocation> &InlineContextStack,
+    bool IncludeLeaf) const {
+  Probe->getInlineContext(InlineContextStack, GUID2FuncDescMap);
   if (!IncludeLeaf)
     return;
   // Note that the context from probe doesn't include leaf frame,
   // hence we need to retrieve and prepend leaf if requested.
   const auto *FuncDesc = getFuncDescForGUID(Probe->getGuid());
-  InlineContextStack.emplace_back(FuncDesc->FuncName + ":" +
-                                  Twine(Probe->getIndex()).str());
+  InlineContextStack.emplace_back(
+      MCPseduoProbeFrameLocation(FuncDesc->FuncName, Probe->getIndex()));
 }
 
 const MCPseudoProbeFuncDesc *MCPseudoProbeDecoder::getInlinerDescForProbe(
     const MCDecodedPseudoProbe *Probe) const {
   MCDecodedPseudoProbeInlineTree *InlinerNode = Probe->getInlineTreeNode();
   if (!InlinerNode->hasInlineSite())
     return nullptr;
   return getFuncDescForGUID(std::get<0>(InlinerNode->ISite));
 }

llvm/tools/llvm-profgen/CSPreInliner.h

 // Pre-compilation inliner based on context-sensitive profile.
 // The PreInliner estimates inline decision using hotness from profile
 // and cost estimation from machine code size. It helps merges context
 // profile globally and achieves better post-inine profile quality, which
 // otherwise won't be possible for ThinLTO. It also reduce context profile
 // size by only keep context that is estimated to be inlined.
 class CSPreInliner {
 public:
-  CSPreInliner(StringMap<FunctionSamples> &Profiles, uint64_t HotThreshold,
+  CSPreInliner(SampleProfileMap &Profiles, uint64_t HotThreshold,
                uint64_t ColdThreshold);
   void run();
 
 private:
   bool getInlineCandidates(ProfiledCandidateQueue &CQueue,
                            const FunctionSamples *FCallerContextSamples);
   std::vector<StringRef> buildTopDownOrder();
   void processFunction(StringRef Name);
   bool shouldInline(ProfiledInlineCandidate &Candidate);
   SampleContextTracker ContextTracker;
-  StringMap<FunctionSamples> &ProfileMap;
+  SampleProfileMap &ProfileMap;
 
   // Count thresholds to answer isHotCount and isColdCount queries.
   // Mirrors the threshold in ProfileSummaryInfo.
   uint64_t HotCountThreshold;
   uint64_t ColdCountThreshold;
 };
 
 } // end namespace sampleprof
 } // end namespace llvm
 
 #endif

llvm/tools/llvm-profgen/CSPreInliner.cpp

 extern cl::opt<int> ProfileInlineLimitMin;
 extern cl::opt<int> ProfileInlineLimitMax;
 
 static cl::opt<bool> SamplePreInlineReplay(
     "csspgo-replay-preinline", cl::Hidden, cl::init(false),
     cl::desc(
         "Replay previous inlining and adjust context profile accordingly"));
 
-CSPreInliner::CSPreInliner(StringMap<FunctionSamples> &Profiles,
-                           uint64_t HotThreshold, uint64_t ColdThreshold)
+CSPreInliner::CSPreInliner(SampleProfileMap &Profiles, uint64_t HotThreshold,
+                           uint64_t ColdThreshold)
     : ContextTracker(Profiles), ProfileMap(Profiles),
       HotCountThreshold(HotThreshold), ColdCountThreshold(ColdThreshold) {}
 
 std::vector<StringRef> CSPreInliner::buildTopDownOrder() {
   std::vector<StringRef> Order;
   ProfiledCallGraph ProfiledCG(ContextTracker);
 
   // Now that we have a profiled call graph, construct top-down order
if ((ShouldInline = shouldInline(Candidate))) {
       // We mark context as inlined as the corresponding context profile
       // won't be merged into that function's base profile.
       ContextTracker.markContextSamplesInlined(Candidate.CalleeSamples);
       Candidate.CalleeSamples->getContext().setAttribute(
           ContextShouldBeInlined);
       FuncFinalSize += Candidate.SizeCost;
       getInlineCandidates(CQueue, Candidate.CalleeSamples);
     }
-    LLVM_DEBUG(dbgs() << (ShouldInline ? "  Inlined" : "  Outlined")
-                      << " context profile for: "
-                      << Candidate.CalleeSamples->getNameWithContext()
-                      << " (callee size: " << Candidate.SizeCost
-                      << ", call count:" << Candidate.CallsiteCount << ")\n");
+    LLVM_DEBUG(
+        dbgs() << (ShouldInline ? "  Inlined" : "  Outlined")
+               << Candidate.CalleeSamples->getContext().getContextString()
   }
 
   LLVM_DEBUG({
     if (!CQueue.empty())
       dbgs() << "  Inline candidates ignored due to size limit (inliner "
                 "original size: "
              << FuncSize << ", inliner final size: " << FuncFinalSize
              << ", size limit: " << SizeLimit << ")\n";
 
     while (!CQueue.empty()) {
       ProfiledInlineCandidate Candidate = CQueue.top();
       CQueue.pop();
       bool WasInlined =
           Candidate.CalleeSamples->getContext().hasAttribute(ContextWasInlined);
-      dbgs() << "    " << Candidate.CalleeSamples->getNameWithContext()
+      dbgs() << "    "
+             << Candidate.CalleeSamples->getContext().getContextString()
              << " (candidate size:" << Candidate.SizeCost
              << ", call count: " << Candidate.CallsiteCount << ", previously "
              << (WasInlined ? "inlined)\n" : "not inlined)\n");
     }
   });
 }
 
 void CSPreInliner::run() {
 #ifndef NDEBUG
-  auto printProfileNames = [](StringMap<FunctionSamples> &Profiles,
-                              bool IsInput) {
+  auto printProfileNames = [](SampleProfileMap &Profiles, bool IsInput) {
     dbgs() << (IsInput ? "Input" : "Output") << " context-sensitive profiles ("
            << Profiles.size() << " total):\n";
     for (auto &It : Profiles) {
       const FunctionSamples &Samples = It.second;
-      dbgs() << "  [" << Samples.getNameWithContext() << "] "
+      dbgs() << "  [" << Samples.getContext().getContextString() << "] "
              << Samples.getTotalSamples() << ":" << Samples.getHeadSamples()
              << "\n";
     }
   };
 #endif
 
   LLVM_DEBUG(printProfileNames(ProfileMap, true));
 
   // Execute global pre-inliner to estimate a global top-down inline
   // decision and merge profiles accordingly. This helps with profile
   // merge for ThinLTO otherwise we won't be able to merge profiles back
   // to base profile across module/thin-backend boundaries.
   // It also helps better compress context profile to control profile
   // size, as we now only need context profile for functions going to
   // be inlined.
   for (StringRef FuncName : buildTopDownOrder()) {
     processFunction(FuncName);
   }
 
   // Not inlined context profiles are merged into its base, so we can
   // trim out such profiles from the output.
-  std::vector<StringRef> ProfilesToBeRemoved;
+  std::vector<SampleContext> ProfilesToBeRemoved;
   for (auto &It : ProfileMap) {
     SampleContext Context = It.second.getContext();
     if (!Context.isBaseContext() && !Context.hasState(InlinedContext)) {
       assert(Context.hasState(MergedContext) &&
              "Not inlined context profile should be merged already");
-      ProfilesToBeRemoved.push_back(It.first());
+      ProfilesToBeRemoved.push_back(It.first);
     }
   }
 
-  for (StringRef ContextName : ProfilesToBeRemoved) {
+  for (auto &ContextName : ProfilesToBeRemoved) {
     ProfileMap.erase(ContextName);
   }
 
   // Make sure ProfileMap's key is consistent with FunctionSamples' name.
   SampleContextTrimmer(ProfileMap).canonicalizeContextProfiles();
 
   LLVM_DEBUG(printProfileNames(ProfileMap, false));
 }

llvm/tools/llvm-profgen/CallContext.h

 #include <sstream>
 #include <string>
 #include <vector>
 
 namespace llvm {
 namespace sampleprof {
 
 // Function name, LineLocation
-typedef std::pair<std::string, LineLocation> FrameLocation;
+using SampleContextStorageType = SampleContextStorageType;
 
-typedef SmallVector<FrameLocation, 4> FrameLocationStack;
-
-inline std::string getCallSite(const FrameLocation &Callsite) {
-  std::string CallsiteStr = Callsite.first;
+inline std::string getCallSite(const SampleCallSiteType &Callsite) {
+  std::string CallsiteStr = Callsite.CallerName.str();
   CallsiteStr += ":";
-  CallsiteStr += Twine(Callsite.second.LineOffset).str();
-  if (Callsite.second.Discriminator > 0) {
+  CallsiteStr += Twine(Callsite.Callsite.LineOffset).str();
+  if (Callsite.Callsite.Discriminator > 0) {
     CallsiteStr += ".";
-    CallsiteStr += Twine(Callsite.second.Discriminator).str();
+    CallsiteStr += Twine(Callsite.Callsite.Discriminator).str();
   }
   return CallsiteStr;
 }
 
 // TODO: This operation is expansive. If it ever gets called multiple times we
 // may think of making a class wrapper with internal states for it.
-inline std::string getLocWithContext(const FrameLocationStack &Context) {
+inline std::string getLocWithContext(const SampleContextStorageType &Context) {
   std::ostringstream OContextStr;
   for (const auto &Callsite : Context) {
     if (OContextStr.str().size())
       OContextStr << " @ ";
     OContextStr << getCallSite(Callsite);
   }
   return OContextStr.str();
 }
 
 // Reverse call context, i.e., in the order of callee frames to caller frames,
 // is useful during instruction printing or pseudo probe printing.
 inline std::string
-getReversedLocWithContext(const FrameLocationStack &Context) {
+getReversedLocWithContext(const SampleContextStorageType &Context) {
   std::ostringstream OContextStr;
   for (const auto &Callsite : reverse(Context)) {
     if (OContextStr.str().size())
       OContextStr << " @ ";
     OContextStr << getCallSite(Callsite);
   }
   return OContextStr.str();
 }
 
 } // end namespace sampleprof
 } // end namespace llvm
 
 #endif

llvm/tools/llvm-profgen/PerfReader.h

struct ContextKey {
   enum ContextKind { CK_StringBased, CK_ProbeBased };
   const ContextKind Kind;
   ContextKind getKind() const { return Kind; }
   ContextKey(ContextKind K) : Kind(K){};
 };
 
 // String based context id
 struct StringBasedCtxKey : public ContextKey {
-  std::string Context;
+  SampleContextStorageType Context;
+
   bool WasLeafInlined;
   StringBasedCtxKey() : ContextKey(CK_StringBased), WasLeafInlined(false){};
   static bool classof(const ContextKey *K) {
     return K->getKind() == CK_StringBased;
   }
 
   bool isEqual(const ContextKey *K) const override {
     const StringBasedCtxKey *Other = dyn_cast<StringBasedCtxKey>(K);
     return Context == Other->Context;
   }
 
-  void genHashCode() { HashCode = hash_value(Context); }
+  void genHashCode() { HashCode = hash_value(SampleContextRefType(Context)); }
 };
 
 // Probe based context key as the intermediate key of context
 // String based context key will introduce redundant string handling
 // since the callee context is inferred from the context string which
 // need to be splitted by '@' to get the last location frame, so we
 // can just use probe instead and generate the string in the end.
 struct ProbeBasedCtxKey : public ContextKey {

llvm/tools/llvm-profgen/PerfReader.cpp

void VirtualUnwinder::unwindBranchWithinFrame(UnwindState &State) {
   uint64_t Source = State.getCurrentLBRSource();
   State.switchToFrame(Source);
   State.InstPtr.update(Source);
 }
 
 std::shared_ptr<StringBasedCtxKey> FrameStack::getContextKey() {
   std::shared_ptr<StringBasedCtxKey> KeyStr =
       std::make_shared<StringBasedCtxKey>();
-  KeyStr->Context =
-      Binary->getExpandedContextStr(Stack, KeyStr->WasLeafInlined);
+  KeyStr->Context = Binary->getExpandedContext(Stack, KeyStr->WasLeafInlined);
   if (KeyStr->Context.empty())
     return nullptr;
   KeyStr->genHashCode();
   return KeyStr;
 }
 
 std::shared_ptr<ProbeBasedCtxKey> ProbeStack::getContextKey() {
   std::shared_ptr<ProbeBasedCtxKey> ProbeBasedKey =
for (auto I : Range.second) {
       outs() << "  (" << format("%" PRIx64, I.first.first) << ", "
              << format("%" PRIx64, I.first.second) << "): " << I.second << "\n";
     }
   }
 }
 
 static std::string getContextKeyStr(ContextKey *K,
                                     const ProfiledBinary *Binary) {
-  std::string ContextStr;
   if (const auto *CtxKey = dyn_cast<StringBasedCtxKey>(K)) {
-    return CtxKey->Context;
+    return SampleContext::getContextString(CtxKey->Context);
   } else if (const auto *CtxKey = dyn_cast<ProbeBasedCtxKey>(K)) {
-    SmallVector<std::string, 16> ContextStack;
+    SampleContextStorageType ContextStack;
     for (const auto *Probe : CtxKey->Probes) {
       Binary->getInlineContextForProbe(Probe, ContextStack, true);
     }
-    for (const auto &Context : ContextStack) {
-      if (ContextStr.size())
-        ContextStr += " @ ";
-      ContextStr += Context;
-    }
+    return SampleContext::getContextString(ContextStack);
+  } else {
+    llvm_unreachable("NYI");
   }
-  return ContextStr;
 }
 
 static void printRangeCounter(ContextSampleCounterMap &Counter,
                               const ProfiledBinary *Binary) {
   OrderedCounterForPrint OrderedCounter;
   for (auto &CI : Counter) {
     OrderedCounter[getContextKeyStr(CI.first.getPtr(), Binary)] =
         CI.second.RangeCounter;
   }

llvm/tools/llvm-profgen/ProfileGenerator.h

public:
   ProfileGenerator(){};
   virtual ~ProfileGenerator() = default;
   static std::unique_ptr<ProfileGenerator>
   create(const BinarySampleCounterMap &BinarySampleCounters,
          enum PerfScriptType SampleType);
   virtual void generateProfile() = 0;
   // Use SampleProfileWriter to serialize profile map
   virtual void write(std::unique_ptr<SampleProfileWriter> Writer,
-                     StringMap<FunctionSamples> &ProfileMap);
+                     SampleProfileMap &ProfileMap);
   void write();
 
 protected:
   /*
   For each region boundary point, mark if it is begin or end (or both) of
   the region. Boundary points are inclusive. Log the sample count as well
   so we can use it when we compute the sample count of each disjoint region
   later. Note that there might be multiple ranges with different sample
   count that share same begin/end point. We need to accumulate the sample
   count for the boundary point for such case, because for the example
   below,
 
   |<--100-->|
   |<------200------>|
   A         B       C
 
   sample count for disjoint region [A,B] would be 300.
   */
   void findDisjointRanges(RangeSample &DisjointRanges,
                           const RangeSample &Ranges);
 
   // Used by SampleProfileWriter
-  StringMap<FunctionSamples> ProfileMap;
+  SampleProfileMap ProfileMap;
 };
 
 class CSProfileGenerator : public ProfileGenerator {
 protected:
   const BinarySampleCounterMap &BinarySampleCounters;
 
 public:
   CSProfileGenerator(const BinarySampleCounterMap &Counters)
while (I <= MaxDedupSize) {
       I++;
       Context.resize(End);
       MaxDedupSize = std::min(static_cast<uint32_t>(End / 2), MaxDedupSize);
     }
   }
 
 protected:
   // Lookup or create FunctionSamples for the context
-  FunctionSamples &getFunctionProfileForContext(StringRef ContextId,
+  FunctionSamples &getFunctionProfileForContext(SampleContextRefType ContextId,
                                                 bool WasLeafInlined = false);
   // Post processing for profiles before writing out, such as mermining
   // and trimming cold profiles, running preinliner on profiles.
   void postProcessProfiles();
   void computeSummaryAndThreshold();
   void write(std::unique_ptr<SampleProfileWriter> Writer,
-             StringMap<FunctionSamples> &ProfileMap) override;
+             SampleProfileMap &ProfileMap) override;
 
   // Thresholds from profile summary to answer isHotCount/isColdCount queries.
   uint64_t HotCountThreshold;
   uint64_t ColdCountThreshold;
 
-  // String table owning context strings created from profile generation.
-  std::unordered_set<std::string> ContextStrings;
-
 private:
   // Helper function for updating body sample for a leaf location in
   // FunctionProfile
   void updateBodySamplesforFunctionProfile(FunctionSamples &FunctionProfile,
-                                           const FrameLocation &LeafLoc,
+                                           const SampleCallSiteType &LeafLoc,
                                            uint64_t Count);
   void populateFunctionBodySamples(FunctionSamples &FunctionProfile,
                                    const RangeSample &RangeCounters,
                                    ProfiledBinary *Binary);
-  void populateFunctionBoundarySamples(StringRef ContextId,
+  void populateFunctionBoundarySamples(SampleContextRefType ContextId,
                                        FunctionSamples &FunctionProfile,
                                        const BranchSample &BranchCounters,
                                        ProfiledBinary *Binary);
   void populateInferredFunctionSamples();
 
 public:
   // Deduplicate adjacent repeated context sequences up to a given sequence
   // length. -1 means no size limit.

 
 private:
   // Go through each address from range to extract the top frame probe by
   // looking up in the Address2ProbeMap
   void extractProbesFromRange(const RangeSample &RangeCounter,
                               ProbeCounterMap &ProbeCounter,
                               ProfiledBinary *Binary);
   // Fill in function body samples from probes
-  void
-  populateBodySamplesWithProbes(const RangeSample &RangeCounter,
-                                SmallVectorImpl<std::string> &ContextStrStack,
+  void populateBodySamplesWithProbes(const RangeSample &RangeCounter,
+                                     SampleContextRefType ContextStack,
                                      ProfiledBinary *Binary);
   // Fill in boundary samples for a call probe
-  void populateBoundarySamplesWithProbes(
-      const BranchSample &BranchCounter,
-      SmallVectorImpl<std::string> &ContextStrStack, ProfiledBinary *Binary);
-  // Helper function to get FunctionSamples for the leaf inlined context
-  FunctionSamples &
-  getFunctionProfileForLeafProbe(SmallVectorImpl<std::string> &ContextStrStack,
-                                 const MCPseudoProbeFuncDesc *LeafFuncDesc,
-                                 bool WasLeafInlined);
+  void populateBoundarySamplesWithProbes(const BranchSample &BranchCounter,
+                                         SampleContextRefType ContextStack,
+                                         ProfiledBinary *Binary);
   // Helper function to get FunctionSamples for the leaf probe
   FunctionSamples &
-  getFunctionProfileForLeafProbe(SmallVectorImpl<std::string> &ContextStrStack,
+  getFunctionProfileForLeafProbe(SampleContextRefType ContextStrStack,
                                  const MCDecodedPseudoProbe *LeafProbe,
                                  ProfiledBinary *Binary);
+
+  // Underlying context table serves for sample profile writer.
+  std::unordered_set<SampleContextStorageType, SampleContextStorageHash>
+      Contexts;
 };
 
 } // end namespace sampleprof
 } // end namespace llvm
 
 #endif

llvm/tools/llvm-profgen/ProfileGenerator.cpp

if (SampleType == PERF_LBR_STACK) {
     // TODO:
     llvm_unreachable("Unsupported perfscript!");
   }
 
   return ProfileGenerator;
 }
 
 void ProfileGenerator::write(std::unique_ptr<SampleProfileWriter> Writer,
-                             StringMap<FunctionSamples> &ProfileMap) {
+                             SampleProfileMap &ProfileMap) {
   if (std::error_code EC = Writer->write(ProfileMap))
     exitWithError(std::move(EC));
 }
 
 void ProfileGenerator::write() {
   auto WriterOrErr = SampleProfileWriter::create(OutputFilename, OutputFormat);
   if (std::error_code EC = WriterOrErr.getError())
     exitWithError(EC, OutputFilename);
if (Point.EndCount) {
       DisjointRanges[{BeginAddress, Address}] = Count;
       Count -= Point.EndCount;
       BeginAddress = Address + 1;
     }
   }
 }
 
 FunctionSamples &
-CSProfileGenerator::getFunctionProfileForContext(StringRef ContextStr,
+CSProfileGenerator::getFunctionProfileForContext(SampleContextRefType Context,
                                                  bool WasLeafInlined) {
-  auto Ret = ProfileMap.try_emplace(ContextStr, FunctionSamples());
+  SampleContext FContext(Context);
+  auto Ret = ProfileMap.emplace(Context, FunctionSamples());
   if (Ret.second) {
-    // Make a copy of the underlying context string in string table
-    // before StringRef wrapper is used for context.
-    auto It = ContextStrings.insert(ContextStr.str());
-    SampleContext FContext(*It.first, RawContext);
+    SampleContext FContext(Context, RawContext);
     if (WasLeafInlined)
       FContext.setAttribute(ContextWasInlined);
     FunctionSamples &FProfile = Ret.first->second;
     FProfile.setContext(FContext);
-    FProfile.setName(FContext.getNameWithoutContext());
   }
   return Ret.first->second;
 }
 
 void CSProfileGenerator::generateProfile() {
   FunctionSamples::ProfileIsCS = true;
   for (const auto &BI : BinarySampleCounters) {
     ProfiledBinary *Binary = BI.first;
     for (const auto &CI : BI.second) {
       const StringBasedCtxKey *CtxKey =
           dyn_cast<StringBasedCtxKey>(CI.first.getPtr());
-      StringRef ContextId(CtxKey->Context);
       // Get or create function profile for the range
       FunctionSamples &FunctionProfile =
-          getFunctionProfileForContext(ContextId, CtxKey->WasLeafInlined);
+          getFunctionProfileForContext(CtxKey->Context, CtxKey->WasLeafInlined);
 
       // Fill in function body samples
       populateFunctionBodySamples(FunctionProfile, CI.second.RangeCounter,
                                   Binary);
       // Fill in boundary sample counts as well as call site samples for calls
-      populateFunctionBoundarySamples(ContextId, FunctionProfile,
+      populateFunctionBoundarySamples(CtxKey->Context, FunctionProfile,
                                       CI.second.BranchCounter, Binary);
     }
   }
   // 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();
 
   postProcessProfiles();
 }
 
 void CSProfileGenerator::updateBodySamplesforFunctionProfile(
-    FunctionSamples &FunctionProfile, const FrameLocation &LeafLoc,
+    FunctionSamples &FunctionProfile, const SampleCallSiteType &LeafLoc,
     uint64_t Count) {
   // Filter out invalid negative(int type) lineOffset
-  if (LeafLoc.second.LineOffset & 0x80000000)
+  if (LeafLoc.Callsite.LineOffset & 0x80000000)
     return;
   // Use the maximum count of samples with same line location
   ErrorOr<uint64_t> R = FunctionProfile.findSamplesAt(
-      LeafLoc.second.LineOffset, LeafLoc.second.Discriminator);
+      LeafLoc.Callsite.LineOffset, LeafLoc.Callsite.Discriminator);
   uint64_t PreviousCount = R ? R.get() : 0;
   if (PreviousCount < Count) {
-    FunctionProfile.addBodySamples(LeafLoc.second.LineOffset,
-                                   LeafLoc.second.Discriminator,
+    FunctionProfile.addBodySamples(LeafLoc.Callsite.LineOffset,
+                                   LeafLoc.Callsite.Discriminator,
                                    Count - PreviousCount);
   }
 }
 
 void CSProfileGenerator::populateFunctionBodySamples(
     FunctionSamples &FunctionProfile, const RangeSample &RangeCounter,
     ProfiledBinary *Binary) {
   // Compute disjoint ranges first, so we can use MAX
while (IP.Address <= RangeEnd) {
       FunctionProfile.addTotalSamples(Count);
       // Move to next IP within the range
       IP.advance();
     }
   }
 }
 
 void CSProfileGenerator::populateFunctionBoundarySamples(
-    StringRef ContextId, FunctionSamples &FunctionProfile,
+    SampleContextRefType ContextId, FunctionSamples &FunctionProfile,
     const BranchSample &BranchCounters, ProfiledBinary *Binary) {
 
   for (auto Entry : BranchCounters) {
     uint64_t SourceOffset = Entry.first.first;
     uint64_t TargetOffset = Entry.first.second;
     uint64_t Count = Entry.second;
     // Get the callee name by branch target if it's a call branch
     StringRef CalleeName = FunctionSamples::getCanonicalFnName(
         Binary->getFuncFromStartOffset(TargetOffset));
     if (CalleeName.size() == 0)
       continue;
 
     // Record called target sample and its count
     auto LeafLoc = Binary->getInlineLeafFrameLoc(SourceOffset);
     if (!LeafLoc.hasValue())
       continue;
-    FunctionProfile.addCalledTargetSamples(LeafLoc->second.LineOffset,
-                                           LeafLoc->second.Discriminator,
+    FunctionProfile.addCalledTargetSamples(LeafLoc->Callsite.LineOffset,
+                                           LeafLoc->Callsite.Discriminator,
                                            CalleeName, Count);
 
     // Record head sample for called target(callee)
-    std::ostringstream OCalleeCtxStr;
-    if (ContextId.find(" @ ") != StringRef::npos) {
-      OCalleeCtxStr << ContextId.rsplit(" @ ").first.str();
-      OCalleeCtxStr << " @ ";
-    }
-    OCalleeCtxStr << getCallSite(*LeafLoc) << " @ " << CalleeName.str();
-
-    FunctionSamples &CalleeProfile =
-        getFunctionProfileForContext(OCalleeCtxStr.str());
+    SampleContextStorageType CalleeCtx(ContextId.begin(), ContextId.end());
+    assert(CalleeCtx.back().CallerName == LeafLoc->CallerName);
+    CalleeCtx.back() = *LeafLoc;
+    CalleeCtx.emplace_back(CalleeName, LineLocation(0, 0));
+    FunctionSamples &CalleeProfile = getFunctionProfileForContext(CalleeCtx);
     assert(Count != 0 && "Unexpected zero weight branch");
     CalleeProfile.addHeadSamples(Count);
   }
 }
 
-static FrameLocation getCallerContext(StringRef CalleeContext,
-                                      StringRef &CallerNameWithContext) {
-  StringRef CallerContext = CalleeContext.rsplit(" @ ").first;
-  CallerNameWithContext = CallerContext.rsplit(':').first;
-  auto ContextSplit = CallerContext.rsplit(" @ ");
-  StringRef CallerFrameStr = ContextSplit.second.size() == 0
-                                 ? ContextSplit.first
-                                 : ContextSplit.second;
-  FrameLocation LeafFrameLoc = {"", {0, 0}};
-  StringRef Funcname;
-  SampleContext::decodeContextString(CallerFrameStr, Funcname,
-                                     LeafFrameLoc.second);
-  LeafFrameLoc.first = Funcname.str();
-  return LeafFrameLoc;
+static SampleCallSiteType
+getCallerContext(SampleContextRefType CalleeContext,
+                 SampleContextStorageType &CallerContext) {
+  assert(CalleeContext.size() > 1);
+  CalleeContext = CalleeContext.drop_back();
+  CallerContext.assign(CalleeContext.begin(), CalleeContext.end());
+  SampleCallSiteType CallerFrame = CallerContext.back();
+  CallerContext.back().Callsite = LineLocation(0, 0);
+  return CallerFrame;
 }
 
 void CSProfileGenerator::populateInferredFunctionSamples() {
   for (const auto &Item : ProfileMap) {
-    const StringRef CalleeContext = Item.first();
+    const auto &CalleeContext = Item.first;
     const FunctionSamples &CalleeProfile = Item.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.find(" @ ") == StringRef::npos)
+    if (CalleeContext.isBaseContext())
       continue;
 
     // Infer Caller's frame loc and context ID through string splitting
-    StringRef CallerContextId;
-    FrameLocation &&CallerLeafFrameLoc =
-        getCallerContext(CalleeContext, CallerContextId);
+    SampleContextStorageType CallerContextId;
+    SampleCallSiteType &&CallerLeafFrameLoc =
+        getCallerContext(CalleeContext.getFullContext(), 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(CallerContextId) == ProfileMap.end())
       continue;
     FunctionSamples &CallerProfile = ProfileMap[CallerContextId];
 
     // 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.second.LineOffset,
-        CallerLeafFrameLoc.second.Discriminator,
+        CallerLeafFrameLoc.Callsite.LineOffset,
+        CallerLeafFrameLoc.Callsite.Discriminator,
         CalleeProfile.getContext().getNameWithoutContext(), EstimatedCallCount);
-    CallerProfile.addBodySamples(CallerLeafFrameLoc.second.LineOffset,
-                                 CallerLeafFrameLoc.second.Discriminator,
+    CallerProfile.addBodySamples(CallerLeafFrameLoc.Callsite.LineOffset,
+                                 CallerLeafFrameLoc.Callsite.Discriminator,
                                  EstimatedCallCount);
     CallerProfile.addTotalSamples(EstimatedCallCount);
   }
 }
 
 void CSProfileGenerator::postProcessProfiles() {
   // Compute hot/cold threshold based on profile. This will be used for cold
   // context profile merging/trimming.
void CSProfileGenerator::computeSummaryAndThreshold() {
   auto Summary = Builder.computeSummaryForProfiles(ProfileMap);
   HotCountThreshold = ProfileSummaryBuilder::getHotCountThreshold(
       (Summary->getDetailedSummary()));
   ColdCountThreshold = ProfileSummaryBuilder::getColdCountThreshold(
       (Summary->getDetailedSummary()));
 }
 
 void CSProfileGenerator::write(std::unique_ptr<SampleProfileWriter> Writer,
-                               StringMap<FunctionSamples> &ProfileMap) {
+                               SampleProfileMap &ProfileMap) {
   if (std::error_code EC = Writer->write(ProfileMap))
     exitWithError(std::move(EC));
 }
 
 // Helper function to extract context prefix string stack
 // Extract context stack for reusing, leaf context stack will
 // be added compressed while looking up function profile
-static void
-extractPrefixContextStack(SmallVectorImpl<std::string> &ContextStrStack,
+static void extractPrefixContextStack(
+    SampleContextStorageType &ContextStack,
     const SmallVectorImpl<const MCDecodedPseudoProbe *> &Probes,
     ProfiledBinary *Binary) {
   for (const auto *P : Probes) {
-    Binary->getInlineContextForProbe(P, ContextStrStack, true);
+    Binary->getInlineContextForProbe(P, ContextStack, true);
   }
 }
 
 void PseudoProbeCSProfileGenerator::generateProfile() {
   // Enable pseudo probe functionalities in SampleProf
   FunctionSamples::ProfileIsProbeBased = true;
   FunctionSamples::ProfileIsCS = true;
   for (const auto &BI : BinarySampleCounters) {
     ProfiledBinary *Binary = BI.first;
     for (const auto &CI : BI.second) {
       const ProbeBasedCtxKey *CtxKey =
           dyn_cast<ProbeBasedCtxKey>(CI.first.getPtr());
-      SmallVector<std::string, 16> ContextStrStack;
-      extractPrefixContextStack(ContextStrStack, CtxKey->Probes, Binary);
+      SampleContextStorageType ContextStack;
+      extractPrefixContextStack(ContextStack, CtxKey->Probes, Binary);
       // Fill in function body samples from probes, also infer caller's samples
       // from callee's probe
-      populateBodySamplesWithProbes(CI.second.RangeCounter, ContextStrStack,
+      populateBodySamplesWithProbes(CI.second.RangeCounter, ContextStack,
                                     Binary);
       // Fill in boundary samples for a call probe
-      populateBoundarySamplesWithProbes(CI.second.BranchCounter,
-                                        ContextStrStack, Binary);
+      populateBoundarySamplesWithProbes(CI.second.BranchCounter, ContextStack,
+                                        Binary);
     }
   }
 
   postProcessProfiles();
 }
 
 void PseudoProbeCSProfileGenerator::extractProbesFromRange(
     const RangeSample &RangeCounter, ProbeCounterMap &ProbeCounter,
while (IP.Address <= RangeEnd) {
       }
 
       IP.advance();
     }
   }
 }
 
 void PseudoProbeCSProfileGenerator::populateBodySamplesWithProbes(
-    const RangeSample &RangeCounter,
-    SmallVectorImpl<std::string> &ContextStrStack, ProfiledBinary *Binary) {
+    const RangeSample &RangeCounter, SampleContextRefType ContextStack,
+    ProfiledBinary *Binary) {
   ProbeCounterMap ProbeCounter;
   // Extract the top frame probes by looking up each address among the range in
   // the Address2ProbeMap
   extractProbesFromRange(RangeCounter, ProbeCounter, Binary);
   std::unordered_map<MCDecodedPseudoProbeInlineTree *,
                      std::unordered_set<FunctionSamples *>>
       FrameSamples;
   for (auto PI : ProbeCounter) {
     const MCDecodedPseudoProbe *Probe = PI.first;
     uint64_t Count = PI.second;
     FunctionSamples &FunctionProfile =
-        getFunctionProfileForLeafProbe(ContextStrStack, Probe, Binary);
+        getFunctionProfileForLeafProbe(ContextStack, Probe, Binary);
     // 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);
       if (InlinerDesc != nullptr) {
         // 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.
-        StringRef CalleeContextId =
-            FunctionProfile.getContext().getNameWithContext();
-        StringRef CallerContextId;
-        FrameLocation &&CallerLeafFrameLoc =
+        SampleContextRefType CalleeContextId =
+            FunctionProfile.getContext().getFullContext();
+        SampleContextStorageType CallerContextId;
+        SampleCallSiteType &&CallerLeafFrameLoc =
             getCallerContext(CalleeContextId, CallerContextId);
-        uint64_t CallerIndex = CallerLeafFrameLoc.second.LineOffset;
+        uint64_t CallerIndex = CallerLeafFrameLoc.Callsite.LineOffset;
         assert(CallerIndex &&
                "Inferred caller's location index shouldn't be zero!");
+        // Save the new context for future references.
+        SampleContextRefType CallerContextRef =
+            *Contexts.insert(CallerContextId).first;
         FunctionSamples &CallerProfile =
-            getFunctionProfileForContext(CallerContextId);
+            getFunctionProfileForContext(CallerContextRef);
         CallerProfile.setFunctionHash(InlinerDesc->FuncHash);
         CallerProfile.addBodySamples(CallerIndex, 0, Count);
         CallerProfile.addTotalSamples(Count);
         CallerProfile.addCalledTargetSamples(
             CallerIndex, 0,
             FunctionProfile.getContext().getNameWithoutContext(), 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 (auto *FunctionProfile : I.second) {
       for (auto *Probe : I.first->getProbes()) {
         FunctionProfile->addBodySamplesForProbe(Probe->getIndex(), 0);
       }
     }
   }
 }
 
 void PseudoProbeCSProfileGenerator::populateBoundarySamplesWithProbes(
-    const BranchSample &BranchCounter,
-    SmallVectorImpl<std::string> &ContextStrStack, ProfiledBinary *Binary) {
+    const BranchSample &BranchCounter, SampleContextRefType ContextStack,
+    ProfiledBinary *Binary) {
   for (auto BI : BranchCounter) {
     uint64_t SourceOffset = BI.first.first;
     uint64_t TargetOffset = BI.first.second;
     uint64_t Count = BI.second;
     uint64_t SourceAddress = Binary->offsetToVirtualAddr(SourceOffset);
     const MCDecodedPseudoProbe *CallProbe =
         Binary->getCallProbeForAddr(SourceAddress);
     if (CallProbe == nullptr)
       continue;
     FunctionSamples &FunctionProfile =
-        getFunctionProfileForLeafProbe(ContextStrStack, CallProbe, Binary);
+        getFunctionProfileForLeafProbe(ContextStack, CallProbe, Binary);
     FunctionProfile.addBodySamples(CallProbe->getIndex(), 0, Count);
     FunctionProfile.addTotalSamples(Count);
     StringRef CalleeName = FunctionSamples::getCanonicalFnName(
         Binary->getFuncFromStartOffset(TargetOffset));
     if (CalleeName.size() == 0)
       continue;
     FunctionProfile.addCalledTargetSamples(CallProbe->getIndex(), 0, CalleeName,
                                            Count);
   }
 }
 
 FunctionSamples &PseudoProbeCSProfileGenerator::getFunctionProfileForLeafProbe(
-    SmallVectorImpl<std::string> &ContextStrStack,
-    const MCPseudoProbeFuncDesc *LeafFuncDesc, bool WasLeafInlined) {
-  assert(ContextStrStack.size() && "Profile context must have the leaf frame");
-  // Compress the context string except for the leaf frame
-  std::string LeafFrame = ContextStrStack.back();
-  ContextStrStack.pop_back();
-  CSProfileGenerator::compressRecursionContext(ContextStrStack);
-  CSProfileGenerator::trimContext(ContextStrStack);
-
-  std::ostringstream OContextStr;
-  for (uint32_t I = 0; I < ContextStrStack.size(); I++) {
-    if (OContextStr.str().size())
-      OContextStr << " @ ";
-    OContextStr << ContextStrStack[I];
-  }
+    SampleContextRefType ContextStack, const MCDecodedPseudoProbe *LeafProbe,
+    ProfiledBinary *Binary) {
+
+  // Explicitly copy the context for appending the leaf context
+  SampleContextStorageType 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"
-  if (OContextStr.str().size())
-    OContextStr << " @ ";
-  OContextStr << StringRef(LeafFrame).split(":").first.str();
-
-  FunctionSamples &FunctionProile =
-      getFunctionProfileForContext(OContextStr.str(), WasLeafInlined);
-  FunctionProile.setFunctionHash(LeafFuncDesc->FuncHash);
-  return FunctionProile;
-}
-
-FunctionSamples &PseudoProbeCSProfileGenerator::getFunctionProfileForLeafProbe(
-    SmallVectorImpl<std::string> &ContextStrStack,
-    const MCDecodedPseudoProbe *LeafProbe, ProfiledBinary *Binary) {
+  auto LeafFrame = NewContextStack.back();
+  LeafFrame.Callsite = 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);
+  // Save the new context for future references.
+  SampleContextRefType NewContextRef = *Contexts.insert(NewContextStack).first;
 
-  // Explicitly copy the context for appending the leaf context
-  SmallVector<std::string, 16> ContextStrStackCopy(ContextStrStack.begin(),
-                                                   ContextStrStack.end());
-  Binary->getInlineContextForProbe(LeafProbe, ContextStrStackCopy, true);
   const auto *FuncDesc = Binary->getFuncDescForGUID(LeafProbe->getGuid());
   bool WasLeafInlined = LeafProbe->getInlineTreeNode()->hasInlineSite();
-  return getFunctionProfileForLeafProbe(ContextStrStackCopy, FuncDesc,
-                                        WasLeafInlined);
+  FunctionSamples &FunctionProile =
+      getFunctionProfileForContext(NewContextRef, WasLeafInlined);
+  FunctionProile.setFunctionHash(FuncDesc->FuncHash);
+  return FunctionProile;
 }
 
 } // end namespace sampleprof
 } // end namespace llvm

llvm/tools/llvm-profgen/ProfiledBinary.h

class ProfiledBinary {
   std::unique_ptr<const MCInstrAnalysis> MIA;
   std::unique_ptr<MCInstPrinter> IPrinter;
   // A list of text sections sorted by start RVA and size. Used to check
   // if a given RVA is a valid code address.
   std::set<std::pair<uint64_t, uint64_t>> TextSections;
   // Function offset to name mapping.
   std::unordered_map<uint64_t, std::string> FuncStartAddrMap;
   // Offset to context location map. Used to expand the context.
-  std::unordered_map<uint64_t, FrameLocationStack> Offset2LocStackMap;
+  std::unordered_map<uint64_t, SampleContextStorageType> Offset2LocStackMap;
   // An array of offsets of all instructions sorted in increasing order. The
   // sorting is needed to fast advance to the next forward/backward instruction.
   std::vector<uint64_t> CodeAddrs;
   // A set of call instruction offsets. Used by virtual unwinding.
   std::unordered_set<uint64_t> CallAddrs;
   // A set of return instruction offsets. Used by virtual unwinding.
   std::unordered_set<uint64_t> RetAddrs;
 
   PrologEpilogTracker ProEpilogTracker;
 
   // The symbolizer used to get inline context for an instruction.
   std::unique_ptr<symbolize::LLVMSymbolizer> Symbolizer;
 
+  // String table owning function name strings created from the symbolizer.
+  std::unordered_set<std::string> NameStrings;
+
   // Pseudo probe decoder
   MCPseudoProbeDecoder ProbeDecoder;
 
   bool UsePseudoProbes = false;
 
   // Indicate if the base loading address is parsed from the mmap event or uses
   // the preferred address
   bool IsLoadedByMMap = false;
class ProfiledBinary {
 
   /// Dissassemble the text section and build various address maps.
   void disassemble(const ELFObjectFileBase *O);
 
   /// Helper function to dissassemble the symbol and extract info for unwinding
   bool dissassembleSymbol(std::size_t SI, ArrayRef<uint8_t> Bytes,
                           SectionSymbolsTy &Symbols, const SectionRef &Section);
   /// Symbolize a given instruction pointer and return a full call context.
-  FrameLocationStack symbolize(const InstructionPointer &IP,
+  SampleContextStorageType symbolize(const InstructionPointer &IP,
                                      bool UseCanonicalFnName = false);
 
   /// Decode the interesting parts of the binary and build internal data
   /// structures. On high level, the parts of interest are:
   ///   1. Text sections, including the main code section and the PLT
   ///   entries that will be used to handle cross-module call transitions.
   ///   2. The .debug_line section, used by Dwarf-based profile generation.
   ///   3. Pseudo probe related sections, used by probe-based profile
   ///   generation.
   void load();
-  const FrameLocationStack &getFrameLocationStack(uint64_t Offset) const {
+  const SampleContextStorageType &getFrameLocationStack(uint64_t Offset) const {
     auto I = Offset2LocStackMap.find(Offset);
     assert(I != Offset2LocStackMap.end() &&
            "Can't find location for offset in the binary");
     return I->second;
   }
 
 public:
   ProfiledBinary(StringRef Path) : Path(Path), ProEpilogTracker(this) {
public:
   uint64_t getCallAddrFromFrameAddr(uint64_t FrameAddr) const {
     return getAddressforIndex(getIndexForAddr(FrameAddr) - 1);
   }
 
   StringRef getFuncFromStartOffset(uint64_t Offset) {
     return FuncStartAddrMap[Offset];
   }
 
-  Optional<FrameLocation> getInlineLeafFrameLoc(uint64_t Offset) {
+  Optional<SampleCallSiteType> getInlineLeafFrameLoc(uint64_t Offset) {
     const auto &Stack = getFrameLocationStack(Offset);
     if (Stack.empty())
       return {};
     return Stack.back();
   }
 
   // Compare two addresses' inline context
   bool inlineContextEqual(uint64_t Add1, uint64_t Add2) const;
 
-  // Get the context string of the current stack with inline context filled in.
+  // Get the full context of the current stack with inline context filled in.
   // It will search the disassembling info stored in Offset2LocStackMap. This is
   // used as the key of function sample map
-  std::string getExpandedContextStr(const SmallVectorImpl<uint64_t> &Stack,
-                                    bool &WasLeafInlined) const;
+  SampleContextStorageType
+  getExpandedContext(const SmallVectorImpl<uint64_t> &Stack,
 
   const MCDecodedPseudoProbe *getCallProbeForAddr(uint64_t Address) const {
     return ProbeDecoder.getCallProbeForAddr(Address);
   }
 
-  void
-  getInlineContextForProbe(const MCDecodedPseudoProbe *Probe,
-                           SmallVectorImpl<std::string> &InlineContextStack,
+  void getInlineContextForProbe(const MCDecodedPseudoProbe *Probe,
+                                SampleContextStorageType &InlineContextStack,
                                 bool IncludeLeaf = false) const {
-    return ProbeDecoder.getInlineContextForProbe(Probe, InlineContextStack,
-                                                 IncludeLeaf);
+    SmallVector<MCPseduoProbeFrameLocation, 16> ProbeInlineContext;
+    ProbeDecoder.getInlineContextForProbe(Probe, ProbeInlineContext,
+    for (auto &Callsite : ProbeInlineContext) {
+      InlineContextStack.push_back(
+          SampleCallSiteType(Callsite.first, LineLocation(Callsite.second, 0)));
+    }
   }
   const AddressProbesMap &getAddress2ProbesMap() const {
     return ProbeDecoder.getAddress2ProbesMap();
   }
   const MCPseudoProbeFuncDesc *getFuncDescForGUID(uint64_t GUID) {
     return ProbeDecoder.getFuncDescForGUID(GUID);
   }
 

llvm/tools/llvm-profgen/ProfiledBinary.cpp

void ProfiledBinary::load() {
 
   // TODO: decode other sections.
 }
 
 bool ProfiledBinary::inlineContextEqual(uint64_t Address1,
                                         uint64_t Address2) const {
   uint64_t Offset1 = virtualAddrToOffset(Address1);
   uint64_t Offset2 = virtualAddrToOffset(Address2);
-  const FrameLocationStack &Context1 = getFrameLocationStack(Offset1);
-  const FrameLocationStack &Context2 = getFrameLocationStack(Offset2);
+  const SampleContextStorageType &Context1 = getFrameLocationStack(Offset1);
+  const SampleContextStorageType &Context2 = getFrameLocationStack(Offset2);
   if (Context1.size() != Context2.size())
     return false;
   if (Context1.empty())
     return false;
   // The leaf frame contains location within the leaf, and it
   // needs to be remove that as it's not part of the calling context
   return std::equal(Context1.begin(), Context1.begin() + Context1.size() - 1,
                     Context2.begin(), Context2.begin() + Context2.size() - 1);
 }
 
-std::string
-ProfiledBinary::getExpandedContextStr(const SmallVectorImpl<uint64_t> &Stack,
-                                      bool &WasLeafInlined) const {
-  std::string ContextStr;
-  SmallVector<std::string, 16> ContextVec;
+SampleContextStorageType
+ProfiledBinary::getExpandedContext(const SmallVectorImpl<uint64_t> &Stack,
+  SampleContextStorageType ContextVec;
   // Process from frame root to leaf
   for (auto Address : Stack) {
     uint64_t Offset = virtualAddrToOffset(Address);
-    const FrameLocationStack &ExpandedContext = getFrameLocationStack(Offset);
+    const SampleContextStorageType &ExpandedContext =
+        getFrameLocationStack(Offset);
     // An instruction without a valid debug line will be ignored by sample
     // processing
     if (ExpandedContext.empty())
-      return std::string();
+      return SampleContextStorageType();
     // Set WasLeafInlined to the size of inlined frame count for the last
     // address which is leaf
     WasLeafInlined = (ExpandedContext.size() > 1);
-    for (const auto &Loc : ExpandedContext) {
-      ContextVec.push_back(getCallSite(Loc));
-    }
+    ContextVec.append(ExpandedContext);
   }
 
-  assert(ContextVec.size() && "Context length should be at least 1");
   // Compress the context string except for the leaf frame
-  std::string LeafFrame = ContextVec.back();
+  auto LeafFrame = ContextVec.back();
+  LeafFrame.Callsite = LineLocation(0, 0);
   ContextVec.pop_back();
-  CSProfileGenerator::compressRecursionContext<std::string>(ContextVec);
-  CSProfileGenerator::trimContext<std::string>(ContextVec);
-
-  std::ostringstream OContextStr;
-  for (uint32_t I = 0; I < (uint32_t)ContextVec.size(); I++) {
-    if (OContextStr.str().size()) {
-      OContextStr << " @ ";
-    }
-    OContextStr << ContextVec[I];
-  }
-  // Only keep the function name for the leaf frame
-  if (OContextStr.str().size())
-    OContextStr << " @ ";
-  OContextStr << StringRef(LeafFrame).split(":").first.str();
-  return OContextStr.str();
+  assert(ContextVec.size() && "Context length should be at least 1");
+  CSProfileGenerator::compressRecursionContext(ContextVec);
+  CSProfileGenerator::trimContext(ContextVec);
+  ContextVec.push_back(LeafFrame);
+  return ContextVec;
 }
 
 template <class ELFT>
 void ProfiledBinary::setPreferredTextSegmentAddresses(const ELFFile<ELFT> &Obj, StringRef FileName) {
   const auto &PhdrRange = unwrapOrError(Obj.program_headers(), FileName);
   for (const typename ELFT::Phdr &Phdr : PhdrRange) {
     if ((Phdr.p_type == ELF::PT_LOAD) && (Phdr.p_flags & ELF::PF_X)) {
         // Segments will always be loaded at a page boundary.
if (ShowDisassemblyOnly) {
       outs() << "\n";
     }
 
     if (Disassembled) {
       const MCInstrDesc &MCDesc = MII->get(Inst.getOpcode());
       // Populate a vector of the symbolized callsite at this location
       // We don't need symbolized info for probe-based profile, just use an
       // empty stack as an entry to indicate a valid binary offset
-      FrameLocationStack SymbolizedCallStack;
+      SampleContextStorageType SymbolizedCallStack;
       if (!UsePseudoProbes) {
         InstructionPointer IP(this, Offset);
         SymbolizedCallStack = symbolize(IP, true);
       }
       Offset2LocStackMap[Offset] = SymbolizedCallStack;
       // Populate address maps.
       CodeAddrs.push_back(Offset);
       if (MCDesc.isCall())
SymbolizerOpts.PrintFunctions =
       DILineInfoSpecifier::FunctionNameKind::LinkageName;
   SymbolizerOpts.Demangle = false;
   SymbolizerOpts.DefaultArch = TheTriple.getArchName().str();
   SymbolizerOpts.UseSymbolTable = false;
   SymbolizerOpts.RelativeAddresses = false;
   Symbolizer = std::make_unique<symbolize::LLVMSymbolizer>(SymbolizerOpts);
 }
 
-FrameLocationStack ProfiledBinary::symbolize(const InstructionPointer &IP,
+SampleContextStorageType ProfiledBinary::symbolize(const InstructionPointer &IP,
                                                    bool UseCanonicalFnName) {
   assert(this == IP.Binary &&
          "Binary should only symbolize its own instruction");
   auto Addr = object::SectionedAddress{IP.Offset + getPreferredBaseAddress(),
                                        object::SectionedAddress::UndefSection};
   DIInliningInfo InlineStack =
       unwrapOrError(Symbolizer->symbolizeInlinedCode(Path, Addr), getName());
 
-  FrameLocationStack CallStack;
+  SampleContextStorageType CallStack;
 
   for (int32_t I = InlineStack.getNumberOfFrames() - 1; I >= 0; I--) {
     const auto &CallerFrame = InlineStack.getFrame(I);
     if (CallerFrame.FunctionName == "<invalid>")
       break;
     StringRef FunctionName(CallerFrame.FunctionName);
     if (UseCanonicalFnName)
       FunctionName = FunctionSamples::getCanonicalFnName(FunctionName);
     LineLocation Line(CallerFrame.Line - CallerFrame.StartLine,
                       DILocation::getBaseDiscriminatorFromDiscriminator(
                           CallerFrame.Discriminator,
                           /* IsFSDiscriminator */ false));
-    FrameLocation Callsite(FunctionName.str(), Line);
+    auto It = NameStrings.insert(FunctionName.str());
+    SampleCallSiteType Callsite(*It.first, Line);
     CallStack.push_back(Callsite);
   }
 
   return CallStack;
 }
 
 InstructionPointer::InstructionPointer(ProfiledBinary *Binary, uint64_t Address,
                                        bool RoundToNext)
     : Binary(Binary), Address(Address) {