[CSSPGO] Support of CS profiles in extended binary format.

[hoy]

Summary

This change brings up support of context-sensitive profiles in the format of extended binary. Existing sample profile reader/writer/merger code is being tweaked to reflect the fact of bracketed input contexts, like ([...]). The paired brackets are also needed in extbinary profiles because we don't yet have an otherwise good way to tell calling contexts apart from regular function names since the context delimiter @ can somehow serve as a part of the C++ mangled names.

[wenlei]

llvm/include/llvm/ProfileData/SampleProf.h

445

What about using getNameWithContext(bool WithBracket = false), also change InputContext to FullContextWithBracket.

760

same here, merge with getNameWithContext(bool WithBracket=false)?

llvm/lib/ProfileData/SampleProfReader.cpp

323

Use only one counter, and assert(CSProfileCount == 0 || CSProfileCount == Profiles.size()) ?

llvm/lib/ProfileData/SampleProfWriter.cpp

150

If we keep StringRef for context with bracket, we can use it for text profile writer as well.

As we discussed off patch, the context provided by llvm-profgen may not have [], and it's probably better to fix that to make it consistent.

llvm/test/tools/llvm-profdata/Inputs/csspgo-profile.proftext

1 ↗

(On Diff #319623)

Add a function without calling context? We can reuse the test case from our fork?

llvm/test/tools/llvm-profdata/merge-cs-profile.test

1 ↗

(On Diff #319623)

What about adding a roundtrip conversion test for text and extbin format?

[hoy]

llvm/include/llvm/ProfileData/SampleProf.h

445

That looks better. Thanks.

llvm/lib/ProfileData/SampleProfReader.cpp

323

Yeah, that's smart.

llvm/lib/ProfileData/SampleProfWriter.cpp

150

Exactly. The text writer for now has to emit the brackets explicitly like below. @wlei this can be fixed by including brackets in the final context string during profile generation.

std::error_code SampleProfileWriterText::writeSample(const FunctionSamples &S) {
  auto &OS = *OutputStream;
  if (FunctionSamples::ProfileIsCS)
    OS << "[" << S.getNameWithContext() << "]:" << S.getTotalSamples();
  else
    OS << S.getName() << ":" << S.getTotalSamples();

llvm/test/tools/llvm-profdata/Inputs/csspgo-profile.proftext

1 ↗

(On Diff #319623)

Sounds good, will port our internal test.

llvm/test/tools/llvm-profdata/merge-cs-profile.test

1 ↗

(On Diff #319623)

It's actually done by line 4.

[hoy]

Addressing Wenlei's feedbacks.

[wenlei]

lgtm, with one nit.

llvm/include/llvm/ProfileData/SampleProf.h

445

missed the rename of InputContext to FullContextWithBracket?

[hoy]

Addressing nit.

[wlei]

llvm/lib/ProfileData/SampleProfWriter.cpp

150

Got it! So I will fix like the code below:

if (FunctionSamples::ProfileIsCS)
  OS << S.getNameWithContext()  << S.getTotalSamples();
else
  OS << S.getName() << ":" << S.getTotalSamples();

[wlei]

llvm/include/llvm/ProfileData/SampleProf.h

445

I didn't see any where in this patch add the bracket. Do we need to add it in this patch? because the input ContextStr actually doesn't have the bracket, the bracket in the regression test is printed by the text format which added the bracket in the end, but not for extended binary.

[hoy]

llvm/include/llvm/ProfileData/SampleProf.h

445

This patch doesn't add brackets. It relies on bracketed input, like a real text profile, or in-memory profiles fed by llvm-profgen. Currently when llvm-profgen creates a samplecontext , it doesn't add brackets. Instead, it sets FunctionSamples::ProfileIsCS explicitly. However, the context should have brackets as well so that the the contextness can be inferred automatically without setting ProfileIsCS.

llvm/lib/ProfileData/SampleProfWriter.cpp

150

Thanks. Ideally it should be just OS << S.getNameWithContext(true) << S.getTotalSamples(); unconditionally once brackets are added to contexts on the llvm-profgen side.

[wlei]

llvm/include/llvm/ProfileData/SampleProf.h

445

Thanks for your explanation, that makes a lot of sense.

So we need to change the llvm-profgen side, to add the brackets for the input ContextStr, then we don't need to relay on the text format, also that's the right way for extended format.

454

Should we merge it to the condition in line 462. because as its name "WithBracket", it should always have bracket or we can add a check for the HasContext.

if (HasContext) {
  FullContextWithBracket = ContextStr;
}

[hoy]

llvm/include/llvm/ProfileData/SampleProf.h

454

Good point. It really refers to the original input context so that we can use it unconditionally for non-CS profile as well. @wenlei thinking about renaming it OrigContext, what do you think?

[wenlei]

llvm/include/llvm/ProfileData/SampleProf.h

454

Good catch, in this case OrigContext or InputContext makes more sense.. Sorry for the confusion.

[wmi]

llvm/include/llvm/ProfileData/SampleProf.h

442

Nit: I feel getNameWithoutContext is more straightforward.

685–687

Is it possible for the two FunctionSamples candidates to have different non empty contexts?

llvm/lib/ProfileData/SampleProfReader.cpp

547

This is unnecessary. The next statement will create the default entry if the key doesn't exist in the map.

[hoy]

llvm/include/llvm/ProfileData/SampleProf.h

442

you mean for the getName function?

454

Sound good. Renamed to InputContext.

685–687

This merge function is called under the assumption that Other has the same context with this or this is empty, which basically means we are initializing this with Other.

llvm/lib/ProfileData/SampleProfReader.cpp

547

Good catch. Will remove.

[wmi]

llvm/include/llvm/ProfileData/SampleProf.h

442

Yes.

[hoy]

llvm/include/llvm/ProfileData/SampleProf.h

442

Sounds good.

[hoy]

Addressing Wei's feedbacks.

[wmi]

LGTM.

llvm/include/llvm/ProfileData/SampleProf.h

public:
   }
 
   operator StringRef() const { return FullContext; }
   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 getName() const { return Name; }
   StringRef getCallingContext() const { return CallingContext; }
-  StringRef getNameWithContext() const { return FullContext; }
+  StringRef getNameWithContext(bool WithBracket = false) const {
+    return WithBracket ? InputContext : 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())
+      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
       // units (without using -unique-internal-linkage-names), or
       // 2. They are really the same function but from different compilations.
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() const {
-    return FunctionSamples::ProfileIsCS ? Context.getNameWithContext() : Name;
+  StringRef getNameWithContext(bool WithBracket = false) const {
+    return FunctionSamples::ProfileIsCS
+               ? Context.getNameWithContext(WithBracket)
+               : 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; }
 
   /// Return the canonical name for a function, taking into account
   /// suffix elision policy attributes.
   static StringRef getCanonicalFnName(const Function &F) {

llvm/include/llvm/ProfileData/SampleProfReader.h

protected:
   /// Compute summary for this profile.
   void computeSummary();
 
   std::unique_ptr<SampleProfileReaderItaniumRemapper> Remapper;
 
   /// \brief Whether samples are collected based on pseudo probes.
   bool ProfileIsProbeBased = false;
 
+  /// Whether function profiles are context-sensitive.
   bool ProfileIsCS = false;
 
+  /// Number of context-sensitive profiles.
+  int CSProfileCount = 0;
+
   /// \brief The format of sample.
   SampleProfileFormat Format = SPF_None;
 };
 
 class SampleProfileReaderText : public SampleProfileReader {
 public:
   SampleProfileReaderText(std::unique_ptr<MemoryBuffer> B, LLVMContext &C)
       : SampleProfileReader(std::move(B), C, SPF_Text) {}

llvm/lib/ProfileData/SampleProfReader.cpp

 /// 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;
-  int CSProfileCount = 0;
-  int RegularProfileCount = 0;
   uint32_t ProbeProfileCount = 0;
 
   // SeenMetadata tracks whether we have processed metadata for the current
   // top-level function profile.
   bool SeenMetadata = false;
 
   for (; !LineIt.is_at_eof(); ++LineIt) {
     if ((*LineIt)[(*LineIt).find_first_not_of(' ')] == '#')
if ((*LineIt)[0] != ' ') {
         reportError(LineIt.line_number(),
                     "Expected 'mangled_name:NUM:NUM', found " + *LineIt);
         return sampleprof_error::malformed;
       }
       SeenMetadata = false;
       SampleContext FContext(FName);
       if (FContext.hasContext())
         ++CSProfileCount;
-      else
-        ++RegularProfileCount;
       Profiles[FContext] = FunctionSamples();
       FunctionSamples &FProfile = Profiles[FContext];
       FProfile.setName(FContext.getName());
       FProfile.setContext(FContext);
       MergeResult(Result, FProfile.addTotalSamples(NumSamples));
       MergeResult(Result, FProfile.addHeadSamples(NumHeadSamples));
       InlineStack.clear();
       InlineStack.push_back(&FProfile);
if ((*LineIt)[0] != ' ') {
         ++ProbeProfileCount;
         SeenMetadata = true;
         break;
       }
       }
     }
   }
 
-  assert((RegularProfileCount == 0 || CSProfileCount == 0) &&
+  assert((CSProfileCount == 0 || CSProfileCount == Profiles.size()) &&
          "Cannot have both context-sensitive and regular profile");
   ProfileIsCS = (CSProfileCount > 0);
   assert((ProbeProfileCount == 0 || ProbeProfileCount == Profiles.size()) &&
          "Cannot have both probe-based profiles and regular profiles");
   ProfileIsProbeBased = (ProbeProfileCount > 0);
   FunctionSamples::ProfileIsProbeBased = ProfileIsProbeBased;
+  FunctionSamples::ProfileIsCS = ProfileIsCS;
 
   if (Result == sampleprof_error::success)
     computeSummary();
 
   return Result;
 }
 
 bool SampleProfileReaderText::hasFormat(const MemoryBuffer &Buffer) {
SampleProfileReaderBinary::readFuncProfile(const uint8_t *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())
     return EC;
 
-  Profiles[*FName] = FunctionSamples();
-  FunctionSamples &FProfile = Profiles[*FName];
-  FProfile.setName(*FName);
-
+  SampleContext FContext(*FName);
+  Profiles[FContext] = FunctionSamples();
+  FunctionSamples &FProfile = Profiles[FContext];
+  FProfile.setName(FContext.getName());
+  FProfile.setContext(FContext);
   FProfile.addHeadSamples(*NumHeadSamples);
 
+  if (FContext.hasContext())
+    CSProfileCount++;
+
   if (std::error_code EC = readProfile(FProfile))
     return EC;
   return sampleprof_error::success;
 }
 
 std::error_code SampleProfileReaderBinary::readImpl() {
   while (!at_eof()) {
     if (std::error_code EC = readFuncProfile(Data))

 std::error_code SampleProfileReaderExtBinaryBase::readFuncProfiles() {
   const uint8_t *Start = Data;
   if (UseAllFuncs) {
     while (Data < End) {
       if (std::error_code EC = readFuncProfile(Data))
         return EC;
     }
     assert(Data == End && "More data is read than expected");
-    return sampleprof_error::success;
-  }
-
+  } else {
     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 {
       for (auto NameOffset : FuncOffsetTable) {
-      auto FuncName = NameOffset.first;
+        SampleContext FContext(NameOffset.first);
+        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;
       }
     }
-
     Data = End;
+  }
+
+  assert((CSProfileCount == 0 || CSProfileCount == Profiles.size()) &&
+         "Cannot have both context-sensitive and regular profile");
+  ProfileIsCS = (CSProfileCount > 0);
+  FunctionSamples::ProfileIsCS = ProfileIsCS;
   return sampleprof_error::success;
 }
 
 std::error_code SampleProfileReaderExtBinaryBase::readProfileSymbolList() {
   if (!ProfSymList)
     ProfSymList = std::make_unique<ProfileSymbolList>();
 
   if (std::error_code EC = ProfSymList->read(Data, End - Data))
for (unsigned I = 0; I < Profiles.size(); ++I) {
     auto FName(readStringFromTable());
     if (std::error_code EC = FName.getError())
       return EC;
 
     auto Checksum = readNumber<uint64_t>();
     if (std::error_code EC = Checksum.getError())
       return EC;
 
-    Profiles[*FName].setFunctionHash(*Checksum);
+    SampleContext FContext(*FName);
+    Profiles[FContext].setFunctionHash(*Checksum);
   }
   return sampleprof_error::success;
 }
 
 std::error_code SampleProfileReaderCompactBinary::readNameTable() {
   auto Size = readNumber<uint64_t>();
   if (std::error_code EC = Size.getError())
     return EC;

llvm/lib/ProfileData/SampleProfWriter.cpp

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.getName();
+  StringRef Name = S.getNameWithContext(true);
   FuncOffsetTable[Name] = Offset - SecLBRProfileStart;
   encodeULEB128(S.getHeadSamples(), *OutputStream);
   return writeBody(S);
 }
 
 std::error_code SampleProfileWriterExtBinaryBase::writeFuncOffsetTable() {
   auto &OS = *OutputStream;
 
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.getName()))
+  if (std::error_code EC = writeNameIdx(S.getNameWithContext(true)))
     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/Transforms/SampleProfile/profile-context-tracker.ll

 ; Test for CSSPGO's SampleContextTracker to make sure context profile tree is promoted and merged properly
 ; based on inline decision, so post inline counts are accurate.
 
+; RUN: llvm-profdata merge --sample --extbinary %S/Inputs/profile-context-tracker.prof -o %t
+
 ; Note that we need new pass manager to enable top-down processing for sample profile loader
 ; Testwe we inlined the following in top-down order and entry counts accurate reflects post-inline base profile
 ;   main:3 @ _Z5funcAi
 ;   main:3 @ _Z5funcAi:1 @ _Z8funcLeafi
 ;   _Z5funcBi:1 @ _Z8funcLeafi
 ; RUN: opt < %s -passes=sample-profile -sample-profile-file=%S/Inputs/profile-context-tracker.prof -sample-profile-inline-size -profile-sample-accurate -S | FileCheck %s --check-prefix=INLINE-ALL
+; RUN: opt < %s -passes=sample-profile -sample-profile-file=%t -sample-profile-inline-size -profile-sample-accurate -S | FileCheck %s --check-prefix=INLINE-ALL
 
 ; Testwe we inlined the following in top-down order and entry counts accurate reflects post-inline base profile
 ;   main:3 @ _Z5funcAi
 ;   _Z5funcAi:1 @ _Z8funcLeafi
 ;   _Z5funcBi:1 @ _Z8funcLeafi
 ; RUN: opt < %s -passes=sample-profile -sample-profile-file=%S/Inputs/profile-context-tracker.prof -profile-sample-accurate -S | FileCheck %s --check-prefix=INLINE-HOT
+; RUN: opt < %s -passes=sample-profile -sample-profile-file=%t -profile-sample-accurate -S | FileCheck %s --check-prefix=INLINE-HOT
 
 
 @factor = dso_local global i32 3, align 4, !dbg !0
 
 define dso_local i32 @main() local_unnamed_addr #0 !dbg !18 {
 ; INLINE-ALL: @main{{.*}}!prof ![[MAIN_PROF:[0-9]+]]
 ; INLINE-HOT: @main{{.*}}!prof ![[MAIN_PROF:[0-9]+]]
 entry:

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
+[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
+[main]:154:0
+ 2: 12
+ 3: 18 _Z5funcAi:11
+ 3.1: 18 _Z5funcBi:19
+[external:12 @ main]:154:12
+ 2: 12
+ 3: 10 _Z5funcAi:7
+ 3.1: 10 _Z5funcBi:11
+[main:3.1 @ _Z5funcBi]:120:19
+ 0: 19
+ 1: 19 _Z8funcLeafi:20
+ 3: 12
+[externalA:17 @ _Z5funcBi]:120:3
+ 0: 3
+ 1: 3
+[external:10 @ _Z5funcBi]:120:10
+ 0: 10
+ 1: 10
+[main:3 @ _Z5funcAi]:99:11
+ 0: 10
+ 1: 10 _Z8funcLeafi:11
+ 3: 24

llvm/test/tools/llvm-profdata/cs-sample-profile.test

+RUN: llvm-profdata merge --sample --text -output=%t.proftext %S/Inputs/cs-sample.proftext
+RUN: diff -b %t.proftext %S/Inputs/cs-sample.proftext
+RUN: llvm-profdata merge --sample --extbinary %p/Inputs/cs-sample.proftext -o %t.prof && llvm-profdata merge --sample --text %t.prof -o %t1.proftext
+RUN: diff -b %t1.proftext %S/Inputs/cs-sample.proftext

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.getName();
+      StringRef FName = Samples.getNameWithContext(true);
       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 =