[llvm-profgen] Generating probe-based non-CS profile.

[hoy]

Summary

I'm bring up the support of pseudo-probe-based non-CS profile generation. The approach is quite similar to generating dwarf-based non-CS profile. The main difference is for a given linear instruction range, instead of each disassembled instruction, pseudo probes that are covered by the range are processed. The pseudo probe extraction code is shared with CS probe profile generation.

I'm seeing 0.7% performance win for one of our internal large benchmark compared to using non-CS dwarf-based profile, and 0.5% win for another large benchmark when combined with profi.

[wlei]

Thanks for adding support for probe-based non-CS profile.

0.5% win for another large benchmark when combined with profi.

is this because with probe, it now have a good support for the unknown(dangling) block, so PROFI can work well for this?

llvm/test/tools/llvm-profgen/inline-pseudoprobe.test

4

Out of curiosity, for the text profile, how do we know the profile is a CS nested profile or a non-CS nested profile?

llvm/tools/llvm-profgen/ProfileGenerator.cpp

444

So here we don't use the way like CS-profile to generate the zero-count(in ProfileGenerator.cpp: 981), instead we reuse the preprocessRangeCounter to initialize all function range with zero, the probe inside the function will naturally be added with zero count. I guess this is the same to the way of CS-profile, right?

[hoy]

In D120335#3339225, @wlei wrote:

Thanks for adding support for probe-based non-CS profile.

0.5% win for another large benchmark when combined with profi.

is this because with probe, it now have a good support for the unknown(dangling) block, so PROFI can work well for this?

I think so. Profi should help probes more than line-based profile.

llvm/test/tools/llvm-profgen/inline-pseudoprobe.test

4

CS nested profile comes with a "shouldInline" attribute for each nested profile, see test/tools/llvm-profdata/cs-sample-nested-profile.test.

llvm/tools/llvm-profgen/ProfileGenerator.cpp

444

Yes, we are using the non-CS way of reporting zero counted probes. There is a difference between CS and non-CS in that for CS, the non-executed probes are reported for its owner frame only, while for non-CS, such probes are reported for the whole inline nest.

[wlei]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

944

Probe.isBlock() check is hoisted from extractProbesFromRange to this for CS profile but Why doesn't add this check back for non-CS profile?

[hoy]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

944

So we don't count callsite probes here since they are handled in populateBoundarySamplesWithProbes. Otherwise they will be double counted.

We don't do this for non-CS probe profile. This is to be consistent with non-CS dwarf implementation where we report zero count in populateBoundarySamplesForAllFunctions. I guess I can also do some refactoring to unify the two implementations.

[wlei]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

944

Yeah, I understand for count==0. Here I'm asking for whether we should add Probe->isBlock() for non-CS probe profile. My understanding is for body sample, we only accumulate total sample for Block profile? otherwise, the count from callsite probes will affect the total samples.

[hoy]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

944

For non-CS, I'm using count==0 for probe to be consistent with dwarf, i.e, both block probes and call probes are counted by populateBodySamples.

For CS, only block probes are counted by populateBodySamples, call probes are counted by populateBoundarySamples.

That's the inconsistency we should probably fix.

[wenlei]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

444

Both preprocessRangeCounter and extractProbesFromRange calls findDisjointRanges, which is duplicated.

for CS, the non-executed probes are reported for its owner frame only, while for non-CS, such probes are reported for the whole inline nest.

what does the term "owner frame" refer to? can you elaborate the above?

497

Perhaps we don't need a parameter for this, just use Binary->usePseudoProbes() instead.

901

nit: move this closer to the functions definitions for ProfileGeneratorBase.

[hoy]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

444

By owner frame I mean the inlinee frame that originally directly contains the probe. Eg., given function A, A inlines B and one original probe of B is sampled, for CS, all of other original probes of B will be reported. None of the original probes of A will be reported if none of A's probe is sampled. But for non-CS, all of A's and B's probes will be reported even if only one B's probe is sampled.

A real example is in the attached inline-pseudoprobe.test where we have

