[CSSPGO][llvm-profgen] Pseudo probe based CS profile generation

[wlei]

Summary

This change implements profile generation infra for pseudo probe in llvm-profgen. During virtual unwinding, the raw profile is extracted into range counter and branch counter and aggregated to sample counter map indexed by the call stack context. This change introduces the last step and produces the eventual profile. Specifically, the body of function sample is recorded by going through each probe among the range and callsite target sample is recorded by extracting the callsite probe from branch's source.

Please refer https://groups.google.com/g/llvm-dev/c/1p1rdYbL93s and https://reviews.llvm.org/D89707 for more context about CSSPGO and llvm-profgen.

Implementation

• Extended PseudoProbeProfileGenerator for pseudo probe based profile generation.
• populateBodySamplesWithProbes reading range counter is responsible for recording function body samples and inferring caller's body samples.
• populateBoundarySamplesWithProbes reading branch counter is responsible for recording call site target samples.
• Each sample is recorded with its calling context(named ContextId). Remind that the probe based context key doesn't include the leaf frame probe info, so the ContextId string is created from two part: one from the probe stack strings' concatenation and other one from the leaf frame probe.
• Added regression test

Test Plan:

ninja & ninja check-llvm

[hoy]

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

16

If you rebase on the latest, a checksum will be in presence here.

llvm/tools/llvm-profgen/ProfileGenerator.cpp

412

Please add a comment about why a dangling probe gets a zero count.

419

Nit: function

470

You mean foo:2 @ bar?

478

Move this out of the loop to save the check?

[wlei]

Address Hongtao's feedback: support profile checksum and some other NFC work

[wlei]

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

16

Thanks for reminding this, code to print checksum is added:
FunctionProile.setFunctionHash(LeafFuncDesc->FuncHash);

llvm/tools/llvm-profgen/ProfileGenerator.cpp

412

comments added

470

typo fixed

478

Yeah, moved the last elements out of the loop.

[hoy]

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

5

Nit: use CHECK-NEXT

llvm/tools/llvm-profgen/ProfileGenerator.cpp

42

The setting of the two flags should not be necessary on the profile generation side. They are used on the loader side. Do you see any issue without setting them?

[wlei]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

42

Yeah, if not explicitly setting here, the value will be false.

you see in llvm/lib/Transforms/IPO/SampleProfile.cpp:

// Apply tweaks if context-sensitive profile is available.
if (Reader->profileIsCS()) {
  ProfileIsCS = true;
  FunctionSamples::ProfileIsCS = true;

  // Tracker for profiles under different context
  ContextTracker =
      std::make_unique<SampleContextTracker>(Reader->getProfiles());
}

It's set when the reader know it's a CS profile. But for llvm-profgen side, it doesn't set this.

[hoy]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

42

I see. So ContextTracker used in profile generation?

[wlei]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

42

Not ContextTracker, it's the SampleProfileWriter to write the text format of profile, like:

if (FunctionSamples::ProfileIsCS)
  OS << "[" << S.getNameWithContext() << "]:" << S.getTotalSamples();
else
  OS << S.getName() << ":" << S.getTotalSamples();

if (FunctionSamples::ProfileIsProbeBased) {
  OS.indent(Indent + 1);
  OS << "!CFGChecksum: " << S.getFunctionHash() << "\n";
}

[hoy]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

43

I see. Sounds like the initializations should be moved into ProfileGenerator::generateProfile or ProfileGenerator::write. What do you think?

[wlei]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

43

Sounds good! So since it's only used for the writer, I will move them to ProfileGenerator::write

[wlei]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

43

Just confirm, it's also used in getEntrySamples, so ProfileGenerator::generateProfile might be the right place.

