[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, 48 insertions, 20 deletions

diff --git a/llvm/lib/ProfileData/SampleProfWriter.cpp b/llvm/lib/ProfileData/SampleProfWriter.cpp
index d32fcbf..43cfb93 100644
--- a/llvm/lib/ProfileData/SampleProfWriter.cpp
+++ b/llvm/lib/ProfileData/SampleProfWriter.cpp
@@ -46,9 +46,11 @@ std::error_code SampleProfileWriter::writeFuncProfiles(
   // 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())
@@ -147,7 +149,7 @@ std::error_code SampleProfileWriterExtBinaryBase::write(
 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);
@@ -160,9 +162,11 @@ std::error_code SampleProfileWriterExtBinaryBase::writeFuncOffsetTable() {
   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;
@@ -174,7 +178,9 @@ std::error_code SampleProfileWriterExtBinaryBase::writeFuncMetadata(
     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)
@@ -204,7 +210,9 @@ std::error_code SampleProfileWriterExtBinaryBase::writeNameTable() {
 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);
   }
 
@@ -378,7 +386,11 @@ std::error_code SampleProfileWriterCompactBinary::write(
 /// 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";
@@ -431,15 +443,26 @@ std::error_code SampleProfileWriterText::writeSample(const FunctionSamples &S) {
   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));
 }
 
@@ -500,9 +523,11 @@ std::error_code SampleProfileWriterCompactBinary::writeFuncOffsetTable() {
   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;
 }
@@ -539,7 +564,9 @@ std::error_code SampleProfileWriterBinary::writeHeader(
 
   // 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);
   }
 
@@ -654,7 +681,8 @@ std::error_code SampleProfileWriterBinary::writeSummary() {
 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);