; CHECK:     main:88:0
; CHECK-NEXT: 1: 0
; CHECK-NEXT: 2: foo:88
; CHECK-NEXT:  1: 0
; CHECK-NEXT:  2: 15

the corresponding CS profile is in inline-cs-pseudoprobe.test where there is no profile generated for the main function.

497

Sounds good.

901

done.

[hoy]

Addressing comments.

[wenlei]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

444

Ok, thanks for clarification - that makes sense. Suggestion: use canonical terms like "inlinee frame" or "leaf inlinee frame" instead of nebulous terms like "owner frame" to avoid confusion.

Both preprocessRangeCounter and extractProbesFromRange calls findDisjointRanges, which is duplicated.

Something we can do to avoid redundant findDisjointRanges?

[hoy]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

444

Something we can do to avoid redundant findDisjointRanges?

Good point. Sorry for missing this earliser. I made it conditional.

[hoy]

Updating D120335: [llvm-profgen] Generating probe-based non-CS profile.

[davidxl]

What is the size overhead on binary with probe? What is the impact on source drifting tolerance level?

[hoy]

In D120335#3352701, @davidxl wrote:

What is the size overhead on binary with probe? What is the impact on source drifting tolerance level?

Pseudo probes have negligible impact on code size. When it comes to binary size, the two encoded probe sections can cause the binary 10% - 15% bigger as of now.

As for source drifting, probe-based profile should be resilient to source changes that do not affect CFG structure and the number of callsites.

[wenlei]

lgtm, thanks.

llvm/test/tools/llvm-profgen/inline-pseudoprobe.test

+; RUN: llvm-profgen --format=text --ignore-stack-samples --perfscript=%S/Inputs/inline-cs-pseudoprobe.perfscript --binary=%S/Inputs/inline-cs-pseudoprobe.perfbin --output=%t --profile-summary-cold-count=0
+; RUN: FileCheck %s --input-file %t
+
+; CHECK:     main:88:0
+; CHECK-NEXT: 1: 0
+; CHECK-NEXT: 2: foo:88
+; CHECK-NEXT:  1: 0
+; CHECK-NEXT:  2: 15
+; CHECK-NEXT:  3: 15
+; CHECK-NEXT:  4: 14
+; CHECK-NEXT:  5: 1
+; CHECK-NEXT:  6: 15
+; CHECK-NEXT:  7: 0
+; CHECK-NEXT:  9: 0
+; CHECK-NEXT:  8: bar:28
+; CHECK-NEXT:   1: 14
+; CHECK-NEXT:   4: 14
+; CHECK-NEXT:   !CFGChecksum: 72617220756
+; CHECK-NEXT:  !CFGChecksum: 563088904013236
+; CHECK-NEXT: !CFGChecksum: 281479271677951
+
+
+; 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
+
+#include <stdio.h>
+
+int bar(int x, int y) {
+  if (x % 3) {
+    return x - y;
+  }
+  return x + y;
+}
+
+void foo() {
+  int s, i = 0;
+  while (i++ < 4000 * 4000)
+    if (i % 91) s = bar(i, s); else s += 30;
+  printf("sum is %d\n", s);
+}
+
+int main() {
+  foo();
+  return 0;
+}

llvm/test/tools/llvm-profgen/noinline-pseudoprobe.test

