[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

271

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.

272

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

271

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

271

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

271

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

271

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.

272–280

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.

[hoy]

llvm/tools/llvm-profgen/CSPreInliner.cpp

272–280

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

272–280

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

272–280

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

 };
 
 /// SampleContextTrimmer impelements helper functions to trim, merge cold
 /// context profiles. It also supports context profile canonicalization to make
 /// 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.
+  // Trim and merge cold context profile when requested. TrimBaseProfileOnly
+  // should only be effective when TrimColdContext is true. On top of
+  // TrimColdContext, TrimBaseProfileOnly can be used to specify to trim all
+  // cold profiles or only cold base profiles.
   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

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