[CSSPGO] split context string II - reader/writer changes

[hoy]

Summary

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

[wmi]

llvm/include/llvm/ProfileData/SampleProfReader.h

743

Same as SampleProfileReaderText, can we use std::list instead of std::vector so the storage underlying will not be moved after more elements are inserted?

llvm/lib/ProfileData/SampleProfReader.cpp

238

Looks like it can be a utility function in SampleProfile.h and may be reused somewhere else.

269–276

Can we make it a utility function?

462

This param is unused. Write it as readNameFromTable(bool /* IsContextName */) to make it more explicit.

[hoy]

llvm/include/llvm/ProfileData/SampleProfReader.h

743

Std::vector is good for looking up with an offset, which is needed during the load of function profiles. Also, std::vector is more efficient for the extbinary reader, since number of contexts is known beforehand. Once the CS name table section is read, the std::vector should be frozen and no more names should be appended.

llvm/lib/ProfileData/SampleProfReader.cpp

238

It's actually moved out of SampleProfile.h since it's only used by the text reader. Also the text reader should be the only part of the toolchain that needs to parse full context strings. Other parts just need to deal with ArrayRef form of context. What do you think?

462

Sounds good.

[hoy]

llvm/lib/ProfileData/SampleProfReader.cpp

462

Actually I can only do that in the header file for the function declaration. The function definition needs a real name for param.

[hoy]

Addressing Wei's comment.

[wmi]

llvm/include/llvm/ProfileData/SampleProfReader.h

743

Yeah, right, std::vector is good for looking up. Now once cs name table is read, there is no other code appending names to the table. Just want to know whether we can have some additional mechanism to enforce that. Can we save the size of the table after reading cs name table and check that the table size is the same before visiting any element from the table?

llvm/lib/ProfileData/SampleProfReader.cpp

238

Ok, if it is less likely to be used by others, it is fine to keep as it is.

[hoy]

llvm/include/llvm/ProfileData/SampleProfReader.h

743

Good point. Setting the underlying object to be constant so that it is immutable once populated.

[hoy]

Addressing Wei's feedback.

[wmi]

LGTM.

[hoy]

Updating D108435: [CSSPGO] split context string II - reader/writer changes

[hoy]

Updating D108435: [CSSPGO] split context string II - reader/writer changes

[wenlei]

llvm/include/llvm/ProfileData/SampleProfReader.h

135

In this section, we need to add documentation/spec for CSNameTable how context is stored, and how it interacts with NameTable.

507

nit: remove the blank line

578

typo: given

llvm/lib/ProfileData/SampleProfReader.cpp

282–283

Even though we only need to deal with string context in profile reader/writer for text profile, it's probably still cleaner to keep all string context related parsing into SampleContext.

createContext is more like a ctor. I'd prefer keep string decoding, createContext in its original place in SampleContext. That way, we can construct a context from string, SampleContext::setContext remain a private helper too, and the logic here can be simpler, just like before.

SampleContext does still have getContextString and toString, so it's not really isolated from string representation, might as well keep all string stuff together there for consistency.

457

I think getNameTable exist as a virtual function only because of the use in sample loader for profile-sample-accuarate and profile-accurate-for-symsinlist. Here can we still access the NameTable directly without going through virtual call? Same for other places where NameTable is directly accessible. For writer, before this change, we indeed only have NameTable for SampleProfileWriterBinary without virtual getter, and we access it without going through virtual function call there.

462

Can we avoid this parameter altogether? For one profile, it's either going to be CS or non-CS, so the dispatch can be based on state/member of the reader instead of relying on a param for each invocation.

592

Same here, avoid IsContextName as param, but dispatch based on SampleProfileReader::ProfileIsCS

597–598

nit: this can be confusing, readNameFromTable can mislead people to think FContext is string (or vector of strings) type. the auto also isn't helping. Either spell out the type name, or rename readNameFromTable to something like readSampleContextFromTable

758

Since we have SampleProfileReader::ProfileIsCS, within reader it probable makes more sense to use the reader instance flag as opposed to the global FunctionSamples::ProfileIsCS flag.

We were not very consistent in past, might be good to clean up as you're expanding the use of ProfileIsCS.

784

OrderedNames -> OrderedContexts

1051

Let emplace_back take variadic arguments and forward to the constructor directly instead of doing a move copy?

1055

How about we emplace an empty slot before going into the inner loop, then we can operate on the vector directly (PNameVec.back().emplace_back(...)) and avoid copying temporary Context onto PNameVec?

1071–1072

nit: peal/hoist the get() so we don't have to call it for every use.

[wenlei]

llvm/include/llvm/ProfileData/SampleProfWriter.h

138–139

Now that the parameter is no longer a string name, rename the function as well, e.g. writeContextIdx?

Same for addName(const SampleContext &Context) -> addContext

llvm/lib/ProfileData/SampleProfReader.cpp

785–803

Which operation here is the most costly and visible for e2e build time? The insert/sort/find or the prefix check? What's the % of time here for both prelink and postlink?

llvm/lib/ProfileData/SampleProfWriter.cpp

249

nit: this contains leaf frame, so it's Frames instead of Callsites.

482

assert !Context.hasContext() ?

llvm/tools/llvm-profdata/llvm-profdata.cpp

524–526

I suggest rename these FName (key of the SampleProfileMap) to be FContext accordingly given this is no longer a string map. Same for other places.

[hoy]

llvm/include/llvm/ProfileData/SampleProfReader.h

135

CSNameTable is a concept of the extensible binary format, not the basic binary format, like other concepts such as func name offset table, func metadata table. We could comment about CSNameTable in the SampleProfileReaderExtBinary/SampleProfileWriterExtBinary, but I didn't find a good centralized place to do that. Seems that descriptions about other sections are scattered around the implementation code.

507

done.

578

fixed.

llvm/include/llvm/ProfileData/SampleProfWriter.h

138–139

Sounds good.

llvm/lib/ProfileData/SampleProfReader.cpp

282–283

Makes sense. Moved back to SampleContext.

457

Yes, we can just use NameTable directly here. Actually all changes here are unnecessary. I just undid them.

462

Good point. Checking FunctionSamples::ProfileIsCS should be enough.

597–598

Fixed by using explicit type.

758

That's reasonable. Replaced all FunctionSamples::ProfileIsCS usage in the reader with the reader ProfileIsCS flag.

784

fixed.

785–803

The insert/sort/find operations are the most expansive. For thinlto postlink, they count to 15% to 20% of the whole backend running time. I haven't measured for prelink but given the similarity of prelink and postlink, they should be expansive there too.

I'm actually working on a sample writer change that emits the func offset table in the order of contexts so that we don't need the set operations here. That turns out very effective. Will send out a separate diff for it.

1055

Sounds good, that should be faster.

1071–1072

You mean save FContext.get() in a temp and use it in the loop? Thought the compiler would do it.

llvm/lib/ProfileData/SampleProfWriter.cpp

249

Frames sounds good.

482

Added the assert.

llvm/tools/llvm-profdata/llvm-profdata.cpp

524–526

Sounds good.

[hoy]

Addressing Wenlei's comment.

[wenlei]

llvm/include/llvm/ProfileData/SampleProfReader.h

135

Ok, didn't realize we don't use header comment for extbin. That's fine. Yeah, it'd be good to have some comment for CSNameTable somewhere to explain the nesting structure of context and raw names.

llvm/lib/ProfileData/SampleProfReader.cpp

282–283

Thanks. Can you move it back in D107299, so we don't see the change back and forth, just for review?

Can we also fold FName.startswith("[") in to SampleContext?

Additionally, why we do need a createContextFromString instead of using overload ctors? This is also inconsistent between how SampleContext is created from text profile (createContextFromString) vs from binary profile (ctor). What I was thinking is have SampleContext decide how to create the object, so there's no changed needed here, just like before.

597–598

Ok, but on 2nd thought, why do we call it readName while it's actually returning a context? Especially given that we've renamed addName to AddContext, writeNameIdx to writeContextIdx.

731

FName to FContext as well?

785–803

I'm actually working on a sample writer change that emits the func offset table in the order of contexts so that we don't need the set operations here. That turns out very effective.

Great. I was thinking about exactly that too. The order of binary profile is opaque to users, so we can order them in the file to save sorting on reading.

For thinlto postlink, they count to 15% to 20% of the whole backend running time.

What was the % before this work when we were all using StringRef?

1071–1072

Yeah, it's less about optimization but for readability to make the code less verbose.

[hoy]

llvm/include/llvm/ProfileData/SampleProfReader.h

135

Comment added to the definition of SampleProfileReaderExtBinaryBase::readCSNameTableSec.

llvm/lib/ProfileData/SampleProfReader.cpp

282–283

Constructing a CS context will require additional parameter than non-CS profile, especially the underlying context vector. That is causing the inconsistency with the context split work and I'm separating the construction of CS and non-CS contexts. Currently CS context is only constructible from a SampleContextFrameVector.

Another reason is that I was hoping to construct SampleContext in a quick manner from StringRef to favor non-CS profile.

731

Oops, this was fixed in my other unified patch, somehow got dropped when rebasing. Will send an update to the other patch as you suggested.

785–803

What was the % before this work when we were all using StringRef?

It's not noticeable. Actually I didn't see the previous sorting as a hot routine in the profile. With the presorted func offset stable, we are able to achieve similar performance with the previous approach.

[wenlei]

llvm/lib/ProfileData/SampleProfReader.cpp

282–283

Can you move it back in D107299, so we don't see the change back and forth, just for review?

Sorry my bad, I meant to say D108433.

I think the key blocker for everything to be taken care from within ctor is that you need reader to own the context created from the string. How about having a ctor with StringRef and CSNameTable as parameter - it puts the context into CSNameTable (owned by reader) for CS profile, and the CSNameTable would be ignored for non-CS profile.

The current implementation works, but reader has to be aware of the actual string representation of context. I thought it'd be cleaner if such representation is all dealt with from within SampleContext. Currently, it's indeed mostly handled by SampleContext, except it's bleeding into reader here.

597–598

I think it'd be good to establish a convention as to when we call things Name vs Context. My thought is it goes with the type, what do you think?

1023

Comment added to the definition of SampleProfileReaderExtBinaryBase::readCSNameTableSec.

Not seeing comments in the latest update. Did I miss anything?

[hoy]

llvm/lib/ProfileData/SampleProfReader.cpp

282–283

Yeah, that's why I originally put everything related to context string parsing in the reader. I thought they were related to profile-specific representation that SampleContext shouldn't care.

I guess moving the parsing code back is better for extension and code sharing, when we have new module doing the same thing in the future. Adding a new constructor with StringRef and CSNameTable as parameter can work but feel like it's easier and more clear for the reader to know what it constructs. Hiding that from the reader is fine, but asking the reader to provide an underlying CS name table which may not be needed for non-CS sounds a bit confusing. That also kinds of exposes reader-specific implementation to SampleContext.

I think having the reader be aware of the specific representation might be reasonable, since context representation is a part of the profile format. What do you think?

597–598

Yeah, I've been using Name as an identifier of of the context and function name string, and using Context for CS profile and String for function names. But sometimes Name and String as mixed.

We are using SampleContext for both CS and non-CS. And we are also using the word context specifically for CS. Sounds like we need a more general name (instead of Name) in the reader for both of them. How about readIdFromTable?

1071–1072

Changed to using *FContext for readability, like other places. Look good?

Alternatively, we can have two variables such as FContextError and FContext.

[hoy]

Updating D108435: [CSSPGO] split context string II - reader/writer changes

[wenlei]

llvm/lib/ProfileData/SampleProfReader.cpp

282–283

asking the reader to provide an underlying CS name table which may not be needed for non-CS sounds a bit confusing. That also kinds of exposes reader-specific implementation to SampleContext.

I don't see this as exposing reader-specific stuff to sample context. It'd just be the expectation of the sample context API that requires a buffer to hold newly created context.

This is in essence no different from how string buffer is passed in getRepInFormat (and only used for MD5), and I don't see it as coupling between reader and context.

I think having the reader be aware of the specific representation might be reasonable, since context representation is a part of the profile format.

I think the actual string presentation of context is not tied to the format, but rather the format is using the string representation directly. Also if string representation is considered part of format, then all of the string related stuff should be part of reader (like your earlier change). The possible reusing of string functions you mentioned also shows that the string representation is not format-specific.

597–598

I've been using Name as an identifier of of the context and function name string, and using Context for CS profile and String for function names.

The problem is you can't do that cleanly because there're cases that covers both CS and non-CS. readNameFromTable is one example.

I think that name it according to the type makes things clearer. For non-CS profile we're still using SampleContext as key in the profile map anyways. Adding a new notion of Id in addition to name and context seem unnecessary (we have Idx already which can confuse people if we add Id). So I think readSampleContextFromTable/readContextFromTable is better.

1071–1072

either works, thanks.

[hoy]

llvm/lib/ProfileData/SampleProfReader.cpp

282–283

This is in essence no different from how string buffer is passed in getRepInFormat (and only used for MD5), and I don't see it as coupling between reader and context.

getRepInFormat takes a string buffer instead of a string table which is not coupled with the reader. However, CSNameTable is an implementation detail of the reader. Having that populated and updated by SampleContext sounds like a coupling with the reader. E.g, from SampleContext point of view, it is questionable why it has to update a std::list but not a std::vector or std::set.

We can pass in a SampleContextFrameVector to the new ctor, and similar with the callsites of getRepInFormat, the reader should know where to place it, just like the current code. But then the reader has to redo some checks in the ctor. Anyway, feels that the reader will need to check if it is constructing a CS context, if we don't expose CSNameTable to SampleContext.

Since the reader is currently the only consumer, maybe just keep it in the reader for now? When we have a new user we can then decide what kind of new ctor to make?

597–598

