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

This patch fixed the following 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.

Differential Revision: https://reviews.llvm.org/D100090

1 files changed, 79 insertions, 0 deletions

diff --git a/llvm/lib/ProfileData/SampleProf.cpp b/llvm/lib/ProfileData/SampleProf.cpp
index 6647b26..565d67e 100644
--- a/llvm/lib/ProfileData/SampleProf.cpp
+++ b/llvm/lib/ProfileData/SampleProf.cpp
@@ -316,6 +316,85 @@ std::error_code ProfileSymbolList::read(const uint8_t *Data,
   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.