[CSSPGO] Sort function offset table to speed up profile loading.

[hoy]

Summary

With the context split work, the context-based (an array of strings) sorting performed at profile load time is way more expansive than single-string-based sorting. This is likely due to auxiliary operations done on each array element, such as indirect references, std::min operations, also likely cache misses. In this change I'm presorting profiles during profile generation time to avoid sorting at compile time.

Compared to the previous context-split work, this effectively cuts down compile time by 20% for one of our large services and brings us closer to non-CS build, with still a small gap in build time.

[hoy]

Updating D109036: [CSSPGO] Sort function offset table to speed up profile loading.

[wenlei]

This effectively cut down compile time by 20% for one of our large services.

Does this gets us back to where we were? or is there still small regression?

llvm/lib/ProfileData/SampleProfReader.cpp

731–732

FName -> FContext

742

assert on the order comparing to OrderedFuncOffsets->back()?

llvm/lib/ProfileData/SampleProfWriter.cpp

44

Why do we need a flag? Can we simplify this by always order the function offset table for CSSPGO profile?

[hoy]

In D109036#2975978, @wenlei wrote:

This effectively cut down compile time by 20% for one of our large services.

Does this gets us back to where we were? or is there still small regression?

We still have a small regression compared to non-CS build, but we are getting closer.

[hoy]

llvm/lib/ProfileData/SampleProfReader.cpp

731–732

Fixed.

742

The order may not be strictly the lexical order of SampleContext for md5 profiles. Md5 profiles may be converted from a string-based profile and the order of the func offset table in the md5 profile will be based on the original string order. However, the caller-callee context order, i.e, the trie tree path order, is still fulfilled. That order is what we want to load profile based on.

llvm/lib/ProfileData/SampleProfWriter.cpp

44

This is mostly for debugging and verification of the unordered path that is kept for being compatible with legacy unordered profiles.

[hoy]

Addressing Wenlei's comment.

[wenlei]

llvm/lib/ProfileData/SampleProfWriter.cpp

44

I don't think we have a backward-compatibility issue, there's no need to keep the legacy unordered profile. Can we remove it? Testing just need to make sure this ordering change does not affecting importing.

[hoy]

llvm/lib/ProfileData/SampleProfWriter.cpp

44

Makes sense. Removed the switch.

[hoy]

Updating D109036: [CSSPGO] Sort function offset table to speed up profile loading.

[wenlei]

llvm/lib/ProfileData/SampleProfReader.cpp

815

When suggesting the removal of ordered switch, I was thinking about always have ordered offset table for extbin CSSPGO profile. The simplification in implementation is the main reason I think removing that switch is better.

The idea is that we don't need unnecessary flexibility, so we can make assumptions about CSSPGO profile and avoid the complexity of handling the combinations of different input. With that, we either assert SecFlagOrdered is always on for extbin CSSPGO profile, or remove SecFlagOrdered completely and make that an implicit assumption for extbin CSSPGO profile.

[wenlei]

llvm/lib/ProfileData/SampleProfWriter.cpp

176

In the comment here, it would be good to point out the need for sorting due to how importing works in prelink. Basically part of the comment from readFuncProfiles.

[hoy]

llvm/lib/ProfileData/SampleProfReader.cpp

815

Sounds good. This makes it cleaner.

llvm/lib/ProfileData/SampleProfWriter.cpp

176

Comment added.

[hoy]

Addressing Wenlei's comment.

[wenlei]

lgtm, thanks.

[hoy]

Rebasing

[wmi]

LGTM.

llvm/lib/ProfileData/SampleProfReader.cpp

790

can it be a const?

791

const?

[hoy]

llvm/lib/ProfileData/SampleProfReader.cpp

790

Yes, const should work.

[hoy]

Addressing Wei's comment.

llvm/include/llvm/ProfileData/SampleProf.h

 };
 
 enum class SecFuncMetadataFlags : uint32_t {
   SecFlagInvalid = 0,
   SecFlagIsProbeBased = (1 << 0),
   SecFlagHasAttribute = (1 << 1)
 };
 
