[CSSPGO][llvm-profgen] Use profile summary based threshold for context trimming and merging

[wenlei]

Summary

Switch to use cold threshold from profile summary for cold context merging and trimming, instead of relying on hard coded values. Minor refactoring included for switch names, etc.

[wenlei]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

231

We need profile name otherwise UseContextLessSummary won't work as the merge of context profile needs name key.

[wenlei]

undo bad linter

[hoy]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

47

I'm wondering if the original switches can be moved into a common place, such as ProfileSummaryBuilder.cpp, so that they can be reused here. Some code in CSProfileGenerator::computeSummaryAndThreshold may also be sharable.

[wmi]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

424–433

Similar as Hongtao's question. Can ProfileSummaryInfo::computeThresholds be used to set the hot/cold threshold?

[wenlei]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

47

That would be nice, however for normal path of ProfileSummaryInfo, the actual summary is retrieved from module metadata which is read from binary profile directly instead of being built from ProfileSummaryBuilder on the spot.

I think it probably makes sense to just use ProfileSummaryInfo here, as Wei suggested. However, looks like the switches from ProfileSummaryInfo.cpp doesn't get exposed that way.. I guess that's because ProfileData is linked into llvm-profogen, instead of being part of the tool?

424–433

Sounds good, changed to use PSI directly here, so we don't need keep thresholds separately.

[wenlei]

Use PSI directly.

[hoy]

Using PSI looks good, thanks. Can you please add a test for -csprof-merge-cold-context? Otherwise lgtm.

[wenlei]

Add test coverage for --csprof-merge-cold-context.

[wenlei]

rebase

[wmi]

LGTM.

[wlei]

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

50

How about we add:

; CHECK-UNMERGED-NOT: [fa]
; CHECK-UNMERGED-NOT: [fb]

to make sure the merging not happen.

llvm/tools/llvm-profgen/ProfileGenerator.cpp

264

There is two places to call this function: for probe and non-probe.
How about to move this to CSProfileGenerator::write(...) so that we only keep one call for this?

[wenlei]

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

50

Sounds good, added.

llvm/tools/llvm-profgen/ProfileGenerator.cpp

264

I think it doesn't belong to writing, though moving there may save one call from the code. The thresholds will be needed for estimating inlining, which I will send up next, so I'd like to keep it this way for now if it's ok, and try to refactor in the upcoming change.

[wenlei]

Address Lei's comment.

[wlei]

LGTM, thanks!

llvm/tools/llvm-profgen/ProfileGenerator.cpp

264

Sounds good, thanks for the clarification.

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-pseudoprobe.perfscript --binary=%S/Inputs/recursion-compression-pseudoprobe.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-pseudoprobe.perfscript --binary=%S/Inputs/recursion-compression-pseudoprobe.perfbin --output=%t --compress-recursion=-1 --csprof-cold-thres=100 --csprof-keep-cold
+; RUN: llvm-profgen --perfscript=%S/Inputs/recursion-compression-pseudoprobe.perfscript --binary=%S/Inputs/recursion-compression-pseudoprobe.perfbin --output=%t --compress-recursion=-1 --csprof-cold-thres=100 --csprof-trim-cold-context=0
 ; 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

 
 
 ; 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);

llvm/tools/llvm-profgen/ProfileGenerator.h

 
 protected:
   // Lookup or create FunctionSamples for the context
   FunctionSamples &getFunctionProfileForContext(StringRef ContextId,
                                                 bool WasLeafInlined = false);
   // Merge cold context profile whose total sample is below threshold
   // into base profile.
   void mergeAndTrimColdProfile(StringMap<FunctionSamples> &ProfileMap);
+  void computeSummaryAndThreshold();
   void write(std::unique_ptr<SampleProfileWriter> Writer,
              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,
                                    const RangeSample &RangeCounters,
                                    ProfiledBinary *Binary);
   void populateFunctionBoundarySamples(StringRef ContextId,
                                        FunctionSamples &FunctionProfile,
                                        const BranchSample &BranchCounters,
                                        ProfiledBinary *Binary);
   void populateInferredFunctionSamples();
 
+  // Count thresholds to answer isHotCount and isColdCount queries.
+  // Mirrors the threshold in ProfileSummaryInfo.
+  uint64_t HotCountThreshold;
+  uint64_t ColdCountThreshold;
+
 public:
   // Deduplicate adjacent repeated context sequences up to a given sequence
   // length. -1 means no size limit.
   static int32_t MaxCompressionSize;
 };
 
 using ProbeCounterMap = std::unordered_map<const PseudoProbe *, uint64_t>;
 

