[CSSPGO][llvm-profdata] Support trimming cold context when merging profiles

[wenlei]

Summary

The change adds support for triming and merging cold context when mergine CSSPGO profiles using llvm-profdata. This is similar to the context profile trimming in llvm-profgen, however the flexibility to trim cold context after profile is generated can be useful.

A bug fix is also included to make sure PSI thresholds are always calcuated before query.

[wenlei]

llvm/tools/llvm-profdata/CMakeLists.txt

2

This is because we now use ProfileSummaryInfo.cpp, so we can get all consistent hot/cold threshold computation and related switches, though pulling in entire Analysis lib is less than ideal. I'm thinking about moving ProfileSummaryInfo into ProfileData/ProfileSummaryBuildder.cpo from Analysis/ProfileSummaryInfo.cpp. But I'd like to get feedback before I move things around.

[wmi]

llvm/tools/llvm-profdata/CMakeLists.txt

2

Extract ProfileSummaryInfo from Analysis/ directory sounds reasonable to me. How about move it to Core -- lib/IR/ProfileSummary.cpp for example? The reason is ProfileSummaryInfo interface is used widely. If someday we want to use the interface in Core or Support which ProfileData depends on, it will introduce cyclic dependence. Putting it in Core will clear the problem.

[wenlei]

llvm/tools/llvm-profdata/CMakeLists.txt

2

Sounds good. Thanks for the suggestion.

[wenlei]

llvm/tools/llvm-profdata/CMakeLists.txt

2

Looks like extracting ProfileSummaryInfo to Core isn't a trivial move. ProfileSummaryInfo has a bunch of interfaces taking BFI as parameter, and also offers invalidation, so it depends on Analysis. It's fine when ProfileSummaryInfo itself is part of Analysis, but if we move it to Core, we will have cyclic dependency. For the same reason, we can't move it to ProfileData either.

If we want to move it, we will need to break it up, so the part that does the threshold computation can be moved, and the part that does the query (relies on Analysis) stays.

Now what I'm thinking about is only moving the threshold/cutoff switches in ProfileSummaryInfo.cpp along with part of ProfileSummaryInfo::computeThresholds into ProfileData/ProfileSummaryBuilder. That won't change dependencies as ProfileSummaryInfo uses ProfileSummaryBuilder already. But it lets tools that only uses ProfileData access consistent threshold computation too.

[wenlei]

Move threshold knobs into ProfileSummaryBuilder, avoid pulling Analysis lib into tools.

[wenlei]

Cleanup.

[wmi]

llvm/tools/llvm-profdata/CMakeLists.txt

2

Ok, that sounds good to me. Thanks for trying to find a good solution.

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

887–891

What will happen if these two flags are enabled for non-CSSPGO sample profile?

llvm/tools/llvm-profgen/CMakeLists.txt

6

After the PSI switch and threshold are moved to ProfData, no need to depend on Analysis anymore?

[wenlei]

The dependency on Analysis should be removed in the latest version.

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

887–891

Good catch, they should be ignored. Added a check.

[wenlei]

Ignore trim/merge flag for non-CS profile.

[wmi]

LGTM, thanks.

llvm/include/llvm/Analysis/ProfileSummaryInfo.h

private:
   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(std::unique_ptr<ProfileSummary> PSI)
-      : M(nullptr), Summary(std::move(PSI)) {}
+      : M(nullptr), Summary(std::move(PSI)) {
+    computeThresholds();
+  }
 
   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/test/tools/llvm-profdata/cs-sample-trimmer.test