+enum class SecFuncOffsetFlags : uint32_t {
+  SecFlagInvalid = 0,
+  // Store function offsets in an order of contexts. The order ensures that
+  // callee contexts of a given context laid out next to it.
+  SecFlagOrdered = (1 << 0),
+};
+
 // Verify section specific flag is used for the correct section.
 template <class SecFlagType>
 static inline void verifySecFlag(SecType Type, SecFlagType Flag) {
   // No verification is needed for common flags.
   if (std::is_same<SecCommonFlags, SecFlagType>())
     return;
 
   // Verification starts here for section specific flag.
   bool IsFlagLegal = false;
   switch (Type) {
   case SecNameTable:
     IsFlagLegal = std::is_same<SecNameTableFlags, SecFlagType>();
     break;
   case SecProfSummary:
     IsFlagLegal = std::is_same<SecProfSummaryFlags, SecFlagType>();
     break;
   case SecFuncMetadata:
     IsFlagLegal = std::is_same<SecFuncMetadataFlags, SecFlagType>();
     break;
   default:
+  case SecFuncOffsetTable:
+    IsFlagLegal = std::is_same<SecFuncOffsetFlags, SecFlagType>();
     break;
   }
   if (!IsFlagLegal)
     llvm_unreachable("Misuse of a flag in an incompatible section");
 }
 
 template <class SecFlagType>
 static inline void addSecFlag(SecHdrTableEntry &Entry, SecFlagType Flag) {

llvm/include/llvm/ProfileData/SampleProfReader.h

protected:
   virtual ErrorOr<SampleContext> readSampleContextFromTable() override;
   ErrorOr<SampleContextFrames> readContextFromTable();
 
   std::unique_ptr<ProfileSymbolList> ProfSymList;
 
   /// The table mapping from function context to the offset of its
   /// FunctionSample towards file start.
   DenseMap<SampleContext, uint64_t> FuncOffsetTable;
+
+  /// Function offset mapping ordered by contexts.
+  std::unique_ptr<std::vector<std::pair<SampleContext, uint64_t>>>
+      OrderedFuncOffsets;
+
   /// 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;
protected:
   /// CSNameTable is used to save full context vectors. This serves as an
   /// underlying immutable buffer for all clients.
   std::unique_ptr<const std::vector<SampleContextFrameVector>> CSNameTable;
 
   /// If SkipFlatProf is true, skip the sections with
   /// SecFlagFlat flag.
   bool SkipFlatProf = false;
 
+  bool FuncOffsetsOrdered = false;
+
 public:
   SampleProfileReaderExtBinaryBase(std::unique_ptr<MemoryBuffer> B,
                                    LLVMContext &C, SampleProfileFormat Format)
       : SampleProfileReaderBinary(std::move(B), C, Format) {}
 
   /// Read sample profiles in extensible format from the associated file.
   std::error_code readImpl() override;
 

llvm/lib/ProfileData/SampleProfReader.cpp

if (std::error_code EC = readCSNameTableSec())
       return EC;
     break;
   }
   case SecLBRProfile:
     if (std::error_code EC = readFuncProfiles())
       return EC;
     break;
   case SecFuncOffsetTable:
