[CSSPGO] Split context string to deduplicate function name used in the context.

[hoy]

Summary

Currently context strings contain a lot of duplicated function names and that significantly increase the profile size. This change split the context into a series of {name, offset, discriminator} tuples so function names used in the context can be replaced by the index into the name table and that significantly reduce the size consumed by context.

A follow-up improvement made in the compiler and profiling tools is to avoid reconstructing full context strings which is time- and memory- consuming. Instead a context vector of StringRef is adopted to represent the full context in all scenarios. As a result, the previous prevalent profile map which was implemented as a StringRef is now engineered as an unordered map keyed by SampleContext. SampleContext is reshaped to using an ArrayRef to represent a full context for CS profile. For non-CS profile, it falls back to use StringRef to represent a contextless function name. Both the ArrayRef and StringRef objects are underpinned by real array and string objects that are stored in producer buffers. For compiler, they are maintained by the sample reader. For llvm-profgen, they are maintained in ProfiledBinary and ProfileGenerator. Full context strings can be generated only in those cases of debugging and printing.

When it comes to profile format, nothing has changed to the text format, though internally CS context is implemented as a vector. Extbinary format is only changed for CS profile, with an additional SecCSNameTable section which stores all full contexts logically in the form of vector<int>, which each element as an offset points to SecNameTable. All occurrences of contexts elsewhere are redirected to using the offset of SecCSNameTable.

Testing
This is no-diff change in terms of code quality and profile content (for text profile).

For our internal large service (aka ads), the profile generation is cut to half, with a 20x smaller string-based extbinary format generated.

The compile time of ads is dropped by 25%.

[hoy]

Updating D107299: [SampleFDO] split context string to deduplicate function name used in the context.

[wmi]

Thanks Hongtao. The compile time result looks great. The patch is large. Is it possible to split it into smaller ones, SampleContext related change and SecCSNameTable related change for example?

[hoy]

In D107299#2955079, @wmi wrote:

Thanks Hongtao. The compile time result looks great. The patch is large. Is it possible to split it into smaller ones, SampleContext related change and SecCSNameTable related change for example?

Thanks for the suggestion. The separation may not be complete. The point is that the previous StringRef context is used in both reader/writer, and now we are changing it to be an ArrayRef which relies on an underlying array buffer (i.e., the CSName the reader builds up). Separating SecCSNameTable from this patch will still require the reader to build the underlying CSName table to bridge full string contexts in profile to ArrayRef contexts used elsewhere. Does this sound good to you?

[wmi]

In D107299#2955596, @hoy wrote:

In D107299#2955079, @wmi wrote:

Thanks Hongtao. The compile time result looks great. The patch is large. Is it possible to split it into smaller ones, SampleContext related change and SecCSNameTable related change for example?

Thanks for the suggestion. The separation may not be complete. The point is that the previous StringRef context is used in both reader/writer, and now we are changing it to be an ArrayRef which relies on an underlying array buffer (i.e., the CSName the reader builds up). Separating SecCSNameTable from this patch will still require the reader to build the underlying CSName table to bridge full string contexts in profile to ArrayRef contexts used elsewhere. Does this sound good to you?

I understand it is not easy to split them into two patches to be committed separately. It needs the bridge you described above which is only useful to make the temporary first part of the patch work and will be immediately thrown away after the rest of the patch is committed. That is unnecessary. Maybe just split the patches for easier review. After all the parts are ok, still commit all the parts in one shot.

[hoy]

This comment was removed by hoy.

[hoy]

Updating D107299: [CSSPGO] Split context string to deduplicate function name used in the context.

[hoy]

In D107299#2955913, @wmi wrote:

In D107299#2955596, @hoy wrote:

In D107299#2955079, @wmi wrote:

Thanks Hongtao. The compile time result looks great. The patch is large. Is it possible to split it into smaller ones, SampleContext related change and SecCSNameTable related change for example?

Thanks for the suggestion. The separation may not be complete. The point is that the previous StringRef context is used in both reader/writer, and now we are changing it to be an ArrayRef which relies on an underlying array buffer (i.e., the CSName the reader builds up). Separating SecCSNameTable from this patch will still require the reader to build the underlying CSName table to bridge full string contexts in profile to ArrayRef contexts used elsewhere. Does this sound good to you?