+RUN: llvm-profdata merge --sample --text -output=%t.proftext %S/Inputs/cs-sample.proftext
+RUN: diff -b %t.proftext %S/Inputs/cs-sample.proftext
+
+RUN: llvm-profdata merge --sample --text -output=%t.proftext %S/Inputs/cs-sample.proftext -sample-merge-cold-context -profile-summary-cold-count=500
+RUN: FileCheck %s --input-file %t.proftext --check-prefixes=CHECK-TRIM,CHECK-MERGE
+RUN: llvm-profdata merge --sample --text -output=%t.proftext %S/Inputs/cs-sample.proftext -sample-merge-cold-context -sample-trim-cold-context -profile-summary-cold-count=500
+RUN: FileCheck %s --input-file %t.proftext --check-prefixes=CHECK-TRIM,CHECK-END
+RUN: llvm-profdata merge --sample --text -output=%t.proftext %S/Inputs/cs-sample.proftext -sample-merge-cold-context -sample-trim-cold-context -profile-summary-cutoff-cold=990000
+RUN: FileCheck %s --input-file %t.proftext --check-prefixes=CHECK-TRIM,CHECK-END
+
+CHECK-TRIM:      [main:3 @ _Z5funcAi:1 @ _Z8funcLeafi]:1467299:11
+CHECK-TRIM-NEXT:  0: 6
+CHECK-TRIM-NEXT:  1: 6
+CHECK-TRIM-NEXT:  3: 287884
+CHECK-TRIM-NEXT:  4: 287864 _Z3fibi:315608
+CHECK-TRIM-NEXT:  15: 23
+CHECK-TRIM-NEXT:  !Attributes: 0
+CHECK-TRIM-NEXT: [main:3.1 @ _Z5funcBi:1 @ _Z8funcLeafi]:500853:20
+CHECK-TRIM-NEXT:  0: 15
+CHECK-TRIM-NEXT:  1: 15
+CHECK-TRIM-NEXT:  3: 74946
+CHECK-TRIM-NEXT:  4: 74941 _Z3fibi:82359
+CHECK-TRIM-NEXT:  10: 23324
+CHECK-TRIM-NEXT:  11: 23327 _Z3fibi:25228
+CHECK-TRIM-NEXT:  15: 11
+CHECK-TRIM-NEXT:  !Attributes: 1
+CHECK-END-NOT:   [
+CHECK-MERGE:     [_Z5funcBi]:360:32
+CHECK-MERGE-NEXT: 0: 32
+CHECK-MERGE-NEXT: 1: 32 _Z8funcLeafi:20
+CHECK-MERGE-NEXT: 3: 12
+CHECK-MERGE-NEXT: !Attributes: 0
+CHECK-MERGE-NEXT:[main]:308:12
+CHECK-MERGE-NEXT: 2: 24
+CHECK-MERGE-NEXT: 3: 28 _Z5funcAi:18
+CHECK-MERGE-NEXT: 3.1: 28 _Z5funcBi:30
+CHECK-MERGE-NEXT: !Attributes: 0
+CHECK-MERGE-NEXT:[_Z5funcAi]:99:11
+CHECK-MERGE-NEXT: 0: 10
+CHECK-MERGE-NEXT: 1: 10 _Z8funcLeafi:11
+CHECK-MERGE-NEXT: 3: 24
+CHECK-MERGE-NEXT: !Attributes: 0

llvm/test/tools/llvm-profgen/cs-preinline.test

 ; Test default llvm-profgen with preinline off
 ; RUN: llvm-profgen --perfscript=%S/Inputs/inline-cs-noprobe.perfscript --binary=%S/Inputs/inline-cs-noprobe.perfbin --output=%t
 ; RUN: FileCheck %s --input-file %t --check-prefix=CHECK-DEFAULT
 
 ; Test llvm-profgen with preinliner on will merge not inlinable profile into base profile.
 ; RUN: llvm-profgen --perfscript=%S/Inputs/inline-cs-noprobe.perfscript --binary=%S/Inputs/inline-cs-noprobe.perfbin --output=%t --csspgo-preinliner=1
 ; RUN: FileCheck %s --input-file %t --check-prefix=CHECK-PREINL
 
 ; Test preinliner threshold that prevents all possible inlining and merges everything into base profile.
-; RUN: llvm-profgen --perfscript=%S/Inputs/inline-cs-noprobe.perfscript --binary=%S/Inputs/inline-cs-noprobe.perfbin --output=%t --csspgo-preinliner=1  -sample-profile-hot-inline-threshold=0
+; RUN: llvm-profgen --perfscript=%S/Inputs/inline-cs-noprobe.perfscript --binary=%S/Inputs/inline-cs-noprobe.perfbin --output=%t --csspgo-preinliner=1  -sample-profile-cold-inline-threshold=0
 ; RUN: FileCheck %s --input-file %t --check-prefix=CHECK-NO-PREINL
 
 ; CHECK-DEFAULT:     [main:1 @ foo]:309:0
 ; CHECK-DEFAULT-NEXT: 2.1: 14
 ; CHECK-DEFAULT-NEXT: 3: 15
 ; CHECK-DEFAULT-NEXT: 3.1: 14 bar:14
 ; CHECK-DEFAULT-NEXT: 3.2: 1
 ; CHECK-DEFAULT-NEXT: !Attributes: 1

llvm/tools/llvm-profdata/CMakeLists.txt

 set(LLVM_LINK_COMPONENTS
+  Analysis
   Core
   ProfileData
   Support
   )
 
 add_llvm_tool(llvm-profdata
   llvm-profdata.cpp
 
   DEPENDS
   intrinsics_gen
   )

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

 //===- llvm-profdata.cpp - LLVM profile data tool -------------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // llvm-profdata merges .profdata files.
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/Analysis/ProfileSummaryInfo.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/ProfileData/InstrProfReader.h"
 #include "llvm/ProfileData/InstrProfWriter.h"
 #include "llvm/ProfileData/ProfileCommon.h"
 #include "llvm/ProfileData/SampleProfReader.h"
 #include "llvm/ProfileData/SampleProfWriter.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Errc.h"
else
       Writer.setPartialProfile();
   }
 }
 
 static void
 mergeSampleProfile(const WeightedFileVector &Inputs, SymbolRemapper *Remapper,
                    StringRef OutputFilename, ProfileFormat OutputFormat,
                    StringRef ProfileSymbolListFile, bool CompressAllSections,
-                   bool UseMD5, bool GenPartialProfile, FailureMode FailMode) {
+                   bool UseMD5, bool GenPartialProfile,
+                   bool SampleMergeColdContext, bool SampleTrimColdContext,
+                   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) {
for (StringMap<FunctionSamples>::iterator I = Profiles.begin(),
       }
     }
 
     std::unique_ptr<sampleprof::ProfileSymbolList> ReaderList =
         Reader->getProfileSymbolList();
     if (ReaderList)
       WriterList.merge(*ReaderList);
   }
