[CSSPGO][llvm-profgen] Merge and trim profile for cold context to reduce profile size

[wlei]

Summary

This change allows merging and trimming cold context profile in llvm-profgen to solve profile size bloat problem. Currently when the profile's total sample is below threshold(supported by a switch), it will be considered cold and merged into a base context-less profile, which will at least keep the profile quality as good as the baseline(non-cs).

For example, two input profiles:
[main @ foo @ bar]:60
[main @ bar]:50
Under threshold = 100, the two profiles will be merge into one with the base context, get result:
[bar]:110

Added two switches:
--csprof-cold-thres=<value>: Specified the total samples threshold for a context profile to be considered cold, with 100 being the default. Any cold context profiles will be merged into context-less base profile by default.
--csprof-keep-cold: Force profile generation to keep cold context profiles instead of dropping them. By default, any cold context will not be written to output profile.

Results:
Though not yet evaluating it with the latest CSSPGO, our internal branch shows neutral on performance but significantly reduce the profile size. Detailed evaluation on llvm-profgen with CSSPGO will come later.

[wmi]

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

7

I see "[main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb]:9:4" is trimmed but [fa]:14:4 is kept, so the threshold should be 10?

llvm/tools/llvm-profgen/ProfileGenerator.cpp

349

emplace_back(I.getKey(), &I.second)?

[wlei]

Rebase and addressing Wei's feedback

[wlei]

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

7

The second test is for --csprof-keep-cold, which means we only merge the cold profile but don't remove it.
So I intended to use a very high threshold(100) to let all the profiles merge together and keep it.

llvm/tools/llvm-profgen/ProfileGenerator.cpp

349

Good catch, fixed

[wmi]

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

7

Ok, I see. Thanks for explaining it.

It means if the total count of a profile with context is smaller than the csprof-cold-thres, it will be merged into its base profile. If the total count of the base profile after all merge is done is still smaller than the csprof-cold-thres, it will be trimmed unless csprof-keep-cold flag is specified. Is it correct? Better document it more clearly in the flag desc.

[wlei]

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

7

Yes, that's exactly what I mean. Will update the flag desc, thanks for your feedback.

[wlei]

Refine the cl::desc for CSProfColdThres and CSProfKeepCold

[wmi]

LGTM. Thanks.

[hoy]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

33

Nit: Specify

I was wondering in the future if we'd like a consistent way of computing cold threshold with SampleProfileSummaryBuilder, i.e., counts below 99.9% (aka cut-off) are considered cold, so that the profile generator and the profile consumer will be on the same page.

[wlei]

fix typo and rebase

[wlei]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

33

Thanks for sharing this. I'm not sure how we get the 100 threshold, I guess this is the result of many iterations. Let me dig into the SampleProfileSummaryBuilder to see the chance.

[hoy]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

33

Yes, it is a number from experiments. The logic in the profile summary builder is a bit different in that it sorts all samples by weight and treats the samples below a given percentage (say 0.01%) as cold. In order words, if there are more than 99.99% samples that are hotter than the current one, then the current is treated as cold. That's the logic used in the compiler.

[wenlei]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

33

Yeah, +1 for what Hongtao said, I wanted to use %-based cutoff, but didn't have access to SampleProfileSummaryBuilder from early prototype.

I think create_llvm_profile tool have both sample_threshold_frac and sample_threshold_absolute? Similarly both could be useful here. We could refine them as we broaden the adoption and continue tuning - I don't have an opinion whether we make all these improvements here or in another patch.

[wlei]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

33

Cool, thanks for your informative suggestions! will have a try and add more experiments on SPEC.

[wlei]

rebase

llvm/test/tools/llvm-profgen/inline-cs-noprobe.test

-; RUN: llvm-profgen --perfscript=%S/Inputs/inline-cs-noprobe.perfscript --binary=%S/Inputs/inline-cs-noprobe.perfbin --output=%t --show-unwinder-output | FileCheck %s --check-prefix=CHECK-UNWINDER
+; RUN: llvm-profgen --perfscript=%S/Inputs/inline-cs-noprobe.perfscript --binary=%S/Inputs/inline-cs-noprobe.perfbin --output=%t --show-unwinder-output --csprof-cold-thres=0 | FileCheck %s --check-prefix=CHECK-UNWINDER
 ; RUN: FileCheck %s --input-file %t
 
 ; CHECK:[main:1 @ foo]:44:0
 ; CHECK: 2.2: 14
 ; CHECK: 3: 15
 ; CHECK: 3.2: 14 bar:14
 ; CHECK: 3.4: 1
 ; CHECK:[main:1 @ foo:3.2 @ bar]:14:0