uint64_t getEntrySamples() const {
  if (FunctionSamples::ProfileIsCS && getHeadSamples()) {
    // For CS profile, if we already have more accurate head samples
    // counted by branch sample from caller, use them as entry samples.
    return getHeadSamples();
  }

[hoy]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

43

Yeah, generateProfile sounds the right place.

[wlei]

added CHECK-NEXT and other refactoring work

[wlei]

fix lint

[wlei]

rebase and fix clang-tidy

[wlei]

rebase

[wmi]

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

13–14

In [main:2 @ foo], probes with 0 count are not dumped. In which case probe with 0 count will be dumped?

llvm/tools/llvm-profgen/ProfileGenerator.cpp

347–349

Move it to header comment of extractPrefixContextId in case extractPrefixContextId is used elsewhere.

llvm/tools/llvm-profgen/ProfileGenerator.h

139

Make it an overload function of getFunctionProfileForLeafProbe so it is known to be used for probe?

[wlei]

Addressing Wei's feedback

[wlei]

fix typo

[wlei]

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

13–14

The 0 count is for a dangling probe case to distinguish a missing count, see the ProfileGenerator.cpp:420

Copied the comments for you to refer:

// Drop the samples collected for a dangling probe since it's misleading.
// We still report the probe but with a special zero count. The compiler
// won't trust the zero count and will rely on the counts inference
// algorithm to get the probe a reasonable count. Note that a zero count is
// different from a missing count, where the latter really tells the
// compiler that a probe is never executed.

llvm/tools/llvm-profgen/ProfileGenerator.cpp

347–349

Good catch, moved.

llvm/tools/llvm-profgen/ProfileGenerator.h

139

Good suggestion!

[wmi]

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

13–14

Note that a zero count is different from a missing count, where the latter really tells the compiler that a probe is never executed.

That is contrary to the debug info based profile where zero count for a line means the line is never executed. Missing count for a line means compiler has to infer the count. Why the probe based profile is implemented in such way?

[hoy]

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

13–14

Good question. The current counts inference algorithm in sample profile loader doesn't seem to differentiate a missing sample from a zero sample. For example, in SampleProfileLoader::propagateThroughEdges, accesses to BlockWeights are not preceded with a check. Therefore a missing entry in BlockWeights will be created as a zero entry.

With pseudo probes, a missing sample really means the probe is not executed. A zero sample, on the other hand, wouldn't be created by the profile generator. We are leveraging zero sample to represent a special type of probe, called dangling probe, which will need the compiler to infer its count. We were thinking about using UINT64_MAX instead, but UINT64_MAX literally is a legal number of samples.

The compiler change that dangles probes haven't been sent out yet. @wlei I think we should exclude this detail from this change for now. But it's better to discuss here if this approach makes sense.

[wmi]

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

13–14

For example, in SampleProfileLoader::propagateThroughEdges, accesses to BlockWeights are not preceded with a check. Therefore a missing entry in BlockWeights will be created as a zero entry.

IIRC, VisitedBlocks is used to differentiate missing count with zero count BB.

With pseudo probes, a missing sample really means the probe is not executed.
We are leveraging zero sample to represent a special type of probe, called dangling probe, which will need the compiler to infer its count.

I think we discussed it in another patch so I understand for probe, a missing sample means the probe is not executed. It is a little weird to leverage zero sample to tell compiler to infer the count. UINT64_MAX looks better to me because there cannot be any real sample counter with value being UINT64_MAX. In addition, using zero is misleading when we read the profile -- for a location which has zero count, it is actually not zero from compiler's perspective.

[hoy]

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

13–14

IIRC, VisitedBlocks is used to differentiate missing count with zero count BB.

Thanks for pointing that out. I misread that code.

Sounds like in order to reuse the existing inference algorithm, we will need to give probes without any samples an explicit zero count in BlockWeights.

Agreed that using UINT64_MAX looks more clear than using zero.

[wenlei]

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

13–14

It is a little weird to leverage zero sample to tell compiler to infer the count.

Agreed.

Sounds like in order to reuse the existing inference algorithm, we will need to give probes without any samples an explicit zero count in BlockWeights.
Agreed that using UINT64_MAX looks more clear than using zero.

This is counter-intuitive, regardless of the inference algorithm, it'd be good if we can avoid using zero to mark dangling. This is actually one of the TODOs on internal patch, I think it's time to take care of it now.

[wlei]

Removed the dangling probe related code, solution will come in another patch

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

13–14

Thanks for your conversation to let me have a good understanding. So I just removed the zero count related code and the dangling probe issue will be addressed in other diff.

[wlei]

Fix one bug with context id: the inferred callsite context id should be generated from
string splitting since the callee's context id might be compressed they should share the
same context prefix

[wlei]

@wmi Hi, could you take a further look at this patch? another three accepted patches depend on this, thank you!

[wmi]

Thanks for pinging me. LGTM.

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

 ; RUN: llvm-profgen --perfscript=%S/Inputs/inline-cs-pseudoprobe.perfscript --binary=%S/Inputs/inline-cs-pseudoprobe.perfbin --output=%t --show-unwinder-output | FileCheck %s --check-prefix=CHECK-UNWINDER
+; RUN: FileCheck %s --input-file %t
+
+; CHECK:     [main:2 @ foo]:74:0
+; CHECK-NEXT: 2: 15
+; CHECK-NEXT: 3: 15
+; CHECK-NEXT: 4: 14
+; CHECK-NEXT: 5: 1
+; CHECK-NEXT: 6: 15
+; CHECK-NEXT: 8: 14 bar:14
+; CHECK-NEXT: !CFGChecksum: 138950591924
+; CHECK-NEXT:[main:2 @ foo:8 @ bar]:56:14
+; CHECK-NEXT: 1: 14
+; CHECK-NEXT: 2: 14
+; CHECK-NEXT: 3: 14
+; CHECK-NEXT: 4: 14
+; CHECK-NEXT: !CFGChecksum: 72617220756
+
 
 ; CHECK-UNWINDER:      Binary(inline-cs-pseudoprobe.perfbin)'s Range Counter:
 ; CHECK-UNWINDER-EMPTY:
 ; CHECK-UNWINDER-NEXT:   (800, 858): 1
 ; CHECK-UNWINDER-NEXT:   (80e, 82b): 1
 ; CHECK-UNWINDER-NEXT:   (80e, 858): 13
 
 ; CHECK-UNWINDER:      Binary(inline-cs-pseudoprobe.perfbin)'s Branch Counter:

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

 ; RUN: llvm-profgen --perfscript=%S/Inputs/noinline-cs-pseudoprobe.perfscript --binary=%S/Inputs/noinline-cs-pseudoprobe.perfbin --output=%t --show-unwinder-output | FileCheck %s --check-prefix=CHECK-UNWINDER
+; RUN: FileCheck %s --input-file %t
+
+; CHECK:     [main:2 @ foo]:75:0
+; CHECK-NEXT: 2: 15
+; CHECK-NEXT: 3: 15
+; CHECK-NEXT: 4: 15
+; CHECK-NEXT: 6: 15
+; CHECK-NEXT: 8: 15 bar:15
+; CHECK-NEXT: !CFGChecksum: 138950591924
+; CHECK-NEXT:[main:2 @ foo:8 @ bar]:60:15
+; CHECK-NEXT: 1: 15
+; CHECK-NEXT: 2: 15
+; CHECK-NEXT: 3: 15
+; CHECK-NEXT: 4: 15
+; CHECK-NEXT: !CFGChecksum: 72617220756
+
 
 ; CHECK-UNWINDER:      Binary(noinline-cs-pseudoprobe.perfbin)'s Range Counter:
 ; CHECK-UNWINDER-NEXT: main:2
 ; CHECK-UNWINDER-NEXT:   (79e, 7bf): 15
 ; CHECK-UNWINDER-NEXT:   (7c4, 7cf): 15
 ; CHECK-UNWINDER-NEXT: main:2 @ foo:8
 ; CHECK-UNWINDER-NEXT:   (760, 77f): 15
 

llvm/tools/llvm-profgen/PerfReader.cpp

void PerfReader::checkAndSetPerfType(
   for (auto FileName : PerfTraceFilenames) {
     if (!isHybridPerfScript(FileName)) {
       HasHybridPerf = false;
       break;
     }
   }
 
   if (HasHybridPerf) {
-    // Set up ProfileIsCS to enable context-sensitive functionalities
-    // in SampleProf
-    FunctionSamples::ProfileIsCS = true;
     PerfType = PERF_LBR_STACK;
-
   } else {
     // TODO: Support other type of perf script
     PerfType = PERF_INVILID;
   }
 
   if (BinaryTable.size() > 1) {
     // TODO: remove this if everything is ready to support multiple binaries.
     exitWithError("Currently only support one input binary, multiple binaries' "

llvm/tools/llvm-profgen/ProfileGenerator.h

 namespace sampleprof {
 
 class ProfileGenerator {
 
 public:
   ProfileGenerator(){};
   virtual ~ProfileGenerator() = default;
   static std::unique_ptr<ProfileGenerator>
-  create(const BinarySampleCounterMap &SampleCounters,
+  create(const BinarySampleCounterMap &BinarySampleCounters,
          enum PerfScriptType SampleType);
   virtual void generateProfile() = 0;
 
   // Use SampleProfileWriter to serialize profile map
   void write();
 
 protected:
   /*
   For each region boundary point, mark if it is begin or end (or both) of
   the region. Boundary points are inclusive. Log the sample count as well
   so we can use it when we compute the sample count of each disjoint region
   later. Note that there might be multiple ranges with different sample
   count that share same begin/end point. We need to accumulate the sample
   count for the boundary point for such case, because for the example
   below,
 
   |<--100-->|
   |<------200------>|
   A         B       C
 
   sample count for disjoint region [A,B] would be 300.
   */
   void findDisjointRanges(RangeSample &DisjointRanges,
                           const RangeSample &Ranges);
-
   // Used by SampleProfileWriter
   StringMap<FunctionSamples> ProfileMap;
 };
 
 class CSProfileGenerator : public ProfileGenerator {
 protected:
   const BinarySampleCounterMap &BinarySampleCounters;
 
 public:
   CSProfileGenerator(const BinarySampleCounterMap &Counters)
       : BinarySampleCounters(Counters){};
 
 public:
   void generateProfile() override {
+    // Enable context-sensitive functionalities in SampleProf
+    FunctionSamples::ProfileIsCS = true;
     for (const auto &BI : BinarySampleCounters) {
       ProfiledBinary *Binary = BI.first;
       for (const auto &CI : BI.second) {
         const StringBasedCtxKey *CtxKey =
             dyn_cast<StringBasedCtxKey>(CI.first.getPtr());
         StringRef ContextId(CtxKey->Context);
         // Get or create function profile for the range
         FunctionSamples &FunctionProfile =
void generateProfile() override {
     }
     // Fill in call site value sample for inlined calls and also use context to
     // infer missing samples. Since we don't have call count for inlined
     // functions, we estimate it from inlinee's profile using the entry of the
     // body sample.
     populateInferredFunctionSamples();
   }
 
+protected:
+  // Lookup or create FunctionSamples for the context
+  FunctionSamples &getFunctionProfileForContext(StringRef ContextId);
+
 private:
   // Helper function for updating body sample for a leaf location in
   // FunctionProfile
   void updateBodySamplesforFunctionProfile(FunctionSamples &FunctionProfile,
                                            const FrameLocation &LeafLoc,
                                            uint64_t Count);
-  // Lookup or create FunctionSamples for the context
-  FunctionSamples &getFunctionProfileForContext(StringRef ContextId);
   void populateFunctionBodySamples(FunctionSamples &FunctionProfile,
                                    const RangeSample &RangeCounters,
                                    ProfiledBinary *Binary);
   void populateFunctionBoundarySamples(StringRef ContextId,
                                        FunctionSamples &FunctionProfile,
                                        const BranchSample &BranchCounters,
                                        ProfiledBinary *Binary);
   void populateInferredFunctionSamples();
 };
 
+using ProbeCounterMap = std::unordered_map<const PseudoProbe *, uint64_t>;
+
 class PseudoProbeCSProfileGenerator : public CSProfileGenerator {
 
 public:
   PseudoProbeCSProfileGenerator(const BinarySampleCounterMap &Counters)
       : CSProfileGenerator(Counters) {}
-  void generateProfile() override {
-    // TODO
-  }
+  void generateProfile() override;
+
+private:
+  // Go throught each address from range to extract the top frame probe by
+  // looking up in the Address2ProbeMap
+  void extractProbesFromRange(const RangeSample &RangeCounter,
+                              ProbeCounterMap &ProbeCounter,
+                              ProfiledBinary *Binary);
+  // Fill in function body samples from probes
+  void populateBodySamplesWithProbes(const RangeSample &RangeCounter,
+                                     StringRef PrefixContextId,
+                                     ProfiledBinary *Binary);
+  // Fill in boundary samples for a call probe
+  void populateBoundarySamplesWithProbes(const BranchSample &BranchCounter,
+                                         StringRef PrefixContextId,
+                                         ProfiledBinary *Binary);
+  // Helper function to get FunctionSamples for the leaf inlined context
+  FunctionSamples &getFunctionProfileForLeafProbe(
+      StringRef PrefixContextId,
+      SmallVector<std::string, 16> &LeafInlinedContext,
+      const PseudoProbeFuncDesc *LeafFuncDesc);
+  // Helper function to get FunctionSamples for the leaf probe
+  FunctionSamples &getFunctionProfileForLeafProbe(StringRef PrefixContextId,
+                                                  const PseudoProbe *LeafProbe,
+                                                  ProfiledBinary *Binary);
 };
 
 } // end namespace sampleprof
 } // end namespace llvm
 
 #endif

llvm/tools/llvm-profgen/ProfileGenerator.cpp

return BinarySampleCounters.size() &&
          BinarySampleCounters.begin()->first->usePseudoProbes();
 }
 
 std::unique_ptr<ProfileGenerator>
 ProfileGenerator::create(const BinarySampleCounterMap &BinarySampleCounters,
                          enum PerfScriptType SampleType) {
   std::unique_ptr<ProfileGenerator> ProfileGenerator;
   if (SampleType == PERF_LBR_STACK) {
     if (usePseudoProbes(BinarySampleCounters)) {
       ProfileGenerator.reset(
           new PseudoProbeCSProfileGenerator(BinarySampleCounters));
     } else {
       ProfileGenerator.reset(new CSProfileGenerator(BinarySampleCounters));
     }
   } else {
     // TODO:
     llvm_unreachable("Unsupported perfscript!");
   }

 }
 
 void ProfileGenerator::findDisjointRanges(RangeSample &DisjointRanges,
                                           const RangeSample &Ranges) {
 
   /*
   Regions may overlap with each other. Using the boundary info, find all
   disjoint ranges and their sample count. BoundaryPoint contains the count
-  mutiple samples begin/end at this points.
+  multiple samples begin/end at this points.
 
   |<--100-->|           Sample1
   |<------200------>|   Sample2
   A         B       C
 
   In the example above,
   Sample1 begins at A, ends at B, its value is 100.
   Sample2 beings at A, ends at C, its value is 200.
CallerProfile.addCalledTargetSamples(
         EstimatedCallCount);
     CallerProfile.addBodySamples(CallerLeafFrameLoc.second.LineOffset,
                                  CallerLeafFrameLoc.second.Discriminator,
                                  EstimatedCallCount);
     CallerProfile.addTotalSamples(EstimatedCallCount);
   }
 }
 
+// Helper function to extract context prefix
+// PrefixContextId is the context id string except for the leaf probe's
+// context, the final ContextId will be:
+// ContextId =  PrefixContextId + LeafContextId;
+// Remind that the string in ContextStrStack is in callee-caller order
+// So process the string vector reversely
+static std::string
+extractPrefixContextId(const SmallVector<const PseudoProbe *, 16> &Probes,
+                       ProfiledBinary *Binary) {
+  SmallVector<std::string, 16> ContextStrStack;
+  for (const auto *P : Probes) {
+    Binary->getInlineContextForProbe(P, ContextStrStack, true);
+  }
+  std::ostringstream OContextStr;
+  for (auto &CxtStr : ContextStrStack) {
+    if (OContextStr.str().size())
+      OContextStr << " @ ";
+    OContextStr << CxtStr;
+  }
+  return OContextStr.str();
+}
+
+void PseudoProbeCSProfileGenerator::generateProfile() {
+  // Enable CS and pseudo probe functionalities in SampleProf
+  FunctionSamples::ProfileIsCS = true;
+  FunctionSamples::ProfileIsProbeBased = true;
+  for (const auto &BI : BinarySampleCounters) {
+    ProfiledBinary *Binary = BI.first;
+    for (const auto &CI : BI.second) {
+      const ProbeBasedCtxKey *CtxKey =
+          dyn_cast<ProbeBasedCtxKey>(CI.first.getPtr());
+      std::string PrefixContextId =
+          extractPrefixContextId(CtxKey->Probes, Binary);
+      // Fill in function body samples from probes, also infer caller's samples
+      // from callee's probe
+      populateBodySamplesWithProbes(CI.second.RangeCounter, PrefixContextId,
+                                    Binary);
+      // Fill in boundary samples for a call probe
+      populateBoundarySamplesWithProbes(CI.second.BranchCounter,
+                                        PrefixContextId, Binary);
+    }
+  }
+}
+
+void PseudoProbeCSProfileGenerator::extractProbesFromRange(
+    const RangeSample &RangeCounter, ProbeCounterMap &ProbeCounter,
+    ProfiledBinary *Binary) {
+  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;
+
+    while (IP.Address <= RangeEnd) {
+      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;
+        }
+      }
+
+      IP.advance();
+    }
+  }
+}
+
+void PseudoProbeCSProfileGenerator::populateBodySamplesWithProbes(
+    const RangeSample &RangeCounter, StringRef PrefixContextId,
+    ProfiledBinary *Binary) {
+  ProbeCounterMap ProbeCounter;
+  // Extract the top frame probes by looking up each address among the range in
+  // the Address2ProbeMap
+  extractProbesFromRange(RangeCounter, ProbeCounter, Binary);
+  for (auto PI : ProbeCounter) {
+    const PseudoProbe *Probe = PI.first;
+    uint64_t Count = PI.second;
+    FunctionSamples &FunctionProfile =
+        getFunctionProfileForLeafProbe(PrefixContextId, Probe, Binary);
+
+    FunctionProfile.addBodySamples(Probe->Index, 0, Count);
+    FunctionProfile.addTotalSamples(Count);
+    if (Probe->isEntry()) {
+      FunctionProfile.addHeadSamples(Count);
+      // Look up for the caller's function profile
+      SmallVector<std::string, 16> LeafInlinedContext;
+      Binary->getInlineContextForProbe(Probe, LeafInlinedContext);
+      const auto *CallerDesc = Binary->getCallerDescForProbe(Probe);
+      if (CallerDesc != nullptr) {
+        StringRef CallerLoc = LeafInlinedContext.back();
+        uint64_t CallerIndex = 0;
+        CallerLoc.split(":").second.getAsInteger(10, CallerIndex);
+        FunctionSamples &CallerProfile = getFunctionProfileForLeafProbe(
+            PrefixContextId, LeafInlinedContext, CallerDesc);
+        CallerProfile.addBodySamples(CallerIndex, 0, Count);
+        CallerProfile.addTotalSamples(Count);
+        CallerProfile.addCalledTargetSamples(CallerIndex, 0,
+                                             FunctionProfile.getName(), Count);
+      }
+    }
+  }
+}
+
+void PseudoProbeCSProfileGenerator::populateBoundarySamplesWithProbes(
+    const BranchSample &BranchCounter, StringRef PrefixContextId,
+    ProfiledBinary *Binary) {
+  for (auto BI : BranchCounter) {
+    uint64_t SourceOffset = BI.first.first;
+    uint64_t TargetOffset = BI.first.second;
+    uint64_t Count = BI.second;
+    uint64_t SourceAddress = Binary->offsetToVirtualAddr(SourceOffset);
+    const PseudoProbe *CallProbe = Binary->getCallProbeForAddr(SourceAddress);
+    if (CallProbe == nullptr)
+      continue;
+    FunctionSamples &FunctionProfile =
+        getFunctionProfileForLeafProbe(PrefixContextId, CallProbe, Binary);
+    FunctionProfile.addBodySamples(CallProbe->Index, 0, Count);
+    FunctionProfile.addTotalSamples(Count);
+    StringRef CalleeName = FunctionSamples::getCanonicalFnName(
+        Binary->getFuncFromStartOffset(TargetOffset));
+    if (CalleeName.size() == 0)
+      continue;
+    FunctionProfile.addCalledTargetSamples(CallProbe->Index, 0, CalleeName,
+                                           Count);
+  }
+}
+
+FunctionSamples &PseudoProbeCSProfileGenerator::getFunctionProfileForLeafProbe(
+    StringRef PrefixContextId, SmallVector<std::string, 16> &LeafInlinedContext,
+    const PseudoProbeFuncDesc *LeafFuncDesc) {
+  assert(LeafInlinedContext.size() &&
+         "Profile context must have the leaf frame");
+  std::ostringstream OContextStr;
+  OContextStr << PrefixContextId.str();
+
+  for (uint32_t I = 0; I < LeafInlinedContext.size() - 1; I++) {
+    if (OContextStr.str().size())
+      OContextStr << " @ ";
+    OContextStr << LeafInlinedContext[I];
+  }
+  // For leaf inlined context with the top frame, we should strip off the top
+  // frame's probe id, like:
+  // Inlined stack: [foo:1, bar:2], the ContextId will be "foo:1 @ bar"
+  if (OContextStr.str().size())
+    OContextStr << " @ ";
+  StringRef LeafLoc = LeafInlinedContext.back();
+  OContextStr << LeafLoc.split(":").first.str();
+
+  FunctionSamples &FunctionProile =
+      getFunctionProfileForContext(OContextStr.str());
+  FunctionProile.setFunctionHash(LeafFuncDesc->FuncHash);
+  return FunctionProile;
+}
+
+FunctionSamples &PseudoProbeCSProfileGenerator::getFunctionProfileForLeafProbe(
+    StringRef PrefixContextId, const PseudoProbe *LeafProbe,
+    ProfiledBinary *Binary) {
+  SmallVector<std::string, 16> LeafInlinedContext;
+  Binary->getInlineContextForProbe(LeafProbe, LeafInlinedContext);
+  // Note that the context from probe doesn't include leaf frame,
+  // hence we need to retrieve and append the leaf frame.
+  const auto *FuncDesc = Binary->getFuncDescForGUID(LeafProbe->GUID);
+  LeafInlinedContext.emplace_back(FuncDesc->FuncName + ":" +
+                                  Twine(LeafProbe->Index).str());
+  return getFunctionProfileForLeafProbe(PrefixContextId, LeafInlinedContext,
+                                        FuncDesc);
+}
+
 } // end namespace sampleprof
 } // end namespace llvm

llvm/tools/llvm-profgen/ProfiledBinary.h

public:
   std::string getExpandedContextStr(const std::list<uint64_t> &stack) const;
 
   const PseudoProbe *getCallProbeForAddr(uint64_t Address) const {
     return ProbeDecoder.getCallProbeForAddr(Address);
   }
   void
   getInlineContextForProbe(const PseudoProbe *Probe,
                            SmallVector<std::string, 16> &InlineContextStack,
-                           bool IncludeLeaf) const {
+                           bool IncludeLeaf = false) const {
     return ProbeDecoder.getInlineContextForProbe(Probe, InlineContextStack,
                                                  IncludeLeaf);
   }
+  const AddressProbesMap &getAddress2ProbesMap() const {
+    return ProbeDecoder.getAddress2ProbesMap();
+  }
+  const PseudoProbeFuncDesc *getFuncDescForGUID(uint64_t GUID) {
+    return ProbeDecoder.getFuncDescForGUID(GUID);
+  }
+  const PseudoProbeFuncDesc *getCallerDescForProbe(const PseudoProbe *Probe) {
+    return ProbeDecoder.getCallerDescForProbe(Probe);
+  }
 };
 
 } // end namespace sampleprof
 } // end namespace llvm
 
 #endif

llvm/tools/llvm-profgen/PseudoProbe.h

public:
   void printGUID2FuncDescMap(raw_ostream &OS);
 
   // Print pseudo_probe section info, used along with show-disassembly
   void printProbeForAddress(raw_ostream &OS, uint64_t Address);
 
   // Look up the probe of a call for the input address
   const PseudoProbe *getCallProbeForAddr(uint64_t Address) const;
 
+  const PseudoProbeFuncDesc *getFuncDescForGUID(uint64_t GUID) const;
+
   // Helper function to populate one probe's inline stack into
   // \p InlineContextStack.
   // Current leaf location info will be added if IncludeLeaf is true
   // Example:
   //  Current probe(bar:3) inlined at foo:2 then inlined at main:1
   //  IncludeLeaf = true,  Output: [main:1, foo:2, bar:3]
-  //  IncludeLeaf = false, OUtput: [main:1, foo:2]
+  //  IncludeLeaf = false, Output: [main:1, foo:2]
   void
   getInlineContextForProbe(const PseudoProbe *Probe,
                            SmallVector<std::string, 16> &InlineContextStack,
                            bool IncludeLeaf) const;
+
+  const AddressProbesMap &getAddress2ProbesMap() const {
+    return Address2ProbesMap;
+  }
+
+  const PseudoProbeFuncDesc *
+  getCallerDescForProbe(const PseudoProbe *Probe) const;
 };
 
 } // end namespace sampleprof
 } // end namespace llvm
 
 #endif

llvm/tools/llvm-profgen/PseudoProbe.cpp

if (Probe.isCall()) {
              "There should be only one call probe corresponding to address "
              "which is a callsite.");
       CallProbe = &Probe;
     }
   }
   return CallProbe;
 }
 
+const PseudoProbeFuncDesc *
+PseudoProbeDecoder::getFuncDescForGUID(uint64_t GUID) const {
+  auto It = GUID2FuncDescMap.find(GUID);
+  assert(It != GUID2FuncDescMap.end() && "Function descriptor doesn't exist");
+  return &It->second;
+}
+
 void PseudoProbeDecoder::getInlineContextForProbe(
     const PseudoProbe *Probe, SmallVector<std::string, 16> &InlineContextStack,
     bool IncludeLeaf) const {
-  if (IncludeLeaf) {
-    // Note that the context from probe doesn't include leaf frame,
-    // hence we need to retrieve and prepend leaf if requested.
-    auto It = GUID2FuncDescMap.find(Probe->GUID);
-    assert(It != GUID2FuncDescMap.end() &&
-           "Should have function descriptor for a valid GUID");
-    StringRef FuncName = It->second.FuncName;
-    // InlineContextStack is in callee-caller order, so push leaf in the front
-    InlineContextStack.emplace_back(FuncName.str() + ":" +
-                                    Twine(Probe->Index).str());
-  }
+  Probe->getInlineContext(InlineContextStack, GUID2FuncDescMap, true);
+  if (!IncludeLeaf)
+    return;
+  const auto *FuncDesc = getFuncDescForGUID(Probe->GUID);
+  InlineContextStack.emplace_back(FuncDesc->FuncName + ":" +
 
-  Probe->getInlineContext(InlineContextStack, GUID2FuncDescMap, true);
+const PseudoProbeFuncDesc *
+PseudoProbeDecoder::getCallerDescForProbe(const PseudoProbe *Probe) const {
+  PseudoProbeInlineTree *CallerNode = Probe->InlineTree;
+  if (std::get<0>(CallerNode->ISite) == 0) {
+    return nullptr;
+  }
+  return getFuncDescForGUID(std::get<0>(CallerNode->ISite));
 }
 
 } // end namespace sampleprof
 } // end namespace llvm