+; RUN: llvm-profgen --format=text --perfscript=%S/Inputs/noinline-cs-pseudoprobe.perfscript --binary=%S/Inputs/noinline-cs-pseudoprobe.perfbin --output=%t1 --ignore-stack-samples
+; RUN: FileCheck %s --input-file %t1
+; RUN: llvm-profgen --format=text --perfscript=%S/Inputs/noinline-cs-pseudoprobe.aggperfscript --binary=%S/Inputs/noinline-cs-pseudoprobe.perfbin --output=%t2 --ignore-stack-samples
+; RUN: FileCheck %s --input-file %t2
+
+
+; CHECK:     foo:75:0
+; CHECK-NEXT:  1: 0
+; CHECK-NEXT:  2: 15
+; CHECK-NEXT:  3: 15
+; CHECK-NEXT:  4: 15
+; CHECK-NEXT:  5: 0
+; CHECK-NEXT:  6: 15
+; CHECK-NEXT:  7: 0
+; CHECK-NEXT:  8: 15 bar:15
+; CHECK-NEXT:  9: 0
+; CHECK-NEXT:  !CFGChecksum: 563088904013236
+; CHECK-NEXT: bar:30:15
+; CHECK-NEXT:  1: 15
+; CHECK-NEXT:  4: 15
+; CHECK-NEXT:  !CFGChecksum: 72617220756
+
+
+
+; 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
+; -fno-inline-functions -g test.c  -o a.out
+
+#include <stdio.h>
+
+int bar(int x, int y) {
+  if (x % 3) {
+    return x - y;
+  }
+  return x + y;
+}
+
+void foo() {
+  int s, i = 0;
+  while (i++ < 4000 * 4000)
+    if (i % 91) s = bar(i, s); else s += 30;
+  printf("sum is %d\n", s);
+}
+
+int main() {
+  foo();
+  return 0;
+}