Sounds good. Will use readSampleContextFromTable.

[hoy]

Renaming readNameFromTable to readSampleContextFromTable.

[wenlei]

llvm/lib/ProfileData/SampleProfReader.cpp

282–283

Ok, I think we've probably spent too much time on this. But bear with me a bit more, still hope we can reach a consensus. :) That said, whatever we choose to do, I don't agree with your reasoning above.

However, CSNameTable is an implementation detail of the reader.

Look at it as buffer, and in this case reader happens to own and provide that buffer.

Having that populated and updated by SampleContext sounds like a coupling with the reader. E.g, from SampleContext point of view, it is questionable why it has to update a std::list but not a std::vector or std::set.

I don't think this is really a concern, nor is it a coupling tbh. Regardless of what reader does, what makes sense for a buffer type? I think std::list simply makes sense as it avoid reallocating. Yes, reader uses std::list, but this is just a choice that makes sense in general for a buffer. If we go down this level, sure every single API call is a coupling because every param has a type.

The way I look at this - sample context abstracts away the string representation (again it's independent of format, and text format just uses that string representation directly), and reader resort to sample context for anything related to that. It's probably not so important to have that logical layering enforced, but if the cost of doing that is small, why not? Clear layering makes it easy to maintain and reason about.

[hoy]

llvm/lib/ProfileData/SampleProfReader.cpp

282–283

I don't think this is really a concern, nor is it a coupling tbh. Regardless of what reader does, what makes sense for a buffer type? I think std::list simply makes sense as it avoid reallocating. Yes, reader uses std::list, but this is just a choice that makes sense in general for a buffer. If we go down this level, sure every single API call is a coupling because every param has a type.

Yeah, we probably spent too much time on this. I will make a new ctor with CSNameTable as a parameter.

[hoy]

Updating D108435: [CSSPGO] split context string II - reader/writer changes

[wenlei]

llvm/include/llvm/ProfileData/SampleProf.h

532

This can be removed now?

[hoy]

Updating D108435: [CSSPGO] split context string II - reader/writer changes

llvm/include/llvm/ProfileData/SampleProf.h

532

yes, it's removed from the other patch but somehow not updated here.

[wenlei]

lgtm, thanks for working through the comments!

llvm/include/llvm/ProfileData/ProfileCommon.h

 
 class SampleProfileSummaryBuilder final : public ProfileSummaryBuilder {
 public:
   SampleProfileSummaryBuilder(std::vector<uint32_t> Cutoffs)
       : ProfileSummaryBuilder(std::move(Cutoffs)) {}
 
   void addRecord(const sampleprof::FunctionSamples &FS,
                  bool isCallsiteSample = false);
-  std::unique_ptr<ProfileSummary> computeSummaryForProfiles(
-      const StringMap<sampleprof::FunctionSamples> &Profiles);
+  std::unique_ptr<ProfileSummary>
+  computeSummaryForProfiles(const sampleprof::SampleProfileMap &Profiles);
   std::unique_ptr<ProfileSummary> getSummary();
 };
 
 /// This is called when a count is seen in the profile.
 void ProfileSummaryBuilder::addCount(uint64_t Count) {
   TotalCount += Count;
   if (Count > MaxCount)
     MaxCount = Count;
   NumCounts++;
   CountFrequencies[Count]++;
 }
 
 } // end namespace llvm
 
 #endif // LLVM_PROFILEDATA_PROFILECOMMON_H

llvm/include/llvm/ProfileData/SampleProf.h

 // 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.
public:
   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 FullContext.size() == 1; }
   StringRef getName() const { return Name; }
   StringRef getNameWithoutContext() const { return Name; }
   SampleContextRefType getFullContext() const { return FullContext; }
 
-  static std::string getContextString(SampleContextRefType Context) {
-    std::string ContextStr;
+  // Decode context string for a frame to get function name and location.
+  // `ContextStr` is in the form of `FuncName:StartLine.Discriminator`.
+  static void decodeContextString(StringRef ContextStr, StringRef &FName,
+                                  LineLocation &LineLoc) {
+    // Get function name
+    auto EntrySplit = ContextStr.split(':');
+    FName = EntrySplit.first;
+
+    LineLoc = {0, 0};
+    if (!EntrySplit.second.empty()) {
+      // Get line offset, use signed int for getAsInteger so string will
+      // be parsed as signed.
+      int LineOffset = 0;
+      auto LocSplit = EntrySplit.second.split('.');
+      LocSplit.first.getAsInteger(10, LineOffset);
+      LineLoc.LineOffset = LineOffset;
+
+      // Get discriminator
+      if (!LocSplit.second.empty())
+        LocSplit.second.getAsInteger(10, LineLoc.Discriminator);
+    }
+  }
+
+  /// Create a context vector from a given context string and save it in
+  /// `Context`.
+  static void createContextFromString(StringRef ContextStr,
+                                      SampleContextFrameVector &Context) {
+    Context.clear();
+    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);
+    }
+  }
+
+  static std::string getContextString(SampleContextFrames Context,
+                                      bool IncludeLeafLineLocation = false) {
+    std::ostringstream OContextStr;
     for (uint32_t I = 0; I < Context.size(); I++) {
       if (ContextStr.size()) {
         ContextStr += " @ ";
       }
       ContextStr += Context[I].str();
     }
     return ContextStr;
   }

llvm/include/llvm/ProfileData/SampleProfReader.h

 // We support the following types of metadata:
 //
 // a. CFG Checksum (a.k.a. function hash):
 //   !CFGChecksum: 12345
 // b. CFG Checksum (see ContextAttributeMask):
 //   !Atribute: 1
 //
 //
 // Binary format
 // -------------
 //
 // This is a more compact encoding. Numbers are encoded as ULEB128 values
 // and all strings are encoded in a name table. The file is organized in
 // the following sections:
 //
 // MAGIC (uint64_t)
 //    File identifier computed by function SPMagic() (0x5350524f463432ff)

 #include "llvm/ProfileData/SampleProf.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/Discriminator.h"
 #include "llvm/Support/ErrorOr.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/SymbolRemappingReader.h"
 #include <algorithm>
 #include <cstdint>
+#include <list>
 #include <memory>
 #include <string>
 #include <system_error>
 #include <unordered_set>
 #include <vector>
 
 namespace llvm {
 
if (Remapper)
       Remapper->applyRemapping(Ctx);
     FunctionSamples::UseMD5 = useMD5();
     return sampleprof_error::success;
   }
 
   /// The implementaion to read sample profiles from the associated file.
   virtual std::error_code readImpl() = 0;
 
-  /// Print the profile for \p FName on stream \p OS.
-  void dumpFunctionProfile(StringRef FName, raw_ostream &OS = dbgs());
+  /// Print the profile for \p FContext on stream \p OS.
+  void dumpFunctionProfile(SampleContext FContext, raw_ostream &OS = dbgs());
 
   /// Collect functions with definitions in Module M. For reader which
   /// support loading function profiles on demand, return true when the
   /// reader has been given a module. Always return false for reader
   /// which doesn't support loading function profiles on demand.
   virtual bool collectFuncsFromModule() { return false; }
 
   /// Print all the profiles on stream \p OS.
if (Remapper) {
         if (It != Profiles.end())
           return &It->second;
       }
     }
     return nullptr;
   }
 
   /// Return all the profiles.
-  StringMap<FunctionSamples> &getProfiles() { return Profiles; }
+  SampleProfileMap &getProfiles() { return Profiles; }
 
   /// Report a parse error message.
   void reportError(int64_t LineNumber, const Twine &Msg) const {
     Ctx.diagnose(DiagnosticInfoSampleProfile(Buffer->getBufferIdentifier(),
                                              LineNumber, Msg));
   }
 
   /// Create a sample profile reader appropriate to the file format.
public:
   void setModule(const Module *Mod) { M = Mod; }
 
 protected:
   /// Map every function to its associated profile.
   ///
   /// The profile of every function executed at runtime is collected
   /// in the structure FunctionSamples. This maps function objects
   /// to their corresponding profiles.
-  StringMap<FunctionSamples> Profiles;
+  SampleProfileMap Profiles;
 
   /// LLVM context used to emit diagnostics.
   LLVMContext &Ctx;
 
   /// Memory buffer holding the profile file.
   std::unique_ptr<MemoryBuffer> Buffer;
 
   /// Extra name buffer holding names created on demand.
public:
   /// Read and validate the file header.
   std::error_code readHeader() override { return sampleprof_error::success; }
 
   /// Read sample profiles from the associated file.
   std::error_code readImpl() override;
 
   /// Return true if \p Buffer is in the format supported by this class.
   static bool hasFormat(const MemoryBuffer &Buffer);
