[CSSPGO] Trim cold profiles with the CS preinliner.

[hoy]

Summary

Adding support to the CS preinliner to trim cold base profiles. This makes trimming consistent with the inline decision made by the preinliner. Also disable the existing profile merger/trimmer when preinliner is on.

llvm/include/llvm/ProfileData/SampleProf.h

 /// sure ProfileMap's key is consistent with FunctionSample's name/context.
 class SampleContextTrimmer {
 public:
   SampleContextTrimmer(SampleProfileMap &Profiles) : ProfileMap(Profiles){};
   // Trim and merge cold context profile when requested.
   void trimAndMergeColdContextProfiles(uint64_t ColdCountThreshold,
                                        bool TrimColdContext,
                                        bool MergeColdContext,
-                                       uint32_t ColdContextFrameLength);
+                                       uint32_t ColdContextFrameLength,
+                                       bool TrimBaseProfileOnly);
   // Canonicalize context profile name and attributes.
   void canonicalizeContextProfiles();
 
 private:
   SampleProfileMap &ProfileMap;
 };
 
 /// ProfileSymbolList records the list of function symbols shown up

llvm/lib/ProfileData/SampleProf.cpp

while (Size < ListSize && StrNum < ProfileSymbolListCutOff) {
     Size += Str.size() + 1;
     StrNum++;
   }
   if (Size != ListSize && StrNum != ProfileSymbolListCutOff)
     return sampleprof_error::malformed;
   return sampleprof_error::success;
 }
 