I understand it is not easy to split them into two patches to be committed separately. It needs the bridge you described above which is only useful to make the temporary first part of the patch work and will be immediately thrown away after the rest of the patch is committed. That is unnecessary. Maybe just split the patches for easier review. After all the parts are ok, still commit all the parts in one shot.

Sounds good. I separated this patch into three D108433 D108435 D108437. Please let me know if they are good for review.

[hoy]

Updating D107299: [CSSPGO] Split context string to deduplicate function name used in the context.

[hoy]

Updating D107299: [CSSPGO] Split context string to deduplicate function name used in the context.

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.
if (!EntrySplit.second.empty()) {
       LineLoc.LineOffset = LineOffset;
 
       // Get discriminator
       if (!LocSplit.second.empty())
         LocSplit.second.getAsInteger(10, LineLoc.Discriminator);
     }
   }
 
+  static std::vector<std::pair<StringRef, LineLocation>>
+  decodeFullContextString(StringRef ContextStr) {
+    // If the context string is enclosed in a pair of brackets, strip the
+    // brackets before decoding.
+    if (ContextStr.startswith("[")) {
+      assert(ContextStr.endswith("]") && "Unpaired [] in context string");
+      ContextStr = ContextStr.substr(1, ContextStr.size() - 2);
+    }
+
+    std::vector<std::pair<StringRef, LineLocation>> decodeResult;
+    StringRef ContextRemain = ContextStr;
+    StringRef ChildContext;
+    LineLocation CallSiteLoc(0, 0);
+    while (!ContextRemain.empty()) {
+      auto ContextSplit = SampleContext::splitContextString(ContextRemain);
+      StringRef CurrentContext = ContextSplit.first;
+      ContextRemain = ContextSplit.second;
+
+      StringRef CalleeName;
+      LineLocation CallSiteLoc(0, 0);
+      SampleContext::decodeContextString(CurrentContext, CalleeName,
+                                         CallSiteLoc);
+      decodeResult.emplace_back(CalleeName, CallSiteLoc);
+    }
+    return decodeResult;
+  }
+
   operator StringRef() const { return FullContext; }
   bool hasAttribute(ContextAttributeMask A) { return Attributes & (uint32_t)A; }
   void setAttribute(ContextAttributeMask A) { Attributes |= (uint32_t)A; }
   uint32_t getAllAttributes() { return Attributes; }
   void setAllAttributes(uint32_t A) { Attributes = A; }
   bool hasState(ContextStateMask S) { return State & (uint32_t)S; }
   void setState(ContextStateMask S) { State |= (uint32_t)S; }
   void clearState(ContextStateMask S) { State &= (uint32_t)~S; }

llvm/include/llvm/ProfileData/SampleProfReader.h

public:
   bool profileIsCS() const { return ProfileIsCS; }
 
   virtual std::unique_ptr<ProfileSymbolList> getProfileSymbolList() {
     return nullptr;
   };
 
   /// It includes all the names that have samples either in outline instance
   /// or inline instance.
-  virtual std::vector<StringRef> *getNameTable() { return nullptr; }
+  virtual std::vector<StringRef> *getNameTable(bool IsContextName = false) {
+    return nullptr;
+  }
   virtual bool dumpSectionInfo(raw_ostream &OS = dbgs()) { return false; };
 
   /// Return whether names in the profile are all MD5 numbers.
   virtual bool useMD5() { return false; }
 
   /// Don't read profile without context if the flag is set. This is only meaningful
   /// for ExtBinary format.
   virtual void setSkipFlatProf(bool Skip) {}
public:
   /// Read and validate the file header.
   virtual std::error_code readHeader() override;
 
   /// Read sample profiles from the associated file.
   std::error_code readImpl() override;
 
   /// It includes all the names that have samples either in outline instance
   /// or inline instance.
-  virtual std::vector<StringRef> *getNameTable() override { return &NameTable; }
+  virtual std::vector<StringRef> *
+  getNameTable(bool IsContextName = false) override {
+    return &NameTable;
+  }
 
 protected:
   /// Read a numeric value of type T from the profile.
   ///
   /// If an error occurs during decoding, a diagnostic message is emitted and
   /// EC is set.
   ///
   /// \returns the read value.
protected:
 
   /// Read the string index and check whether it overflows the table.
   template <typename T> inline ErrorOr<uint32_t> readStringIndex(T &Table);
 
   /// Return true if we've reached the end of file.
   bool at_eof() const { return Data >= End; }
 
   /// Read the next function profile instance.