+    FuncOffsetsOrdered = hasSecFlag(Entry, SecFuncOffsetFlags::SecFlagOrdered);
     if (std::error_code EC = readFuncOffsetTable())
       return EC;
     break;
   case SecFuncMetadata: {
     ProfileIsProbeBased =
         hasSecFlag(Entry, SecFuncMetadataFlags::SecFlagIsProbeBased);
     FunctionSamples::ProfileIsProbeBased = ProfileIsProbeBased;
     bool HasAttribute =
std::error_code SampleProfileReaderExtBinaryBase::readFuncOffsetTable() {
   // is read.
   FuncOffsetTable.clear();
 
   auto Size = readNumber<uint64_t>();
   if (std::error_code EC = Size.getError())
     return EC;
 
   FuncOffsetTable.reserve(*Size);
+
+  if (FuncOffsetsOrdered) {
+    OrderedFuncOffsets =
+        std::make_unique<std::vector<std::pair<SampleContext, uint64_t>>>();
+    OrderedFuncOffsets->reserve(*Size);
+  }
+
   for (uint32_t I = 0; I < *Size; ++I) {
-    auto FName(readSampleContextFromTable());
-    if (std::error_code EC = FName.getError())
+    auto FContext(readSampleContextFromTable());
+    if (std::error_code EC = FContext.getError())
       return EC;
 
     auto Offset = readNumber<uint64_t>();
     if (std::error_code EC = Offset.getError())
       return EC;
 
-    FuncOffsetTable[*FName] = *Offset;
+    FuncOffsetTable[*FContext] = *Offset;
+    if (FuncOffsetsOrdered)
+      OrderedFuncOffsets->emplace_back(*FContext, *Offset);
   }
+
   return sampleprof_error::success;
 }
 
 std::error_code SampleProfileReaderExtBinaryBase::readFuncProfiles() {
   // Collect functions used by current module if the Reader has been
   // given a module.
   // collectFuncsFromModule uses FunctionSamples::getCanonicalFnName
   // which will query FunctionSamples::HasUniqSuffix, so it has to be
if (!LoadFuncsToBeUsed) {
     // Load function profiles on demand.
     if (Remapper) {
       for (auto Name : FuncsToUse) {
         Remapper->insert(Name);
       }
     }
 
     if (ProfileIsCS) {
-      // Compute the ordered set of names, so we can
-      // get all context profiles under a subtree by
-      // iterating through the ordered names.
-      std::set<SampleContext> OrderedContexts;
-      for (auto Name : FuncOffsetTable) {
-        OrderedContexts.insert(Name.first);
-      }
-
       DenseSet<uint64_t> FuncGuidsToUse;
       if (useMD5()) {
         for (auto Name : FuncsToUse)
           FuncGuidsToUse.insert(Function::getGUID(Name));
       }
 
-      // For each function in current module, load all
-      // context profiles for the function.
-      for (auto NameOffset : FuncOffsetTable) {
-        SampleContext FContext = NameOffset.first;
-        auto FuncName = FContext.getName();
-        if ((useMD5() && !FuncGuidsToUse.count(std::stoull(FuncName.data()))) ||
-            (!useMD5() && !FuncsToUse.count(FuncName) &&
-             (!Remapper || !Remapper->exist(FuncName))))
-          continue;
-
-        // For each context profile we need, try to load
-        // all context profile in the subtree. This can
-        // help profile guided importing for ThinLTO.
-        auto It = OrderedContexts.find(FContext);
-        while (It != OrderedContexts.end() && FContext.IsPrefixOf(*It)) {
-          const uint8_t *FuncProfileAddr = Start + FuncOffsetTable[*It];
+      // For each function in current module, load all context profiles for
+      // the function as well as their callee contexts which can help profile
+      // guided importing for ThinLTO. This can be achieved by walking
+      // through an ordered context container, where contexts are laid out
+      // as if they were walked in preorder of a context trie. While
+      // traversing the trie, a link to the highest common ancestor node is
+      // kept so that all of its decendants will be loaded.
+      assert(OrderedFuncOffsets.get() &&
+             "func offset table should always be sorted in CS profile");
+      const SampleContext *CommonContext = nullptr;
+      for (const auto &NameOffset : *OrderedFuncOffsets) {
+        const auto &FContext = NameOffset.first;
+        auto FName = FContext.getName();
+        // For function in the current module, keep its farthest ancestor
+        // context. This can be used to load itself and its child and
+        // sibling contexts.
+        if ((useMD5() && FuncGuidsToUse.count(std::stoull(FName.data()))) ||
+            (!useMD5() && (FuncsToUse.count(FName) ||
+                           (Remapper && Remapper->exist(FName))))) {
+          if (!CommonContext || !CommonContext->IsPrefixOf(FContext))
+            CommonContext = &FContext;
+        }
+
+        if (CommonContext == &FContext ||
+            (CommonContext && CommonContext->IsPrefixOf(FContext))) {
+          // Load profile for the current context which originated from
+          // the common ancestor.
+          const uint8_t *FuncProfileAddr = Start + NameOffset.second;
           assert(FuncProfileAddr < End && "out of LBRProfile section");
           if (std::error_code EC = readFuncProfile(FuncProfileAddr))
             return EC;
-          // Remove loaded context profile so we won't
-          // load it repeatedly.
-          It = OrderedContexts.erase(It);
         }
       }
     } else {
       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;
static std::string getSecFlagsStr(const SecHdrTableEntry &Entry) {
   case SecProfSummary:
     if (hasSecFlag(Entry, SecProfSummaryFlags::SecFlagPartial))
       Flags.append("partial,");
     if (hasSecFlag(Entry, SecProfSummaryFlags::SecFlagFullContext))
       Flags.append("context,");
     if (hasSecFlag(Entry, SecProfSummaryFlags::SecFlagFSDiscriminator))
       Flags.append("fs-discriminator,");
     break;
+  case SecFuncOffsetTable:
+    if (hasSecFlag(Entry, SecFuncOffsetFlags::SecFlagOrdered))
+      Flags.append("ordered,");
+    break;
   default:
     break;
   }
   char &last = Flags.back();
   if (last == ',')
     last = '}';
   else
     Flags.append("}");

llvm/lib/ProfileData/SampleProfWriter.cpp

 #include <set>
 #include <system_error>
 #include <utility>
 #include <vector>
 
 using namespace llvm;
 using namespace sampleprof;
 
 std::error_code
 SampleProfileWriter::writeFuncProfiles(const SampleProfileMap &ProfileMap) {
   std::vector<NameFunctionSamples> V;
   sortFuncProfiles(ProfileMap, V);
   for (const auto &I : V) {
     if (std::error_code EC = writeSample(*I.second))
       return EC;
   }
   return sampleprof_error::success;

 
 std::error_code SampleProfileWriterExtBinaryBase::writeFuncOffsetTable() {
   auto &OS = *OutputStream;
 
   // Write out the table size.
   encodeULEB128(FuncOffsetTable.size(), OS);
 
   // Write out FuncOffsetTable.
-  for (auto Entry : FuncOffsetTable) {
-    if (std::error_code EC = writeContextIdx(Entry.first))
+  auto WriteItem = [&](const SampleContext &Context, uint64_t Offset) {
+    if (std::error_code EC = writeContextIdx(Context))
+      return EC;
+    encodeULEB128(Offset, OS);
+    return (std::error_code)sampleprof_error::success;
+  };
+
+  if (FunctionSamples::ProfileIsCS) {
+    // Sort the contexts before writing them out. This is to help fast load all
+    // context profiles for a function as well as their callee contexts which
+    // can help profile-guided importing for ThinLTO.
+    std::map<SampleContext, uint64_t> OrderedFuncOffsetTable(
+        FuncOffsetTable.begin(), FuncOffsetTable.end());
+    for (const auto &Entry : OrderedFuncOffsetTable) {
+      if (std::error_code EC = WriteItem(Entry.first, Entry.second))
         return EC;
-    encodeULEB128(Entry.second, OS);
     }
+    addSectionFlag(SecFuncOffsetTable, SecFuncOffsetFlags::SecFlagOrdered);
+  } else {
+    for (const auto &Entry : FuncOffsetTable) {
+      if (std::error_code EC = WriteItem(Entry.first, Entry.second))
+        return EC;
+    }
+  }
+
   FuncOffsetTable.clear();
   return sampleprof_error::success;
 }
 
 std::error_code SampleProfileWriterExtBinaryBase::writeFuncMetadata(
     const SampleProfileMap &Profiles) {
   if (!FunctionSamples::ProfileIsProbeBased && !FunctionSamples::ProfileIsCS)
     return sampleprof_error::success;

llvm/test/Transforms/SampleProfile/csspgo-import-list.ll

 ; Make sure Import GUID list for ThinLTO properly set for CSSPGO
 ; RUN: opt < %s -passes='thinlto-pre-link<O2>' -pgo-kind=pgo-sample-use-pipeline -sample-profile-file=%S/Inputs/csspgo-import-list.prof -S | FileCheck %s
 ; RUN: llvm-profdata merge --sample --extbinary %S/Inputs/csspgo-import-list.prof -o %t.prof
 ; RUN: opt < %s -passes='thinlto-pre-link<O2>' -pgo-kind=pgo-sample-use-pipeline -sample-profile-file=%t.prof -S | FileCheck %s
+; RUN: llvm-profdata show --sample -show-sec-info-only %t.prof | FileCheck %s --check-prefix=CHECK-ORDERED
 ; RUN: llvm-profdata merge --sample --extbinary --use-md5 %S/Inputs/csspgo-import-list.prof -o %t.md5
 ; RUN: opt < %s -passes='thinlto-pre-link<O2>' -pgo-kind=pgo-sample-use-pipeline -sample-profile-file=%t.md5 -S | FileCheck %s
+; RUN: llvm-profdata show --sample -show-sec-info-only %t.md5 | FileCheck %s --check-prefix=CHECK-ORDERED
+
 
 declare i32 @_Z5funcBi(i32 %x)
 declare i32 @_Z5funcAi(i32 %x)
 
 define dso_local i32 @main() local_unnamed_addr #0 !dbg !18 {
 entry:
   br label %for.body, !dbg !25
 
for.body:                                         ; preds = %for.body, %entry
   %cmp = icmp eq i32 %x.011, 0, !dbg !38
   br i1 %cmp, label %for.cond.cleanup, label %for.body, !dbg !25
 }
 
 ; Make sure the ImportGUID stays with entry count metadata for ThinLTO-PreLink
 ; CHECK: distinct !DISubprogram(name: "main"
 ; CHECK: !{!"function_entry_count", i64 3, i64 446061515086924981, i64 3815895320998406042, i64 7102633082150537521, i64 -2862076748587597320}
 
+; CHECK-ORDERED: FuncOffsetTableSection {{.*}} {ordered}
 
 attributes #0 = { nofree noinline norecurse nounwind uwtable "correctly-rounded-divide-sqrt-fp-math"="false" "disable-tail-calls"="false" "frame-pointer"="none" "less-precise-fpmad"="false" "min-legal-vector-width"="0" "no-infs-fp-math"="false" "no-jump-tables"="false" "no-nans-fp-math"="false" "no-signed-zeros-fp-math"="false" "no-trapping-math"="false" "stack-protector-buffer-size"="8" "target-cpu"="x86-64" "target-features"="+cx8,+fxsr,+mmx,+sse,+sse2,+x87" "unsafe-fp-math"="false" "use-soft-float"="false" "use-sample-profile" }
 
 !llvm.dbg.cu = !{!2}
 !llvm.module.flags = !{!14, !15, !16}
 !llvm.ident = !{!17}
 
 !0 = !DIGlobalVariableExpression(var: !1, expr: !DIExpression())