+
+private:
+  /// CSNameTable is used to save full context vectors. This serves as an
+  /// underlying immutable buffer for all clients.
+  std::list<SampleContextFrameVector> CSNameTable;
 };
 
 class SampleProfileReaderBinary : public SampleProfileReader {
 public:
   SampleProfileReaderBinary(std::unique_ptr<MemoryBuffer> B, LLVMContext &C,
                             SampleProfileFormat Format = SPF_None)
       : SampleProfileReader(std::move(B), C, Format) {}
 
   /// Read and validate the file header.
   virtual std::error_code readHeader() override;
 
protected:
   /// Points to the end of the buffer.
   const uint8_t *End = nullptr;
 
   /// Function name table.
   std::vector<StringRef> NameTable;
 
   /// Read a string indirectly via the name table.
   virtual ErrorOr<StringRef> readStringFromTable();
+  virtual ErrorOr<SampleContext> readNameFromTable();
 
 private:
   std::error_code readSummaryEntry(std::vector<ProfileSummaryEntry> &Entries);
   virtual std::error_code verifySPMagic(uint64_t Magic) = 0;
 };
 
 class SampleProfileReaderRawBinary : public SampleProfileReaderBinary {
 private:
protected:
   std::error_code readSecHdrTableEntry(uint32_t Idx);
   std::error_code readSecHdrTable();
 
   std::error_code readFuncMetadata(bool ProfileHasAttribute);
   std::error_code readFuncOffsetTable();
   std::error_code readFuncProfiles();
   std::error_code readMD5NameTable();
   std::error_code readNameTableSec(bool IsMD5);
+  std::error_code readCSNameTableSec();
   std::error_code readProfileSymbolList();
 
   virtual std::error_code readHeader() override;
   virtual std::error_code verifySPMagic(uint64_t Magic) override = 0;
   virtual std::error_code readOneSection(const uint8_t *Start, uint64_t Size,
                                          const SecHdrTableEntry &Entry);
   // placeholder for subclasses to dispatch their own section readers.
   virtual std::error_code readCustomSection(const SecHdrTableEntry &Entry) = 0;
   virtual ErrorOr<StringRef> readStringFromTable() override;
+  virtual ErrorOr<SampleContext> readNameFromTable() override;
+  ErrorOr<SampleContextFrames> readContextFromTable();
 
   std::unique_ptr<ProfileSymbolList> ProfSymList;
 
-  /// The table mapping from function name to the offset of its FunctionSample
-  /// towards file start.
-  DenseMap<StringRef, uint64_t> FuncOffsetTable;
+  /// The table mapping from function context to the offset of its
+  /// FunctionSample towards file start.
+  DenseMap<SampleContext, uint64_t> FuncOffsetTable;
   /// The set containing the functions to use when compiling a module.
   DenseSet<StringRef> FuncsToUse;
 
   /// Use fixed length MD5 instead of ULEB128 encoding so NameTable doesn't
   /// need to be read in up front and can be directly accessed using index.
   bool FixedLengthMD5 = false;
   /// The starting address of NameTable containing fixed length MD5.
   const uint8_t *MD5NameMemStart = nullptr;
 
   /// If MD5 is used in NameTable section, the section saves uint64_t data.
   /// The uint64_t data has to be converted to a string and then the string
   /// will be used to initialize StringRef in NameTable.
   /// Note NameTable contains StringRef so it needs another buffer to own
   /// the string data. MD5StringBuf serves as the string buffer that is
   /// referenced by NameTable (vector of StringRef). We make sure
   /// the lifetime of MD5StringBuf is not shorter than that of NameTable.
   std::unique_ptr<std::vector<std::string>> MD5StringBuf;
 
+  /// CSNameTable is used to save full context vectors. This serves as an
+  /// underlying immutable buffer for all clients.
+  std::unique_ptr<const std::vector<SampleContextFrameVector>> CSNameTable;
+
   /// If SkipFlatProf is true, skip the sections with
   /// SecFlagFlat flag.
   bool SkipFlatProf = false;
 
 public:
   SampleProfileReaderExtBinaryBase(std::unique_ptr<MemoryBuffer> B,
                                    LLVMContext &C, SampleProfileFormat Format)
       : SampleProfileReaderBinary(std::move(B), C, Format) {}

llvm/include/llvm/ProfileData/SampleProfWriter.h

public:
   /// Write sample profiles in \p S.
   ///
   /// \returns status code of the file update operation.
   virtual std::error_code writeSample(const FunctionSamples &S) = 0;
 
   /// Write all the sample profiles in the given map of samples.
   ///
   /// \returns status code of the file update operation.
-  virtual std::error_code write(const StringMap<FunctionSamples> &ProfileMap);
+  virtual std::error_code write(const SampleProfileMap &ProfileMap);
 
   raw_ostream &getOutputStream() { return *OutputStream; }
 
   /// Profile writer factory.
   ///
   /// Create a new file writer based on the value of \p Format.
   static ErrorOr<std::unique_ptr<SampleProfileWriter>>
   create(StringRef Filename, SampleProfileFormat Format);
public:
   virtual void setPartialProfile() {}
   virtual void resetSecLayout(SectionLayout SL) {}
 
 protected:
   SampleProfileWriter(std::unique_ptr<raw_ostream> &OS)
       : OutputStream(std::move(OS)) {}
 
   /// Write a file header for the profile file.
-  virtual std::error_code
-  writeHeader(const StringMap<FunctionSamples> &ProfileMap) = 0;
+  virtual std::error_code writeHeader(const SampleProfileMap &ProfileMap) = 0;
 
   // Write function profiles to the profile file.
-  virtual std::error_code
-  writeFuncProfiles(const StringMap<FunctionSamples> &ProfileMap);
+  virtual std::error_code writeFuncProfiles(const SampleProfileMap &ProfileMap);
 
   /// Output stream where to emit the profile to.
   std::unique_ptr<raw_ostream> OutputStream;
 
   /// Profile summary.
   std::unique_ptr<ProfileSummary> Summary;
 
   /// Compute summary for this profile.
-  void computeSummary(const StringMap<FunctionSamples> &ProfileMap);
+  void computeSummary(const SampleProfileMap &ProfileMap);
 
   /// Profile format.
   SampleProfileFormat Format = SPF_None;
 };
 
 /// Sample-based profile writer (text format).
 class SampleProfileWriterText : public SampleProfileWriter {
 public:
   std::error_code writeSample(const FunctionSamples &S) override;
 
 protected:
   SampleProfileWriterText(std::unique_ptr<raw_ostream> &OS)
       : SampleProfileWriter(OS), Indent(0) {}
 
-  std::error_code
-  writeHeader(const StringMap<FunctionSamples> &ProfileMap) override {
+  std::error_code writeHeader(const SampleProfileMap &ProfileMap) override {
     return sampleprof_error::success;
   }
 
 private:
   /// Indent level to use when writing.
   ///
   /// This is used when printing inlined callees.
   unsigned Indent;
 
   friend ErrorOr<std::unique_ptr<SampleProfileWriter>>
   SampleProfileWriter::create(std::unique_ptr<raw_ostream> &OS,
                               SampleProfileFormat Format);
 };
 
 /// Sample-based profile writer (binary format).
 class SampleProfileWriterBinary : public SampleProfileWriter {
 public:
   SampleProfileWriterBinary(std::unique_ptr<raw_ostream> &OS)
       : SampleProfileWriter(OS) {}
 
   virtual std::error_code writeSample(const FunctionSamples &S) override;
 
 protected:
+  virtual MapVector<StringRef, uint32_t> &getNameTable() { return NameTable; }
   virtual std::error_code writeMagicIdent(SampleProfileFormat Format);
   virtual std::error_code writeNameTable();
   virtual std::error_code
-  writeHeader(const StringMap<FunctionSamples> &ProfileMap) override;
+  writeHeader(const SampleProfileMap &ProfileMap) override;
   std::error_code writeSummary();
-  std::error_code writeNameIdx(StringRef FName, bool IsContextName = false);
+  virtual std::error_code writeContextIdx(const SampleContext &Context);
+  std::error_code writeNameIdx(StringRef FName);
   std::error_code writeBody(const FunctionSamples &S);
-  inline void stablizeNameTable(std::set<StringRef> &V);
+  inline void stablizeNameTable(MapVector<StringRef, uint32_t> &NameTable,
+                                std::set<StringRef> &V);
 
   MapVector<StringRef, uint32_t> NameTable;
-  std::unordered_set<std::string> BracketedContextStr;
 
-  void addName(StringRef FName, bool IsContextName = false);
+  void addName(StringRef FName);
+  virtual void addContext(const SampleContext &Context);
   void addNames(const FunctionSamples &S);
 
 private:
   friend ErrorOr<std::unique_ptr<SampleProfileWriter>>
   SampleProfileWriter::create(std::unique_ptr<raw_ostream> &OS,
                               SampleProfileFormat Format);
 };
 
 class SampleProfileWriterRawBinary : public SampleProfileWriterBinary {
   using SampleProfileWriterBinary::SampleProfileWriterBinary;
 };
 
 const std::array<SmallVector<SecHdrTableEntry, 8>, NumOfLayout>
     ExtBinaryHdrLayoutTable = {
         // Note that SecFuncOffsetTable section is written after SecLBRProfile
         // in the profile, but is put before SecLBRProfile in SectionHdrLayout.
         // This is because sample reader follows the order in SectionHdrLayout
         // to read each section. To read function profiles on demand, sample
         // reader need to get the offset of each function profile first.
         //
         // DefaultLayout
         SmallVector<SecHdrTableEntry, 8>({{SecProfSummary, 0, 0, 0, 0},
                                           {SecNameTable, 0, 0, 0, 0},
+                                          {SecCSNameTable, 0, 0, 0, 0},
                                           {SecFuncOffsetTable, 0, 0, 0, 0},
                                           {SecLBRProfile, 0, 0, 0, 0},
                                           {SecProfileSymbolList, 0, 0, 0, 0},
                                           {SecFuncMetadata, 0, 0, 0, 0}}),
         // CtxSplitLayout
         SmallVector<SecHdrTableEntry, 8>({{SecProfSummary, 0, 0, 0, 0},
                                           {SecNameTable, 0, 0, 0, 0},
                                           // profile with context
                                           // for next two sections
                                           {SecFuncOffsetTable, 0, 0, 0, 0},
                                           {SecLBRProfile, 0, 0, 0, 0},
                                           // profile without context
                                           // for next two sections
                                           {SecFuncOffsetTable, 0, 0, 0, 0},
                                           {SecLBRProfile, 0, 0, 0, 0},
                                           {SecProfileSymbolList, 0, 0, 0, 0},
                                           {SecFuncMetadata, 0, 0, 0, 0}}),
 };
 
 class SampleProfileWriterExtBinaryBase : public SampleProfileWriterBinary {
   using SampleProfileWriterBinary::SampleProfileWriterBinary;
 public:
-  virtual std::error_code
-  write(const StringMap<FunctionSamples> &ProfileMap) override;
+  virtual std::error_code write(const SampleProfileMap &ProfileMap) override;
 
   virtual void setToCompressAllSections() override;
   void setToCompressSection(SecType Type);
   virtual std::error_code writeSample(const FunctionSamples &S) override;
 
   // Set to use MD5 to represent string in NameTable.
   virtual void setUseMD5() override {
     UseMD5 = true;
for (auto &Entry : SectionHdrLayout) {
         addSecFlag(Entry, Flag);
     }
   }
   template <class SecFlagType>
   void addSectionFlag(uint32_t SectionIdx, SecFlagType Flag) {
     addSecFlag(SectionHdrLayout[SectionIdx], Flag);
   }
 
+  virtual void addContext(const SampleContext &Context) override;
+
   // placeholder for subclasses to dispatch their own section writers.
   virtual std::error_code writeCustomSection(SecType Type) = 0;
   // Verify the SecLayout is supported by the format.
   virtual void verifySecLayout(SectionLayout SL) = 0;
 
   // specify the order to write sections.
-  virtual std::error_code
-  writeSections(const StringMap<FunctionSamples> &ProfileMap) = 0;
+  virtual std::error_code writeSections(const SampleProfileMap &ProfileMap) = 0;
 
   // Dispatch section writer for each section. \p LayoutIdx is the sequence
   // number indicating where the section is located in SectionHdrLayout.
-  virtual std::error_code
-  writeOneSection(SecType Type, uint32_t LayoutIdx,
-                  const StringMap<FunctionSamples> &ProfileMap);
+  virtual std::error_code writeOneSection(SecType Type, uint32_t LayoutIdx,
+                                          const SampleProfileMap &ProfileMap);
 
   // Helper function to write name table.
   virtual std::error_code writeNameTable() override;
+  virtual std::error_code
+  writeContextIdx(const SampleContext &Context) override;
+  std::error_code writeCSNameIdx(const SampleContext &Context);
+  std::error_code writeCSNameTableSection();
 
-  std::error_code writeFuncMetadata(const StringMap<FunctionSamples> &Profiles);
+  std::error_code writeFuncMetadata(const SampleProfileMap &Profiles);
 
   // Functions to write various kinds of sections.
-  std::error_code
-  writeNameTableSection(const StringMap<FunctionSamples> &ProfileMap);
+  std::error_code writeNameTableSection(const SampleProfileMap &ProfileMap);
   std::error_code writeFuncOffsetTable();
   std::error_code writeProfileSymbolListSection();
 
   SectionLayout SecLayout = DefaultLayout;
   // Specifiy the order of sections in section header table. Note
   // the order of sections in SecHdrTable may be different that the
   // order in SectionHdrLayout. sample Reader will follow the order
   // in SectionHdrLayout to read each section.
   SmallVector<SecHdrTableEntry, 8> SectionHdrLayout =
       ExtBinaryHdrLayoutTable[DefaultLayout];
 
   // Save the start of SecLBRProfile so we can compute the offset to the
   // start of SecLBRProfile for each Function's Profile and will keep it
   // in FuncOffsetTable.
   uint64_t SecLBRProfileStart = 0;
 
 private:
   void allocSecHdrTable();
   std::error_code writeSecHdrTable();
   virtual std::error_code
-  writeHeader(const StringMap<FunctionSamples> &ProfileMap) override;
+  writeHeader(const SampleProfileMap &ProfileMap) override;
   std::error_code compressAndOutput();
 
   // We will swap the raw_ostream held by LocalBufStream and that
   // held by OutputStream if we try to add a section which needs
   // compression. After the swap, all the data written to output
   // will be temporarily buffered into the underlying raw_string_ostream
   // originally held by LocalBufStream. After the data writing for the
   // section is completed, compress the data in the local buffer,
   // swap the raw_ostream back and write the compressed data to the
   // real output.
   std::unique_ptr<raw_ostream> LocalBufStream;
   // The location where the output stream starts.
   uint64_t FileStart;
   // The location in the output stream where the SecHdrTable should be
   // written to.
   uint64_t SecHdrTableOffset;
   // The table contains SecHdrTableEntry entries in order of how they are
   // populated in the writer. It may be different from the order in
   // SectionHdrLayout which specifies the sequence in which sections will
   // be read.
   std::vector<SecHdrTableEntry> SecHdrTable;
 
-  // FuncOffsetTable maps function name to its profile offset in SecLBRProfile
-  // section. It is used to load function profile on demand.
-  MapVector<StringRef, uint64_t> FuncOffsetTable;
+  // FuncOffsetTable maps function context to its profile offset in
+  // SecLBRProfile section. It is used to load function profile on demand.
+  MapVector<SampleContext, uint64_t> FuncOffsetTable;
   // Whether to use MD5 to represent string.
   bool UseMD5 = false;
 
+  /// CSNameTable maps function context to its offset in SecCSNameTable section.
+  /// The offset will be used everywhere where the context is referenced.
+  MapVector<SampleContext, uint32_t> CSNameTable;
+
   ProfileSymbolList *ProfSymList = nullptr;
 };
 
 class SampleProfileWriterExtBinary : public SampleProfileWriterExtBinaryBase {
 public:
   SampleProfileWriterExtBinary(std::unique_ptr<raw_ostream> &OS)
       : SampleProfileWriterExtBinaryBase(OS) {}
 
 private:
-  std::error_code
-  writeDefaultLayout(const StringMap<FunctionSamples> &ProfileMap);
-  std::error_code
-  writeCtxSplitLayout(const StringMap<FunctionSamples> &ProfileMap);
+  std::error_code writeDefaultLayout(const SampleProfileMap &ProfileMap);
+  std::error_code writeCtxSplitLayout(const SampleProfileMap &ProfileMap);
 
   virtual std::error_code
-  writeSections(const StringMap<FunctionSamples> &ProfileMap) override;
+  writeSections(const SampleProfileMap &ProfileMap) override;
 
   virtual std::error_code writeCustomSection(SecType Type) override {
     return sampleprof_error::success;
   };
 
   virtual void verifySecLayout(SectionLayout SL) override {
     assert((SL == DefaultLayout || SL == CtxSplitLayout) &&
            "Unsupported layout");

 //
 // We need Part2 because profile reader can use it to find out and read
 // function offset table without reading Part3 first.
 class SampleProfileWriterCompactBinary : public SampleProfileWriterBinary {
   using SampleProfileWriterBinary::SampleProfileWriterBinary;
 
 public:
   virtual std::error_code writeSample(const FunctionSamples &S) override;
-  virtual std::error_code
-  write(const StringMap<FunctionSamples> &ProfileMap) override;
+  virtual std::error_code write(const SampleProfileMap &ProfileMap) override;
 
 protected:
   /// The table mapping from function name to the offset of its FunctionSample
   /// towards profile start.
   MapVector<StringRef, uint64_t> FuncOffsetTable;
   /// The offset of the slot to be filled with the offset of FuncOffsetTable
   /// towards profile start.
   uint64_t TableOffset;
   virtual std::error_code writeNameTable() override;
   virtual std::error_code
-  writeHeader(const StringMap<FunctionSamples> &ProfileMap) override;
+  writeHeader(const SampleProfileMap &ProfileMap) override;
   std::error_code writeFuncOffsetTable();
 };
 
 } // end namespace sampleprof
 } // end namespace llvm
 
 #endif // LLVM_PROFILEDATA_SAMPLEPROFWRITER_H

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,
+    const SampleProfileMap &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));
+    assert(I.first == I.second.getContext() && "Inconsistent profile map");
+    SortedProfiles.push_back(std::make_pair(I.second.getContext(), &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;
 }
if (!TrimColdContext && !MergeColdContext)
     return;
 
   // Nothing to merge if sample threshold is zero
   if (ColdCountThreshold == 0)
     return;
 
   // Filter the cold profiles from ProfileMap and move them into a tmp
   // container
-  std::vector<std::pair<StringRef, const FunctionSamples *>> ColdProfiles;
+  std::vector<std::pair<SampleContext, 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);
+    ColdProfiles.emplace_back(I.first, &I.second);
   }
 
   // Remove the cold profile from ProfileMap and merge them into
   // MergedProfileMap by the last K frames of context
