[CSSPGO] Aggregation by the last K context frames for cold profiles

[wlei]

Summary

This change provides the option to merge and aggregate cold context by the last k frames instead of context-less name. By default K = 1 means the context-less one.

This is for better perf tuning. The more selective merging and trimming will rely on llvm-profgen's preinliner.

[wenlei]

Thanks for the patch. I think this can be useful for tuning. However as we discussed, it'd be nice if the level of merging (number of frames to keep after merging) can be adaptive based on hotness. But to do that, the preinliner in llvm-profgen is the way to go.

llvm/include/llvm/ProfileData/SampleProf.h

441

We need to find @ instead of @ as mangle name can contain @. And how about using StringRef::rfind ?

llvm/lib/ProfileData/SampleProf.cpp

346–348

Now rename this to MergedProfileMap?

[wlei]

Addressing Wenlei's comments

[wlei]

change to use uint32_t

[wenlei]

llvm/include/llvm/ProfileData/SampleProf.h

1014

Looks like all call site have explicit argument, do we still need a default one?

llvm/tools/llvm-profgen/ProfileGenerator.cpp

47

nit: how about name it frame-depth-for-cold-context? similar name for llvm-profdata flag and all variables.

[wlei]

Addressing Wenlei's comments

[wenlei]

lgtm, thanks!

[hoy]

LGTM.

llvm/include/llvm/ProfileData/SampleProf.h

public:
   }
 
   // Split the top context frame (left-most substr) from context.
   static std::pair<StringRef, StringRef>
   splitContextString(StringRef ContextStr) {
     return ContextStr.split(" @ ");
   }
 
+  // Reconstruct a new context with the last k frames, return the context-less
+  // name if K = 1
+  StringRef getContextWithLastKFrames(uint32_t K) {
+    if (K == 1)
+      return getNameWithoutContext();
+
+    size_t I = FullContext.size();
+    while (K--) {
+      I = FullContext.find_last_of(" @ ", I);
+      if (I == StringRef::npos)
+        return FullContext;
+      I -= 2;
+    }
+    return FullContext.slice(I + 3, StringRef::npos);
+  }
+
   // Decode context string for a frame to get function name and location.
   // `ContextStr` is in the form of `FuncName:StartLine.Discriminator`.
   static void decodeContextString(StringRef ContextStr, StringRef &FName,
                                   LineLocation &LineLoc) {
     // Get function name
     auto EntrySplit = ContextStr.split(':');
     FName = EntrySplit.first;
 

 /// 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);
+                                       bool MergeColdContext = true,
+                                       uint32_t KeepLastKFrames = 1);
   // Canonicalize context profile name and attributes.
   void canonicalizeContextProfiles();
 
 private:
   StringMap<FunctionSamples> &ProfileMap;
 };
 
 /// ProfileSymbolList records the list of function symbols shown up

llvm/lib/ProfileData/SampleProf.cpp

while (Size < ListSize && StrNum < ProfileSymbolListCutOff) {
     StrNum++;
   }
   if (Size != ListSize && StrNum != ProfileSymbolListCutOff)
     return sampleprof_error::malformed;
   return sampleprof_error::success;
 }
 
 void SampleContextTrimmer::trimAndMergeColdContextProfiles(
-    uint64_t ColdCountThreshold, bool TrimColdContext, bool MergeColdContext) {
+    uint64_t ColdCountThreshold, bool TrimColdContext, bool MergeColdContext,
+    uint32_t KeepLastKFrames) {
   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;
+  // Remove the cold profile from ProfileMap and merge them into
+  // MergedProfileMap by the last K frames of context
+  StringMap<FunctionSamples> MergedProfileMap;
   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);
+      auto Ret = MergedProfileMap.try_emplace(
+          I.second->getContext().getContextWithLastKFrames(KeepLastKFrames),
+          FunctionSamples());
+      FunctionSamples &MergedProfile = Ret.first->second;
+      MergedProfile.merge(*I.second);
     }
     ProfileMap.erase(I.first);
   }
 