llvm/tools/llvm-profgen/ProfileGenerator.h

 #include <unordered_set>
 
 using namespace llvm;
 using namespace sampleprof;
 
 namespace llvm {
 namespace sampleprof {
 
+using ProbeCounterMap =
+    std::unordered_map<const MCDecodedPseudoProbe *, uint64_t>;
+
 // This base class for profile generation of sample-based PGO. We reuse all
 // structures relating to function profiles and profile writers as seen in
 // /ProfileData/SampleProf.h.
 class ProfileGeneratorBase {
 
 public:
   ProfileGeneratorBase(ProfiledBinary *Binary,
                        const ContextSampleCounterMap &Counters)
protected:
   |<--100-->|
   |<------200------>|
   A         B       C
 
   sample count for disjoint region [A,B] would be 300.
   */
   void findDisjointRanges(RangeSample &DisjointRanges,
                           const RangeSample &Ranges);
+
+  // Go through each address from range to extract the top frame probe by
+  // looking up in the Address2ProbeMap
+  void extractProbesFromRange(const RangeSample &RangeCounter,
+                              ProbeCounterMap &ProbeCounter);
+
   // Helper function for updating body sample for a leaf location in
   // FunctionProfile
   void updateBodySamplesforFunctionProfile(FunctionSamples &FunctionProfile,
                                            const SampleContextFrame &LeafLoc,
                                            uint64_t Count);
   void updateTotalSamples();
 
   StringRef getCalleeNameForOffset(uint64_t TargetOffset);

 public:
   ProfileGenerator(ProfiledBinary *Binary,
                    const ContextSampleCounterMap &Counters)
       : ProfileGeneratorBase(Binary, Counters){};
   void generateProfile() override;
 
 private:
   void generateLineNumBasedProfile();
+  void generateProbeBasedProfile();
   RangeSample preprocessRangeCounter(const RangeSample &RangeCounter);
   FunctionSamples &getTopLevelFunctionProfile(StringRef FuncName);
   // Helper function to get the leaf frame's FunctionProfile by traversing the
   // inline stack and meanwhile it adds the total samples for each frame's
   // function profile.
   FunctionSamples &
   getLeafProfileAndAddTotalSamples(const SampleContextFrameVector &FrameVec,
-                                   uint64_t Count);
+                                   uint64_t Count, bool WithProbe = false);
   void populateBodySamplesForAllFunctions(const RangeSample &RangeCounter);
   void
   populateBoundarySamplesForAllFunctions(const BranchSample &BranchCounters);
+  void populateBodySamplesWithProbesForAllFunctions(const RangeSample &RangeCounter);
+  void
+  populateBoundarySamplesWithProbesForAllFunctions(const BranchSample &BranchCounters);
   void postProcessProfiles();
   void trimColdProfiles(const SampleProfileMap &Profiles,
                         uint64_t ColdCntThreshold);
 };
 
-using ProbeCounterMap =
-    std::unordered_map<const MCDecodedPseudoProbe *, uint64_t>;
-
 class CSProfileGenerator : public ProfileGeneratorBase {
 public:
   CSProfileGenerator(ProfiledBinary *Binary,
                      const ContextSampleCounterMap &Counters)
       : ProfileGeneratorBase(Binary, Counters){};
 
   void generateProfile() override;
 
private:
   void populateBodySamplesForFunction(FunctionSamples &FunctionProfile,
                                       const RangeSample &RangeCounters);
   void populateBoundarySamplesForFunction(SampleContextFrames ContextId,
                                           FunctionSamples &FunctionProfile,
                                           const BranchSample &BranchCounters);
   void populateInferredFunctionSamples();
 
   void generateProbeBasedProfile();
-  // Go through each address from range to extract the top frame probe by
-  // looking up in the Address2ProbeMap
-  void extractProbesFromRange(const RangeSample &RangeCounter,
-                              ProbeCounterMap &ProbeCounter);
+
   // Fill in function body samples from probes
   void populateBodySamplesWithProbes(const RangeSample &RangeCounter,
                                      SampleContextFrames ContextStack);
   // Fill in boundary samples for a call probe
   void populateBoundarySamplesWithProbes(const BranchSample &BranchCounter,
                                          SampleContextFrames ContextStack);
   // Helper function to get FunctionSamples for the leaf probe
   FunctionSamples &

llvm/tools/llvm-profgen/ProfileGenerator.cpp

if (Ret.second) {
     FunctionSamples &FProfile = Ret.first->second;
     FProfile.setContext(Context);
   }
   return Ret.first->second;
 }
 
 void ProfileGenerator::generateProfile() {
   if (Binary->usePseudoProbes()) {
-    // TODO: Support probe based profile generation
-    exitWithError("Probe based profile generation not supported for AutoFDO, "
-      "consider dropping `--ignore-stack-samples` or adding `--use-dwarf-correlation`.");
+    generateProbeBasedProfile();
   } else {
     generateLineNumBasedProfile();
   }
   postProcessProfiles();
 }
 
 void ProfileGenerator::postProcessProfiles() {
   computeSummaryAndThreshold();
void ProfileGenerator::generateLineNumBasedProfile() {
   // Fill in function body samples
   populateBodySamplesForAllFunctions(SC.RangeCounter);
   // Fill in boundary sample counts as well as call site samples for calls
   populateBoundarySamplesForAllFunctions(SC.BranchCounter);
 
   updateTotalSamples();
 }
 
+void ProfileGenerator::generateProbeBasedProfile() {
+  assert(SampleCounters.size() == 1 &&
+         "Must have one entry for profile generation.");
+  // Enable pseudo probe functionalities in SampleProf
+  FunctionSamples::ProfileIsProbeBased = true;
+  const SampleCounter &SC = SampleCounters.begin()->second;
+  // Fill in function body samples
+  populateBodySamplesWithProbesForAllFunctions(SC.RangeCounter);
+  // Fill in boundary sample counts as well as call site samples for calls
+  populateBoundarySamplesWithProbesForAllFunctions(SC.BranchCounter);
+
+  updateTotalSamples();
+}
+
+void ProfileGenerator::populateBodySamplesWithProbesForAllFunctions(
+    const RangeSample &RangeCounter) {
+  ProbeCounterMap ProbeCounter;
+  extractProbesFromRange(preprocessRangeCounter(RangeCounter), ProbeCounter);
+
+  for (const auto &PI : ProbeCounter) {
+    const MCDecodedPseudoProbe *Probe = PI.first;
+    uint64_t Count = PI.second;
+    SampleContextFrameVector FrameVec;
+    Binary->getInlineContextForProbe(Probe, FrameVec, true);
+    FunctionSamples &FunctionProfile = getLeafProfileAndAddTotalSamples(FrameVec, Count, true);
+    FunctionProfile.addBodySamplesForProbe(Probe->getIndex(), Count);
+    if (Probe->isEntry())
+      FunctionProfile.addHeadSamples(Count);
+  }
+}
+
+void ProfileGenerator::populateBoundarySamplesWithProbesForAllFunctions(
+    const BranchSample &BranchCounters) {
+  for (const auto &Entry : BranchCounters) {
+    uint64_t SourceOffset = Entry.first.first;
+    uint64_t TargetOffset = Entry.first.second;
+    uint64_t Count = Entry.second;
+    assert(Count != 0 && "Unexpected zero weight branch");
+
+    StringRef CalleeName = getCalleeNameForOffset(TargetOffset);
+    if (CalleeName.size() == 0)
+      continue;
+
+    uint64_t SourceAddress = Binary->offsetToVirtualAddr(SourceOffset);
+    const MCDecodedPseudoProbe *CallProbe =
+        Binary->getCallProbeForAddr(SourceAddress);
+    if (CallProbe == nullptr)
+      continue;
+
+    // Record called target sample and its count.
+    SampleContextFrameVector FrameVec;
+    Binary->getInlineContextForProbe(CallProbe, FrameVec, true);
+
+    if (!FrameVec.empty()) {
+      FunctionSamples &FunctionProfile =
+          getLeafProfileAndAddTotalSamples(FrameVec, 0, true);
+      FunctionProfile.addCalledTargetSamples(
+          FrameVec.back().Location.LineOffset,
+          0,
+          CalleeName, Count);
+    }
+  }
+}
+
 FunctionSamples &ProfileGenerator::getLeafProfileAndAddTotalSamples(
-    const SampleContextFrameVector &FrameVec, uint64_t Count) {
+    const SampleContextFrameVector &FrameVec, uint64_t Count, bool WithProbe) {
   // Get top level profile
   FunctionSamples *FunctionProfile =
       &getTopLevelFunctionProfile(FrameVec[0].FuncName);
   FunctionProfile->addTotalSamples(Count);
+  if (WithProbe) {
+    const auto *FuncDesc = Binary->getFuncDescForGUID(Function::getGUID(FunctionProfile->getName()));
+    FunctionProfile->setFunctionHash(FuncDesc->FuncHash);
+  }
 
   for (size_t I = 1; I < FrameVec.size(); I++) {
     LineLocation Callsite(
         FrameVec[I - 1].Location.LineOffset,
         getBaseDiscriminator(FrameVec[I - 1].Location.Discriminator));
     FunctionSamplesMap &SamplesMap =
         FunctionProfile->functionSamplesAt(Callsite);
     auto Ret =
         SamplesMap.emplace(FrameVec[I].FuncName.str(), FunctionSamples());
     if (Ret.second) {
       SampleContext Context(FrameVec[I].FuncName);
       Ret.first->second.setContext(Context);
     }
     FunctionProfile = &Ret.first->second;
     FunctionProfile->addTotalSamples(Count);
+    if (WithProbe) {
+      const auto *FuncDesc = Binary->getFuncDescForGUID(Function::getGUID(FunctionProfile->getName()));
+      FunctionProfile->setFunctionHash(FuncDesc->FuncHash);
+    }
   }
 
   return *FunctionProfile;
 }
 
 RangeSample
 ProfileGenerator::preprocessRangeCounter(const RangeSample &RangeCounter) {
   RangeSample Ranges(RangeCounter.begin(), RangeCounter.end());

 
 void CSProfileGenerator::generateProfile() {
   FunctionSamples::ProfileIsCSFlat = true;
 
   if (Binary->getTrackFuncContextSize())
     computeSizeForProfiledFunctions();
 
   if (Binary->usePseudoProbes()) {
-    // Enable pseudo probe functionalities in SampleProf
-    FunctionSamples::ProfileIsProbeBased = true;
     generateProbeBasedProfile();
   } else {
     generateLineNumBasedProfile();
   }
   postProcessProfiles();
 }
 
 void CSProfileGenerator::computeSizeForProfiledFunctions() {
static void extractPrefixContextStack(
     const SmallVectorImpl<const MCDecodedPseudoProbe *> &Probes,
     ProfiledBinary *Binary) {
   for (const auto *P : Probes) {
     Binary->getInlineContextForProbe(P, ContextStack, true);
   }
 }
 
 void CSProfileGenerator::generateProbeBasedProfile() {
+  // Enable pseudo probe functionalities in SampleProf
+  FunctionSamples::ProfileIsProbeBased = true;
   for (const auto &CI : SampleCounters) {
     const auto *CtxKey = cast<ProbeBasedCtxKey>(CI.first.getPtr());
     SampleContextFrameVector ContextStack;
     extractPrefixContextStack(ContextStack, CtxKey->Probes, Binary);
     // Fill in function body samples from probes, also infer caller's samples
     // from callee's probe
     populateBodySamplesWithProbes(CI.second.RangeCounter, ContextStack);
     // Fill in boundary samples for a call probe
     populateBoundarySamplesWithProbes(CI.second.BranchCounter, ContextStack);
   }
 }
 
-void CSProfileGenerator::extractProbesFromRange(const RangeSample &RangeCounter,
+void ProfileGeneratorBase::extractProbesFromRange(const RangeSample &RangeCounter,
                                                 ProbeCounterMap &ProbeCounter) {
   RangeSample Ranges;
   findDisjointRanges(Ranges, RangeCounter);
   for (const auto &Range : Ranges) {
     uint64_t RangeBegin = Binary->offsetToVirtualAddr(Range.first.first);
     uint64_t RangeEnd = Binary->offsetToVirtualAddr(Range.first.second);
     uint64_t Count = Range.second;
-    // Disjoint ranges have introduce zero-filled gap that
-    // doesn't belong to current context, filter them out.
-    if (Count == 0)
-      continue;
 
     InstructionPointer IP(Binary, RangeBegin, true);
     // Disjoint ranges may have range in the middle of two instr,
     // e.g. If Instr1 at Addr1, and Instr2 at Addr2, disjoint range
     // can be Addr1+1 to Addr2-1. We should ignore such range.
     if (IP.Address > RangeEnd)
       continue;
 
     do {
       const AddressProbesMap &Address2ProbesMap =
           Binary->getAddress2ProbesMap();
       auto It = Address2ProbesMap.find(IP.Address);
       if (It != Address2ProbesMap.end()) {
         for (const auto &Probe : It->second) {
-          if (!Probe.isBlock())
-            continue;
           ProbeCounter[&Probe] += Count;
         }
       }
     } while (IP.advance() && IP.Address <= RangeEnd);
   }
 }
 
 void CSProfileGenerator::populateBodySamplesWithProbes(
     const RangeSample &RangeCounter, SampleContextFrames ContextStack) {
   ProbeCounterMap ProbeCounter;
   // Extract the top frame probes by looking up each address among the range in
   // the Address2ProbeMap
   extractProbesFromRange(RangeCounter, ProbeCounter);
   std::unordered_map<MCDecodedPseudoProbeInlineTree *,
                      std::unordered_set<FunctionSamples *>>
       FrameSamples;
   for (const auto &PI : ProbeCounter) {
     const MCDecodedPseudoProbe *Probe = PI.first;
     uint64_t Count = PI.second;
+    // Disjoint ranges have introduce zero-filled gap that
+    // doesn't belong to current context, filter them out.
+    if (!Probe->isBlock() || Count == 0)
+      continue;
     FunctionSamples &FunctionProfile =
         getFunctionProfileForLeafProbe(ContextStack, Probe);
     // Record the current frame and FunctionProfile whenever samples are
     // collected for non-danglie probes. This is for reporting all of the
     // zero count probes of the frame later.
     FrameSamples[Probe->getInlineTreeNode()].insert(&FunctionProfile);
     FunctionProfile.addBodySamplesForProbe(Probe->getIndex(), Count);
     FunctionProfile.addTotalSamples(Count);