[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

206

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

399–408

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?

399–408

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

54

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

239

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

54

Sounds good, added.

llvm/tools/llvm-profgen/ProfileGenerator.cpp

239

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

239

Sounds good, thanks for the clarification.

llvm/include/llvm/Analysis/ProfileSummaryInfo.h

 /// isColdCount that tells whether a given profile count is considered hot/cold
 /// based on the profile summary. This also provides convenience methods to
 /// check whether a function is hot or cold.
 
 // FIXME: Provide convenience methods to determine hotness/coldness of other IR
 // units. This would require making this depend on BFI.
 class ProfileSummaryInfo {
 private:
-  const Module &M;
+  const Module *M;
   std::unique_ptr<ProfileSummary> Summary;
   void computeThresholds();
   // Count thresholds to answer isHotCount and isColdCount queries.
   Optional<uint64_t> HotCountThreshold, ColdCountThreshold;
   // True if the working set size of the code is considered huge,
   // because the number of profile counts required to reach the hot
   // percentile is above a huge threshold.
   Optional<bool> HasHugeWorkingSetSize;
   // True if the working set size of the code is considered large,
   // because the number of profile counts required to reach the hot
   // percentile is above a large threshold.
   Optional<bool> HasLargeWorkingSetSize;
   // Compute the threshold for a given cutoff.
   Optional<uint64_t> computeThreshold(int PercentileCutoff) const;
   // The map that caches the threshold values. The keys are the percentile
   // cutoff values and the values are the corresponding threshold values.
   mutable DenseMap<int, uint64_t> ThresholdCache;
 
 public:
-  ProfileSummaryInfo(const Module &M) : M(M) { refresh(); }
+  ProfileSummaryInfo(const Module &M) : M(&M) { refresh(); }
+  ProfileSummaryInfo(std::unique_ptr<ProfileSummary> PSI)
+      : M(nullptr), Summary(std::move(PSI)) {}
 
   ProfileSummaryInfo(ProfileSummaryInfo &&Arg) = default;
 
   /// If no summary is present, attempt to refresh.
   void refresh();
 
   /// Returns true if profile summary is available.
   bool hasProfileSummary() const { return Summary != nullptr; }

llvm/lib/Analysis/ProfileSummaryInfo.cpp

 // The profile summary metadata may be attached either by the frontend or by
 // any backend passes (IR level instrumentation, for example). This method
 // checks if the Summary is null and if so checks if the summary metadata is now
 // available in the module and parses it to get the Summary object.
 void ProfileSummaryInfo::refresh() {
   if (hasProfileSummary())
     return;
   // First try to get context sensitive ProfileSummary.
-  auto *SummaryMD = M.getProfileSummary(/* IsCS */ true);
+  auto *SummaryMD = M->getProfileSummary(/* IsCS */ true);
   if (SummaryMD)
     Summary.reset(ProfileSummary::getFromMD(SummaryMD));
 
   if (!hasProfileSummary()) {
     // This will actually return PSK_Instr or PSK_Sample summary.
-    SummaryMD = M.getProfileSummary(/* IsCS */ false);
+    SummaryMD = M->getProfileSummary(/* IsCS */ false);
     if (SummaryMD)
       Summary.reset(ProfileSummary::getFromMD(SummaryMD));
   }
   if (!hasProfileSummary())
     return;
   computeThresholds();
 }
 

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
+; Test --csprof-trim-cold-context=0
+; 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
 
+; Test --csprof-merge-cold-context=0
+; 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=10 --csprof-merge-cold-context=0
+; RUN: FileCheck %s --input-file %t --check-prefix=CHECK-UNMERGED
+
 ; 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-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
 
+; CHECK-UNMERGED:     [main:2 @ foo:5 @ fa:8 @ fa:7 @ fb:5 @ fb]:13:4
+; CHECK-UNMERGED-NEXT: 1: 4
+; 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]
+
 
 ; 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-profgen/ProfileGenerator.h

 //===-- ProfileGenerator.h - 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_TOOLS_LLVM_PROGEN_PROFILEGENERATOR_H
 #define LLVM_TOOLS_LLVM_PROGEN_PROFILEGENERATOR_H
 #include "ErrorHandling.h"
 #include "PerfReader.h"
 #include "ProfiledBinary.h"
+#include "llvm/Analysis/ProfileSummaryInfo.h"
 #include "llvm/ProfileData/SampleProfWriter.h"
+#include <memory>
 
 using namespace llvm;
 using namespace sampleprof;
 
 namespace llvm {
 namespace sampleprof {
 
 class ProfileGenerator {

 
 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();
 
+  // Profile summary to answer isHotCount and isColdCount queries.
+  std::unique_ptr<ProfileSummaryInfo> PSI;
+
 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."));
 
 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);
+  PSI.reset(new ProfileSummaryInfo(std::move(Summary)));
+}
+
 void CSProfileGenerator::mergeAndTrimColdProfile(
     StringMap<FunctionSamples> &ProfileMap) {
+  if (!CSProfMergeColdContext && !CSProfTrimColdContext)
+    return;
+
+  // Use threshold calculated from profile summary unless specified
+  uint64_t ColdThreshold = PSI->getColdCountThreshold();
+  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) {