[CSSPGO] Fix dangling context strings and improve profile order consistency and error handling

[wenlei]

Summary

This patch fixed the follow issues along side with some refactoring:

1. Fix bugs where StringRef for context string out live the underlying std::string. We now keep string table in profile generator to hold std::strings. We also do the same for bracketed context strings in profile writer.
2. Make sure profile output strictly follow (total sample, name) order. Previously, there's inconsistency between ProfileMap's key and FunctionSamples's name, leading to inconsistent ordering. This is now fixed by introducing context profile canonicalization. Assertions are also added to make sure ProfileMap's key and FunctionSamples's name are always consistent.
3. Enhanced error handling for profile writing to make sure we bubble up errors properly for both llvm-profgen and llvm-profdata when string table is not populated correctly for extended binary profile.
4. Keep all internal context representation bracket free. This avoids creating new strings for context trimming, merging and preinline. getNameWithContext API is now simplied accordingly.
5. Factor out the code for context trimming and merging into SampleContextTrimmer in SampleProf.cpp. This enables llvm-profdata to use the trimmer when merging profiles. Changes in llvm-profgen will be in separate patch.

[hoy]

Thanks for fixing the name volatileness issue which was overlooked previously.

llvm/lib/ProfileData/SampleProf.cpp

385

For a profile whose key is different from the profile name string, are there other profiles merged into it, or is it merged into other profiles? A comment could be helpful.

llvm/lib/ProfileData/SampleProfWriter.cpp

503

Can brackets be output here instead of adding it to name table and looking it up there?

[wenlei]

llvm/lib/ProfileData/SampleProf.cpp

385

It could be a profile from context merging, see line 358. For pre-inliner, it could also be promoted context profile without merging. Will add comment.

llvm/lib/ProfileData/SampleProfWriter.cpp

503

Not every name needs bracket. One example is the call target name for call sites. If we do it here, we can't differentiate between such cases and the full context names used in profile header.

[wenlei]

Update comments.

[hoy]

llvm/lib/ProfileData/SampleProf.cpp

338

Nit: cold profile

388

Still trying to understand. So for cold profile trim, what kind of a profile can be here. In trimAndMergeColdContextProfiles all cold profiles are removed and all based profiles are added into ProfileMap.

llvm/lib/ProfileData/SampleProfWriter.cpp

503

I see. Yeah, unfortunately we cannot tell which one needs brackets or not.

[wenlei]

llvm/lib/ProfileData/SampleProf.cpp

388

Cold profile trimming is simply removing cold profiles, so it won't lead to inconsistency between profile map key and actual profile's context. Actually now looking at the merging part again, the currently implementation doesn't require canonicalization either as the new base profile has the correct key.

I'm removing canonicalization from trimAndMergeColdContextProfiles, but pre-inliner still needs it.

[wenlei]

Address Hongtao's comments.

[hoy]

LGTM, thanks.

llvm/include/llvm/ProfileData/SampleProf.h

public:
   void setAllAttributes(uint32_t A) { Attributes = A; }
   bool hasState(ContextStateMask S) { return State & (uint32_t)S; }
   void setState(ContextStateMask S) { State |= (uint32_t)S; }
   void clearState(ContextStateMask S) { State &= (uint32_t)~S; }
   bool hasContext() const { return State != UnknownContext; }
   bool isBaseContext() const { return CallingContext.empty(); }
   StringRef getNameWithoutContext() const { return Name; }
   StringRef getCallingContext() const { return CallingContext; }
-  StringRef getNameWithContext(bool WithBracket = false) const {
-    return WithBracket ? InputContext : FullContext;
-  }
+  StringRef getNameWithContext() const { return FullContext; }
 
 private:
   // Give a context string, decode and populate internal states like
   // Function name, Calling context and context state. Example of input
   // `ContextStr`: `[main:3 @ _Z5funcAi:1 @ _Z8funcLeafi]`
   void setContext(StringRef ContextStr, ContextStateMask CState) {
     assert(!ContextStr.empty());
-    InputContext = ContextStr;
     // Note that `[]` wrapped input indicates a full context string, otherwise
     // it's treated as context-less function name only.
     bool HasContext = ContextStr.startswith("[");
     if (!HasContext && CState == UnknownContext) {
       State = UnknownContext;
       Name = FullContext = ContextStr;
     } else {
       // Assume raw context profile if unspecified
if (!HasContext && CState == UnknownContext) {
         CallingContext = NameContext.second;
       } else {
         Name = NameContext.second;
         CallingContext = NameContext.first;
       }
     }
   }
 
-  // Input context string including bracketed calling context and leaf function
-  // name
-  StringRef InputContext;
   // Full context string including calling context and leaf function name
   StringRef FullContext;
   // Function name for the associated sample profile
   StringRef Name;
   // Calling context (leaf function excluded) for the associated sample profile
   StringRef CallingContext;
   // State of the associated sample profile
   uint32_t State;
