[CSSPGO] Trim cold base profiles for 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 when preinliner is on unless explicitly specified.

[hoy]

Updating D112489: [CSSPGO] Trim cold profiles with the CS preinliner.

[wenlei]

llvm/tools/llvm-profgen/CSPreInliner.cpp

286

I think it's cleaner if we defer all trimming to SampleContextTrimmer. We can add TrimBaseProfileOnly flag to trimAndMergeColdContextProfiles interface, or have a separate interface (will still need to make sure we reuse code). We create SampleContextTrimmer right above anyways.

With this change, CSProfMergeColdContext will always be ignored when preinliner is on, but that's not the case before. So alternatively, if TrimBaseProfileOnly is added to the interface, we can call the trimmer outside of preinliner like before, but set TrimBaseProfileOnly to true when preinliner is on.

287

This logging is a bit inconsistent - we don't have logging for this level of activity from SampleContextTrimmer. However if we unify as suggested above, we won't have such inconsistency. I'm not sure if we need such logging though.

[hoy]

llvm/tools/llvm-profgen/CSPreInliner.cpp

286

Sounds good. I'll defer SampleContextTrimmer to do the trimming.

There is a gray area with previous implementation when the preinliner and the trim are both on. In that case the previous check below will accidentally cause the profile merger to run. The profile merger defaulted to on but did not actually run when preilniner was on and neither merger/trimmer switch was given. Turning on the trimmer via the switch caused the merger to run which otherwise would not run.

if (!EnableCSPreInliner || CSProfTrimColdContext.getNumOccurrences() ||
      CSProfMergeColdContext.getNumOccurrences())

I think we want to run the trimmer based on the command line switch, regardless of whether preilniner is on. How about the merger?

[wenlei]

llvm/tools/llvm-profgen/CSPreInliner.cpp

286

Maybe we should do this below. That makes sure we don't accidentally run the merging if only trimming is asked for. That is - make preinliner and merging either-or unless explicitly specified, but leave trimming orthogonal.

if (EnableCSPreInliner && !CSProfMergeColdContext.getNumOccurrences())
   CSProfMergeColdContext = false;

[hoy]

llvm/tools/llvm-profgen/CSPreInliner.cpp

286

That should work. I was thinking if we should completely disable separate merging for preinliner, since we can tune the hotness cutoff to have preinliner do more merge.

[hoy]

Deferring to SampleContextTrimmer to trim cold profiles..

[wenlei]

llvm/tools/llvm-profgen/CSPreInliner.cpp

275–283

Sorry if I wasn't clear.

So alternatively, if TrimBaseProfileOnly is added to the interface, we can call the trimmer outside of preinliner like before, but set TrimBaseProfileOnly to true when preinliner is on.

What I meant by the above to do the trimming outside of preinliner. So preinliner does strictly preinlining, and trimming is done at the original place, with TrimBaseProfileOnly set to EnableCSPreInliner. This way we don't have trimming functionality scattered.

The setting of CSProfMergeColdContext also better be done outside of preinliner.

286

Yeah, that would be the default behavior. But if user asks for merging on top of preinliner, it rarely makes sense but I assume they know what they're doing.

[hoy]

llvm/tools/llvm-profgen/CSPreInliner.cpp

275–283

I found that separating the logic for the preinlining and non-preinlining path is easier to read and less error prone. Don't have a strong preference. Can do what you suggested.

[hoy]

Addressing Wenlei's comment.

llvm/tools/llvm-profgen/CSPreInliner.cpp

275–283

Actually the logic looks better after moving this out.

[wlei]

llvm/lib/ProfileData/SampleProf.cpp

335

what if TrimBaseProfileOnly is true but those two are false, shall we skip the trim?

maybe change to if (!TrimColdContext && !MergeColdContext && !TrimBaseProfileOnly)

357

if(MergeColdContext && !TrimBaseProfileOnly)

So that MergedProfileMap will be empty for TrimBaseProfileOnly

367

Then we can add an assertion: TrimBaseProfileOnly && MergedProfileMap.empty()

[hoy]

llvm/lib/ProfileData/SampleProf.cpp

335

TrimBaseProfileOnly should only be effective when TrimColdContext is true. Trimming means to trim cold profile. On top of that we can choose to trim all cold profiles or only base cold profiles. Does that make sense? I can add a comment for that.

[wenlei]

lgtm, thanks.

llvm/lib/ProfileData/SampleProf.cpp

335

Yeah, sounds good to me. A comment would be helpful.

llvm/tools/llvm-profgen/CSPreInliner.cpp

275–283

Yes, I think separating them out is cleaner.

[wlei]

llvm/lib/ProfileData/SampleProf.cpp

335

Makes sense.

I was thinking if we could reuse this for non-CS cold profile trimming(we have that in our internal patch). in that case, it's the same logic as TrimBaseProfileOnly = true. Anyway, that should be another patch. thanks!

[hoy]

Adding comment to trimAndMergeColdContextProfiles

[wlei]

llvm/lib/ProfileData/SampleProf.cpp

357

How about this comments? is MergeColdContext always false when TrimBaseProfileOnly is true?

[hoy]

llvm/lib/ProfileData/SampleProf.cpp

357

Sorry for overlooking this comment. MergeColdContext and TrimBaseProfileOnly can be both true when user specifies the switches for both merging and trimming. We were talking about whether to turn off merging for preinliner completely and we agree to let user decide. In that case, I think we should ignore TrimBaseProfileOnly since merging is meant to not honor the preinliner decision.