-  // Merge the base profiles into ProfileMap;
-  for (const auto &I : BaseProfileMap) {
-    // Filter the cold base profile
+  // Move the merged profiles into ProfileMap;
+  for (const auto &I : MergedProfileMap) {
+    // Filter the cold merged 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);

llvm/test/tools/llvm-profgen/merge-cold-profile.test

 ; Used the data from recursion-compression.test, refer it for the unmerged output
 ; RUN: llvm-profgen --format=text --perfscript=%S/Inputs/recursion-compression-pseudoprobe.perfscript --binary=%S/Inputs/recursion-compression-pseudoprobe.perfbin --output=%t1 --compress-recursion=-1 --profile-summary-cold-count=8
 ; RUN: FileCheck %s --input-file %t1
 
 ; Test --csprof-trim-cold-context=0
 ; RUN: llvm-profgen --format=text --perfscript=%S/Inputs/recursion-compression-pseudoprobe.perfscript --binary=%S/Inputs/recursion-compression-pseudoprobe.perfbin --output=%t2 --compress-recursion=-1 --profile-summary-cold-count=100 --csprof-trim-cold-context=0
 ; RUN: FileCheck %s --input-file %t2 --check-prefix=CHECK-KEEP-COLD
 
 ; Test --csprof-merge-cold-context=0
 ; RUN: llvm-profgen --format=text --perfscript=%S/Inputs/recursion-compression-pseudoprobe.perfscript --binary=%S/Inputs/recursion-compression-pseudoprobe.perfbin --output=%t3 --compress-recursion=-1 --profile-summary-cold-count=10 --csprof-merge-cold-context=0
 ; RUN: FileCheck %s --input-file %t3 --check-prefix=CHECK-UNMERGED
 
+; Test --csprof-keep-last-k-frames
+; RUN: llvm-profgen --format=text --perfscript=%S/Inputs/recursion-compression-pseudoprobe.perfscript --binary=%S/Inputs/recursion-compression-pseudoprobe.perfbin --output=%t2 --compress-recursion=-1 --profile-summary-cold-count=100 --csprof-trim-cold-context=0 --csprof-keep-last-k-frames=2
+; RUN: FileCheck %s --input-file %t2 --check-prefix=CHECK-KEEP-K-COLD
+
+
+
 ; CHECK:     [fa]:14:4
 ; CHECK-NEXT: 1: 4
 ; CHECK-NEXT: 2: 18446744073709551615
 ; CHECK-NEXT: 3: 4
 ; CHECK-NEXT: 4: 2
 ; CHECK-NEXT: 5: 1
 ; CHECK-NEXT: 7: 2 fb:2
 ; CHECK-NEXT: 8: 1 fa:1

 ; CHECK-UNMERGED-NEXT: 2: 3
 ; CHECK-UNMERGED-NEXT: 3: 1
 ; CHECK-UNMERGED-NEXT: 5: 4 fb:4
 ; CHECK-UNMERGED-NEXT: 6: 1 fa:1
 ; CHECK-UNMERGED-NEXT: !CFGChecksum: 72617220756
 ; CHECK-UNMERGED-NOT: [fa]
 ; CHECK-UNMERGED-NOT: [fb]
 
+; CHECK-KEEP-K-COLD:     [fb:5 @ fb]:13:4
+; CHECK-KEEP-K-COLD-NEXT: 1: 4
+; CHECK-KEEP-K-COLD-NEXT: 2: 3
+; CHECK-KEEP-K-COLD-NEXT: 3: 1
+; CHECK-KEEP-K-COLD-NEXT: 5: 4 fb:4
+; CHECK-KEEP-K-COLD-NEXT: 6: 1 fa:1
+; CHECK-KEEP-K-COLD-NEXT: !CFGChecksum: 72617220756
+; CHECK-KEEP-K-COLD-NEXT: !Attributes: 0
+; CHECK-KEEP-K-COLD-NEXT:[fb:6 @ fa]:10:3
+; CHECK-KEEP-K-COLD-NEXT: 1: 3
+; CHECK-KEEP-K-COLD-NEXT: 2: 18446744073709551615
+; CHECK-KEEP-K-COLD-NEXT: 3: 3
+; CHECK-KEEP-K-COLD-NEXT: 4: 1
+; CHECK-KEEP-K-COLD-NEXT: 5: 1
+; CHECK-KEEP-K-COLD-NEXT: 7: 1 fb:1
+; CHECK-KEEP-K-COLD-NEXT: 8: 1 fa:1
+; CHECK-KEEP-K-COLD-NEXT: !CFGChecksum: 120515930909
+; CHECK-KEEP-K-COLD-NEXT: !Attributes: 0
+; CHECK-KEEP-K-COLD-NEXT:[fa:7 @ fb]:6:2
+; CHECK-KEEP-K-COLD-NEXT: 1: 2
+; CHECK-KEEP-K-COLD-NEXT: 3: 2
+; CHECK-KEEP-K-COLD-NEXT: 6: 2 fa:2
+; CHECK-KEEP-K-COLD-NEXT: !CFGChecksum: 72617220756
+; CHECK-KEEP-K-COLD-NEXT: !Attributes: 0
+; CHECK-KEEP-K-COLD-NEXT:[fa:8 @ fa]:4:1
+; CHECK-KEEP-K-COLD-NEXT: 1: 1
+; CHECK-KEEP-K-COLD-NEXT: 2: 18446744073709551615
+; CHECK-KEEP-K-COLD-NEXT: 3: 1
+; CHECK-KEEP-K-COLD-NEXT: 4: 1
+; CHECK-KEEP-K-COLD-NEXT: 7: 1 fb:1
+; CHECK-KEEP-K-COLD-NEXT: !CFGChecksum: 120515930909
+; CHECK-KEEP-K-COLD-NEXT: !Attributes: 0
 
 ; clang -O3 -fexperimental-new-pass-manager -fuse-ld=lld -fpseudo-probe-for-profiling
 ; -fno-omit-frame-pointer -mno-omit-leaf-frame-pointer -Xclang -mdisable-tail-calls
 ; -g test.c  -o a.out
 
 ; Copied from recursion-compression.test
 #include <stdio.h>
 

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

 }
 
 static void
 mergeSampleProfile(const WeightedFileVector &Inputs, SymbolRemapper *Remapper,
                    StringRef OutputFilename, ProfileFormat OutputFormat,
                    StringRef ProfileSymbolListFile, bool CompressAllSections,
                    bool UseMD5, bool GenPartialProfile,
                    bool SampleMergeColdContext, bool SampleTrimColdContext,
-                   FailureMode FailMode) {
+                   bool SampleKeepLastKFrames, FailureMode FailMode) {
   using namespace sampleprof;
   StringMap<FunctionSamples> ProfileMap;
   SmallVector<std::unique_ptr<sampleprof::SampleProfileReader>, 5> Readers;
   LLVMContext Context;
   sampleprof::ProfileSymbolList WriterList;
   Optional<bool> ProfileIsProbeBased;
   Optional<bool> ProfileIsCS;
   for (const auto &Input : Inputs) {
if (ProfileIsCS && (SampleMergeColdContext || SampleTrimColdContext)) {
     SampleProfileSummaryBuilder Builder(ProfileSummaryBuilder::DefaultCutoffs);
     auto Summary = Builder.computeSummaryForProfiles(ProfileMap);
     uint64_t SampleProfColdThreshold =
         ProfileSummaryBuilder::getColdCountThreshold(
             (Summary->getDetailedSummary()));
 
     // Trim and merge cold context profile using cold threshold above;
     SampleContextTrimmer(ProfileMap)
-        .trimAndMergeColdContextProfiles(SampleProfColdThreshold,
-                                         SampleTrimColdContext,
-                                         SampleMergeColdContext);
+        .trimAndMergeColdContextProfiles(
+            SampleProfColdThreshold, SampleTrimColdContext,
+            SampleMergeColdContext, SampleKeepLastKFrames);
   }
 
   auto WriterOrErr =
       SampleProfileWriter::create(OutputFilename, FormatMap[OutputFormat]);
   if (std::error_code EC = WriterOrErr.getError())
     exitWithErrorCode(EC, OutputFilename);
 
   auto Writer = std::move(WriterOrErr.get());
static int merge_main(int argc, const char *argv[]) {
   cl::opt<bool> SampleMergeColdContext(
       "sample-merge-cold-context", cl::init(false), cl::Hidden,
       cl::desc(
           "Merge context sample profiles whose count is below cold threshold"));
   cl::opt<bool> SampleTrimColdContext(
       "sample-trim-cold-context", cl::init(false), cl::Hidden,
       cl::desc(
           "Trim context sample profiles whose count is below cold threshold"));
+  cl::opt<uint32_t> SampleKeepLastKFrames(
+      "sample-keep-last-k-frames", cl::init(1), cl::ZeroOrMore,
+      cl::desc("Keep the last K frames while merging cold profile. 1 means the "
+               "context-less base profile"));
   cl::opt<bool> GenPartialProfile(
       "gen-partial-profile", cl::init(false), cl::Hidden,
       cl::desc("Generate a partial profile (only meaningful for -extbinary)"));
   cl::opt<std::string> SupplInstrWithSample(
       "supplement-instr-with-sample", cl::init(""), cl::Hidden,
       cl::desc("Supplement an instr profile with sample profile, to correct "
                "the profile unrepresentativeness issue. The sample "
                "profile is the input of the flag. Output will be in instr "
static int merge_main(int argc, const char *argv[]) {
 
   if (ProfileKind == instr)
     mergeInstrProfile(WeightedInputs, Remapper.get(), OutputFilename,
                       OutputFormat, OutputSparse, NumThreads, FailureMode);
   else
     mergeSampleProfile(WeightedInputs, Remapper.get(), OutputFilename,
                        OutputFormat, ProfileSymbolListFile, CompressAllSections,
                        UseMD5, GenPartialProfile, SampleMergeColdContext,
-                       SampleTrimColdContext, FailureMode);
+                       SampleTrimColdContext, SampleKeepLastKFrames,
+                       FailureMode);
 
   return 0;
 }
 
 /// Computer the overlap b/w profile BaseFilename and profile TestFilename.
 static void overlapInstrProfile(const std::string &BaseFilename,
                                 const std::string &TestFilename,
                                 const OverlapFuncFilters &FuncFilter,

llvm/tools/llvm-profgen/ProfileGenerator.cpp

cl::desc("If the total count of context profile is smaller than "
              "the threshold, it will be merged into context-less base "
              "profile."));
 
 static cl::opt<bool> CSProfTrimColdContext(
     "csprof-trim-cold-context", cl::init(true), cl::ZeroOrMore,
     cl::desc("If the total count of the profile after all merge is done "
              "is still smaller than threshold, it will be trimmed."));
 
+static cl::opt<uint32_t> CSProfKeepLastKFrames(
+    "csprof-keep-last-k-frames", cl::init(1), cl::ZeroOrMore,
+    cl::desc("Keep the last K frames while merging cold profile. 1 means the "
+             "context-less base profile"));
+
 extern cl::opt<int> ProfileSummaryCutoffCold;
 
 using namespace llvm;
 using namespace sampleprof;
 
 namespace llvm {
 namespace sampleprof {
 
void CSProfileGenerator::postProcessProfiles() {
 
   // Run global pre-inliner to adjust/merge context profile based on estimated
   // inline decisions.
   CSPreInliner(ProfileMap, HotCountThreshold, ColdCountThreshold).run();
 
   // Trim and merge cold context profile using cold threshold above;
   SampleContextTrimmer(ProfileMap)
       .trimAndMergeColdContextProfiles(
-          ColdCountThreshold, CSProfTrimColdContext, CSProfMergeColdContext);
+          ColdCountThreshold, CSProfTrimColdContext, CSProfMergeColdContext,
+          CSProfKeepLastKFrames);
 }
 
 void CSProfileGenerator::computeSummaryAndThreshold() {
   // Update the default value of cold cutoff for llvm-profgen.
   // Do it here because we don't want to change the global default,
   // which would lead CS profile size too large.
   if (!ProfileSummaryCutoffCold.getNumOccurrences())
     ProfileSummaryCutoffCold = 999000;