public:
 
   /// Merge the samples in \p Other into this one.
   /// Optionally scale samples by \p Weight.
   sampleprof_error merge(const FunctionSamples &Other, uint64_t Weight = 1) {
     sampleprof_error Result = sampleprof_error::success;
     Name = Other.getName();
     if (!GUIDToFuncNameMap)
       GUIDToFuncNameMap = Other.GUIDToFuncNameMap;
-    if (Context.getNameWithContext(true).empty())
+    if (Context.getNameWithContext().empty())
       Context = Other.getContext();
     if (FunctionHash == 0) {
       // Set the function hash code for the target profile.
       FunctionHash = Other.getFunctionHash();
     } else if (FunctionHash != Other.getFunctionHash()) {
       // The two profiles coming with different valid hash codes indicates
       // either:
       // 1. They are same-named static functions from different compilation
public:
 
   /// Set the name of the function.
   void setName(StringRef FunctionName) { Name = FunctionName; }
 
   /// Return the function name.
   StringRef getName() const { return Name; }
 
   /// Return function name with context.
-  StringRef getNameWithContext(bool WithBracket = false) const {
-    return FunctionSamples::ProfileIsCS
-               ? Context.getNameWithContext(WithBracket)
-               : Name;
+  StringRef getNameWithContext() const {
+    return FunctionSamples::ProfileIsCS ? Context.getNameWithContext() : Name;
   }
 
   /// Return the original function name.
   StringRef getFuncName() const { return getFuncName(Name); }
 
   void setFunctionHash(uint64_t Hash) { FunctionHash = Hash; }
 
   uint64_t getFunctionHash() const { return FunctionHash; }
public:
   }
 
   const SamplesWithLocList &get() const { return V; }
 
 private:
   SamplesWithLocList V;
 };
 
+/// SampleContextTrimmer impelements helper functions to trim, merge cold
+/// context profiles. It also supports context profile canonicalization to make
+/// sure ProfileMap's key is consistent with FunctionSample's name/context.
+class SampleContextTrimmer {
+public:
+  SampleContextTrimmer(StringMap<FunctionSamples> &Profiles)
+      : ProfileMap(Profiles){};
+  // Trim and merge cold context profile when requested.
+  void trimAndMergeColdContextProfiles(uint64_t ColdCountThreshold,
+                                       bool TrimColdContext = true,
+                                       bool MergeColdContext = true);
+  // Canonicalize context profile name and attributes.
+  void canonicalizeContextProfiles();
+
+private:
+  StringMap<FunctionSamples> &ProfileMap;
+};
+
 /// ProfileSymbolList records the list of function symbols shown up
 /// in the binary used to generate the profile. It is useful to
 /// to discriminate a function being so cold as not to shown up
 /// in the profile and a function newly added.
 class ProfileSymbolList {
 public:
   /// copy indicates whether we need to copy the underlying memory
   /// for the input Name.

llvm/include/llvm/ProfileData/SampleProfWriter.h

 #include "llvm/ProfileData/SampleProf.h"
 #include "llvm/Support/ErrorOr.h"
 #include "llvm/Support/raw_ostream.h"
 #include <algorithm>
 #include <cstdint>
 #include <memory>
 #include <set>
 #include <system_error>
+#include <unordered_set>
 
 namespace llvm {
 namespace sampleprof {
 
 enum SectionLayout {
   DefaultLayout,
   // The layout splits profile with context information from profile without
   // context information. When Thinlto is enabled, ThinLTO postlink phase only
public:
   virtual std::error_code writeSample(const FunctionSamples &S) override;
 
 protected:
   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);
+  std::error_code writeNameIdx(StringRef FName, bool IsContextName = false);
   std::error_code writeBody(const FunctionSamples &S);
   inline void stablizeNameTable(std::set<StringRef> &V);
 
   MapVector<StringRef, uint32_t> NameTable;
+  std::unordered_set<std::string> BracketedContextStr;
 
-  void addName(StringRef FName);
+  void addName(StringRef FName, bool IsContextName = false);
   void addNames(const FunctionSamples &S);
 
 private:
   friend ErrorOr<std::unique_ptr<SampleProfileWriter>>
   SampleProfileWriter::create(std::unique_ptr<raw_ostream> &OS,
                               SampleProfileFormat Format);
 };
 

llvm/lib/ProfileData/SampleProf.cpp

while (Size < ListSize && StrNum < ProfileSymbolListCutOff) {
     Size += Str.size() + 1;
     StrNum++;
   }
   if (Size != ListSize && StrNum != ProfileSymbolListCutOff)
     return sampleprof_error::malformed;
   return sampleprof_error::success;
 }
 
+void SampleContextTrimmer::trimAndMergeColdContextProfiles(
+    uint64_t ColdCountThreshold, bool TrimColdContext, bool MergeColdContext) {
+  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;
+  for (const auto &I : ProfileMap) {
+    const FunctionSamples &FunctionProfile = I.second;
+    if (FunctionProfile.getTotalSamples() >= ColdCountThreshold)
+      continue;
+    ColdProfiles.emplace_back(I.getKey(), &I.second);
+  }
+
+  // Remove the cold profile from ProfileMap and merge them into BaseProileMap
+  StringMap<FunctionSamples> BaseProfileMap;
+  for (const auto &I : ColdProfiles) {
+    if (MergeColdContext) {
+      auto Ret = BaseProfileMap.try_emplace(
+          I.second->getContext().getNameWithoutContext(), FunctionSamples());
+      FunctionSamples &BaseProfile = Ret.first->second;
+      BaseProfile.merge(*I.second);
+    }
+    ProfileMap.erase(I.first);
+  }
+
+  // Merge the base profiles into ProfileMap;
+  for (const auto &I : BaseProfileMap) {
+    // Filter the cold base profile
+    if (TrimColdContext && I.second.getTotalSamples() < ColdCountThreshold &&
+        ProfileMap.find(I.getKey()) == ProfileMap.end())
+      continue;
+    // Merge the profile if the original profile exists, otherwise just insert
+    // as a new profile
+    auto Ret = ProfileMap.try_emplace(I.getKey(), FunctionSamples());
+    if (Ret.second) {
+      SampleContext FContext(Ret.first->first(), RawContext);
+      FunctionSamples &FProfile = Ret.first->second;
+      FProfile.setContext(FContext);
+      FProfile.setName(FContext.getNameWithoutContext());
+    }
+    FunctionSamples &OrigProfile = Ret.first->second;
+    OrigProfile.merge(I.second);
+  }
+}
+
+void SampleContextTrimmer::canonicalizeContextProfiles() {
+  StringSet<> ProfilesToBeRemoved;
+  // Note that StringMap order is guaranteed to be top-down order,
+  // this makes sure we make room for promoted/merged context in the
+  // map, before we move profiles in the map.
+  for (auto &I : ProfileMap) {
+    FunctionSamples &FProfile = I.second;
+    StringRef ContextStr = FProfile.getNameWithContext();
+    if (I.first() == ContextStr)
+      continue;
+
+    // Use the context string from FunctionSamples to update the keys of
+    // ProfileMap. They can get out of sync after context profile promotion
+    // through pre-inliner.
+    auto Ret = ProfileMap.try_emplace(ContextStr, FProfile);
+    assert(Ret.second && "Conext conflict during canonicalization");
+    FProfile = Ret.first->second;
+
+    // Track the context profile to remove
+    ProfilesToBeRemoved.erase(ContextStr);
+    ProfilesToBeRemoved.insert(I.first());
+  }
+
+  for (auto &I : ProfilesToBeRemoved) {
+    ProfileMap.erase(I.first());
+  }
+}
+
 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);
 
   std::string OutputString;
   for (auto &Sym : SortedList) {

llvm/lib/ProfileData/SampleProfWriter.cpp

 using namespace llvm;
 using namespace sampleprof;
 
 std::error_code SampleProfileWriter::writeFuncProfiles(
     const StringMap<FunctionSamples> &ProfileMap) {
   // Sort the ProfileMap by total samples.
   typedef std::pair<StringRef, const FunctionSamples *> NameFunctionSamples;
   std::vector<NameFunctionSamples> V;
-  for (const auto &I : ProfileMap)
-    V.push_back(std::make_pair(I.getKey(), &I.second));
-
+  for (const auto &I : ProfileMap) {
+    assert(I.getKey() == I.second.getNameWithContext() &&
+           "Inconsistent profile map");
+    V.push_back(std::make_pair(I.second.getNameWithContext(), &I.second));
+  }
   llvm::stable_sort(
       V, [](const NameFunctionSamples &A, const NameFunctionSamples &B) {
         if (A.second->getTotalSamples() == B.second->getTotalSamples())
           return A.first > B.first;
         return A.second->getTotalSamples() > B.second->getTotalSamples();
       });
 
   for (const auto &I : V) {
if (std::error_code EC = writeSecHdrTable())
     return EC;
 
   return sampleprof_error::success;
 }
 
 std::error_code
 SampleProfileWriterExtBinaryBase::writeSample(const FunctionSamples &S) {
   uint64_t Offset = OutputStream->tell();
-  StringRef Name = S.getNameWithContext(true);
+  StringRef Name = S.getNameWithContext();
   FuncOffsetTable[Name] = 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) {
-    writeNameIdx(entry.first);
-    encodeULEB128(entry.second, OS);
+  for (auto Entry : FuncOffsetTable) {
+    if (std::error_code EC =
+            writeNameIdx(Entry.first, FunctionSamples::ProfileIsCS))
+      return EC;
+    encodeULEB128(Entry.second, OS);
   }
   FuncOffsetTable.clear();
   return sampleprof_error::success;
 }
 
 std::error_code SampleProfileWriterExtBinaryBase::writeFuncMetadata(
     const StringMap<FunctionSamples> &Profiles) {
   if (!FunctionSamples::ProfileIsProbeBased && !FunctionSamples::ProfileIsCS)
     return sampleprof_error::success;
   auto &OS = *OutputStream;
   for (const auto &Entry : Profiles) {
-    writeNameIdx(Entry.first());
+    if (std::error_code EC = writeNameIdx(Entry.second.getNameWithContext(),
+                                          FunctionSamples::ProfileIsCS))
+      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() {
   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) {
-    addName(I.first());
+    assert(I.first() == I.second.getNameWithContext() &&
+           "Inconsistent profile map");
+    addName(I.second.getNameWithContext(), FunctionSamples::ProfileIsCS);
     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) {

 /// 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;
-  OS << S.getNameWithContext(true) << ":" << S.getTotalSamples();
+  if (FunctionSamples::ProfileIsCS)
+    OS << "[" << S.getNameWithContext() << "]:" << S.getTotalSamples();
+  else
+    OS << S.getName() << ":" << S.getTotalSamples();
+
   if (Indent == 0)
     OS << ":" << S.getHeadSamples();
   OS << "\n";
 
   SampleSorter<LineLocation, SampleRecord> SortedSamples(S.getBodySamples());
   for (const auto &I : SortedSamples.get()) {
     LineLocation Loc = I->first;
     const SampleRecord &Sample = I->second;
if (FunctionSamples::ProfileIsCS) {
       OS.indent(Indent + 1);
       OS << "!Attributes: " << S.getContext().getAllAttributes() << "\n";
     }
   }
 
   return sampleprof_error::success;
 }
 
-std::error_code SampleProfileWriterBinary::writeNameIdx(StringRef FName) {
-  const auto &ret = NameTable.find(FName);
-  if (ret == NameTable.end())
+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())
     return sampleprof_error::truncated_name_table;
-  encodeULEB128(ret->second, *OutputStream);
+  encodeULEB128(Ret->second, *OutputStream);
   return sampleprof_error::success;
 }
 
-void SampleProfileWriterBinary::addName(StringRef FName) {
+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::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())

 std::error_code SampleProfileWriterBinary::writeNameTable() {
   auto &OS = *OutputStream;
   std::set<StringRef> V;
   stablizeNameTable(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() {
   auto &OS = *OutputStream;
 
   // Fill the slot remembered by TableOffset with the offset of FuncOffsetTable.
   auto &OFS = static_cast<raw_fd_ostream &>(OS);
   uint64_t FuncOffsetTableStart = OS.tell();
   if (OFS.seek(TableOffset) == (uint64_t)-1)
     return sampleprof_error::ostream_seek_unsupported;
   support::endian::Writer Writer(*OutputStream, support::little);
   Writer.write(FuncOffsetTableStart);
   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) {
-    writeNameIdx(entry.first);
-    encodeULEB128(entry.second, OS);
+  for (auto Entry : FuncOffsetTable) {
+    if (std::error_code EC =
+            writeNameIdx(Entry.first, FunctionSamples::ProfileIsCS))
+      return EC;
+    encodeULEB128(Entry.second, OS);
   }
   return sampleprof_error::success;
 }
 
 std::error_code SampleProfileWriterCompactBinary::writeNameTable() {
   auto &OS = *OutputStream;
   std::set<StringRef> V;
   stablizeNameTable(V);
std::error_code SampleProfileWriterBinary::writeHeader(
   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) {
-    addName(I.first());
+    assert(I.first() == I.second.getNameWithContext() &&
+           "Inconsistent profile map");
+    addName(I.first(), FunctionSamples::ProfileIsCS);
     addNames(I.second);
   }
 
   writeNameTable();
   return sampleprof_error::success;
 }
 
 void SampleProfileWriterExtBinaryBase::setToCompressAllSections() {
for (auto Entry : Entries) {
     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(true)))
+  if (std::error_code EC =
+          writeNameIdx(S.getNameWithContext(), FunctionSamples::ProfileIsCS))
     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;

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

 ; Firstly test uncompression(--compress-recursion=0)
 ; RUN: llvm-profgen --perfscript=%S/Inputs/recursion-compression-noprobe.perfscript --binary=%S/Inputs/recursion-compression-noprobe.perfbin --output=%t --compress-recursion=0 --csprof-cold-thres=0
 ; RUN: FileCheck %s --input-file %t -check-prefix=CHECK-UNCOMPRESS
 ; RUN: llvm-profgen --perfscript=%S/Inputs/recursion-compression-noprobe.perfscript --binary=%S/Inputs/recursion-compression-noprobe.perfbin --output=%t --csprof-cold-thres=0
 ; RUN: FileCheck %s --input-file %t
 
 ; CHECK-UNCOMPRESS:[main:1 @ foo:3 @ fa:2 @ fb]:48:0
 ; CHECK-UNCOMPRESS: 1: 11
 ; CHECK-UNCOMPRESS: 2: 1 fa:1
 ; CHECK-UNCOMPRESS:[main:1 @ foo:3 @ fa]:24:0
 ; CHECK-UNCOMPRESS: 1: 1
 ; CHECK-UNCOMPRESS: 2: 13 fb:11
-; CHECK-UNCOMPRESS:[main:1 @ foo]:7:0
-; CHECK-UNCOMPRESS: 2: 1
-; CHECK-UNCOMPRESS: 3: 2 fa:1
 ; CHECK-UNCOMPRESS:[main:1 @ foo:3 @ fa:2 @ fb:2 @ fa]:7:0
 ; CHECK-UNCOMPRESS: 1: 1
 ; CHECK-UNCOMPRESS: 2: 2 fb:1
-; CHECK-UNCOMPRESS:[main:1 @ foo:3 @ fa:2 @ fb:2 @ fa:2 @ fb]:2:0
-; CHECK-UNCOMPRESS: 2: 1 fa:1
+; CHECK-UNCOMPRESS:[main:1 @ foo]:7:0
+; CHECK-UNCOMPRESS: 2: 1
+; CHECK-UNCOMPRESS: 3: 2 fa:1
 ; CHECK-UNCOMPRESS:[main:1 @ foo:3 @ fa:2 @ fb:2 @ fa:2 @ fb:2 @ fa]:2:0
 ; CHECK-UNCOMPRESS: 4: 1
+; CHECK-UNCOMPRESS:[main:1 @ foo:3 @ fa:2 @ fb:2 @ fa:2 @ fb]:2:0
+; CHECK-UNCOMPRESS: 2: 1 fa:1
 
 
 ; CHECK: [main:1 @ foo:3 @ fa:2 @ fb]:48:0
 ; CHECK:  1: 11
 ; CHECK:  2: 1 fa:1
 ; CHECK: [main:1 @ foo:3 @ fa]:24:0
 ; CHECK:  1: 1
 ; CHECK:  2: 13 fb:11

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

 ; Firstly test uncompression(--compress-recursion=0)
 ; RUN: llvm-profgen --perfscript=%S/Inputs/recursion-compression-pseudoprobe.perfscript --binary=%S/Inputs/recursion-compression-pseudoprobe.perfbin --output=%t --compress-recursion=0 --csprof-cold-thres=0
 ; RUN: FileCheck %s --input-file %t -check-prefix=CHECK-UNCOMPRESS
 ; RUN: llvm-profgen --perfscript=%S/Inputs/recursion-compression-pseudoprobe.perfscript --binary=%S/Inputs/recursion-compression-pseudoprobe.perfbin --output=%t --show-unwinder-output --csprof-cold-thres=0 | FileCheck %s --check-prefix=CHECK-UNWINDER
 ; RUN: FileCheck %s --input-file %t
 
-; CHECK-UNCOMPRESS: [main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:6 @ fa]:4:1
+; CHECK-UNCOMPRESS: [main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:6 @ fa:8 @ fa:7 @ fb:6 @ fa]:4:1
 ; CHECK-UNCOMPRESS:  1: 1
 ; CHECK-UNCOMPRESS:  3: 1
-; CHECK-UNCOMPRESS:  5: 1
-; CHECK-UNCOMPRESS:  8: 1 fa:1
+; CHECK-UNCOMPRESS:  4: 1
+; CHECK-UNCOMPRESS:  7: 1 fb:1
 ; CHECK-UNCOMPRESS:  !CFGChecksum: 120515930909
 ; CHECK-UNCOMPRESS: [main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:6 @ fa:8 @ fa]:4:1
 ; CHECK-UNCOMPRESS:  1: 1
 ; CHECK-UNCOMPRESS:  3: 1
 ; CHECK-UNCOMPRESS:  4: 1
 ; CHECK-UNCOMPRESS:  7: 1 fb:1
 ; CHECK-UNCOMPRESS:  !CFGChecksum: 120515930909
-; CHECK-UNCOMPRESS: [main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:6 @ fa:8 @ fa:7 @ fb:6 @ fa]:4:1
+; CHECK-UNCOMPRESS: [main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:6 @ fa]:4:1
 ; CHECK-UNCOMPRESS:  1: 1
 ; CHECK-UNCOMPRESS:  3: 1
-; CHECK-UNCOMPRESS:  4: 1
-; CHECK-UNCOMPRESS:  7: 1 fb:1
+; CHECK-UNCOMPRESS:  5: 1
+; CHECK-UNCOMPRESS:  8: 1 fa:1
 ; CHECK-UNCOMPRESS:  !CFGChecksum: 120515930909
-; CHECK-UNCOMPRESS: [main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb]:3:1
-; CHECK-UNCOMPRESS:  1: 1
-; CHECK-UNCOMPRESS:  2: 1
-; CHECK-UNCOMPRESS:  5: 1 fb:1
-; CHECK-UNCOMPRESS:  !CFGChecksum: 72617220756
-; CHECK-UNCOMPRESS: [main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb]:3:1
-; CHECK-UNCOMPRESS:  1: 1
-; CHECK-UNCOMPRESS:  2: 1
-; CHECK-UNCOMPRESS:  5: 1 fb:1
-; CHECK-UNCOMPRESS:  !CFGChecksum: 72617220756
-; CHECK-UNCOMPRESS: [main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb]:3:1
-; CHECK-UNCOMPRESS:  1: 1
-; CHECK-UNCOMPRESS:  2: 1
-; CHECK-UNCOMPRESS:  5: 1 fb:1
-; CHECK-UNCOMPRESS:  !CFGChecksum: 72617220756
-; CHECK-UNCOMPRESS: [main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb]:3:1
+; CHECK-UNCOMPRESS: [main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:6 @ fa:8 @ fa:7 @ fb:6 @ fa:7 @ fb]:3:1
 ; CHECK-UNCOMPRESS:  1: 1
 ; CHECK-UNCOMPRESS:  3: 1
 ; CHECK-UNCOMPRESS:  6: 1 fa:1
 ; CHECK-UNCOMPRESS:  !CFGChecksum: 72617220756
 ; CHECK-UNCOMPRESS: [main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:6 @ fa:8 @ fa:7 @ fb]:3:1
 ; CHECK-UNCOMPRESS:  1: 1
 ; CHECK-UNCOMPRESS:  3: 1
 ; CHECK-UNCOMPRESS:  6: 1 fa:1
 ; CHECK-UNCOMPRESS:  !CFGChecksum: 72617220756
-; CHECK-UNCOMPRESS: [main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:6 @ fa:8 @ fa:7 @ fb:6 @ fa:7 @ fb]:3:1
+; CHECK-UNCOMPRESS: [main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb]:3:1
 ; CHECK-UNCOMPRESS:  1: 1
 ; CHECK-UNCOMPRESS:  3: 1
 ; CHECK-UNCOMPRESS:  6: 1 fa:1
 ; CHECK-UNCOMPRESS:  !CFGChecksum: 72617220756
+; CHECK-UNCOMPRESS: [main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb]:3:1
+; CHECK-UNCOMPRESS:  1: 1
+; CHECK-UNCOMPRESS:  2: 1
+; CHECK-UNCOMPRESS:  5: 1 fb:1
+; CHECK-UNCOMPRESS:  !CFGChecksum: 72617220756
+; CHECK-UNCOMPRESS: [main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb]:3:1
+; CHECK-UNCOMPRESS:  1: 1
+; CHECK-UNCOMPRESS:  2: 1
+; CHECK-UNCOMPRESS:  5: 1 fb:1
+; CHECK-UNCOMPRESS:  !CFGChecksum: 72617220756
+; CHECK-UNCOMPRESS: [main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb]:3:1
+; CHECK-UNCOMPRESS:  1: 1
+; CHECK-UNCOMPRESS:  2: 1
+; CHECK-UNCOMPRESS:  5: 1 fb:1
+; CHECK-UNCOMPRESS:  !CFGChecksum: 72617220756
 ; CHECK-UNCOMPRESS: [main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:6 @ fa:8 @ fa:7 @ fb:6 @ fa:7 @ fb:6 @ fa]:2:1
 ; CHECK-UNCOMPRESS:  1: 1
 ; CHECK-UNCOMPRESS:  3: 1
 ; CHECK-UNCOMPRESS:  !CFGChecksum: 120515930909
 ; CHECK-UNCOMPRESS: [main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb]:1:0
 ; CHECK-UNCOMPRESS:  5: 1 fb:1
 ; CHECK-UNCOMPRESS:  !CFGChecksum: 72617220756
 
 
 ; CHECK: [main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb]:13:4
 ; CHECK:  1: 4
 ; CHECK:  2: 3
 ; CHECK:  3: 1
 ; CEHCK:  5: 4 fb:4
 ; CHECK:  6: 1 fa:1
 ; CHECK !CFGChecksum: 72617220756
 ; CHECK: [main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb:6 @ fa:8 @ fa:7 @ fb:6 @ fa]:6:2
 ; CHECK:  1: 2
 ; CHECK:  3: 2
 ; CHECK:  4: 1
 ; CHECK:  7: 1 fb:1
 ; CHECK:  !CFGChecksum: 120515930909
-; CHECK: [main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb:6 @ fa]:4:1
-; CHECK:  1: 1
-; CHECK:  3: 1
-; CHECK:  5: 1
-; CHECK:  8: 1 fa:1
-; CHECK:  !CFGChecksum: 120515930909
 ; CHECK: [main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb:6 @ fa:8 @ fa]:4:1
 ; CHECK:  1: 1
 ; CHECK:  3: 1
 ; CHECK:  4: 1
 ; CHECK:  7: 1 fb:1
 ; CHECK:  !CFGChecksum: 120515930909
-; CHECK: [main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb:6 @ fa:8 @ fa:7 @ fb]:3:1
+; CHECK: [main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb:6 @ fa]:4:1
+; CHECK:  1: 1
+; CHECK:  3: 1
+; CHECK:  5: 1
+; CHECK:  8: 1 fa:1
+; CHECK:  !CFGChecksum: 120515930909
+; CHECK: [main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb:6 @ fa:8 @ fa:7 @ fb:6 @ fa:7 @ fb]:3:1
 ; CHECK:  1: 1
 ; CHECK:  3: 1
 ; CHECK:  6: 1 fa:1
 ; CHECK:  !CFGChecksum: 72617220756
-; CHECK: [main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb:6 @ fa:8 @ fa:7 @ fb:6 @ fa:7 @ fb]:3:1
+; CHECK: [main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb:6 @ fa:8 @ fa:7 @ fb]:3:1
 ; CHECK:  1: 1
 ; CHECK:  3: 1
 ; CHECK:  6: 1 fa:1
 ; CHECK:  !CFGChecksum: 72617220756
 
 
 
+
 ; CHECK-UNWINDER: Binary(recursion-compression-pseudoprobe.perfbin)'s Range Counter:
 ; CHECK-UNWINDER: main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5
 ; CHECK-UNWINDER:   (7a0, 7a7): 1
 ; CHECK-UNWINDER:   (7a0, 7ab): 3
 ; CHECK-UNWINDER:   (7b2, 7b5): 1
 ; CHECK-UNWINDER: main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb:6
 ; CHECK-UNWINDER:   (7c0, 7d4): 1
 ; CHECK-UNWINDER: main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb:6 @ fa:8

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

for (const auto &Input : Inputs) {
     for (StringMap<FunctionSamples>::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(true);
+      StringRef FName = Samples.getNameWithContext();
       MergeResult(Result, ProfileMap[FName].merge(Samples, Input.Weight));
       if (Result != sampleprof_error::success) {
         std::error_code EC = make_error_code(Result);
         handleMergeWriterError(errorCodeToError(EC), Input.Filename, FName);
       }
     }
 
     std::unique_ptr<sampleprof::ProfileSymbolList> ReaderList =
         Reader->getProfileSymbolList();
     if (ReaderList)
       WriterList.merge(*ReaderList);
   }
   auto WriterOrErr =
       SampleProfileWriter::create(OutputFilename, FormatMap[OutputFormat]);
   if (std::error_code EC = WriterOrErr.getError())
     exitWithErrorCode(EC, OutputFilename);
 
   auto Writer = std::move(WriterOrErr.get());
   // WriterList will have StringRef refering to string in Buffer.
   // Make sure Buffer lives as long as WriterList.
   auto Buffer = getInputFileBuf(ProfileSymbolListFile);
   handleExtBinaryWriter(*Writer, OutputFormat, Buffer.get(), WriterList,
                         CompressAllSections, UseMD5, GenPartialProfile);
-  Writer->write(ProfileMap);
+  if (std::error_code EC = Writer->write(ProfileMap))
+    exitWithErrorCode(std::move(EC));
 }
 
 static WeightedFile parseWeightedFile(const StringRef &WeightedFilename) {
   StringRef WeightStr, FileName;
   std::tie(WeightStr, FileName) = WeightedFilename.split(',');
 
   uint64_t Weight;
   if (WeightStr.getAsInteger(10, Weight) || Weight < 1)

llvm/tools/llvm-profgen/CSPreInliner.cpp

if (!Context.isBaseContext() && !Context.hasState(InlinedContext)) {
       ProfilesToBeRemoved.push_back(It.first());
     }
   }
 
   for (StringRef ContextName : ProfilesToBeRemoved) {
     ProfileMap.erase(ContextName);
   }
 
+  // Make sure ProfileMap's key is consistent with FunctionSamples' name.
+  SampleContextTrimmer(ProfileMap).canonicalizeContextProfiles();
+
   LLVM_DEBUG(printProfileNames(ProfileMap, false));
 }

llvm/tools/llvm-profgen/ProfileGenerator.h

 #define LLVM_TOOLS_LLVM_PROGEN_PROFILEGENERATOR_H
 #include "CSPreInliner.h"
 #include "ErrorHandling.h"
 #include "PerfReader.h"
 #include "ProfiledBinary.h"
 #include "llvm/Analysis/ProfileSummaryInfo.h"
 #include "llvm/ProfileData/SampleProfWriter.h"
 #include <memory>
+#include <unordered_set>
 
 using namespace llvm;
 using namespace sampleprof;
 
 namespace llvm {
 namespace sampleprof {
 
 class ProfileGenerator {

 
 protected:
   // Lookup or create FunctionSamples for the context
   FunctionSamples &getFunctionProfileForContext(StringRef ContextId,
                                                 bool WasLeafInlined = false);
   // Post processing for profiles before writing out, such as mermining
   // and trimming cold profiles, running preinliner on profiles.
   void postProcessProfiles();
-  // Merge cold context profile whose total sample is below threshold
-  // into base profile.
-  void mergeAndTrimColdProfile(StringMap<FunctionSamples> &ProfileMap);
   void computeSummaryAndThreshold();
   void write(std::unique_ptr<SampleProfileWriter> Writer,
              StringMap<FunctionSamples> &ProfileMap) override;
 
   // Profile summary to answer isHotCount and isColdCount queries.
   std::unique_ptr<ProfileSummaryInfo> PSI;
 
+  // String table owning context strings created from profile generation.
+  std::unordered_set<std::string> ContextStrings;
+
 private:
   // Helper function for updating body sample for a leaf location in
   // FunctionProfile
   void updateBodySamplesforFunctionProfile(FunctionSamples &FunctionProfile,
                                            const FrameLocation &LeafLoc,
                                            uint64_t Count);
   void populateFunctionBodySamples(FunctionSamples &FunctionProfile,
                                    const RangeSample &RangeCounters,

llvm/tools/llvm-profgen/ProfileGenerator.cpp

if (SampleType == PERF_LBR_STACK) {
     llvm_unreachable("Unsupported perfscript!");
   }
 
   return ProfileGenerator;
 }
 
 void ProfileGenerator::write(std::unique_ptr<SampleProfileWriter> Writer,
                              StringMap<FunctionSamples> &ProfileMap) {
-  Writer->write(ProfileMap);
+  if (std::error_code EC = Writer->write(ProfileMap))
+    exitWithError(std::move(EC));
 }
 
 void ProfileGenerator::write() {
   auto WriterOrErr = SampleProfileWriter::create(OutputFilename, OutputFormat);
   if (std::error_code EC = WriterOrErr.getError())
     exitWithError(EC, OutputFilename);
   write(std::move(WriterOrErr.get()), ProfileMap);
 }
void ProfileGenerator::findDisjointRanges(RangeSample &DisjointRanges,
   }
 }
 
 FunctionSamples &
 CSProfileGenerator::getFunctionProfileForContext(StringRef ContextStr,
                                                  bool WasLeafInlined) {
   auto Ret = ProfileMap.try_emplace(ContextStr, FunctionSamples());
   if (Ret.second) {
-    SampleContext FContext(Ret.first->first(), RawContext);
+    // Make a copy of the underlying context string in string table
+    // before StringRef wrapper is used for context.
+    auto It = ContextStrings.insert(ContextStr.str());
+    SampleContext FContext(*It.first, RawContext);
     if (WasLeafInlined)
       FContext.setAttribute(ContextWasInlined);
     FunctionSamples &FProfile = Ret.first->second;
     FProfile.setContext(FContext);
     FProfile.setName(FContext.getNameWithoutContext());
   }
   return Ret.first->second;
 }
void CSProfileGenerator::postProcessProfiles() {
   computeSummaryAndThreshold();
 
   // Run global pre-inliner to adjust/merge context profile based on estimated
   // inline decisions.
   CSPreInliner(ProfileMap, PSI->getHotCountThreshold(),
                PSI->getColdCountThreshold())
       .run();
 
-  mergeAndTrimColdProfile(ProfileMap);
+  // Trim and merge cold context profile using cold threshold above;
+  SampleContextTrimmer(ProfileMap)
+      .trimAndMergeColdContextProfiles(
+          CSProfColdThreshold, CSProfTrimColdContext, CSProfMergeColdContext);
 }
 
 void CSProfileGenerator::computeSummaryAndThreshold() {
   SampleProfileSummaryBuilder Builder(ProfileSummaryBuilder::DefaultCutoffs);
   auto Summary = Builder.computeSummaryForProfiles(ProfileMap);
   PSI.reset(new ProfileSummaryInfo(std::move(Summary)));
-}
-
-void CSProfileGenerator::mergeAndTrimColdProfile(
-    StringMap<FunctionSamples> &ProfileMap) {
-  if (!CSProfMergeColdContext && !CSProfTrimColdContext)
-    return;
 
-  // Use threshold calculated from profile summary unless specified
-  uint64_t ColdThreshold = PSI->getColdCountThreshold();
-  if (CSProfColdThreshold.getNumOccurrences()) {
-    ColdThreshold = CSProfColdThreshold;
-  }
-
-  // Nothing to merge if sample threshold is zero
-  if (ColdThreshold == 0)
-    return;
-
-  // Filter the cold profiles from ProfileMap and move them into a tmp
-  // container
-  std::vector<std::pair<StringRef, const FunctionSamples *>> ColdProfiles;
-  for (const auto &I : ProfileMap) {
-    const FunctionSamples &FunctionProfile = I.second;
-    if (FunctionProfile.getTotalSamples() >= ColdThreshold)
-      continue;
-    ColdProfiles.emplace_back(I.getKey(), &I.second);
-  }
-
-  // Remove the code profile from ProfileMap and merge them into BaseProileMap
-  StringMap<FunctionSamples> BaseProfileMap;
-  for (const auto &I : ColdProfiles) {
-    if (CSProfMergeColdContext) {
-      auto Ret = BaseProfileMap.try_emplace(
-          I.second->getContext().getNameWithoutContext(), FunctionSamples());
-      FunctionSamples &BaseProfile = Ret.first->second;
-      BaseProfile.merge(*I.second);
-    }
-    ProfileMap.erase(I.first);
-  }
-
-  // Merge the base profiles into ProfileMap;
-  for (const auto &I : BaseProfileMap) {
-    // Filter the cold base profile
-    if (CSProfTrimColdContext &&
-        I.second.getTotalSamples() < CSProfColdThreshold &&
-        ProfileMap.find(I.getKey()) == ProfileMap.end())
-      continue;
-    // Merge the profile if the original profile exists, otherwise just insert
-    // as a new profile
-    FunctionSamples &OrigProfile = getFunctionProfileForContext(I.getKey());
-    OrigProfile.merge(I.second);
+  // Use threshold calculated from profile summary unless specified.
+  if (!CSProfColdThreshold.getNumOccurrences()) {
+    CSProfColdThreshold = PSI->getColdCountThreshold();
   }
 }
 
 void CSProfileGenerator::write(std::unique_ptr<SampleProfileWriter> Writer,
                                StringMap<FunctionSamples> &ProfileMap) {
-  // Add bracket for context key to support different profile binary format
-  StringMap<FunctionSamples> CxtWithBracketPMap;
-  for (const auto &Item : ProfileMap) {
-    // After CSPreInliner the key of ProfileMap is no longer accurate for
-    // context, use the context attached to function samples instead.
-    std::string ContextWithBracket =
-        "[" + Item.second.getNameWithContext().str() + "]";
-    auto Ret = CxtWithBracketPMap.try_emplace(ContextWithBracket, Item.second);
-    assert(Ret.second && "Must be a unique context");
-    SampleContext FContext(Ret.first->first(), RawContext);
-    FunctionSamples &FProfile = Ret.first->second;
-    FContext.setAllAttributes(FProfile.getContext().getAllAttributes());
-    FProfile.setName(FContext.getNameWithoutContext());
-    FProfile.setContext(FContext);
-  }
-  Writer->write(CxtWithBracketPMap);
+  if (std::error_code EC = Writer->write(ProfileMap))
+    exitWithError(std::move(EC));
 }
 
 // Helper function to extract context prefix string stack
 // Extract context stack for reusing, leaf context stack will
 // be added compressed while looking up function profile
 static void
 extractPrefixContextStack(SmallVectorImpl<std::string> &ContextStrStack,
                           const SmallVectorImpl<const PseudoProbe *> &Probes,
if (Probe->isEntry()) {
       FunctionProfile.addHeadSamples(Count);
       // Look up for the caller's function profile
       const auto *InlinerDesc = Binary->getInlinerDescForProbe(Probe);
       if (InlinerDesc != nullptr) {
         // Since the context id will be compressed, we have to use callee's
         // context id to infer caller's context id to ensure they share the
         // same context prefix.
         StringRef CalleeContextId =
-            FunctionProfile.getContext().getNameWithContext(true);
+            FunctionProfile.getContext().getNameWithContext();
         StringRef CallerContextId;
         FrameLocation &&CallerLeafFrameLoc =
             getCallerContext(CalleeContextId, CallerContextId);
         uint64_t CallerIndex = CallerLeafFrameLoc.second.LineOffset;
         assert(CallerIndex &&
                "Inferred caller's location index shouldn't be zero!");
         FunctionSamples &CallerProfile =
             getFunctionProfileForContext(CallerContextId);