It is a good catch, thanks!

[hoy]

Addressing Lei's comment.

[wlei]

LGTM.

llvm/lib/ProfileData/SampleProf.cpp

357

Thanks for the clarification!

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

std::error_code ProfileSymbolList::read(const uint8_t *Data,
   }
   if (Size != ListSize && StrNum != ProfileSymbolListCutOff)
     return sampleprof_error::malformed;
   return sampleprof_error::success;
 }
 
 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();
       if (ColdContextFrameLength < MergedContext.size())
         MergedContext = MergedContext.take_back(ColdContextFrameLength);
       auto Ret = MergedProfileMap.emplace(MergedContext, FunctionSamples());
       FunctionSamples &MergedProfile = Ret.first->second;
       MergedProfile.merge(*I.second);
     }
     ProfileMap.erase(I.first);
   }
 
   // Move the merged profiles into ProfileMap;
   for (const auto &I : MergedProfileMap) {
     // Filter the cold merged profile
     if (TrimColdContext && I.second.getTotalSamples() < ColdCountThreshold &&
         ProfileMap.find(I.first) == ProfileMap.end())
       continue;
     // Merge the profile if the original profile exists, otherwise just insert
     // as a new profile

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

 // The switches specify inline thresholds used in SampleProfileLoader inlining.
 // TODO: the actual threshold to be tuned here because the size here is based
 // on machine code not LLVM IR.
 extern cl::opt<int> SampleHotCallSiteThreshold;
 extern cl::opt<int> SampleColdCallSiteThreshold;
 extern cl::opt<int> ProfileInlineGrowthLimit;
 extern cl::opt<int> ProfileInlineLimitMin;
 extern cl::opt<int> ProfileInlineLimitMax;
+extern cl::opt<bool> CSProfMergeColdContext;
+extern cl::opt<bool> CSProfTrimColdContext;
+extern cl::opt<uint32_t> CSProfMaxColdContextDepth;
 
 cl::opt<bool> EnableCSPreInliner(
     "csspgo-preinliner", cl::Hidden, cl::init(false),
     cl::desc("Run a global pre-inliner to merge context profile based on "
              "estimated global top-down inline decisions"));
 
 cl::opt<bool> UseContextCostForPreInliner(
     "use-context-cost-for-preinliner", cl::Hidden, cl::init(true),
#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();
+  SampleContextTrimmer Trimmer(ProfileMap);
+  Trimmer.canonicalizeContextProfiles();
+
+  // Turn off the profile merger by default unless it is explicitly enabled.
+  if (!CSProfMergeColdContext.getNumOccurrences())
+    CSProfMergeColdContext = false;
+
+  // Trim or merge cold profiles from ProfileMap.
+  if (CSProfTrimColdContext || CSProfMergeColdContext) {
+    Trimmer.trimAndMergeColdContextProfiles(
+        HotCountThreshold, CSProfTrimColdContext, CSProfMergeColdContext,
+        CSProfMaxColdContextDepth, true);
+  }
 
   LLVM_DEBUG(printProfileNames(ProfileMap, false));
 }

llvm/tools/llvm-profgen/ProfileGenerator.cpp

 
 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<bool> CSProfMergeColdContext(
+cl::opt<bool> CSProfMergeColdContext(
     "csprof-merge-cold-context", cl::init(true), cl::ZeroOrMore,
     cl::desc("If the total count of context profile is smaller than "
              "the threshold, it will be merged into context-less base "
              "profile."));
 
-static cl::opt<bool> CSProfTrimColdContext(
+cl::opt<bool> CSProfTrimColdContext(
     "csprof-trim-cold-context", cl::init(false), cl::ZeroOrMore,
     cl::desc("If the total count of the profile after all merge is done "
              "is still smaller than threshold, it will be trimmed."));
 
-static cl::opt<uint32_t> CSProfMaxColdContextDepth(
+cl::opt<uint32_t> CSProfMaxColdContextDepth(
     "csprof-max-cold-context-depth", cl::init(1), cl::ZeroOrMore,
     cl::desc("Keep the last K contexts while merging cold profile. 1 means the "
              "context-less base profile"));
 
 static cl::opt<int, true> CSProfMaxContextDepth(
     "csprof-max-context-depth", cl::ZeroOrMore,
     cl::desc("Keep the last K contexts while merging profile. -1 means no "
              "depth limit."),
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();
-  }
-
-  // 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()) {
+  } else {
+    // Trim and merge cold context profile using cold threshold above. By
+    // default, we skip such merging and trimming when preinliner is on.
+    if (CSProfTrimColdContext || CSProfMergeColdContext) {
       SampleContextTrimmer(ProfileMap)
           .trimAndMergeColdContextProfiles(
               HotCountThreshold, CSProfTrimColdContext, CSProfMergeColdContext,
-            CSProfMaxColdContextDepth);
+              CSProfMaxColdContextDepth, false);
+    }
   }
 }
 
 void CSProfileGenerator::computeSummaryAndThreshold() {
   SampleProfileSummaryBuilder Builder(ProfileSummaryBuilder::DefaultCutoffs);
   auto Summary = Builder.computeSummaryForProfiles(ProfileMap);
   HotCountThreshold = ProfileSummaryBuilder::getHotCountThreshold(
       (Summary->getDetailedSummary()));