-  std::error_code readFuncProfile(const uint8_t *Start);
+  std::error_code readFuncProfile(const uint8_t *Start,
+                                  bool IsContextName = false);
 
   /// Read the contents of the given profile instance.
   std::error_code readProfile(FunctionSamples &FProfile);
 
   /// Read the contents of Magic number and Version number.
   std::error_code readMagicIdent();
 
   /// Read profile summary.
   std::error_code readSummary();
 
   /// Read the whole name table.
   virtual std::error_code readNameTable();
 
   /// Points to the current location in the buffer.
   const uint8_t *Data = nullptr;
 
   /// 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<StringRef> readStringFromTable(bool IsContextName = false);
 
 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<StringRef>
+  readStringFromTable(bool IsContextName = false) override;
 
   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 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;
 
+  /// CSNameBuf buffer is used to save full context names.
+  std::unique_ptr<std::vector<std::string>> CSNameBuf;
+  /// CSNameTable contains StringRefs referring to the names in CSNameBuf.
+  std::unique_ptr<std::vector<StringRef>> 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) {}
public:
   /// Return whether names in the profile are all MD5 numbers.
   virtual bool useMD5() override { return MD5StringBuf.get(); }
 
   virtual std::unique_ptr<ProfileSymbolList> getProfileSymbolList() override {
     return std::move(ProfSymList);
   };
 
   virtual void setSkipFlatProf(bool Skip) override { SkipFlatProf = Skip; }
+
+  virtual std::vector<StringRef> *
+  getNameTable(bool IsContextName = false) override {
+    if (IsContextName && CSNameTable.get())
+      return CSNameTable.get();
+    return &NameTable;
+  }
 };
 
 class SampleProfileReaderExtBinary : public SampleProfileReaderExtBinaryBase {
 private:
   virtual std::error_code verifySPMagic(uint64_t Magic) override;
   virtual std::error_code
   readCustomSection(const SecHdrTableEntry &Entry) override {
     return sampleprof_error::success;
private:
   /// The table mapping from function name to the offset of its FunctionSample
   /// towards file start.
   DenseMap<StringRef, uint64_t> FuncOffsetTable;
   /// The set containing the functions to use when compiling a module.
   DenseSet<StringRef> FuncsToUse;
   virtual std::error_code verifySPMagic(uint64_t Magic) override;
   virtual std::error_code readNameTable() override;
   /// Read a string indirectly via the name table.
-  virtual ErrorOr<StringRef> readStringFromTable() override;
+  virtual ErrorOr<StringRef>
+  readStringFromTable(bool IsContextName = false) override;
   virtual std::error_code readHeader() override;
   std::error_code readFuncOffsetTable();
 
 public:
   SampleProfileReaderCompactBinary(std::unique_ptr<MemoryBuffer> B,
                                    LLVMContext &C)
       : SampleProfileReaderBinary(std::move(B), C, SPF_Compact_Binary) {}
 

llvm/include/llvm/ProfileData/SampleProfWriter.h

 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(bool IsContextName) {
+    return NameTable;
+  }
   virtual std::error_code writeMagicIdent(SampleProfileFormat Format);
   virtual std::error_code writeNameTable();
   virtual std::error_code
   writeHeader(const StringMap<FunctionSamples> &ProfileMap) override;
   std::error_code writeSummary();
   std::error_code writeNameIdx(StringRef FName, bool IsContextName = false);
   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 addNames(const FunctionSamples &S);
 
 private:
ExtBinaryHdrLayoutTable = {
         // 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 (auto &Entry : SectionHdrLayout) {
              "resetSecLayout has to be called before any flag setting");
     }
 #endif
     SecLayout = SL;
     SectionHdrLayout = ExtBinaryHdrLayoutTable[SL];
   }
 
 protected:
+  virtual MapVector<StringRef, uint32_t> &getNameTable(bool IsContextName) {
+    return IsContextName ? CSNameTable : NameTable;
+  }
   uint64_t markSectionStart(SecType Type, uint32_t LayoutIdx);
   std::error_code addNewSection(SecType Sec, uint32_t LayoutIdx,
                                 uint64_t SectionStart);
   template <class SecFlagType>
   void addSectionFlag(SecType Type, SecFlagType Flag) {
     for (auto &Entry : SectionHdrLayout) {
       if (Entry.Type == Type)
         addSecFlag(Entry, Flag);
protected:
   // 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);
 
   // Helper function to write name table.
   virtual std::error_code writeNameTable() override;
+  std::error_code writeCSNameTableSection();
 
   std::error_code writeFuncMetadata(const StringMap<FunctionSamples> &Profiles);
 
   // Functions to write various kinds of sections.
   std::error_code
   writeNameTableSection(const StringMap<FunctionSamples> &ProfileMap);
   std::error_code writeFuncOffsetTable();
   std::error_code writeProfileSymbolListSection();
private:
   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;
   // Whether to use MD5 to represent string.
   bool UseMD5 = false;
 
+  // Maps context name to its index in the context name table.
+  MapVector<StringRef, uint32_t> CSNameTable;
+
   ProfileSymbolList *ProfSymList = nullptr;
 };
 
 class SampleProfileWriterExtBinary : public SampleProfileWriterExtBinaryBase {
 public:
   SampleProfileWriterExtBinary(std::unique_ptr<raw_ostream> &OS)
       : SampleProfileWriterExtBinaryBase(OS) {}
 

llvm/lib/ProfileData/SampleProfReader.cpp

inline ErrorOr<uint32_t> SampleProfileReaderBinary::readStringIndex(T &Table) {
   auto Idx = readNumber<uint32_t>();
   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);
+ErrorOr<StringRef>
+SampleProfileReaderBinary::readStringFromTable(bool IsContextName) {
+  auto NTable = getNameTable(IsContextName);
+  assert(NTable && "NTable should not be nullptr");
+  auto Idx = readStringIndex(*NTable);
   if (std::error_code EC = Idx.getError())
     return EC;
-
-  return NameTable[*Idx];
+  return (*NTable)[*Idx];
 }
 
-ErrorOr<StringRef> SampleProfileReaderExtBinaryBase::readStringFromTable() {
+ErrorOr<StringRef>
+SampleProfileReaderExtBinaryBase::readStringFromTable(bool IsContextName) {
   if (!FixedLengthMD5)
-    return SampleProfileReaderBinary::readStringFromTable();
+    return SampleProfileReaderBinary::readStringFromTable(IsContextName);
 
   // read NameTable index.
   auto Idx = readStringIndex(NameTable);
   if (std::error_code EC = Idx.getError())
     return EC;
 
   // Check whether the name to be accessed has been accessed before,
   // if not, read it from memory directly.
if (SR.empty()) {
     // in MD5StringBuf.
     MD5StringBuf->push_back(std::to_string(*FID));
     SR = MD5StringBuf->back();
     Data = SavedData;
   }
   return SR;
 }
 
-ErrorOr<StringRef> SampleProfileReaderCompactBinary::readStringFromTable() {
+ErrorOr<StringRef>
+SampleProfileReaderCompactBinary::readStringFromTable(bool IsContextName) {
   auto Idx = readStringIndex(NameTable);
   if (std::error_code EC = Idx.getError())
     return EC;
 
   return StringRef(NameTable[*Idx]);
 }
 
 std::error_code
for (uint32_t J = 0; J < *NumCallsites; ++J) {
     CalleeProfile.setName(*FName);
     if (std::error_code EC = readProfile(CalleeProfile))
       return EC;
   }
 
   return sampleprof_error::success;
 }
 
-std::error_code
-SampleProfileReaderBinary::readFuncProfile(const uint8_t *Start) {
+std::error_code SampleProfileReaderBinary::readFuncProfile(const uint8_t *Start,
+                                                           bool IsContextName) {
   Data = Start;
   auto NumHeadSamples = readNumber<uint64_t>();
   if (std::error_code EC = NumHeadSamples.getError())
     return EC;
 
-  auto FName(readStringFromTable());
+  auto FName(readStringFromTable(IsContextName));
   if (std::error_code EC = FName.getError())
     return EC;
 
   SampleContext FContext(*FName);
   Profiles[FContext] = FunctionSamples();
   FunctionSamples &FProfile = Profiles[FContext];
   FProfile.setName(FContext.getNameWithoutContext());
   FProfile.setContext(FContext);
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(readStringFromTable(FunctionSamples::ProfileIsCS));
     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() {
   // called after FunctionSamples::HasUniqSuffix is set, i.e. after
   // 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))
+      if (std::error_code EC =
+              readFuncProfile(Data, FunctionSamples::ProfileIsCS))
         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);

 }
 
 std::error_code SampleProfileReaderExtBinaryBase::readNameTableSec(bool IsMD5) {
   if (IsMD5)
     return readMD5NameTable();
   return SampleProfileReaderBinary::readNameTable();
 }
 
+std::error_code SampleProfileReaderExtBinaryBase::readCSNameTableSec() {
+  auto Size = readNumber<uint32_t>();
+  if (std::error_code EC = Size.getError())
+    return EC;
+
+  CSNameBuf = std::make_unique<std::vector<std::string>>();
+  CSNameBuf->reserve(*Size + CSNameBuf->size());
+  CSNameTable = std::make_unique<std::vector<StringRef>>();
+  CSNameTable->reserve(*Size + CSNameTable->size());
+  for (uint32_t I = 0; I < *Size; ++I) {
+    std::string Callsite;
+    auto CallsiteSize = readNumber<uint32_t>();
+    if (std::error_code EC = CallsiteSize.getError())
+      return EC;
+    for (uint32_t J = 0; J < *CallsiteSize; ++J) {
+      if (J != 0)
+        Callsite = Callsite + " @ ";
+
+      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;
+
+      if (J == (*CallsiteSize) - 1) {
+        Callsite = Callsite + (*FName).str();
+        continue;
+      }
+      Callsite = Callsite + (*FName).str() + ":" + std::to_string(*LineOffset) +
+                 "." + std::to_string(*Discriminator);
+    }
+    Callsite = "[" + Callsite + "]";
+    CSNameBuf->push_back(Callsite);
+    CSNameTable->push_back(CSNameBuf->back());
+  }
+
+  return sampleprof_error::success;
+}
+
 std::error_code
 SampleProfileReaderExtBinaryBase::readFuncMetadata(bool ProfileHasAttribute) {
   while (Data < End) {
-    auto FName(readStringFromTable());
+    auto FName(readStringFromTable(FunctionSamples::ProfileIsCS));
     if (std::error_code EC = FName.getError())
       return EC;
 
     SampleContext FContext(*FName);
     bool ProfileInMap = Profiles.count(FContext);
 
     if (ProfileIsProbeBased) {
       auto Checksum = readNumber<uint64_t>();

llvm/lib/ProfileData/SampleProfWriter.cpp

 }
 
 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) {
   for (const auto &I : ProfileMap) {
     assert(I.first() == I.second.getNameWithContext() &&
            "Inconsistent profile map");
-    addName(I.second.getNameWithContext(), FunctionSamples::ProfileIsCS);
+    if (FunctionSamples::ProfileIsCS) {
+      StringRef CSName = I.second.getNameWithContext();
+      addName(CSName, true);
+      std::vector<std::pair<StringRef, LineLocation>> callsites =
+          SampleContext::decodeFullContextString(CSName);
+      for (const auto NameLoc : callsites)
+        addName(NameLoc.first, false);
+    } else {
+      addName(I.second.getName(), false);
+    }
     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() {
+  auto &OS = *OutputStream;
+  std::set<StringRef> CSNameSet;
+  stablizeNameTable(CSNameTable, CSNameSet);
+
+  encodeULEB128(CSNameSet.size(), OS);
+  support::endian::Writer Writer(OS, support::little);
+  for (auto CSName : CSNameSet) {
+    std::vector<std::pair<StringRef, LineLocation>> callsites =
+        SampleContext::decodeFullContextString(CSName);
+    encodeULEB128(callsites.size(), OS);
+    for (const auto NameLoc : callsites) {
+      if (std::error_code EC = writeNameIdx(NameLoc.first))
+        return EC;
+      encodeULEB128(NameLoc.second.LineOffset, OS);
+      encodeULEB128(NameLoc.second.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;
 }
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;
std::error_code SampleProfileWriterExtBinary::writeDefaultLayout(
   // 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) {

 std::error_code SampleProfileWriterBinary::writeNameIdx(StringRef FName,
                                                         bool IsContextName) {
   std::string BracketedName;
   if (IsContextName) {
     BracketedName = "[" + FName.str() + "]";
     FName = StringRef(BracketedName);
   }
 
-  const auto &Ret = NameTable.find(FName);
-  if (Ret == NameTable.end())
+  auto &NTable = getNameTable(IsContextName);
+  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));
+  auto &NTable = getNameTable(IsContextName);
+  NTable.insert(std::make_pair(FName, 0));
 }
 
 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 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;
for (auto Entry : FuncOffsetTable) {
     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;
 }