llvm/test/tools/llvm-profgen/inline-cs-pseudoprobe.test

-; RUN: llvm-profgen --perfscript=%S/Inputs/inline-cs-pseudoprobe.perfscript --binary=%S/Inputs/inline-cs-pseudoprobe.perfbin --output=%t --show-unwinder-output | FileCheck %s --check-prefix=CHECK-UNWINDER
+; RUN: llvm-profgen --perfscript=%S/Inputs/inline-cs-pseudoprobe.perfscript --binary=%S/Inputs/inline-cs-pseudoprobe.perfbin --output=%t --show-unwinder-output --csprof-cold-thres=0 | FileCheck %s --check-prefix=CHECK-UNWINDER
 ; RUN: FileCheck %s --input-file %t
 
 ; CHECK:     [main:2 @ foo]:74:0
 ; CHECK-NEXT: 2: 15
 ; CHECK-NEXT: 3: 15
 ; CHECK-NEXT: 4: 14
 ; CHECK-NEXT: 5: 1
 ; CHECK-NEXT: 6: 15

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 --perfscript=%S/Inputs/recursion-compression.perfscript --binary=%S/Inputs/recursion-compression.perfbin --output=%t --compress-recursion=-1 --csprof-cold-thres=8
+; RUN: FileCheck %s --input-file %t
+
+; Test --csprof-keep-cold
+; RUN: llvm-profgen --perfscript=%S/Inputs/recursion-compression.perfscript --binary=%S/Inputs/recursion-compression.perfbin --output=%t --compress-recursion=-1 --csprof-cold-thres=100 --csprof-keep-cold
+; RUN: FileCheck %s --input-file %t --check-prefix=CHECK-KEEP-COLD
+
+; CHECK:     [fa]:14:4
+; CHECK-NEXT: 1: 4
+; 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-NEXT: !CFGChecksum: 120515930909
+; CHECK-NEXT:[fb]:10:2
+; CHECK-NEXT: 1: 2
+; CHECK-NEXT: 3: 2
+; CHECK-NEXT: 5: 4 fb:4
+; CHECK-NEXT: 6: 2 fa:2
+; CHECK-NEXT: !CFGChecksum: 72617220756
+; CHECK-NEXT:[main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb]:9:4
+; CHECK-NEXT: 1: 4
+; CHECK-NEXT: 2: 3
+; CHECK-NEXT: 3: 1
+; CHECK-NEXT: 6: 1 fa:1
+; CHECK-NEXT: !CFGChecksum: 72617220756
+
+; CHECK-KEEP-COLD:     [fb]:19:6
+; CHECK-KEEP-COLD-NEXT: 1: 6
+; CHECK-KEEP-COLD-NEXT: 2: 3
+; CHECK-KEEP-COLD-NEXT: 3: 3
+; CHECK-KEEP-COLD-NEXT: 5: 4 fb:4
+; CHECK-KEEP-COLD-NEXT: 6: 3 fa:3
+; CHECK-KEEP-COLD-NEXT: !CFGChecksum: 72617220756
+; CHECK-KEEP-COLD-NEXT:[fa]:14:4
+; CHECK-KEEP-COLD-NEXT: 1: 4
+; CHECK-KEEP-COLD-NEXT: 3: 4
+; CHECK-KEEP-COLD-NEXT: 4: 2
+; CHECK-KEEP-COLD-NEXT: 5: 1
+; CHECK-KEEP-COLD-NEXT: 7: 2 fb:2
+; CHECK-KEEP-COLD-NEXT: 8: 1 fa:1
+; CHECK-KEEP-COLD-NEXT: !CFGChecksum: 120515930909
+
+
+; 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>
+
+int fb(int n) {
+  if(n > 10) return fb(n / 2);
+  return fa(n - 1);
+}
+
+int fa(int n) {
+  if(n < 2) return n;
+  if(n % 2) return fb(n - 1);
+  return fa(n - 1);
+}
+
+void foo() {
+  int s, i = 0;
+  while (i++ < 10000)
+    s += fa(i);
+  printf("sum is %d\n", s);
+}
+
+int main() {
+  foo();
+  return 0;
+}

llvm/test/tools/llvm-profgen/noinline-cs-noprobe.test