-  StringMap<FunctionSamples> MergedProfileMap;
+  SampleProfileMap MergedProfileMap;
   for (const auto &I : ColdProfiles) {
     if (MergeColdContext) {
-      auto Ret = MergedProfileMap.try_emplace(
-          I.second->getContext().getContextWithLastKFrames(
-              ColdContextFrameLength),
-          FunctionSamples());
+      auto MergedContext = I.second->getContext().getContextFrames();
+      if (ColdContextFrameLength < MergedContext.size())
+        MergedContext = MergedContext.take_back(ColdContextFrameLength);
+      auto Ret = MergedProfileMap.emplace(MergedContext, 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())
+        ProfileMap.find(I.first) == 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());
+    auto Ret = ProfileMap.emplace(I.first, FunctionSamples());
     if (Ret.second) {
-      SampleContext FContext(Ret.first->first(), RawContext);
+      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;
+  std::vector<SampleContext> ProfilesToBeRemoved;
+  SampleProfileMap ProfilesToBeAdded;
   for (auto &I : ProfileMap) {
     FunctionSamples &FProfile = I.second;
-    StringRef ContextStr = FProfile.getNameWithContext();
-    if (I.first() == ContextStr)
+    SampleContext &Context = FProfile.getContext();
+    if (I.first == Context)
       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);
+    auto Ret = ProfilesToBeAdded.emplace(Context, FProfile);
     (void)Ret;
     assert(Ret.second && "Context conflict during canonicalization");
-    ProfilesToBeRemoved.push_back(I.first());
+    ProfilesToBeRemoved.push_back(I.first);
   }
 
   for (auto &I : ProfilesToBeRemoved) {
     ProfileMap.erase(I);
   }
 
   for (auto &I : ProfilesToBeAdded) {
-    ProfileMap.try_emplace(I.first(), I.second);
+    ProfileMap.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/ProfileData/SampleProfReader.cpp

 static cl::opt<bool> ProfileIsFSDisciminator(
     "profile-isfs", cl::Hidden, cl::init(false),
     cl::desc("Profile uses flow sensitive discriminators"));
 
 /// Dump the function profile for \p FName.
 ///
 /// \param FName Name of the function to print.
 /// \param OS Stream to emit the output to.
-void SampleProfileReader::dumpFunctionProfile(StringRef FName,
+void SampleProfileReader::dumpFunctionProfile(SampleContext FContext,
                                               raw_ostream &OS) {
-  OS << "Function: " << FName << ": " << Profiles[FName];
+  OS << "Function: " << FContext.toString() << ": " << Profiles[FContext];
 }
 
 /// Dump all the function profiles found on stream \p OS.
 void SampleProfileReader::dump(raw_ostream &OS) {
   std::vector<NameFunctionSamples> V;
   sortFuncProfiles(Profiles, V);
   for (const auto &I : V)
     dumpFunctionProfile(I.first, OS);
if (isDigit(Rest[0])) {
     if (Rest.substr(n3 + 1).getAsInteger(10, NumSamples))
       return false;
   }
   return true;
 }
 
 /// Load samples from a text file.
 ///
 /// See the documentation at the top of the file for an explanation of
 /// the expected format.
 ///
 /// \returns true if the file was loaded successfully, false otherwise.
 std::error_code SampleProfileReaderText::readImpl() {
   line_iterator LineIt(*Buffer, /*SkipBlanks=*/true, '#');
   sampleprof_error Result = sampleprof_error::success;
 
   InlineCallStack InlineStack;
for (; !LineIt.is_at_eof(); ++LineIt) {
     // the compiler decides not to emit the function (e.g., it was inlined
     // and removed). In this case, the binary will not have the linkage
     // name for the function, so the profiler will emit the function's
     // unmangled name, which may contain characters like ':' and '>' in its
     // name (member functions, templates, etc).
     //
     // The only requirement we place on the identifier, then, is that it
     // should not begin with a number.
     if ((*LineIt)[0] != ' ') {
       uint64_t NumSamples, NumHeadSamples;
       StringRef FName;
       if (!ParseHead(*LineIt, FName, NumSamples, NumHeadSamples)) {
         reportError(LineIt.line_number(),
                     "Expected 'mangled_name:NUM:NUM', found " + *LineIt);
         return sampleprof_error::malformed;
       }
       SeenMetadata = false;
       SampleContext FContext(FName);
-      if (FContext.hasContext())
+      if (FName.startswith("[")) {
+        // Note that `[]` wrapped input indicates a full context string,
+        // otherwise it's treated as context-less function name only.
+        // Remove encapsulating '[' and ']' and create a context vector from a
+        // given context string and save it in CSNameTable.
+        FName = FName.substr(1, FName.size() - 2);
+        CSNameTable.emplace_back(SampleContextFrameVector());
+        SampleContextFrameVector &Context = CSNameTable.back();
+        SampleContext::createContextFromString(FName, Context);
+        FContext.setContext(Context);
         ++CSProfileCount;
+      }
       Profiles[FContext] = FunctionSamples();
       FunctionSamples &FProfile = Profiles[FContext];
-      FProfile.setName(FContext.getNameWithoutContext());
       FProfile.setContext(FContext);
       MergeResult(Result, FProfile.addTotalSamples(NumSamples));
       MergeResult(Result, FProfile.addHeadSamples(NumHeadSamples));
       InlineStack.clear();
       InlineStack.push_back(&FProfile);
     } else {
       uint64_t NumSamples;
       StringRef FName;
inline ErrorOr<uint32_t> SampleProfileReaderBinary::readStringIndex(T &Table) {
   if (std::error_code EC = Idx.getError())
     return EC;
   if (*Idx >= Table.size())
     return sampleprof_error::truncated_name_table;
   return *Idx;
 }
 
 ErrorOr<StringRef> SampleProfileReaderBinary::readStringFromTable() {
   auto Idx = readStringIndex(NameTable);
   if (std::error_code EC = Idx.getError())
     return EC;
 
   return NameTable[*Idx];
 }
 
+ErrorOr<SampleContext> SampleProfileReaderBinary::readNameFromTable() {
+  auto FName(readStringFromTable());
+  if (std::error_code EC = FName.getError())
+    return EC;
+  return SampleContext(*FName);
+}
+
 ErrorOr<StringRef> SampleProfileReaderExtBinaryBase::readStringFromTable() {
   if (!FixedLengthMD5)
     return SampleProfileReaderBinary::readStringFromTable();
 
   // read NameTable index.
   auto Idx = readStringIndex(NameTable);
   if (std::error_code EC = Idx.getError())
     return EC;
if (std::error_code EC = readProfile(CalleeProfile))
       return EC;
   }
 
   return sampleprof_error::success;
 }
 
 std::error_code
 SampleProfileReaderBinary::readFuncProfile(const uint8_t *Start) {
   Data = Start;
   auto NumHeadSamples = readNumber<uint64_t>();
   if (std::error_code EC = NumHeadSamples.getError())
     return EC;
 
-  auto FName(readStringFromTable());
-  if (std::error_code EC = FName.getError())
+  ErrorOr<SampleContext> FContext(readNameFromTable());
+  if (std::error_code EC = FContext.getError())
     return EC;
 
-  SampleContext FContext(*FName);
-  Profiles[FContext] = FunctionSamples();
-  FunctionSamples &FProfile = Profiles[FContext];
-  FProfile.setName(FContext.getNameWithoutContext());
-  FProfile.setContext(FContext);
+  Profiles[*FContext] = FunctionSamples();
+  FunctionSamples &FProfile = Profiles[*FContext];
+  FProfile.setContext(*FContext);
   FProfile.addHeadSamples(*NumHeadSamples);
 
-  if (FContext.hasContext())
+  if (FContext->hasContext())
     CSProfileCount++;
 
   if (std::error_code EC = readProfile(FProfile))
     return EC;
   return sampleprof_error::success;
 }
 
 std::error_code SampleProfileReaderBinary::readImpl() {
   ProfileIsFS = ProfileIsFSDisciminator;
   while (!at_eof()) {
     if (std::error_code EC = readFuncProfile(Data))
       return EC;
   }
 
   return sampleprof_error::success;
 }
 
+ErrorOr<SampleContextFrames>
+SampleProfileReaderExtBinaryBase::readContextFromTable() {
+  auto ContextIdx = readNumber<uint32_t>();
+  if (std::error_code EC = ContextIdx.getError())
+    return EC;
+  if (*ContextIdx >= CSNameTable->size())
+    return sampleprof_error::truncated_name_table;
+  return (*CSNameTable)[*ContextIdx];
+}
+
+ErrorOr<SampleContext> SampleProfileReaderExtBinaryBase::readNameFromTable() {
+  if (ProfileIsCS) {
+    auto FContext(readContextFromTable());
+    if (std::error_code EC = FContext.getError())
+      return EC;
+    return SampleContext(*FContext);
+  } else {
+    auto FName(readStringFromTable());
+    if (std::error_code EC = FName.getError())
+      return EC;
+    return SampleContext(*FName);
+  }
+}
+
 std::error_code SampleProfileReaderExtBinaryBase::readOneSection(
     const uint8_t *Start, uint64_t Size, const SecHdrTableEntry &Entry) {
   Data = Start;
   End = Start + Size;
   switch (Entry.Type) {
   case SecProfSummary:
     if (std::error_code EC = readSummary())
       return EC;
case SecNameTable: {
     assert((!FixedLengthMD5 || UseMD5) &&
            "If FixedLengthMD5 is true, UseMD5 has to be true");
     FunctionSamples::HasUniqSuffix =
         hasSecFlag(Entry, SecNameTableFlags::SecFlagUniqSuffix);
     if (std::error_code EC = readNameTableSec(UseMD5))
       return EC;
     break;
   }
+  case SecCSNameTable: {
+    if (std::error_code EC = readCSNameTableSec())
+      return EC;
+    break;
+  }
   case SecLBRProfile:
     if (std::error_code EC = readFuncProfiles())
       return EC;
     break;
   case SecFuncOffsetTable:
     if (std::error_code EC = readFuncOffsetTable())
       return EC;
     break;
std::error_code SampleProfileReaderExtBinaryBase::readFuncOffsetTable() {
   FuncOffsetTable.clear();
 
   auto Size = readNumber<uint64_t>();
   if (std::error_code EC = Size.getError())
     return EC;
 
   FuncOffsetTable.reserve(*Size);
   for (uint32_t I = 0; I < *Size; ++I) {
-    auto FName(readStringFromTable());
+    auto FName(readNameFromTable());
     if (std::error_code EC = FName.getError())
       return EC;
 
     auto Offset = readNumber<uint64_t>();
     if (std::error_code EC = Offset.getError())
       return EC;
 
     FuncOffsetTable[*FName] = *Offset;
std::error_code SampleProfileReaderExtBinaryBase::readFuncProfiles() {
   // NameTable section is read.
   bool LoadFuncsToBeUsed = collectFuncsFromModule();
 
   // When LoadFuncsToBeUsed is false, load all the function profiles.
   const uint8_t *Start = Data;
   if (!LoadFuncsToBeUsed) {
     while (Data < End) {
       if (std::error_code EC = readFuncProfile(Data))
         return EC;
     }
     assert(Data == End && "More data is read than expected");
   } else {
     // Load function profiles on demand.
     if (Remapper) {
       for (auto Name : FuncsToUse) {
         Remapper->insert(Name);
       }
     }
 
     if (useMD5()) {
       for (auto Name : FuncsToUse) {
         auto GUID = std::to_string(MD5Hash(Name));
         auto iter = FuncOffsetTable.find(StringRef(GUID));
         if (iter == FuncOffsetTable.end())
           continue;
         const uint8_t *FuncProfileAddr = Start + iter->second;
         assert(FuncProfileAddr < End && "out of LBRProfile section");
         if (std::error_code EC = readFuncProfile(FuncProfileAddr))
           return EC;
       }
-    } else if (FunctionSamples::ProfileIsCS) {
+    } else if (ProfileIsCS) {
       // Compute the ordered set of names, so we can
       // get all context profiles under a subtree by
       // iterating through the ordered names.
-      struct Comparer {
-        // Ignore the closing ']' when ordering context
-        bool operator()(const StringRef &L, const StringRef &R) const {
-          return L.substr(0, L.size() - 1) < R.substr(0, R.size() - 1);
-        }
-      };
-      std::set<StringRef, Comparer> OrderedNames;
+      std::set<SampleContext> OrderedContexts;
       for (auto Name : FuncOffsetTable) {
-        OrderedNames.insert(Name.first);
+        OrderedContexts.insert(Name.first);
       }
 
       // For each function in current module, load all
       // context profiles for the function.
       for (auto NameOffset : FuncOffsetTable) {
-        StringRef ContextName = NameOffset.first;
-        SampleContext FContext(ContextName);
-        auto FuncName = FContext.getNameWithoutContext();
+        SampleContext FContext = NameOffset.first;
+        auto FuncName = FContext.getName();
         if (!FuncsToUse.count(FuncName) &&
             (!Remapper || !Remapper->exist(FuncName)))
           continue;
 
         // For each context profile we need, try to load
         // all context profile in the subtree. This can
         // help profile guided importing for ThinLTO.
-        auto It = OrderedNames.find(ContextName);
-        while (It != OrderedNames.end() &&
-               It->startswith(ContextName.substr(0, ContextName.size() - 1))) {
+        auto It = OrderedContexts.find(FContext);
+        while (It != OrderedContexts.end() && FContext.IsPrefixOf(*It)) {
           const uint8_t *FuncProfileAddr = Start + FuncOffsetTable[*It];
           assert(FuncProfileAddr < End && "out of LBRProfile section");
           if (std::error_code EC = readFuncProfile(FuncProfileAddr))
             return EC;
           // Remove loaded context profile so we won't
           // load it repeatedly.
-          It = OrderedNames.erase(It);
+          It = OrderedContexts.erase(It);
         }
       }
     } else {
       for (auto NameOffset : FuncOffsetTable) {
         SampleContext FContext(NameOffset.first);
-        auto FuncName = FContext.getNameWithoutContext();
+        auto FuncName = FContext.getName();
         if (!FuncsToUse.count(FuncName) &&
             (!Remapper || !Remapper->exist(FuncName)))
           continue;
         const uint8_t *FuncProfileAddr = Start + NameOffset.second;
         assert(FuncProfileAddr < End && "out of LBRProfile section");
         if (std::error_code EC = readFuncProfile(FuncProfileAddr))
           return EC;
       }

 }
 
 std::error_code SampleProfileReaderExtBinaryBase::readNameTableSec(bool IsMD5) {
   if (IsMD5)
     return readMD5NameTable();
   return SampleProfileReaderBinary::readNameTable();
 }
 
-std::error_code
-SampleProfileReaderExtBinaryBase::readFuncMetadata(bool ProfileHasAttribute) {
-  while (Data < End) {
+// Read in the CS name table section, which basically contains a list of context
+// vectors. Each element of a context vector, aka a frame, refers to the
+// underlying raw function names that are stored in the name table, as well as
+// a callsite identifier that only makes sense for non-leaf frames.
+std::error_code SampleProfileReaderExtBinaryBase::readCSNameTableSec() {
+  auto Size = readNumber<uint32_t>();
+  if (std::error_code EC = Size.getError())
+    return EC;
+
+  std::vector<SampleContextFrameVector> *PNameVec =
+      new std::vector<SampleContextFrameVector>();
+  PNameVec->reserve(*Size);
+  for (uint32_t I = 0; I < *Size; ++I) {
+    PNameVec->emplace_back(SampleContextFrameVector());
+    auto ContextSize = readNumber<uint32_t>();
+    if (std::error_code EC = ContextSize.getError())
+      return EC;
+    for (uint32_t J = 0; J < *ContextSize; ++J) {
       auto FName(readStringFromTable());
       if (std::error_code EC = FName.getError())
         return EC;
+      auto LineOffset = readNumber<uint64_t>();
+      if (std::error_code EC = LineOffset.getError())
+        return EC;
+
+      if (!isOffsetLegal(*LineOffset))
+        return std::error_code();
+
+      auto Discriminator = readNumber<uint64_t>();
+      if (std::error_code EC = Discriminator.getError())
+        return EC;
+
+      PNameVec->back().emplace_back(
+          FName.get(), LineLocation(LineOffset.get(), Discriminator.get()));
+    }
+  }
 
-    SampleContext FContext(*FName);
-    bool ProfileInMap = Profiles.count(FContext);
+  // From this point the underlying object of CSNameTable should be immutable.
+  CSNameTable.reset(PNameVec);
+  return sampleprof_error::success;
+}
+
+std::error_code
+SampleProfileReaderExtBinaryBase::readFuncMetadata(bool ProfileHasAttribute) {
+  while (Data < End) {
+    auto FContext(readNameFromTable());
+    if (std::error_code EC = FContext.getError())
+      return EC;
 
+    bool ProfileInMap = Profiles.count(*FContext);
     if (ProfileIsProbeBased) {
       auto Checksum = readNumber<uint64_t>();
       if (std::error_code EC = Checksum.getError())
         return EC;
       if (ProfileInMap)
-        Profiles[FContext].setFunctionHash(*Checksum);
+        Profiles[*FContext].setFunctionHash(*Checksum);
     }
 
     if (ProfileHasAttribute) {
       auto Attributes = readNumber<uint32_t>();
       if (std::error_code EC = Attributes.getError())
         return EC;
       if (ProfileInMap)
-        Profiles[FContext].getContext().setAllAttributes(*Attributes);
+        Profiles[*FContext].getContext().setAllAttributes(*Attributes);
     }
   }
 
   assert(Data == End && "More data is read than expected");
   return sampleprof_error::success;
 }
 
 std::error_code SampleProfileReaderCompactBinary::readNameTable() {

llvm/lib/ProfileData/SampleProfWriter.cpp

 #include <set>
 #include <system_error>
 #include <utility>
 #include <vector>
 
 using namespace llvm;
 using namespace sampleprof;
 
-std::error_code SampleProfileWriter::writeFuncProfiles(
-    const StringMap<FunctionSamples> &ProfileMap) {
+std::error_code
+SampleProfileWriter::writeFuncProfiles(const SampleProfileMap &ProfileMap) {
   std::vector<NameFunctionSamples> V;
   sortFuncProfiles(ProfileMap, V);
   for (const auto &I : V) {
     if (std::error_code EC = writeSample(*I.second))
       return EC;
   }
   return sampleprof_error::success;
 }
 
-std::error_code
-SampleProfileWriter::write(const StringMap<FunctionSamples> &ProfileMap) {
+std::error_code SampleProfileWriter::write(const SampleProfileMap &ProfileMap) {
   if (std::error_code EC = writeHeader(ProfileMap))
     return EC;
 
   if (std::error_code EC = writeFuncProfiles(ProfileMap))
     return EC;
 
   return sampleprof_error::success;
 }
if (hasSecFlag(Entry, SecCommonFlags::SecFlagCompress)) {
     if (std::error_code EC = compressAndOutput())
       return EC;
   }
   SecHdrTable.push_back({Type, Entry.Flags, SectionStart - FileStart,
                          OutputStream->tell() - SectionStart, LayoutIdx});
   return sampleprof_error::success;
 }
 
-std::error_code SampleProfileWriterExtBinaryBase::write(
-    const StringMap<FunctionSamples> &ProfileMap) {
+std::error_code
+SampleProfileWriterExtBinaryBase::write(const SampleProfileMap &ProfileMap) {
   if (std::error_code EC = writeHeader(ProfileMap))
     return EC;
 
   std::string LocalBuf;
   LocalBufStream = std::make_unique<raw_string_ostream>(LocalBuf);
   if (std::error_code EC = writeSections(ProfileMap))
     return EC;
 
   if (std::error_code EC = writeSecHdrTable())
     return EC;
 
   return sampleprof_error::success;
 }
 
+std::error_code SampleProfileWriterExtBinaryBase::writeContextIdx(
+    const SampleContext &Context) {
+  if (Context.hasContext())
+    return writeCSNameIdx(Context);
+  else
+    return SampleProfileWriterBinary::writeNameIdx(Context.getName());
+}
+
+std::error_code
+SampleProfileWriterExtBinaryBase::writeCSNameIdx(const SampleContext &Context) {
+  const auto &Ret = CSNameTable.find(Context);
+  if (Ret == CSNameTable.end())
+    return sampleprof_error::truncated_name_table;
+  encodeULEB128(Ret->second, *OutputStream);
+  return sampleprof_error::success;
+}
+
 std::error_code
 SampleProfileWriterExtBinaryBase::writeSample(const FunctionSamples &S) {
   uint64_t Offset = OutputStream->tell();
-  StringRef Name = S.getNameWithContext();
-  FuncOffsetTable[Name] = Offset - SecLBRProfileStart;
+  auto &Context = S.getContext();
+  FuncOffsetTable[Context] = Offset - SecLBRProfileStart;
   encodeULEB128(S.getHeadSamples(), *OutputStream);
   return writeBody(S);
 }
 
 std::error_code SampleProfileWriterExtBinaryBase::writeFuncOffsetTable() {
   auto &OS = *OutputStream;
 
   // Write out the table size.
   encodeULEB128(FuncOffsetTable.size(), OS);
 
   // Write out FuncOffsetTable.
   for (auto Entry : FuncOffsetTable) {
-    if (std::error_code EC =
-            writeNameIdx(Entry.first, FunctionSamples::ProfileIsCS))
+    if (std::error_code EC = writeContextIdx(Entry.first))
       return EC;
     encodeULEB128(Entry.second, OS);
   }
   FuncOffsetTable.clear();
   return sampleprof_error::success;
 }
 
 std::error_code SampleProfileWriterExtBinaryBase::writeFuncMetadata(
-    const StringMap<FunctionSamples> &Profiles) {
+    const SampleProfileMap &Profiles) {
   if (!FunctionSamples::ProfileIsProbeBased && !FunctionSamples::ProfileIsCS)
     return sampleprof_error::success;
   auto &OS = *OutputStream;
   for (const auto &Entry : Profiles) {
-    if (std::error_code EC = writeNameIdx(Entry.second.getNameWithContext(),
-                                          FunctionSamples::ProfileIsCS))
+    if (std::error_code EC = writeContextIdx(Entry.second.getContext()))
       return EC;
     if (FunctionSamples::ProfileIsProbeBased)
       encodeULEB128(Entry.second.getFunctionHash(), OS);
     if (FunctionSamples::ProfileIsCS)
       encodeULEB128(Entry.second.getContext().getAllAttributes(), OS);
   }
   return sampleprof_error::success;
 }
 
 std::error_code SampleProfileWriterExtBinaryBase::writeNameTable() {
   if (!UseMD5)
     return SampleProfileWriterBinary::writeNameTable();
 
   auto &OS = *OutputStream;
   std::set<StringRef> V;
-  stablizeNameTable(V);
+  stablizeNameTable(NameTable, V);
 
   // Write out the MD5 name table. We wrote unencoded MD5 so reader can
   // retrieve the name using the name index without having to read the
   // whole name table.
   encodeULEB128(NameTable.size(), OS);
   support::endian::Writer Writer(OS, support::little);
   for (auto N : V)
     Writer.write(MD5Hash(N));
   return sampleprof_error::success;
 }
 
 std::error_code SampleProfileWriterExtBinaryBase::writeNameTableSection(
-    const StringMap<FunctionSamples> &ProfileMap) {
+    const SampleProfileMap &ProfileMap) {
   for (const auto &I : ProfileMap) {
-    assert(I.first() == I.second.getNameWithContext() &&
-           "Inconsistent profile map");
-    addName(I.second.getNameWithContext(), FunctionSamples::ProfileIsCS);
+    assert(I.first == I.second.getContext() && "Inconsistent profile map");
+    addContext(I.second.getContext());
     addNames(I.second);
   }
 
   // If NameTable contains ".__uniq." suffix, set SecFlagUniqSuffix flag
   // so compiler won't strip the suffix during profile matching after
   // seeing the flag in the profile.
   for (const auto &I : NameTable) {
     if (I.first.find(FunctionSamples::UniqSuffix) != StringRef::npos) {
       addSectionFlag(SecNameTable, SecNameTableFlags::SecFlagUniqSuffix);
       break;
     }
   }
 
   if (auto EC = writeNameTable())
     return EC;
   return sampleprof_error::success;
 }
 
+std::error_code SampleProfileWriterExtBinaryBase::writeCSNameTableSection() {
+  // Sort the names to make CSNameTable deterministic.
+  std::set<SampleContext> OrderedContexts;
+  for (const auto &I : CSNameTable)
+    OrderedContexts.insert(I.first);
+  assert(OrderedContexts.size() == CSNameTable.size() &&
+         "Unmatched ordered and unordered contexts");
+  uint64_t I = 0;
+  for (auto &Context : OrderedContexts)
+    CSNameTable[Context] = I++;
+
+  auto &OS = *OutputStream;
+  encodeULEB128(OrderedContexts.size(), OS);
+  support::endian::Writer Writer(OS, support::little);
+  for (auto Context : OrderedContexts) {
+    auto Frames = Context.getContextFrames();
+    encodeULEB128(Frames.size(), OS);
+    for (auto &Callsite : Frames) {
+      if (std::error_code EC = writeNameIdx(Callsite.CallerName))
+        return EC;
+      encodeULEB128(Callsite.Callsite.LineOffset, OS);
+      encodeULEB128(Callsite.Callsite.Discriminator, OS);
+    }
+  }
+
+  return sampleprof_error::success;
+}
+
 std::error_code
 SampleProfileWriterExtBinaryBase::writeProfileSymbolListSection() {
   if (ProfSymList && ProfSymList->size() > 0)
     if (std::error_code EC = ProfSymList->write(*OutputStream))
       return EC;
 
   return sampleprof_error::success;
 }
 
 std::error_code SampleProfileWriterExtBinaryBase::writeOneSection(
-    SecType Type, uint32_t LayoutIdx,
-    const StringMap<FunctionSamples> &ProfileMap) {
+    SecType Type, uint32_t LayoutIdx, const SampleProfileMap &ProfileMap) {
   // The setting of SecFlagCompress should happen before markSectionStart.
   if (Type == SecProfileSymbolList && ProfSymList && ProfSymList->toCompress())
     setToCompressSection(SecProfileSymbolList);
   if (Type == SecFuncMetadata && FunctionSamples::ProfileIsProbeBased)
     addSectionFlag(SecFuncMetadata, SecFuncMetadataFlags::SecFlagIsProbeBased);
   if (Type == SecProfSummary && FunctionSamples::ProfileIsCS)
     addSectionFlag(SecProfSummary, SecProfSummaryFlags::SecFlagFullContext);
   if (Type == SecFuncMetadata && FunctionSamples::ProfileIsCS)
     addSectionFlag(SecFuncMetadata, SecFuncMetadataFlags::SecFlagHasAttribute);
   if (Type == SecProfSummary && FunctionSamples::ProfileIsFS)
     addSectionFlag(SecProfSummary, SecProfSummaryFlags::SecFlagFSDiscriminator);
 
   uint64_t SectionStart = markSectionStart(Type, LayoutIdx);
   switch (Type) {
   case SecProfSummary:
     computeSummary(ProfileMap);
     if (auto EC = writeSummary())
       return EC;
     break;
   case SecNameTable:
     if (auto EC = writeNameTableSection(ProfileMap))
       return EC;
     break;
+  case SecCSNameTable:
+    if (auto EC = writeCSNameTableSection())
+      return EC;
+    break;
   case SecLBRProfile:
     SecLBRProfileStart = OutputStream->tell();
     if (std::error_code EC = writeFuncProfiles(ProfileMap))
       return EC;
     break;
   case SecFuncOffsetTable:
     if (auto EC = writeFuncOffsetTable())
       return EC;
default:
     break;
   }
   if (std::error_code EC = addNewSection(Type, LayoutIdx, SectionStart))
     return EC;
   return sampleprof_error::success;
 }
 
 std::error_code SampleProfileWriterExtBinary::writeDefaultLayout(
-    const StringMap<FunctionSamples> &ProfileMap) {
+    const SampleProfileMap &ProfileMap) {
   // The const indices passed to writeOneSection below are specifying the
   // positions of the sections in SectionHdrLayout. Look at
   // initSectionHdrLayout to find out where each section is located in
   // SectionHdrLayout.
   if (auto EC = writeOneSection(SecProfSummary, 0, ProfileMap))
     return EC;
   if (auto EC = writeOneSection(SecNameTable, 1, ProfileMap))
     return EC;
-  if (auto EC = writeOneSection(SecLBRProfile, 3, ProfileMap))
+  if (auto EC = writeOneSection(SecCSNameTable, 2, ProfileMap))
+    return EC;
+  if (auto EC = writeOneSection(SecLBRProfile, 4, ProfileMap))
     return EC;
-  if (auto EC = writeOneSection(SecProfileSymbolList, 4, ProfileMap))
+  if (auto EC = writeOneSection(SecProfileSymbolList, 5, ProfileMap))
     return EC;
-  if (auto EC = writeOneSection(SecFuncOffsetTable, 2, ProfileMap))
+  if (auto EC = writeOneSection(SecFuncOffsetTable, 3, ProfileMap))
     return EC;
-  if (auto EC = writeOneSection(SecFuncMetadata, 5, ProfileMap))
+  if (auto EC = writeOneSection(SecFuncMetadata, 6, ProfileMap))
     return EC;
   return sampleprof_error::success;
 }
 
-static void
-splitProfileMapToTwo(const StringMap<FunctionSamples> &ProfileMap,
-                     StringMap<FunctionSamples> &ContextProfileMap,
-                     StringMap<FunctionSamples> &NoContextProfileMap) {
+static void splitProfileMapToTwo(const SampleProfileMap &ProfileMap,
+                                 SampleProfileMap &ContextProfileMap,
+                                 SampleProfileMap &NoContextProfileMap) {
   for (const auto &I : ProfileMap) {
     if (I.second.getCallsiteSamples().size())
-      ContextProfileMap.insert({I.first(), I.second});
+      ContextProfileMap.insert({I.first, I.second});
     else
-      NoContextProfileMap.insert({I.first(), I.second});
+      NoContextProfileMap.insert({I.first, I.second});
   }
 }
 
 std::error_code SampleProfileWriterExtBinary::writeCtxSplitLayout(
-    const StringMap<FunctionSamples> &ProfileMap) {
-  StringMap<FunctionSamples> ContextProfileMap, NoContextProfileMap;
+    const SampleProfileMap &ProfileMap) {
+  SampleProfileMap ContextProfileMap, NoContextProfileMap;
   splitProfileMapToTwo(ProfileMap, ContextProfileMap, NoContextProfileMap);
 
   if (auto EC = writeOneSection(SecProfSummary, 0, ProfileMap))
     return EC;
   if (auto EC = writeOneSection(SecNameTable, 1, ProfileMap))
     return EC;
   if (auto EC = writeOneSection(SecLBRProfile, 3, ContextProfileMap))
     return EC;
if (auto EC = writeOneSection(SecProfileSymbolList, 6, ProfileMap))
     return EC;
   if (auto EC = writeOneSection(SecFuncMetadata, 7, ProfileMap))
     return EC;
 
   return sampleprof_error::success;
 }
 
 std::error_code SampleProfileWriterExtBinary::writeSections(
-    const StringMap<FunctionSamples> &ProfileMap) {
+    const SampleProfileMap &ProfileMap) {
   std::error_code EC;
   if (SecLayout == DefaultLayout)
     EC = writeDefaultLayout(ProfileMap);
   else if (SecLayout == CtxSplitLayout)
     EC = writeCtxSplitLayout(ProfileMap);
   else
     llvm_unreachable("Unsupported layout");
   return EC;
 }
 
-std::error_code SampleProfileWriterCompactBinary::write(
-    const StringMap<FunctionSamples> &ProfileMap) {
+std::error_code
+SampleProfileWriterCompactBinary::write(const SampleProfileMap &ProfileMap) {
   if (std::error_code EC = SampleProfileWriter::write(ProfileMap))
     return EC;
   if (std::error_code EC = writeFuncOffsetTable())
     return EC;
   return sampleprof_error::success;
 }
 
 /// Write samples to a text file.
 ///
 /// Note: it may be tempting to implement this in terms of
 /// FunctionSamples::print().  Please don't.  The dump functionality is intended
 /// for debugging and has no specified form.
 ///
 /// The format used here is more structured and deliberate because
 /// it needs to be parsed by the SampleProfileReaderText class.
 std::error_code SampleProfileWriterText::writeSample(const FunctionSamples &S) {
   auto &OS = *OutputStream;
   if (FunctionSamples::ProfileIsCS)
-    OS << "[" << S.getNameWithContext() << "]:" << S.getTotalSamples();
+    OS << "[" << S.getContext().toString() << "]:" << S.getTotalSamples();
   else
     OS << S.getName() << ":" << S.getTotalSamples();
 
   if (Indent == 0)
     OS << ":" << S.getHeadSamples();
   OS << "\n";
 
   SampleSorter<LineLocation, SampleRecord> SortedSamples(S.getBodySamples());
if (FunctionSamples::ProfileIsCS) {
       OS.indent(Indent + 1);
       OS << "!Attributes: " << S.getContext().getAllAttributes() << "\n";
     }
   }
 
   return sampleprof_error::success;
 }
 
-std::error_code SampleProfileWriterBinary::writeNameIdx(StringRef FName,
-                                                        bool IsContextName) {
-  std::string BracketedName;
-  if (IsContextName) {
-    BracketedName = "[" + FName.str() + "]";
-    FName = StringRef(BracketedName);
-  }
+std::error_code
+SampleProfileWriterBinary::writeContextIdx(const SampleContext &Context) {
+  assert(!Context.hasContext() && "cs profile is not supported");
+  return writeNameIdx(Context.getName());
 
-  const auto &Ret = NameTable.find(FName);
-  if (Ret == NameTable.end())
+std::error_code SampleProfileWriterBinary::writeNameIdx(StringRef FName) {
+  auto &NTable = getNameTable();
+  const auto &Ret = NTable.find(FName);
+  if (Ret == NTable.end())
     return sampleprof_error::truncated_name_table;
   encodeULEB128(Ret->second, *OutputStream);
   return sampleprof_error::success;
 }
 
-void SampleProfileWriterBinary::addName(StringRef FName, bool IsContextName) {
-  if (IsContextName) {
-    auto It = BracketedContextStr.insert("[" + FName.str() + "]");
-    FName = StringRef(*It.first);
-  }
-  NameTable.insert(std::make_pair(FName, 0));
+void SampleProfileWriterBinary::addName(StringRef FName) {
+  auto &NTable = getNameTable();
+  NTable.insert(std::make_pair(FName, 0));
+
+void SampleProfileWriterBinary::addContext(const SampleContext &Context) {
+  addName(Context.getName());
 }
 
 void SampleProfileWriterBinary::addNames(const FunctionSamples &S) {
   // Add all the names in indirect call targets.
   for (const auto &I : S.getBodySamples()) {
     const SampleRecord &Sample = I.second;
     for (const auto &J : Sample.getCallTargets())
       addName(J.first());
   }
 
   // Recursively add all the names for inlined callsites.
   for (const auto &J : S.getCallsiteSamples())
     for (const auto &FS : J.second) {
       const FunctionSamples &CalleeSamples = FS.second;
       addName(CalleeSamples.getName());
       addNames(CalleeSamples);
     }
 }
 
-void SampleProfileWriterBinary::stablizeNameTable(std::set<StringRef> &V) {
+void SampleProfileWriterExtBinaryBase::addContext(
+    const SampleContext &Context) {
+  if (Context.hasContext()) {
+    for (auto &Callsite : Context.getContextFrames())
+      SampleProfileWriterBinary::addName(Callsite.CallerName);
+    CSNameTable.insert(std::make_pair(Context, 0));
+  } else {
+    SampleProfileWriterBinary::addName(Context.getName());
+  }
+}
+
+void SampleProfileWriterBinary::stablizeNameTable(
+    MapVector<StringRef, uint32_t> &NameTable, std::set<StringRef> &V) {
   // Sort the names to make NameTable deterministic.
   for (const auto &I : NameTable)
     V.insert(I.first);
   int i = 0;
   for (const StringRef &N : V)
     NameTable[N] = i++;
 }
 
 std::error_code SampleProfileWriterBinary::writeNameTable() {
   auto &OS = *OutputStream;
   std::set<StringRef> V;
-  stablizeNameTable(V);
+  stablizeNameTable(NameTable, V);
 
   // Write out the name table.
   encodeULEB128(NameTable.size(), OS);
   for (auto N : V) {
     OS << N;
     encodeULEB128(0, OS);
   }
   return sampleprof_error::success;
std::error_code SampleProfileWriterCompactBinary::writeFuncOffsetTable() {
   if (OFS.seek(FuncOffsetTableStart) == (uint64_t)-1)
     return sampleprof_error::ostream_seek_unsupported;
 
   // Write out the table size.
   encodeULEB128(FuncOffsetTable.size(), OS);
 
   // Write out FuncOffsetTable.
   for (auto Entry : FuncOffsetTable) {
-    if (std::error_code EC =
-            writeNameIdx(Entry.first, FunctionSamples::ProfileIsCS))
+    if (std::error_code EC = writeNameIdx(Entry.first))
       return EC;
     encodeULEB128(Entry.second, OS);
   }
   return sampleprof_error::success;
 }
 
 std::error_code SampleProfileWriterCompactBinary::writeNameTable() {
   auto &OS = *OutputStream;
   std::set<StringRef> V;
-  stablizeNameTable(V);
+  stablizeNameTable(NameTable, V);
 
   // Write out the name table.
   encodeULEB128(NameTable.size(), OS);
   for (auto N : V) {
     encodeULEB128(MD5Hash(N), OS);
   }
   return sampleprof_error::success;
 }
 
 std::error_code
 SampleProfileWriterBinary::writeMagicIdent(SampleProfileFormat Format) {
   auto &OS = *OutputStream;
   // Write file magic identifier.
   encodeULEB128(SPMagic(Format), OS);
   encodeULEB128(SPVersion(), OS);
   return sampleprof_error::success;
 }
 
-std::error_code SampleProfileWriterBinary::writeHeader(
-    const StringMap<FunctionSamples> &ProfileMap) {
+std::error_code
+SampleProfileWriterBinary::writeHeader(const SampleProfileMap &ProfileMap) {
   writeMagicIdent(Format);
 
   computeSummary(ProfileMap);
   if (auto EC = writeSummary())
     return EC;
 
   // Generate the name table for all the functions referenced in the profile.
   for (const auto &I : ProfileMap) {
-    assert(I.first() == I.second.getNameWithContext() &&
-           "Inconsistent profile map");
-    addName(I.first(), FunctionSamples::ProfileIsCS);
+    assert(I.first == I.second.getContext() && "Inconsistent profile map");
+    addContext(I.first);
     addNames(I.second);
   }
 
   writeNameTable();
   return sampleprof_error::success;
 }
 
 void SampleProfileWriterExtBinaryBase::setToCompressAllSections() {
std::error_code SampleProfileWriterExtBinaryBase::writeSecHdrTable() {
   // Reset OutputStream.
   if (OFS.seek(Saved) == (uint64_t)-1)
     return sampleprof_error::ostream_seek_unsupported;
 
   return sampleprof_error::success;
 }
 
 std::error_code SampleProfileWriterExtBinaryBase::writeHeader(
-    const StringMap<FunctionSamples> &ProfileMap) {
+    const SampleProfileMap &ProfileMap) {
   auto &OS = *OutputStream;
   FileStart = OS.tell();
   writeMagicIdent(Format);
 
   allocSecHdrTable();
   return sampleprof_error::success;
 }
 
 std::error_code SampleProfileWriterCompactBinary::writeHeader(
-    const StringMap<FunctionSamples> &ProfileMap) {
+    const SampleProfileMap &ProfileMap) {
   support::endian::Writer Writer(*OutputStream, support::little);
   if (auto EC = SampleProfileWriterBinary::writeHeader(ProfileMap))
     return EC;
 
   // Reserve a slot for the offset of function offset table. The slot will
   // be populated with the offset of FuncOffsetTable later.
   TableOffset = OutputStream->tell();
   Writer.write(static_cast<uint64_t>(-2));
for (auto Entry : Entries) {
     encodeULEB128(Entry.Cutoff, OS);
     encodeULEB128(Entry.MinCount, OS);
     encodeULEB128(Entry.NumCounts, OS);
   }
   return sampleprof_error::success;
 }
 std::error_code SampleProfileWriterBinary::writeBody(const FunctionSamples &S) {
   auto &OS = *OutputStream;
-
-  if (std::error_code EC =
-          writeNameIdx(S.getNameWithContext(), FunctionSamples::ProfileIsCS))
+  if (std::error_code EC = writeContextIdx(S.getContext()))
     return EC;
 
   encodeULEB128(S.getTotalSamples(), OS);
 
   // Emit all the body samples.
   encodeULEB128(S.getBodySamples().size(), OS);
   for (const auto &I : S.getBodySamples()) {
     LineLocation Loc = I.first;
SampleProfileWriter::create(std::unique_ptr<raw_ostream> &OS,
 
   if (EC)
     return EC;
 
   Writer->Format = Format;
   return std::move(Writer);
 }
 
-void SampleProfileWriter::computeSummary(
-    const StringMap<FunctionSamples> &ProfileMap) {
+void SampleProfileWriter::computeSummary(const SampleProfileMap &ProfileMap) {
   SampleProfileSummaryBuilder Builder(ProfileSummaryBuilder::DefaultCutoffs);
   Summary = Builder.computeSummaryForProfiles(ProfileMap);
 }

llvm/test/tools/llvm-profdata/Inputs/cs-sample.proftext

 [main:3 @ _Z5funcAi:1 @ _Z8funcLeafi]:1467299:11
  0: 6
  1: 6
  3: 287884
  4: 287864 _Z3fibi:315608
  15: 23
  !Attributes: 0
 [main:3.1 @ _Z5funcBi:1 @ _Z8funcLeafi]:500853:20
  0: 15
  1: 15
  3: 74946
  4: 74941 _Z3fibi:82359
  10: 23324
  11: 23327 _Z3fibi:25228
  15: 11
  !Attributes: 1
+[external:12 @ main]:154:12
+ 2: 12
+ 3: 10 _Z5funcAi:7
+ 3.1: 10 _Z5funcBi:11
+ !Attributes: 0
 [main]:154:0
  2: 12
  3: 18 _Z5funcAi:11
  3.1: 18 _Z5funcBi:19
  !Attributes: 0
-[external:12 @ main]:154:12
- 2: 12
- 3: 10 _Z5funcAi:7
- 3.1: 10 _Z5funcBi:11
+[external:10 @ _Z5funcBi]:120:10
+ 0: 10
+ 1: 10
+ !Attributes: 0
+[externalA:17 @ _Z5funcBi]:120:3
+ 0: 3
+ 1: 3
  !Attributes: 0
 [main:3.1 @ _Z5funcBi]:120:19
  0: 19
  1: 19 _Z8funcLeafi:20
  3: 12
  !Attributes: 1
-[externalA:17 @ _Z5funcBi]:120:3
- 0: 3
- 1: 3
- !Attributes: 0
-[external:10 @ _Z5funcBi]:120:10
- 0: 10
- 1: 10
- !Attributes: 0
 [main:3 @ _Z5funcAi]:99:11
  0: 10
  1: 10 _Z8funcLeafi:11
  3: 24
- !Attributes: 0
+ !Attributes: 0
 No newline at end of file

llvm/tools/llvm-profdata/llvm-profdata.cpp

uint64_t ColdInstrThreshold =
           : ProfileSummaryBuilder::getEntryForPercentile(
                 InstrPS.getDetailedSummary(),
                 ProfileSummaryBuilder::DefaultCutoffs[ColdPercentileIdx])
                 .MinCount;
 
   // Find hot/warm functions in sample profile which is cold in instr profile
   // and adjust the profiles of those functions in the instr profile.
   for (const auto &PD : Reader->getProfiles()) {
-    StringRef FName = PD.getKey();
-    const sampleprof::FunctionSamples &FS = PD.getValue();
-    auto It = InstrProfileMap.find(FName);
+    auto &FContext = PD.first;
+    const sampleprof::FunctionSamples &FS = PD.second;
+    auto It = InstrProfileMap.find(FContext.toString());
     if (FS.getHeadSamples() > ColdSampleThreshold &&
         It != InstrProfileMap.end() &&
         It->second.MaxCount <= ColdInstrThreshold &&
         FS.getBodySamples().size() >= SupplMinSizeThreshold) {
       updateInstrProfileEntry(It->second, HotInstrThreshold,
                               ZeroCounterThreshold);
     }
   }

 static void
 mergeSampleProfile(const WeightedFileVector &Inputs, SymbolRemapper *Remapper,
                    StringRef OutputFilename, ProfileFormat OutputFormat,
                    StringRef ProfileSymbolListFile, bool CompressAllSections,
                    bool UseMD5, bool GenPartialProfile,
                    bool SampleMergeColdContext, bool SampleTrimColdContext,
                    bool SampleColdContextFrameDepth, FailureMode FailMode) {
   using namespace sampleprof;
-  StringMap<FunctionSamples> ProfileMap;
+  SampleProfileMap ProfileMap;
   SmallVector<std::unique_ptr<sampleprof::SampleProfileReader>, 5> Readers;
   LLVMContext Context;
   sampleprof::ProfileSymbolList WriterList;
   Optional<bool> ProfileIsProbeBased;
   Optional<bool> ProfileIsCS;
   for (const auto &Input : Inputs) {
     auto ReaderOrErr = SampleProfileReader::create(Input.Filename, Context,
                                                    FSDiscriminatorPassOption);
for (const auto &Input : Inputs) {
     Readers.push_back(std::move(ReaderOrErr.get()));
     const auto Reader = Readers.back().get();
     if (std::error_code EC = Reader->read()) {
       warnOrExitGivenError(FailMode, EC, Input.Filename);
       Readers.pop_back();
       continue;
     }
 
-    StringMap<FunctionSamples> &Profiles = Reader->getProfiles();
+    SampleProfileMap &Profiles = Reader->getProfiles();
     if (ProfileIsProbeBased.hasValue() &&
         ProfileIsProbeBased != FunctionSamples::ProfileIsProbeBased)
       exitWithError(
           "cannot merge probe-based profile with non-probe-based profile");
     ProfileIsProbeBased = FunctionSamples::ProfileIsProbeBased;
     if (ProfileIsCS.hasValue() && ProfileIsCS != FunctionSamples::ProfileIsCS)
       exitWithError("cannot merge CS profile with non-CS profile");
     ProfileIsCS = FunctionSamples::ProfileIsCS;
-    for (StringMap<FunctionSamples>::iterator I = Profiles.begin(),
-                                              E = Profiles.end();
+    for (SampleProfileMap::iterator I = Profiles.begin(), E = Profiles.end();
          I != E; ++I) {
       sampleprof_error Result = sampleprof_error::success;
       FunctionSamples Remapped =
           Remapper ? remapSamples(I->second, *Remapper, Result)
                    : FunctionSamples();
       FunctionSamples &Samples = Remapper ? Remapped : I->second;
-      StringRef FName = Samples.getNameWithContext();
-      MergeResult(Result, ProfileMap[FName].merge(Samples, Input.Weight));
+      SampleContext FContext = Samples.getContext();
+      MergeResult(Result, ProfileMap[FContext].merge(Samples, Input.Weight));
       if (Result != sampleprof_error::success) {
         std::error_code EC = make_error_code(Result);
-        handleMergeWriterError(errorCodeToError(EC), Input.Filename, FName);
+        handleMergeWriterError(errorCodeToError(EC), Input.Filename,
+                               FContext.toString());
       }
     }
 
     std::unique_ptr<sampleprof::ProfileSymbolList> ReaderList =
         Reader->getProfileSymbolList();
     if (ReaderList)
       WriterList.merge(*ReaderList);
   }
static void overlapInstrProfile(const std::string &BaseFilename,
   loadInput(WeightedInput, nullptr, &Context);
   overlapInput(BaseFilename, TestFilename, &Context, Overlap, FuncFilter, OS,
                IsCS);
   Overlap.dump(OS);
 }
 
 namespace {
 struct SampleOverlapStats {
-  StringRef BaseName;
-  StringRef TestName;
+  SampleContext BaseName;
+  SampleContext TestName;
   // Number of overlap units
   uint64_t OverlapCount;
   // Total samples of overlap units
   uint64_t OverlapSample;
   // Number of and total samples of units that only present in base or test
   // profile
   uint64_t BaseUniqueCount;
   uint64_t BaseUniqueSample;
public:
   /// profiles. This function also computes and keeps the sum of samples and
   /// max sample counts of each function in BaseStats and TestStats for later
   /// use to avoid re-computations.
   void initializeSampleProfileOverlap();
 
   /// Load profiles specified by BaseFilename and TestFilename.
   std::error_code loadProfiles();
 
+  using FuncSampleStatsMap =
+      std::unordered_map<SampleContext, FuncSampleStats, SampleContext::Hash>;
+
 private:
   SampleOverlapStats ProfOverlap;
   SampleOverlapStats HotFuncOverlap;
   SampleOverlapStats HotBlockOverlap;
   std::string BaseFilename;
   std::string TestFilename;
   std::unique_ptr<sampleprof::SampleProfileReader> BaseReader;
   std::unique_ptr<sampleprof::SampleProfileReader> TestReader;
   // BaseStats and TestStats hold FuncSampleStats for each function, with
   // function name as the key.
-  StringMap<FuncSampleStats> BaseStats;
-  StringMap<FuncSampleStats> TestStats;
+  FuncSampleStatsMap BaseStats;
+  FuncSampleStatsMap TestStats;
   // Low similarity threshold in floating point number
   double LowSimilarityThreshold;
   // Block samples above BaseHotThreshold or TestHotThreshold are considered hot
   // for tracking hot blocks.
   uint64_t BaseHotThreshold;
   uint64_t TestHotThreshold;
   // A small threshold used to round the results of floating point accumulations
   // to resolve imprecision.
private:
   ///    this function in test profile ST, compute BS(i) = 1.0 - fabs(BB(i)/SB -
   ///    BT(i)/ST), ranging in [0.0f to 1.0f] with 0.0 meaning no-overlap.
   double computeBlockSimilarity(uint64_t BaseSample, uint64_t TestSample,
                                 const SampleOverlapStats &FuncOverlap) const;
 
   void updateHotBlockOverlap(uint64_t BaseSample, uint64_t TestSample,
                              uint64_t HotBlockCount);
 
-  void getHotFunctions(const StringMap<FuncSampleStats> &ProfStats,
-                       StringMap<FuncSampleStats> &HotFunc,
+  void getHotFunctions(const FuncSampleStatsMap &ProfStats,
+                       FuncSampleStatsMap &HotFunc,
                        uint64_t HotThreshold) const;
 
   void computeHotFuncOverlap();
 
   /// This function updates statistics in FuncOverlap, HotBlockOverlap, and
   /// Difference for two sample units in a matched function according to the
   /// given match status.
   void updateOverlapStatsForFunction(uint64_t BaseSample, uint64_t TestSample,
if (IsBaseHot)
     HotBlockOverlap.BaseCount += HotBlockCount;
   if (IsTestHot)
     HotBlockOverlap.TestCount += HotBlockCount;
   if (IsBaseHot && IsTestHot)
     HotBlockOverlap.OverlapCount += HotBlockCount;
 }
 
 void SampleOverlapAggregator::getHotFunctions(
-    const StringMap<FuncSampleStats> &ProfStats,
-    StringMap<FuncSampleStats> &HotFunc, uint64_t HotThreshold) const {
+    const FuncSampleStatsMap &ProfStats, FuncSampleStatsMap &HotFunc,
+    uint64_t HotThreshold) const {
   for (const auto &F : ProfStats) {
     if (isFunctionHot(F.second, HotThreshold))
-      HotFunc.try_emplace(F.first(), F.second);
+      HotFunc.emplace(F.first, F.second);
   }
 }
 
 void SampleOverlapAggregator::computeHotFuncOverlap() {
-  StringMap<FuncSampleStats> BaseHotFunc;
+  FuncSampleStatsMap BaseHotFunc;
   getHotFunctions(BaseStats, BaseHotFunc, BaseHotThreshold);
   HotFuncOverlap.BaseCount = BaseHotFunc.size();
 
-  StringMap<FuncSampleStats> TestHotFunc;
+  FuncSampleStatsMap TestHotFunc;
   getHotFunctions(TestStats, TestHotFunc, TestHotThreshold);
   HotFuncOverlap.TestCount = TestHotFunc.size();
   HotFuncOverlap.UnionCount = HotFuncOverlap.TestCount;
 
   for (const auto &F : BaseHotFunc) {
-    if (TestHotFunc.count(F.first()))
+    if (TestHotFunc.count(F.first))
       ++HotFuncOverlap.OverlapCount;
     else
       ++HotFuncOverlap.UnionCount;
   }
 }
 
 void SampleOverlapAggregator::updateOverlapStatsForFunction(
     uint64_t BaseSample, uint64_t TestSample, uint64_t HotBlockCount,
double SampleOverlapAggregator::computeSampleFunctionOverlap(
   FuncSimilarity = weightForFuncSimilarity(FuncInternalSimilarity,
                                            BaseFuncSample, TestFuncSample);
   return FuncSimilarity;
 }
 
 void SampleOverlapAggregator::computeSampleProfileOverlap(raw_fd_ostream &OS) {
   using namespace sampleprof;
 
-  StringMap<const FunctionSamples *> BaseFuncProf;
+  std::unordered_map<SampleContext, const FunctionSamples *,
+                     SampleContext::Hash>
+      BaseFuncProf;
   const auto &BaseProfiles = BaseReader->getProfiles();
   for (const auto &BaseFunc : BaseProfiles) {
-    BaseFuncProf.try_emplace(BaseFunc.second.getNameWithContext(),
-                             &(BaseFunc.second));
+    BaseFuncProf.emplace(BaseFunc.second.getContext(), &(BaseFunc.second));
   }
   ProfOverlap.UnionCount = BaseFuncProf.size();
 
   const auto &TestProfiles = TestReader->getProfiles();
   for (const auto &TestFunc : TestProfiles) {
     SampleOverlapStats FuncOverlap;
-    FuncOverlap.TestName = TestFunc.second.getNameWithContext();
+    FuncOverlap.TestName = TestFunc.second.getContext();
     assert(TestStats.count(FuncOverlap.TestName) &&
            "TestStats should have records for all functions in test profile "
            "except inlinees");
     FuncOverlap.TestSample = TestStats[FuncOverlap.TestName].SampleSum;
 
     const auto Match = BaseFuncProf.find(FuncOverlap.TestName);
     if (Match == BaseFuncProf.end()) {
       const FuncSampleStats &FuncStats = TestStats[FuncOverlap.TestName];
if (Match == BaseFuncProf.end()) {
 
       ++ProfOverlap.UnionCount;
       ProfOverlap.UnionSample += FuncStats.SampleSum;
     } else {
       ++ProfOverlap.OverlapCount;
 
       // Two functions match with each other. Compute function-level overlap and
       // aggregate them into profile-level overlap.
-      FuncOverlap.BaseName = Match->second->getNameWithContext();
+      FuncOverlap.BaseName = Match->second->getContext();
       assert(BaseStats.count(FuncOverlap.BaseName) &&
              "BaseStats should have records for all functions in base profile "
              "except inlinees");
       FuncOverlap.BaseSample = BaseStats[FuncOverlap.BaseName].SampleSum;
 
       FuncOverlap.Similarity = computeSampleFunctionOverlap(
           Match->second, &TestFunc.second, &FuncOverlap, FuncOverlap.BaseSample,
           FuncOverlap.TestSample);
for (const auto &TestFunc : TestProfiles) {
     // Print function-level similarity information if specified by options.
     assert(TestStats.count(FuncOverlap.TestName) &&
            "TestStats should have records for all functions in test profile "
            "except inlinees");
     if (TestStats[FuncOverlap.TestName].MaxSample >= FuncFilter.ValueCutoff ||
         (Match != BaseFuncProf.end() &&
          FuncOverlap.Similarity < LowSimilarityThreshold) ||
         (Match != BaseFuncProf.end() && !FuncFilter.NameFilter.empty() &&
-         FuncOverlap.BaseName.find(FuncFilter.NameFilter) !=
-             FuncOverlap.BaseName.npos)) {
+         FuncOverlap.BaseName.toString().find(FuncFilter.NameFilter) !=
+             std::string::npos)) {
       assert(ProfOverlap.BaseSample > 0 &&
              "Total samples in base profile should be greater than 0");
       FuncOverlap.BaseWeight =
           static_cast<double>(FuncOverlap.BaseSample) / ProfOverlap.BaseSample;
       assert(ProfOverlap.TestSample > 0 &&
              "Total samples in test profile should be greater than 0");
       FuncOverlap.TestWeight =
           static_cast<double>(FuncOverlap.TestSample) / ProfOverlap.TestSample;
       FuncSimilarityDump.emplace(FuncOverlap.BaseWeight, FuncOverlap);
     }
   }
 
   // Traverse through functions in base profile but not in test profile.
   for (const auto &F : BaseFuncProf) {
-    assert(BaseStats.count(F.second->getNameWithContext()) &&
+    assert(BaseStats.count(F.second->getContext()) &&
            "BaseStats should have records for all functions in base profile "
            "except inlinees");
-    const FuncSampleStats &FuncStats =
-        BaseStats[F.second->getNameWithContext()];
+    const FuncSampleStats &FuncStats = BaseStats[F.second->getContext()];
     ++ProfOverlap.BaseUniqueCount;
     ProfOverlap.BaseUniqueSample += FuncStats.SampleSum;
 
     updateHotBlockOverlap(FuncStats.SampleSum, 0, FuncStats.HotBlockCount);
 
     double FuncSimilarity = computeSampleFunctionOverlap(
         nullptr, nullptr, nullptr, FuncStats.SampleSum, 0);
     ProfOverlap.Similarity +=

 
 void SampleOverlapAggregator::initializeSampleProfileOverlap() {
   const auto &BaseProf = BaseReader->getProfiles();
   for (const auto &I : BaseProf) {
     ++ProfOverlap.BaseCount;
     FuncSampleStats FuncStats;
     getFuncSampleStats(I.second, FuncStats, BaseHotThreshold);
     ProfOverlap.BaseSample += FuncStats.SampleSum;
-    BaseStats.try_emplace(I.second.getNameWithContext(), FuncStats);
+    BaseStats.emplace(I.second.getContext(), FuncStats);
   }
 
   const auto &TestProf = TestReader->getProfiles();
   for (const auto &I : TestProf) {
     ++ProfOverlap.TestCount;
     FuncSampleStats FuncStats;
     getFuncSampleStats(I.second, FuncStats, TestHotThreshold);
     ProfOverlap.TestSample += FuncStats.SampleSum;
-    TestStats.try_emplace(I.second.getNameWithContext(), FuncStats);
+    TestStats.emplace(I.second.getContext(), FuncStats);
   }
 
   ProfOverlap.BaseName = StringRef(BaseFilename);
   ProfOverlap.TestName = StringRef(TestFilename);
 }
 
 void SampleOverlapAggregator::dumpFuncSimilarity(raw_fd_ostream &OS) const {
   using namespace sampleprof;
for (const auto &F : FuncSimilarityDump) {
     FOS << format("%.2f%%", BaseUniquePercent * 100);
     FOS.PadToColumn(TestUniqueCol);
     FOS << format("%.2f%%", TestUniquePercent * 100);
     FOS.PadToColumn(BaseSampleCol);
     FOS << F.second.BaseSample;
     FOS.PadToColumn(TestSampleCol);
     FOS << F.second.TestSample;
     FOS.PadToColumn(FuncNameCol);
-    FOS << F.second.TestName << "\n";
+    FOS << F.second.TestName.toString() << "\n";
   }
 }
 
 void SampleOverlapAggregator::dumpProgramSummary(raw_fd_ostream &OS) const {
-  OS << "Profile overlap infomation for base_profile: " << ProfOverlap.BaseName
-     << " and test_profile: " << ProfOverlap.TestName << "\nProgram level:\n";
+  OS << "Profile overlap infomation for base_profile: "
+     << ProfOverlap.BaseName.toString()
+     << " and test_profile: " << ProfOverlap.TestName.toString()
+     << "\nProgram level:\n";
 
   OS << "  Whole program profile similarity: "
      << format("%.3f%%", ProfOverlap.Similarity * 100) << "\n";
 
   assert(ProfOverlap.UnionSample > 0 &&
          "Total samples in two profile should be greater than 0");
   double OverlapPercent =
       static_cast<double>(ProfOverlap.OverlapSample) / ProfOverlap.UnionSample;
WithColor::warning() << "-show-sec-info-only is only supported for "
                          << "sample profile in extbinary format and is "
                          << "ignored for other formats.\n";
     return;
   }
 }
 
 namespace {
 struct HotFuncInfo {
-  StringRef FuncName;
+  std::string FuncName;
   uint64_t TotalCount;
   double TotalCountPercent;
   uint64_t MaxCount;
   uint64_t EntryCount;
 
   HotFuncInfo()
       : FuncName(), TotalCount(0), TotalCountPercent(0.0f), MaxCount(0),
         EntryCount(0) {}
 
   HotFuncInfo(StringRef FN, uint64_t TS, double TSP, uint64_t MS, uint64_t ES)
-      : FuncName(FN), TotalCount(TS), TotalCountPercent(TSP), MaxCount(MS),
-        EntryCount(ES) {}
+      : FuncName(FN.begin(), FN.end()), TotalCount(TS), TotalCountPercent(TSP),
+        MaxCount(MS), EntryCount(ES) {}
 };
 } // namespace
 
 // Print out detailed information about hot functions in PrintValues vector.
 // Users specify titles and offset of every columns through ColumnTitle and
 // ColumnOffset. The size of ColumnTitle and ColumnOffset need to be the same
 // and at least 4. Besides, users can optionally give a HotFuncMetric string to
 // print out or let it be an empty string.
for (const HotFuncInfo &R : PrintValues) {
     FOS << R.MaxCount;
     FOS.PadToColumn(ColumnOffset[2]);
     FOS << R.EntryCount;
     FOS.PadToColumn(ColumnOffset[3]);
     FOS << R.FuncName << "\n";
   }
 }
 
-static int
-showHotFunctionList(const StringMap<sampleprof::FunctionSamples> &Profiles,
+static int showHotFunctionList(const sampleprof::SampleProfileMap &Profiles,
                                ProfileSummary &PS, raw_fd_ostream &OS) {
   using namespace sampleprof;
 
   const uint32_t HotFuncCutoff = 990000;
   auto &SummaryVector = PS.getDetailedSummary();
   uint64_t MinCountThreshold = 0;
   for (const ProfileSummaryEntry &SummaryEntry : SummaryVector) {
     if (SummaryEntry.Cutoff == HotFuncCutoff) {
       MinCountThreshold = SummaryEntry.MinCount;
static int showHotFunctionList(const sampleprof::SampleProfileMap &Profiles,
   std::vector<HotFuncInfo> PrintValues;
   for (const auto &FuncPair : HotFunc) {
     const FunctionSamples &Func = *FuncPair.second.first;
     double TotalSamplePercent =
         (ProfileTotalSample > 0)
             ? (Func.getTotalSamples() * 100.0) / ProfileTotalSample
             : 0;
     PrintValues.emplace_back(HotFuncInfo(
-        Func.getNameWithContext(), Func.getTotalSamples(), TotalSamplePercent,
-        FuncPair.second.second, Func.getEntrySamples()));
+        Func.getContext().toString(), Func.getTotalSamples(),
+        TotalSamplePercent, FuncPair.second.second, Func.getEntrySamples()));
   }
   dumpHotFunctionList(ColumnTitle, ColumnOffset, PrintValues, HotFuncCount,
                       Profiles.size(), HotFuncSample, ProfileTotalSample,
                       Metric, OS);
 
   return 0;
 }
 
static int showSampleProfile(const std::string &Filename, bool ShowCounts,
   }
 
   if (std::error_code EC = Reader->read())
     exitWithErrorCode(EC, Filename);
 
   if (ShowAllFunctions || ShowFunction.empty())
     Reader->dump(OS);
   else
-    Reader->dumpFunctionProfile(ShowFunction, OS);
+    // TODO: parse context string to support filtering by contexts.
+    Reader->dumpFunctionProfile(StringRef(ShowFunction), OS);
 
   if (ShowProfileSymbolList) {
     std::unique_ptr<sampleprof::ProfileSymbolList> ReaderList =
         Reader->getProfileSymbolList();
     ReaderList->dump(OS);
   }
 
   if (ShowDetailedSummary) {

llvm/unittests/ProfileData/SampleProfTest.cpp

void testRoundTrip(SampleProfileFormat Format, bool Remap, bool UseMD5) {
 
     StringRef BooName("_Z3booi");
     FunctionSamples BooSamples;
     BooSamples.setName(BooName);
     BooSamples.addTotalSamples(1232);
     BooSamples.addHeadSamples(1);
     BooSamples.addBodySamples(1, 0, 1232);
 
-    StringMap<FunctionSamples> Profiles;
+    SampleProfileMap Profiles;
     Profiles[FooName] = std::move(FooSamples);
     Profiles[BarName] = std::move(BarSamples);
     Profiles[BazName] = std::move(BazSamples);
     Profiles[BooName] = std::move(BooSamples);
 
     Module M("my_module", Context);
     FunctionType *fn_type =
         FunctionType::get(Type::getVoidTy(Context), {}, false);
void testRoundTrip(SampleProfileFormat Format, bool Remap, bool UseMD5) {
 
     verifyProfileSummary(Summary, M, false, false);
 
     Summary.setPartialProfile(true);
     Summary.setPartialProfileRatio(0.5);
     verifyProfileSummary(Summary, M, true, true);
   }
 
-  void addFunctionSamples(StringMap<FunctionSamples> *Smap, const char *Fname,
+  void addFunctionSamples(SampleProfileMap *Smap, const char *Fname,
                           uint64_t TotalSamples, uint64_t HeadSamples) {
     StringRef Name(Fname);
     FunctionSamples FcnSamples;
     FcnSamples.setName(Name);
     FcnSamples.addTotalSamples(TotalSamples);
     FcnSamples.addHeadSamples(HeadSamples);
     FcnSamples.addBodySamples(1, 0, HeadSamples);
     (*Smap)[Name] = FcnSamples;
   }
 
-  StringMap<FunctionSamples> setupFcnSamplesForElisionTest(StringRef Policy) {
-    StringMap<FunctionSamples> Smap;
+  SampleProfileMap setupFcnSamplesForElisionTest(StringRef Policy) {
+    SampleProfileMap Smap;
     addFunctionSamples(&Smap, "foo", uint64_t(20301), uint64_t(1437));
     if (Policy == "" || Policy == "all")
       return Smap;
     addFunctionSamples(&Smap, "foo.bar", uint64_t(20303), uint64_t(1439));
     if (Policy == "selected")
       return Smap;
     addFunctionSamples(&Smap, "foo.llvm.2465", uint64_t(20305), uint64_t(1441));
     return Smap;
struct SampleProfTest : ::testing::Test {
   }
 
   void testSuffixElisionPolicy(SampleProfileFormat Format, StringRef Policy,
                                const StringMap<uint64_t> &Expected) {
     TempFile ProfileFile("profile", "", "", /*Unique*/ true);
 
     Module M("my_module", Context);
     setupModuleForElisionTest(&M, Policy);
-    StringMap<FunctionSamples> ProfMap = setupFcnSamplesForElisionTest(Policy);
+    SampleProfileMap ProfMap = setupFcnSamplesForElisionTest(Policy);
 
     // write profile
     createWriter(Format, ProfileFile.path());
     std::error_code EC;
     EC = Writer->write(ProfMap);
     ASSERT_TRUE(NoError(EC));
     Writer->getOutputStream().flush();