+
+  if (SampleMergeColdContext || SampleTrimColdContext) {
+    // Use threshold calculated from profile summary unless specified.
+    SampleProfileSummaryBuilder Builder(ProfileSummaryBuilder::DefaultCutoffs);
+    auto Summary = Builder.computeSummaryForProfiles(ProfileMap);
+    uint64_t SampleProfColdThreshold =
+        ProfileSummaryInfo(std::move(Summary)).getColdCountThreshold();
+
+    // Trim and merge cold context profile using cold threshold above;
+    SampleContextTrimmer(ProfileMap)
+        .trimAndMergeColdContextProfiles(SampleProfColdThreshold,
+                                         SampleTrimColdContext,
+                                         SampleMergeColdContext);
+  }
+
   auto WriterOrErr =
       SampleProfileWriter::create(OutputFilename, FormatMap[OutputFormat]);
   if (std::error_code EC = WriterOrErr.getError())
     exitWithErrorCode(EC, OutputFilename);
 
   auto Writer = std::move(WriterOrErr.get());
   // WriterList will have StringRef refering to string in Buffer.
   // Make sure Buffer lives as long as WriterList.
static int merge_main(int argc, const char *argv[]) {
   cl::opt<bool> CompressAllSections(
       "compress-all-sections", cl::init(false), cl::Hidden,
       cl::desc("Compress all sections when writing the profile (only "
                "meaningful for -extbinary)"));
   cl::opt<bool> UseMD5(
       "use-md5", cl::init(false), cl::Hidden,
       cl::desc("Choose to use MD5 to represent string in name table (only "
                "meaningful for -extbinary)"));
+  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<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, FailureMode);
+                       UseMD5, GenPartialProfile, SampleMergeColdContext,
+                       SampleTrimColdContext, 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/CMakeLists.txt

 
 set(LLVM_LINK_COMPONENTS
   AllTargetsDescs
   AllTargetsDisassemblers
   AllTargetsInfos
+  Analysis
   Core
   MC
   IPO
   MCDisassembler
   Object
   ProfileData
   Support
   Symbolize