+// TrimBaseProfileOnly should only be effective when TrimColdContext is true.
+// Trimming means to trim cold profile. On top of that TrimBaseProfileOnly can
+// be used to specify triming all cold profiles or only cold base profiles.
 void SampleContextTrimmer::trimAndMergeColdContextProfiles(
     uint64_t ColdCountThreshold, bool TrimColdContext, bool MergeColdContext,
-    uint32_t ColdContextFrameLength) {
+    uint32_t ColdContextFrameLength, bool TrimBaseProfileOnly) {
   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<SampleContext, const FunctionSamples *>> ColdProfiles;
   for (const auto &I : ProfileMap) {
+    const SampleContext &Context = I.first;
     const FunctionSamples &FunctionProfile = I.second;
-    if (FunctionProfile.getTotalSamples() >= ColdCountThreshold)
-      continue;
-    ColdProfiles.emplace_back(I.first, &I.second);
+    if (FunctionProfile.getTotalSamples() < ColdCountThreshold &&
+        (!TrimBaseProfileOnly || Context.isBaseContext()))
+      ColdProfiles.emplace_back(Context, &I.second);
   }
 
   // Remove the cold profile from ProfileMap and merge them into
   // MergedProfileMap by the last K frames of context
   SampleProfileMap MergedProfileMap;
   for (const auto &I : ColdProfiles) {
     if (MergeColdContext) {
       auto MergedContext = I.second->getContext().getContextFrames();

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

 ; Test default llvm-profgen with preinline off
 ; RUN: llvm-profgen --format=text --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 --format=text --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 --format=text --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
 
+; Test cold profile trimming. Only base profiles should be dropped.
+; RUN: llvm-profgen --format=text --perfscript=%S/Inputs/inline-cs-noprobe.perfscript --binary=%S/Inputs/inline-cs-noprobe.perfbin --output=%t --csspgo-preinliner=1  --csprof-trim-cold-context=1 --profile-summary-hot-count=250
+
+; RUN: FileCheck %s --input-file %t --check-prefix=CHECK-TRIM
+
 ; CHECK-DEFAULT:     [main:1 @ foo]:225: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
 ; CHECK-DEFAULT-NEXT:[main:1 @ foo:3.1 @ bar]:84:0
 ; CHECK-DEFAULT-NEXT: 1: 14

 ; CHECK-NO-PREINL-NEXT: 2.1: 14
 ; CHECK-NO-PREINL-NEXT: 3: 15
 ; CHECK-NO-PREINL-NEXT: 3.1: 14 bar:14
 ; CHECK-NO-PREINL-NEXT: 3.2: 1
 ; CHECK-NO-PREINL-NEXT: !Attributes: 1
 ; CHECK-NO-PREINL-NEXT:[bar]:84:0
 ; CHECK-NO-PREINL-NEXT: 1: 14
 ; CHECK-NO-PREINL-NEXT: !Attributes: 1
+
+; CHECK-TRIM-NOT: [foo]:225:0
+; CHECK-TRIM:[foo:3.1 @ bar]:84:0
+; CHECK-TRIM-NEXT: 1: 14
+; CHECK-TRIM-NEXT: !Attributes: 3

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

if (ProfileIsCS && (SampleMergeColdContext || SampleTrimColdContext)) {
     uint64_t SampleProfColdThreshold =
         ProfileSummaryBuilder::getColdCountThreshold(
             (Summary->getDetailedSummary()));
 
     // Trim and merge cold context profile using cold threshold above;
     SampleContextTrimmer(ProfileMap)
         .trimAndMergeColdContextProfiles(
             SampleProfColdThreshold, SampleTrimColdContext,
-            SampleMergeColdContext, SampleColdContextFrameDepth);
+            SampleMergeColdContext, SampleColdContextFrameDepth, false);
   }
 
   auto WriterOrErr =
       SampleProfileWriter::create(OutputFilename, FormatMap[OutputFormat]);
   if (std::error_code EC = WriterOrErr.getError())
     exitWithErrorCode(EC, OutputFilename);
 
   auto Writer = std::move(WriterOrErr.get());

llvm/tools/llvm-profgen/CSPreInliner.cpp

#endif
   for (StringRef FuncName : buildTopDownOrder()) {
     processFunction(FuncName);
   }
 
   // Not inlined context profiles are merged into its base, so we can
   // trim out such profiles from the output.
   std::vector<SampleContext> ProfilesToBeRemoved;
   for (auto &It : ProfileMap) {
-    SampleContext Context = It.second.getContext();
+    SampleContext &Context = It.second.getContext();
     if (!Context.isBaseContext() && !Context.hasState(InlinedContext)) {
       assert(Context.hasState(MergedContext) &&
              "Not inlined context profile should be merged already");
       ProfilesToBeRemoved.push_back(It.first);
     }
   }
 
   for (auto &ContextName : ProfilesToBeRemoved) {
     ProfileMap.erase(ContextName);
   }
 
   // Make sure ProfileMap's key is consistent with FunctionSamples' name.
   SampleContextTrimmer(ProfileMap).canonicalizeContextProfiles();
 
   LLVM_DEBUG(printProfileNames(ProfileMap, false));
 }

llvm/tools/llvm-profgen/ProfileGenerator.cpp

void CSProfileGenerator::postProcessProfiles() {
   // context profile merging/trimming.
   computeSummaryAndThreshold();
 
   // Run global pre-inliner to adjust/merge context profile based on estimated
   // inline decisions.
   if (EnableCSPreInliner) {
     CSPreInliner(ProfileMap, *Binary, HotCountThreshold, ColdCountThreshold)
         .run();
+    // Turn off the profile merger by default unless it is explicitly enabled.
+    if (!CSProfMergeColdContext.getNumOccurrences())
+      CSProfMergeColdContext = false;
   }
 
-  // Trim and merge cold context profile using cold threshold above. By default,
-  // we skip such merging and trimming when preinliner is on.
-  if (!EnableCSPreInliner || CSProfTrimColdContext.getNumOccurrences() ||
-      CSProfMergeColdContext.getNumOccurrences()) {
+  // Trim and merge cold context profile using cold threshold above. 
+  if (CSProfTrimColdContext || CSProfMergeColdContext) {
     SampleContextTrimmer(ProfileMap)
         .trimAndMergeColdContextProfiles(
             HotCountThreshold, CSProfTrimColdContext, CSProfMergeColdContext,
-            CSProfMaxColdContextDepth);
+            CSProfMaxColdContextDepth, EnableCSPreInliner);
   }
 }
 
 void CSProfileGenerator::computeSummaryAndThreshold() {
   SampleProfileSummaryBuilder Builder(ProfileSummaryBuilder::DefaultCutoffs);
   auto Summary = Builder.computeSummaryForProfiles(ProfileMap);
   HotCountThreshold = ProfileSummaryBuilder::getHotCountThreshold(
       (Summary->getDetailedSummary()));