-; RUN: llvm-profgen --perfscript=%S/Inputs/noinline-cs-noprobe.perfscript --binary=%S/Inputs/noinline-cs-noprobe.perfbin --output=%t --show-unwinder-output | FileCheck %s --check-prefix=CHECK-UNWINDER
+; RUN: llvm-profgen --perfscript=%S/Inputs/noinline-cs-noprobe.perfscript --binary=%S/Inputs/noinline-cs-noprobe.perfbin --output=%t --show-unwinder-output --csprof-cold-thres=0 | FileCheck %s --check-prefix=CHECK-UNWINDER
 ; RUN: FileCheck %s --input-file %t
 
 ; CHECK:[main:1 @ foo:3 @ bar]:12:3
 ; CHECK: 0: 3
 ; CHECK: 1: 3
 ; CHECK: 2: 2
 ; CHECK: 4: 1
 ; CHECK: 5: 3

llvm/test/tools/llvm-profgen/noinline-cs-pseudoprobe.test

-; RUN: llvm-profgen --perfscript=%S/Inputs/noinline-cs-pseudoprobe.perfscript --binary=%S/Inputs/noinline-cs-pseudoprobe.perfbin --output=%t --show-unwinder-output | FileCheck %s --check-prefix=CHECK-UNWINDER
+; RUN: llvm-profgen --perfscript=%S/Inputs/noinline-cs-pseudoprobe.perfscript --binary=%S/Inputs/noinline-cs-pseudoprobe.perfbin --output=%t --show-unwinder-output --csprof-cold-thres=0 | FileCheck %s --check-prefix=CHECK-UNWINDER
 ; RUN: FileCheck %s --input-file %t
 
 ; CHECK:     [main:2 @ foo]:75:0
 ; CHECK-NEXT: 2: 15
 ; CHECK-NEXT: 3: 15
 ; CHECK-NEXT: 4: 15
 ; CHECK-NEXT: 6: 15
 ; CHECK-NEXT: 8: 15 bar:15

llvm/test/tools/llvm-profgen/recursion-compression.test

 ; Firstly test uncompression(--compress-recursion=0)
-; RUN: llvm-profgen --perfscript=%S/Inputs/recursion-compression.perfscript --binary=%S/Inputs/recursion-compression.perfbin --output=%t --compress-recursion=0
+; RUN: llvm-profgen --perfscript=%S/Inputs/recursion-compression.perfscript --binary=%S/Inputs/recursion-compression.perfbin --output=%t --compress-recursion=0 --csprof-cold-thres=0
 ; RUN: FileCheck %s --input-file %t -check-prefix=CHECK-UNCOMPRESS
-; RUN: llvm-profgen --perfscript=%S/Inputs/recursion-compression.perfscript --binary=%S/Inputs/recursion-compression.perfbin --output=%t --compress-recursion=-1 --show-unwinder-output | FileCheck %s --check-prefix=CHECK-UNWINDER
+; RUN: llvm-profgen --perfscript=%S/Inputs/recursion-compression.perfscript --binary=%S/Inputs/recursion-compression.perfbin --output=%t --compress-recursion=-1 --show-unwinder-output --csprof-cold-thres=0  | FileCheck %s --check-prefix=CHECK-UNWINDER
 ; RUN: FileCheck %s --input-file %t
 
 ; CHECK-UNCOMPRESS: [main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:5 @ fb:6 @ fa:8 @ fa:7 @ fb:6 @ fa]:4:1
 ; CHECK-UNCOMPRESS:  1: 1
 ; CHECK-UNCOMPRESS:  3: 1
 ; CHECK-UNCOMPRESS:  4: 1
 ; CHECK-UNCOMPRESS:  7: 1 fb:1
 ; CHECK-UNCOMPRESS:  !CFGChecksum: 120515930909

llvm/tools/llvm-profgen/ProfileGenerator.h

 
 public:
   ProfileGenerator(){};
   virtual ~ProfileGenerator() = default;
   static std::unique_ptr<ProfileGenerator>
   create(const BinarySampleCounterMap &BinarySampleCounters,
          enum PerfScriptType SampleType);
   virtual void generateProfile() = 0;