llvm/tools/llvm-profgen/ProfileGenerator.cpp

 //===-- ProfileGenerator.cpp - Profile Generator  ---------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "ProfileGenerator.h"
+#include "llvm/ProfileData/ProfileCommon.h"
 
 static cl::opt<std::string> OutputFilename("output", cl::value_desc("output"),
                                            cl::Required,
                                            cl::desc("Output profile file"));
 static cl::alias OutputA("o", cl::desc("Alias for --output"),
                          cl::aliasopt(OutputFilename));
 
 static cl::opt<SampleProfileFormat> OutputFormat(
     "format", cl::desc("Format of output profile"), cl::init(SPF_Text),
     cl::values(
         clEnumValN(SPF_Binary, "binary", "Binary encoding (default)"),
         clEnumValN(SPF_Compact_Binary, "compbinary", "Compact binary encoding"),
         clEnumValN(SPF_Ext_Binary, "extbinary", "Extensible binary encoding"),
         clEnumValN(SPF_Text, "text", "Text encoding"),
         clEnumValN(SPF_GCC, "gcc",
                    "GCC encoding (only meaningful for -sample)")));
 
 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(llvm::sampleprof::CSProfileGenerator::MaxCompressionSize));
 
-static cl::opt<uint64_t> CSProfColdThres(
+static cl::opt<uint64_t> CSProfColdThreshold(
     "csprof-cold-thres", cl::init(100), cl::ZeroOrMore,
     cl::desc("Specify the total samples threshold for a context profile to "
              "be considered cold, any cold profiles will be merged into "
              "context-less base profiles"));
 
-static cl::opt<bool> CSProfKeepCold(
-    "csprof-keep-cold", cl::init(false), cl::ZeroOrMore,
+static cl::opt<bool> CSProfMergeColdContext(
+    "csprof-merge-cold-context", cl::init(true), cl::ZeroOrMore,
     cl::desc("This works together with --csprof-cold-thres. If the total count "
-             "of the 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."));
+             "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("This works together with --csprof-cold-thres. If the total count "
+             "of the profile after all merge is done is still smaller than "
+             "threshold, "
+             "it will be trimmed."));
+
+// The 4 switches below mirrors those in ProfileSummaryInfo.cpp
+// We need them here so the pre-inliner in llvm-profgen is tunable
+// and can better align with compiler.
+static cl::opt<int> ProfileSummaryCutoffHot(
+    "profile-summary-cutoff-hot", cl::Hidden, cl::init(990000), cl::ZeroOrMore,
+    cl::desc("A count is hot if it exceeds the minimum count to"
+             " reach this percentile of total counts."));
+
+static cl::opt<int> ProfileSummaryCutoffCold(
+    "profile-summary-cutoff-cold", cl::Hidden, cl::init(999999), cl::ZeroOrMore,
+    cl::desc("A count is cold if it is below the minimum count"
+             " to reach this percentile of total counts."));
+
+static cl::opt<int> ProfileSummaryHotCount(
+    "profile-summary-hot-count", cl::ReallyHidden, cl::ZeroOrMore,
+    cl::desc("A fixed hot count that overrides the count derived from"
+             " profile-summary-cutoff-hot"));
+
+static cl::opt<int> ProfileSummaryColdCount(
+    "profile-summary-cold-count", cl::ReallyHidden, cl::ZeroOrMore,
+    cl::desc("A fixed cold count that overrides the count derived from"
+             " profile-summary-cutoff-cold"));
 
 using namespace llvm;
 using namespace sampleprof;
 
 namespace llvm {
 namespace sampleprof {
 
 // Initialize the MaxCompressionSize to -1 which means no size limit
CSProfileGenerator::getFunctionProfileForContext(StringRef ContextStr,
                                                  bool WasLeafInlined) {
   auto Ret = ProfileMap.try_emplace(ContextStr, FunctionSamples());
   if (Ret.second) {
     SampleContext FContext(Ret.first->first(), RawContext);
     if (WasLeafInlined)
       FContext.setAttribute(ContextWasInlined);
     FunctionSamples &FProfile = Ret.first->second;
     FProfile.setContext(FContext);
+    FProfile.setName(FContext.getNameWithoutContext());
   }
   return Ret.first->second;
 }
 
 void CSProfileGenerator::generateProfile() {
   FunctionSamples::ProfileIsCS = true;
   for (const auto &BI : BinarySampleCounters) {
     ProfiledBinary *Binary = BI.first;
for (const auto &CI : BI.second) {
                                       CI.second.BranchCounter, Binary);
     }
   }
   // Fill in call site value sample for inlined calls and also use context to
   // infer missing samples. Since we don't have call count for inlined
   // functions, we estimate it from inlinee's profile using the entry of the
   // body sample.
   populateInferredFunctionSamples();
+
+  // Compute hot/cold threshold based on profile. This will be used for cold
+  // context profile merging/trimming.
+  computeSummaryAndThreshold();
 }
 
 void CSProfileGenerator::updateBodySamplesforFunctionProfile(
     FunctionSamples &FunctionProfile, const FrameLocation &LeafLoc,
     uint64_t Count) {
   // Filter out invalid negative(int type) lineOffset
   if (LeafLoc.second.LineOffset & 0x80000000)
     return;
CallerProfile.addCalledTargetSamples(
         CalleeProfile.getContext().getNameWithoutContext(), EstimatedCallCount);
     CallerProfile.addBodySamples(CallerLeafFrameLoc.second.LineOffset,
                                  CallerLeafFrameLoc.second.Discriminator,
                                  EstimatedCallCount);
     CallerProfile.addTotalSamples(EstimatedCallCount);
   }
 }
 