-
+  // Merge and trim profile with cold context before serialization,
+  // only eligible for CS profile
+  virtual void
+  mergeAndTrimColdProfile(StringMap<FunctionSamples> &ProfileMap){};
   // Use SampleProfileWriter to serialize profile map
   void write();
 
 protected:
   /*
   For each region boundary point, mark if it is begin or end (or both) of
   the region. Boundary points are inclusive. Log the sample count as well
   so we can use it when we compute the sample count of each disjoint region
void generateProfile() override {
     // functions, we estimate it from inlinee's profile using the entry of the
     // body sample.
     populateInferredFunctionSamples();
   }
 
 protected:
   // Lookup or create FunctionSamples for the context
   FunctionSamples &getFunctionProfileForContext(StringRef ContextId);
+  // Merge cold context profile whose total sample is below threshold
+  // into base profile.
+  void mergeAndTrimColdProfile(StringMap<FunctionSamples> &ProfileMap) override;
 
 private:
   // Helper function for updating body sample for a leaf location in
   // FunctionProfile
   void updateBodySamplesforFunctionProfile(FunctionSamples &FunctionProfile,
                                            const FrameLocation &LeafLoc,
                                            uint64_t Count);
   void populateFunctionBodySamples(FunctionSamples &FunctionProfile,

llvm/tools/llvm-profgen/ProfileGenerator.cpp

cl::values(
                    "GCC encoding (only meaningful for -sample)")));
 int32_t CompressionSize = -1;
 static cl::opt<int32_t, true> RecursionCompression(
     "compress-recursion",
     cl::desc("Compressing recursion by deduplicating adjacent frame "
              "sequences up to the specified size. -1 means no size limit."),
     cl::Hidden, cl::location(CompressionSize));
 
+static cl::opt<uint64_t> CSProfColdThres(
+    "csprof-cold-thres", cl::init(100), cl::ZeroOrMore,
+    cl::desc(
+        "Total sample threshold for context profile to be considered cold."));
+
+static cl::opt<bool>
+    CSProfKeepCold("csprof-keep-cold", cl::init(false), cl::ZeroOrMore,
+                   cl::desc("Trim cold context profile in output (will be "
+                            "merged to context-less profile by default)."));
+
 using namespace llvm;
 using namespace sampleprof;
 
 namespace llvm {
 namespace sampleprof {
 
 static bool
 usePseudoProbes(const BinarySampleCounterMap &BinarySampleCounters) {
ProfileGenerator::create(const BinarySampleCounterMap &BinarySampleCounters,
   return ProfileGenerator;
 }
 
 void ProfileGenerator::write() {
   auto WriterOrErr = SampleProfileWriter::create(OutputFilename, OutputFormat);
   if (std::error_code EC = WriterOrErr.getError())
     exitWithError(EC, OutputFilename);
   auto Writer = std::move(WriterOrErr.get());
+  mergeAndTrimColdProfile(ProfileMap);
   Writer->write(ProfileMap);
 }
 
 void ProfileGenerator::findDisjointRanges(RangeSample &DisjointRanges,
                                           const RangeSample &Ranges) {
 
   /*
   Regions may overlap with each other. Using the boundary info, find all
CallerProfile.addCalledTargetSamples(
         CallerLeafFrameLoc.second.LineOffset,
         CallerLeafFrameLoc.second.Discriminator, CalleeProfile.getName(),
         EstimatedCallCount);
     updateBodySamplesforFunctionProfile(CallerProfile, CallerLeafFrameLoc,
                                         EstimatedCallCount);
   }
 }
 
+void CSProfileGenerator::mergeAndTrimColdProfile(
+    StringMap<FunctionSamples> &ProfileMap) {
+  // Nothing to merge if sample threshold is zero
+  if (!CSProfColdThres)
+    return;
+
+  // Filter the cold profiles from ProfileMap and move them into a tmp
+  // container
+  std::vector<std::pair<StringRef, const FunctionSamples *>> ToRemoveVec;
+  for (const auto &I : ProfileMap) {
+    const FunctionSamples &FunctionProfile = I.second;
+    if (FunctionProfile.getTotalSamples() >= CSProfColdThres)
+      continue;
+    ToRemoveVec.emplace_back(I.getKey(), &I.second);
+  }
+
+  // Remove the code profile from ProfileMap and merge them into BaseProileMap
+  StringMap<FunctionSamples> BaseProfileMap;
+  for (const auto &I : ToRemoveVec) {
+    auto Ret =
+        BaseProfileMap.try_emplace(I.second->getName(), 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 (!CSProfKeepCold && I.second.getTotalSamples() < CSProfColdThres &&
+        ProfileMap.find(I.getKey()) == ProfileMap.end())
+      continue;
+    // Merge the profile if the original profile exists, otherwise just insert
+    // as a new profile
+    FunctionSamples &OrigProfile = getFunctionProfileForContext(I.getKey());
+    StringRef TmpName = OrigProfile.getName();
+    OrigProfile.merge(I.second);
+    // Should use the name ref from ProfileMap's key to avoid name being freed
+    // from BaseProfileMap
+    OrigProfile.setName(TmpName);
+  }
+}
+
 // Helper function to extract context prefix string stack
 static void
 extractPrefixContextStack(SmallVectorImpl<std::string> &ContextStrStack,
                           const SmallVectorImpl<const PseudoProbe *> &Probes,
                           ProfiledBinary *Binary) {
   for (const auto *P : Probes) {
     Binary->getInlineContextForProbe(P, ContextStrStack, true);
   }