+void CSProfileGenerator::computeSummaryAndThreshold() {
+  SampleProfileSummaryBuilder Builder(ProfileSummaryBuilder::DefaultCutoffs);
+  auto Summary = Builder.computeSummaryForProfiles(ProfileMap);
+  auto &DetailedSummary = Summary->getDetailedSummary();
+  auto &HotEntry = ProfileSummaryBuilder::getEntryForPercentile(
+      DetailedSummary, ProfileSummaryCutoffHot);
+  HotCountThreshold = HotEntry.MinCount;
+  if (ProfileSummaryHotCount.getNumOccurrences() > 0)
+    HotCountThreshold = ProfileSummaryHotCount;
+  auto &ColdEntry = ProfileSummaryBuilder::getEntryForPercentile(
+      DetailedSummary, ProfileSummaryCutoffCold);
+  ColdCountThreshold = ColdEntry.MinCount;
+  if (ProfileSummaryColdCount.getNumOccurrences() > 0)
+    ColdCountThreshold = ProfileSummaryColdCount;
+}
+
 void CSProfileGenerator::mergeAndTrimColdProfile(
     StringMap<FunctionSamples> &ProfileMap) {
+  if (!CSProfMergeColdContext && !CSProfTrimColdContext)
+    return;
+
+  // Use threshold calculated from profile summary unless specified
+  uint64_t ColdThreshold = ColdCountThreshold;
+  if (CSProfColdThreshold.getNumOccurrences()) {
+    ColdThreshold = CSProfColdThreshold;
+  }
+
   // Nothing to merge if sample threshold is zero
-  if (!CSProfColdThres)
+  if (ColdThreshold == 0)
     return;
 
   // Filter the cold profiles from ProfileMap and move them into a tmp
   // container
-  std::vector<std::pair<StringRef, const FunctionSamples *>> ToRemoveVec;
+  std::vector<std::pair<StringRef, const FunctionSamples *>> ColdProfiles;
   for (const auto &I : ProfileMap) {
     const FunctionSamples &FunctionProfile = I.second;
-    if (FunctionProfile.getTotalSamples() >= CSProfColdThres)
+    if (FunctionProfile.getTotalSamples() >= ColdThreshold)
       continue;
-    ToRemoveVec.emplace_back(I.getKey(), &I.second);
+    ColdProfiles.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) {
+  for (const auto &I : ColdProfiles) {
+    if (CSProfMergeColdContext) {
       auto Ret = BaseProfileMap.try_emplace(
           I.second->getContext().getNameWithoutContext(), 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 &&
+    if (CSProfTrimColdContext &&
+        I.second.getTotalSamples() < CSProfColdThreshold &&
         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());
     OrigProfile.merge(I.second);
   }
 }
for (const auto &CI : BI.second) {
       // from callee's probe
       populateBodySamplesWithProbes(CI.second.RangeCounter, ContextStrStack,
                                     Binary);
       // Fill in boundary samples for a call probe
       populateBoundarySamplesWithProbes(CI.second.BranchCounter,
                                         ContextStrStack, Binary);
     }
   }
+
+  // Compute hot/cold threshold based on profile. This will be used for cold
+  // context profile merging/trimming.
+  computeSummaryAndThreshold();
 }
 
 void PseudoProbeCSProfileGenerator::extractProbesFromRange(
     const RangeSample &RangeCounter, ProbeCounterMap &ProbeCounter,
     ProfiledBinary *Binary) {
   RangeSample Ranges;
   findDisjointRanges(Ranges, RangeCounter);
   for (const auto &Range : Ranges) {