[CSSPGO][llvm-profgen] Virtual unwinding with pseudo probe

[wlei]

Summary

This change extends virtual unwinder to support pseudo probe in llvm-profgen. 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

• Added ProbeBasedCtxKey derived from ContextKey for sample counter aggregation. As we need string splitting to infer the profile for callee function, string based context introduces more string handling overhead, here we just use probe pointer based context.
• For linear unwinding, as inline context is encoded in each pseudo probe, we don't need to go through each instruction to extract range sharing same inliner. So just record the range for the context.
• For probe based context, we should ignore the top frame probe since it will be extracted from the address range. we defer the extraction in ProfileGeneration.
• Added PseudoProbeProfileGenerator for pseudo probe based profile generation.
• Some helper function to get pseduo probe info(call probe, inline context) from profiled binary.
• Added regression test for unwinder's output

The pseudo probe based profile generation will be in the upcoming patch.

Test Plan:

ninja & ninja check-llvm

[wlei]

rebase

[hoy]

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

2

Please remove this line here and in other .perfscript files.

llvm/tools/llvm-profgen/PerfReader.cpp

115

for (auto Address : reverse(CallStack)) ...

273–293

This can return StringRef when a ProbeBasedCtxKey holds a string data field.

llvm/tools/llvm-profgen/PerfReader.h

285

How about make a string field here to avoid string copying like in getContextKeyStr? Also the pseudo probe related code in getContextKeyStr can be moved here as an toString() API.

297

use std::equal?

309

assert Probes are empty here?

llvm/tools/llvm-profgen/ProfileGenerator.h

111

Name it PseudoProbeCSProfileGenerator ?

[wlei]

llvm/tools/llvm-profgen/PerfReader.cpp

115

seems it's a bug, it should not be reversed, let me fix.

llvm/tools/llvm-profgen/PerfReader.h

285

Good point. One question I got is we might not use string here directly due to context compression. The compression will implements on the vector of string/probe, but here we miss the top frame probe. So I just want to redesign to use string vector as the output of virtual unwinding. In ProfileGenerator, the top frame context string will be pushed and fed for the compression to generate the string. e.g:

[a:1, b:2, c:3]               virtual unwinding
[a:1, b:2, c:3, b2: d:4]   push back top frame context
[a:1, b:2, d:4]               compression
"a:1 @ b:2 @ d:4"       string joining

309

Same as above, the Probes can be empty since it doesn't include the top frame probe.

llvm/tools/llvm-profgen/ProfileGenerator.h

111

Changed. I see that in our internal tool, we use a switch variable to decide whether to gen CS profile. Here not sure whether we need to merge them in a general class or we should separate in different class. but I guess we can leave it when come to non-cs implementation.

[wlei]

llvm/tools/llvm-profgen/PerfReader.cpp

115

it's not a bug, the call stack is leaf to root order, should be reversed.
But here note that I need to skip the leaf frame. so the code will be like:

bool IsLeaf = true;
for (auto Address : reverse(CallStack)) {
    if (IsLeaf) {
        IsLeaf = false;
        continue;
   }
   ...
}

Btw, I didn't use reverse like this way before, so it will act as a iterator wrapper, won't reverse the vector's inside elements, right?

[wlei]

Address some feedback from hongtao. The string handling part will be redesigned and come later

[hoy]

llvm/tools/llvm-profgen/PerfReader.cpp

115

Yes, see llvm/include/llvm/ADT/STLExtras.h

// Returns an iterator_range over the given container which iterates in reverse.
// Note that the container must have rbegin()/rend() methods for this to work.
template <typename ContainerTy>
auto reverse(ContainerTy &&C,
             std::enable_if_t<has_rbegin<ContainerTy>::value> * = nullptr) {
  return make_range(C.rbegin(), C.rend());
}

llvm/tools/llvm-profgen/PerfReader.h

285

I see. Yeah, compression should be done on vectors. Will the string form computed only once after compression is done? I was wondering if you ever need the string more than once, making a field and fill it out after compression is probably good. But if the sting is only needed once, what you have now is good enough.

309

Sure, I mean to add an assertion for it. Sorry for the confusion.

[wlei]

rebase and fix clang-tidy

[wlei]

llvm/tools/llvm-profgen/PerfReader.cpp

115

As we discussed offline, we need to skip the top frame for this, so just keep using the iterator way.

llvm/tools/llvm-profgen/PerfReader.h

285

Yeah, after compression the string is only used once, so I will keep this.

309

Oh, I mean the probes doesn't include the top frame probe, so the probes will have the change to be empty.
See line 98, it has the assertion to check whether the Hashcode is 0

[wmi]

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

4

Add CHECK for context similar as noinline-cs-pseudoprobe.test below? It is good to make sure we get the context right w/wo inline.

llvm/tools/llvm-profgen/PerfReader.cpp

102

Why it needs to be a shared_ptr?

122

Nit: !CallProbe

138–139

How about move the conditional statement into getOrCreateCounter?

llvm/tools/llvm-profgen/PerfReader.h

292

Add an assertion for O != nullptr?

379

pseduo => pseudo

llvm/tools/llvm-profgen/PseudoProbe.cpp

306

Nit: CallProbe == nullptr --> !CallProbe

322

Add an assertion message.

[wlei]

Address Wei's feedback

[wlei]

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

4

Yeah, this is because the context is empty. added the CHECK-UNWINDER-EMPTY for this.

llvm/tools/llvm-profgen/PerfReader.cpp

102

This is because the Hashable interface is using shared_ptr to allow being the key of unordered_map, I haven't found a workable way for unique_ptr yet, could we leave for later refining? I will keep an eye on this.

122

fixed!

138–139

Good suggestion, moved getOrCreateCounterForProbe into getOrCreateCounter

llvm/tools/llvm-profgen/PerfReader.h

292

assertion added

379

fixed

llvm/tools/llvm-profgen/PseudoProbe.cpp

306

fixed

322

assertion added

[wmi]

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

4

Why the context is empty? To my understanding, for CSSPGO profile, the profile context will be the same no matter how the inlining is happening for the profiled binary, so I expect to see the context just like noinline-cs-pseudoprobe.test. What am I missing?

main:2
...
main:2 @ foo:8
...

llvm/tools/llvm-profgen/PerfReader.cpp

102

I see. Thanks.

[wlei]

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

4

This is not the final context, the final context won't be empty. This is the output for the unwinder, has the context except for the leaf context. For unwinder, we record the leaf frame probe into the ranges, so the leaf context is not generated. In the profile generation patch(https://reviews.llvm.org/D92998), the range will be extract and the leaf context will be generated and appended.

Here for this case, because of inline, the contexts like [main @ foo @ bar] are all the leaf context, so the unwinder's context is empty.

[wmi]

LGTM.

llvm/test/tools/llvm-profgen/Inputs/inline-cs-noprobe.perfscript

-Using perf wrapper that supports hot-text. Try perf.real if you encounter any issues.
 PERF_RECORD_MMAP2 2854748/2854748: [0x400000(0x1000) @ 0 00:1d 123291722 526021]: r-xp /home/inline-cs-noprobe.perfbin
 
 
             40067e
   5541f689495641d7
  0x4006c8/0x40067e/P/-/-/0  0x4006c8/0x40067e/P/-/-/0  0x4006c8/0x40067e/P/-/-/0  0x4006c8/0x40067e/P/-/-/0  0x4006c8/0x40067e/P/-/-/0  0x4006c8/0x40067e/P/-/-/0  0x4006c8/0x40067e/P/-/-/0  0x4006c8/0x40067e/P/-/-/0  0x4006c8/0x40067e/P/-/-/0  0x4006c8/0x40067e/P/-/-/0  0x4006c8/0x40067e/P/-/-/0  0x4006c8/0x40067e/P/-/-/0  0x40069b/0x400670/M/-/-/0  0x4006c8/0x40067e/P/-/-/0  0x4006c8/0x40067e/P/-/-/0  0x4006c8/0x40067e/P/-/-/0

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

+PERF_RECORD_MMAP2 595196/595196: [0x201000(0x1000) @ 0 00:1d 224227621 1042948]: r-xp /home/inline-cs-pseudoprobe.perfbin
+
+            20180e
+  5541f689495641d7
+ 0x201858/0x20180e/P/-/-/0  0x201858/0x20180e/P/-/-/0  0x201858/0x20180e/P/-/-/0  0x201858/0x20180e/P/-/-/0  0x201858/0x20180e/P/-/-/0  0x201858/0x20180e/P/-/-/0  0x201858/0x20180e/P/-/-/0  0x201858/0x20180e/P/-/-/0  0x201858/0x20180e/P/-/-/0  0x201858/0x20180e/P/-/-/0  0x201858/0x20180e/P/-/-/0  0x20182b/0x201800/M/-/-/0  0x201858/0x20180e/P/-/-/0  0x201858/0x20180e/P/-/-/0  0x201858/0x20180e/P/-/-/0  0x201858/0x20180e/P/-/-/0

llvm/test/tools/llvm-profgen/Inputs/noinline-cs-noprobe.perfscript

-Using perf wrapper that supports hot-text. Try perf.real if you encounter any issues.
 PERF_RECORD_MMAP2 2854748/2854748: [0x400000(0x1000) @ 0 00:1d 123291722 526021]: r-xp /home/noinline-cs-noprobe.perfbin
 
             4005dc
             400634
             400684
       7f68c5788793
  0x4005c8/0x4005dc/P/-/-/0  0x40062f/0x4005b0/P/-/-/0  0x400645/0x4005ff/P/-/-/0  0x400637/0x400645/P/-/-/0  0x4005e9/0x400634/P/-/-/0  0x4005d7/0x4005e5/P/-/-/0  0x40062f/0x4005b0/P/-/-/0  0x400645/0x4005ff/P/-/-/0  0x400637/0x400645/P/-/-/0  0x4005e9/0x400634/P/-/-/0  0x4005d7/0x4005e5/P/-/-/0  0x40062f/0x4005b0/P/-/-/0  0x400645/0x4005ff/P/-/-/0  0x400637/0x400645/P/-/-/0  0x4005e9/0x400634/P/-/-/0  0x4005c8/0x4005dc/P/-/-/0
 

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

+PERF_RECORD_MMAP2 1243676/1243676: [0x201000(0x1000) @ 0 00:1d 224517108 1044165]: r-xp /home/noinline-cs-pseudoprobe.perfbin
+
+            20179e
+            2017f9
+      7f83e84e7793
+  5541f689495641d7
+ 0x2017cf/0x20179e/P/-/-/0  0x20177f/0x2017c4/P/-/-/0  0x2017bf/0x201760/P/-/-/0  0x2017cf/0x20179e/P/-/-/0  0x20177f/0x2017c4/P/-/-/0  0x2017bf/0x201760/P/-/-/0  0x2017cf/0x20179e/P/-/-/0  0x20177f/0x2017c4/P/-/-/0  0x2017bf/0x201760/P/-/-/0  0x2017cf/0x20179e/P/-/-/0  0x20177f/0x2017c4/P/-/-/0  0x2017bf/0x201760/P/-/-/0  0x2017cf/0x20179e/P/-/-/0  0x20177f/0x2017c4/P/-/-/0  0x2017bf/0x201760/P/-/-/0  0x2017cf/0x20179e/P/-/-/0
+
+            2017c4
+            2017f9
+      7f83e84e7793
+  5541f689495641d7
+ 0x20177f/0x2017c4/P/-/-/0  0x2017bf/0x201760/P/-/-/0  0x2017cf/0x20179e/P/-/-/0  0x20177f/0x2017c4/P/-/-/0  0x2017bf/0x201760/P/-/-/0  0x2017cf/0x20179e/P/-/-/0  0x20177f/0x2017c4/P/-/-/0  0x2017bf/0x201760/P/-/-/0  0x2017cf/0x20179e/P/-/-/0  0x20177f/0x2017c4/P/-/-/0  0x2017bf/0x201760/P/-/-/0  0x2017cf/0x20179e/P/-/-/0  0x20177f/0x2017c4/P/-/-/0  0x2017bf/0x201760/P/-/-/0  0x2017cf/0x20179e/P/-/-/0  0x20177f/0x2017c4/P/-/-/0
+
+            2017c4
+            2017f9
+      7f83e84e7793
+  5541f689495641d7
+ 0x20177f/0x2017c4/P/-/-/0  0x2017bf/0x201760/P/-/-/0  0x2017cf/0x20179e/P/-/-/0  0x20177f/0x2017c4/P/-/-/0  0x2017bf/0x201760/P/-/-/0  0x2017cf/0x20179e/P/-/-/0  0x20177f/0x2017c4/P/-/-/0  0x2017bf/0x201760/P/-/-/0  0x2017cf/0x20179e/P/-/-/0  0x20177f/0x2017c4/P/-/-/0  0x2017bf/0x201760/P/-/-/0  0x2017cf/0x20179e/P/-/-/0  0x20177f/0x2017c4/P/-/-/0  0x2017bf/0x201760/P/-/-/0  0x2017cf/0x20179e/P/-/-/0  0x20177f/0x2017c4/P/-/-/0

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
+
+; 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:
+; CHECK-UNWINDER-EMPTY:
+; CHECK-UNWINDER-NEXT:   (82b, 800): 1
+; CHECK-UNWINDER-NEXT:   (858, 80e): 15
+
+; 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-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
+
+; 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
+
+; CHECK-UNWINDER:      Binary(noinline-cs-pseudoprobe.perfbin)'s Branch Counter:
+; CHECK-UNWINDER-NEXT: main:2
+; CHECK-UNWINDER-NEXT:   (7bf, 760): 15
+; CHECK-UNWINDER-NEXT:   (7cf, 79e): 16
+; CHECK-UNWINDER-NEXT: main:2 @ foo:8
+; CHECK-UNWINDER-NEXT:   (77f, 7c4): 17
+
+
+; 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/PerfReader.h

struct ContextKey {
   uint64_t HashCode = 0;
   virtual ~ContextKey() = default;
   uint64_t getHashCode() const { return HashCode; }
   virtual bool isEqual(const ContextKey *K) const {
     return HashCode == K->HashCode;
   };
 
   // Utilities for LLVM-style RTTI
-  enum ContextKind { CK_StringBased };
+  enum ContextKind { CK_StringBased, CK_ProbeBased };
   const ContextKind Kind;
   ContextKind getKind() const { return Kind; }
   ContextKey(ContextKind K) : Kind(K){};
 };
 
 // String based context id
 struct StringBasedCtxKey : public ContextKey {
   std::string Context;
   StringBasedCtxKey() : ContextKey(CK_StringBased){};
   static bool classof(const ContextKey *K) {
     return K->getKind() == CK_StringBased;
   }
 
   bool isEqual(const ContextKey *K) const override {
     const StringBasedCtxKey *Other = dyn_cast<StringBasedCtxKey>(K);
     return Context == Other->Context;
   }
 
   void genHashCode() { HashCode = hash_value(Context); }
 };
 
+// Probe based context key as the intermediate key of context
+// String based context key will introduce redundant string handling
+// since the callee context is inferred from the context string which
+// need to be splitted by '@' to get the last location frame, so we
+// can just use probe instead and generate the string in the end.
+struct ProbeBasedCtxKey : public ContextKey {
+  SmallVector<const PseudoProbe *, 16> Probes;
+
+  ProbeBasedCtxKey() : ContextKey(CK_ProbeBased) {}
+  static bool classof(const ContextKey *K) {
+    return K->getKind() == CK_ProbeBased;
+  }
+
+  bool isEqual(const ContextKey *K) const override {
+    const ProbeBasedCtxKey *O = dyn_cast<ProbeBasedCtxKey>(K);
+    assert(O != nullptr && "Probe based key shouldn't be null in isEqual");
+    return std::equal(Probes.begin(), Probes.end(), O->Probes.begin(),
+                      O->Probes.end());
+  }
+
+  void genHashCode() {
+    for (const auto *P : Probes) {
+      HashCode = hash_combine(HashCode, P);
+    }
+    if (HashCode == 0) {
+      // Avoid zero value of HashCode when it's an empty list
+      HashCode = 1;
+    }
+  }
+};
+
 // The counter of branch samples for one function indexed by the branch,
 // which is represented as the source and target offset pair.
 using BranchSample = std::map<std::pair<uint64_t, uint64_t>, uint64_t>;
 // The counter of range samples for one function indexed by the range,
 // which is represented as the start and end offset pair.
 using RangeSample = std::map<std::pair<uint64_t, uint64_t>, uint64_t>;
 // Wrapper for sample counters including range counter and branch counter
 struct SampleCounter {
   RangeSample RangeCounter;
public:
   void unwindLinear(UnwindState &State, uint64_t Repeat);
   void unwindReturn(UnwindState &State);
   void unwindBranchWithinFrame(UnwindState &State);
   bool unwind(const HybridSample *Sample, uint64_t Repeat);
   void recordRangeCount(uint64_t Start, uint64_t End, UnwindState &State,
                         uint64_t Repeat);
   void recordBranchCount(const LBREntry &Branch, UnwindState &State,
                          uint64_t Repeat);
-  SampleCounter &getOrCreateSampleCounter(const ProfiledBinary *Binary,
+  SampleCounter &getOrCreateCounter(const ProfiledBinary *Binary,
+                                    std::list<uint64_t> &CallStack);
+  // Use pseudo probe based context key to get the sample counter
+  // A context stands for a call path from 'main' to an uninlined
+  // callee with all inline frames recovered on that path. The probes
+  // belonging to that call path is the probes either originated from
+  // the callee or from any functions inlined into the callee. Since
+  // pseudo probes are organized in a tri-tree style after decoded,
+  // the tree path from the tri-tree root (which is the uninlined
+  // callee) to the probe node forms an inline context.
+  // Here we use a list of probe(pointer) as the context key to speed up
+  // aggregation and the final context string will be generate in
+  // ProfileGenerator
+  SampleCounter &getOrCreateCounterForProbe(const ProfiledBinary *Binary,
                                             std::list<uint64_t> &CallStack);
 
 private:
   ContextSampleCounterMap *CtxCounterMap;
 };
 
 // Filename to binary map
 using BinaryMap = StringMap<ProfiledBinary>;
 // Address to binary map for fast look-up

llvm/tools/llvm-profgen/PerfReader.cpp

void VirtualUnwinder::unwindCall(UnwindState &State) {
   }
   State.InstPtr.update(Source);
 }
 
 void VirtualUnwinder::unwindLinear(UnwindState &State, uint64_t Repeat) {
   InstructionPointer &IP = State.InstPtr;
   uint64_t Target = State.getCurrentLBRTarget();
   uint64_t End = IP.Address;
+  if (State.getBinary()->usePseudoProbes()) {
+    // The outcome of the virtual unwinding with pseudo probes is a
+    // map from a context key to the address range being unwound.
+    // This means basically linear unwinding is not needed for pseudo
+    // probes. The range will be simply recorded here and will be
+    // converted to a list of pseudo probes to report in ProfileGenerator.
+    recordRangeCount(Target, End, State, Repeat);
+  } else {
     // Unwind linear execution part
     while (IP.Address >= Target) {
       uint64_t PrevIP = IP.Address;
       IP.backward();
       // Break into segments for implicit call/return due to inlining
       bool SameInlinee =
           State.getBinary()->inlineContextEqual(PrevIP, IP.Address);
       if (!SameInlinee || PrevIP == Target) {
         recordRangeCount(PrevIP, End, State, Repeat);
         End = IP.Address;
       }
       State.CallStack.front() = IP.Address;
     }
   }
+}
 
 void VirtualUnwinder::unwindReturn(UnwindState &State) {
   // Add extra frame as we unwind through the return
   const LBREntry &LBR = State.getCurrentLBR();
   uint64_t CallAddr = State.getBinary()->getCallAddrFromFrameAddr(LBR.Target);
   State.CallStack.front() = CallAddr;
   State.CallStack.push_front(LBR.Source);
   State.InstPtr.update(LBR.Source);
 }
 
 void VirtualUnwinder::unwindBranchWithinFrame(UnwindState &State) {
   // TODO: Tolerate tail call for now, as we may see tail call from libraries.
   // This is only for intra function branches, excluding tail calls.
   uint64_t Source = State.getCurrentLBRSource();
   State.CallStack.front() = Source;
   State.InstPtr.update(Source);
 }
 
 SampleCounter &
-VirtualUnwinder::getOrCreateSampleCounter(const ProfiledBinary *Binary,
-                                          std::list<uint64_t> &CallStack) {
+VirtualUnwinder::getOrCreateCounter(const ProfiledBinary *Binary,
+  if (Binary->usePseudoProbes()) {
+    return getOrCreateCounterForProbe(Binary, CallStack);
+  }
   std::shared_ptr<StringBasedCtxKey> KeyStr =
       std::make_shared<StringBasedCtxKey>();
   KeyStr->Context = Binary->getExpandedContextStr(CallStack);
   KeyStr->genHashCode();
   auto Ret =
       CtxCounterMap->emplace(Hashable<ContextKey>(KeyStr), SampleCounter());
   return Ret.first->second;
 }
 
+SampleCounter &
+VirtualUnwinder::getOrCreateCounterForProbe(const ProfiledBinary *Binary,
+                                            std::list<uint64_t> &CallStack) {
+  std::shared_ptr<ProbeBasedCtxKey> ProbeBasedKey =
+      std::make_shared<ProbeBasedCtxKey>();
+  if (CallStack.size() > 1) {
+    // We don't need to top frame probe since it should be extracted
+    // from the range.
+    // The top of stack is an instruction from the function where
+    // the LBR address range physcially resides. Strip it since
+    // the function is not a part of the call context. We also
+    // don't need its inline context since the probes being unwound
+    // come with an inline context all the way back to the uninlined
+    // function in their prefix tree.
+    auto Iter = CallStack.rbegin();
+    auto EndT = std::prev(CallStack.rend());
+    for (; Iter != EndT; Iter++) {
+      uint64_t Address = *Iter;
+      const PseudoProbe *CallProbe = Binary->getCallProbeForAddr(Address);
+      // We may not find a probe for a merged or external callsite.
+      // Callsite merging may cause the loss of original probe IDs.
+      // Cutting off the context from here since the inline will
+      // not know how to consume a context with unknown callsites.
+      if (!CallProbe)
+        break;
+      ProbeBasedKey->Probes.emplace_back(CallProbe);
+    }
+  }
+  ProbeBasedKey->genHashCode();
+  Hashable<ContextKey> ContextId(ProbeBasedKey);
+  auto Ret = CtxCounterMap->emplace(ContextId, SampleCounter());
+  return Ret.first->second;
+}
+
 void VirtualUnwinder::recordRangeCount(uint64_t Start, uint64_t End,
                                        UnwindState &State, uint64_t Repeat) {
   uint64_t StartOffset = State.getBinary()->virtualAddrToOffset(Start);
   uint64_t EndOffset = State.getBinary()->virtualAddrToOffset(End);
   SampleCounter &SCounter =
-      getOrCreateSampleCounter(State.getBinary(), State.CallStack);
+      getOrCreateCounter(State.getBinary(), State.CallStack);
   SCounter.recordRangeCount(StartOffset, EndOffset, Repeat);
 }
 
 void VirtualUnwinder::recordBranchCount(const LBREntry &Branch,
                                         UnwindState &State, uint64_t Repeat) {
   if (Branch.IsArtificial)
     return;
   uint64_t SourceOffset = State.getBinary()->virtualAddrToOffset(Branch.Source);
   uint64_t TargetOffset = State.getBinary()->virtualAddrToOffset(Branch.Target);
   SampleCounter &SCounter =
-      getOrCreateSampleCounter(State.getBinary(), State.CallStack);
+      getOrCreateCounter(State.getBinary(), State.CallStack);
   SCounter.recordBranchCount(SourceOffset, TargetOffset, Repeat);
 }
 
 bool VirtualUnwinder::unwind(const HybridSample *Sample, uint64_t Repeat) {
   // Capture initial state as starting point for unwinding.
   UnwindState State(Sample);
 
   // Sanity check - making sure leaf of LBR aligns with leaf of stack sample
if (Iter == AddrToBinaryMap.end() || Iter->first != Address) {
     if (Iter == AddrToBinaryMap.begin())
       return nullptr;
     Iter--;
   }
   return Iter->second;
 }
 
 // Use ordered map to make the output deterministic
-using OrderedCounterForPrint = std::map<StringRef, RangeSample>;
+using OrderedCounterForPrint = std::map<std::string, RangeSample>;
 
 static void printSampleCounter(OrderedCounterForPrint &OrderedCounter) {
   for (auto Range : OrderedCounter) {
     outs() << Range.first << "\n";
     for (auto I : Range.second) {
       outs() << "  (" << format("%" PRIx64, I.first.first) << ", "
              << format("%" PRIx64, I.first.second) << "): " << I.second << "\n";
     }
   }
 }
 
-static void printRangeCounter(ContextSampleCounterMap &Counter) {
+static std::string getContextKeyStr(ContextKey *K,
+                                    const ProfiledBinary *Binary) {
+  std::string ContextStr;
+  if (const auto *CtxKey = dyn_cast<StringBasedCtxKey>(K)) {
+    return CtxKey->Context;
+  } else if (const auto *CtxKey = dyn_cast<ProbeBasedCtxKey>(K)) {
+    SmallVector<std::string, 16> ContextStack;
+    for (const auto *Probe : CtxKey->Probes) {
+      Binary->getInlineContextForProbe(Probe, ContextStack, true);
+    }
+    for (const auto &Context : ContextStack) {
+      if (ContextStr.size())
+        ContextStr += " @ ";
+      ContextStr += Context;
+    }
+  }
+  return ContextStr;
+}
+
+static void printRangeCounter(ContextSampleCounterMap &Counter,
+                              const ProfiledBinary *Binary) {
   OrderedCounterForPrint OrderedCounter;
   for (auto &CI : Counter) {
-    const StringBasedCtxKey *CtxKey =
-        dyn_cast<StringBasedCtxKey>(CI.first.getPtr());
-    OrderedCounter[CtxKey->Context] = CI.second.RangeCounter;
+    OrderedCounter[getContextKeyStr(CI.first.getPtr(), Binary)] =
+        CI.second.RangeCounter;
   }
   printSampleCounter(OrderedCounter);
 }
 
-static void printBranchCounter(ContextSampleCounterMap &Counter) {
+static void printBranchCounter(ContextSampleCounterMap &Counter,
+                               const ProfiledBinary *Binary) {
   OrderedCounterForPrint OrderedCounter;
   for (auto &CI : Counter) {
-    const StringBasedCtxKey *CtxKey =
-        dyn_cast<StringBasedCtxKey>(CI.first.getPtr());
-    OrderedCounter[CtxKey->Context] = CI.second.BranchCounter;
+    OrderedCounter[getContextKeyStr(CI.first.getPtr(), Binary)] =
+        CI.second.BranchCounter;
   }
   printSampleCounter(OrderedCounter);
 }
 
 void PerfReader::printUnwinderOutput() {
   for (auto I : BinarySampleCounters) {
     const ProfiledBinary *Binary = I.first;
     outs() << "Binary(" << Binary->getName().str() << ")'s Range Counter:\n";
-    printRangeCounter(I.second);
+    printRangeCounter(I.second, Binary);
     outs() << "\nBinary(" << Binary->getName().str() << ")'s Branch Counter:\n";
-    printBranchCounter(I.second);
+    printBranchCounter(I.second, Binary);
   }
 }
 
 void PerfReader::unwindSamples() {
   for (const auto &Item : AggregatedSamples) {
     const HybridSample *Sample = dyn_cast<HybridSample>(Item.first.getPtr());
     VirtualUnwinder Unwinder(&BinarySampleCounters[Sample->Binary]);
     Unwinder.unwind(Sample, Item.second);

llvm/tools/llvm-profgen/ProfileGenerator.h

protected:
   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 {
void populateFunctionBodySamples(FunctionSamples &FunctionProfile,
                                    ProfiledBinary *Binary);
   void populateFunctionBoundarySamples(StringRef ContextId,
                                        FunctionSamples &FunctionProfile,
                                        const BranchSample &BranchCounters,
                                        ProfiledBinary *Binary);
   void populateInferredFunctionSamples();
 };
 
+class PseudoProbeCSProfileGenerator : public CSProfileGenerator {
+
+public:
+  PseudoProbeCSProfileGenerator(const BinarySampleCounterMap &Counters)
+      : CSProfileGenerator(Counters) {}
+  void generateProfile() override {
+    // TODO
+  }
+};
+
 } // end namespace sampleprof
 } // end namespace llvm
 
 #endif

llvm/tools/llvm-profgen/ProfileGenerator.cpp

cl::values(
                    "GCC encoding (only meaningful for -sample)")));
 
 using namespace llvm;
 using namespace sampleprof;
 
 namespace llvm {
 namespace sampleprof {
 
+static bool
+usePseudoProbes(const BinarySampleCounterMap &BinarySampleCounters) {
+  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!");
   }
 
   return ProfileGenerator;
 }
 

llvm/tools/llvm-profgen/ProfiledBinary.h

class ProfiledBinary {
   PrologEpilogTracker ProEpilogTracker;
 
   // The symbolizer used to get inline context for an instruction.
   std::unique_ptr<symbolize::LLVMSymbolizer> Symbolizer;
 
   // Pseudo probe decoder
   PseudoProbeDecoder ProbeDecoder;
 
+  bool UsePseudoProbes = false;
+
   void setPreferredBaseAddress(const ELFObjectFileBase *O);
 
   void decodePseudoProbe(const ELFObjectFileBase *Obj);
 
   // Set up disassembler and related components.
   void setUpDisassembler(const ELFObjectFileBase *Obj);
   void setupSymbolizer();
 
bool addressInPrologEpilog(uint64_t Address) const {
     uint64_t Offset = virtualAddrToOffset(Address);
     return ProEpilogTracker.PrologEpilogSet.count(Offset);
   }
 
   uint64_t getAddressforIndex(uint64_t Index) const {
     return offsetToVirtualAddr(CodeAddrs[Index]);
   }
 
+  bool usePseudoProbes() const { return UsePseudoProbes; }
   // Get the index in CodeAddrs for the address
   // As we might get an address which is not the code
   // here it would round to the next valid code address by
   // using lower bound operation
   uint32_t getIndexForAddr(uint64_t Address) const {
     uint64_t Offset = virtualAddrToOffset(Address);
     auto Low = llvm::lower_bound(CodeAddrs, Offset);
     return Low - CodeAddrs.begin();
public:
 
   // Compare two addresses' inline context
   bool inlineContextEqual(uint64_t Add1, uint64_t Add2) const;
 
   // Get the context string of the current stack with inline context filled in.
   // It will search the disassembling info stored in Offset2LocStackMap. This is
   // used as the key of function sample map
   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 {
+    return ProbeDecoder.getInlineContextForProbe(Probe, InlineContextStack,
+                                                 IncludeLeaf);
+  }
 };
 
 } // end namespace sampleprof
 } // end namespace llvm
 
 #endif

llvm/tools/llvm-profgen/ProfiledBinary.cpp

for (uint32_t I = 0; I < (uint32_t)ContextVec.size(); I++) {
     if (I == ContextVec.size() - 1) {
       // Only keep the function name for the leaf frame
       StringRef Ref(ContextVec[I]);
       OContextStr << Ref.split(":").first.str();
     } else {
       OContextStr << ContextVec[I];
     }
   }
-
   return OContextStr.str();
 }
 
 void ProfiledBinary::setPreferredBaseAddress(const ELFObjectFileBase *Obj) {
   for (section_iterator SI = Obj->section_begin(), SE = Obj->section_end();
        SI != SE; ++SI) {
     const SectionRef &Section = *SI;
     if (Section.isText()) {
for (section_iterator SI = Obj->section_begin(), SE = Obj->section_end();
     if (SectionName == ".pseudo_probe_desc") {
       StringRef Contents = unwrapOrError(Section.getContents(), FileName);
       ProbeDecoder.buildGUID2FuncDescMap(
           reinterpret_cast<const uint8_t *>(Contents.data()), Contents.size());
     } else if (SectionName == ".pseudo_probe") {
       StringRef Contents = unwrapOrError(Section.getContents(), FileName);
       ProbeDecoder.buildAddress2ProbeMap(
           reinterpret_cast<const uint8_t *>(Contents.data()), Contents.size());
+      // set UsePseudoProbes flag, used for PerfReader
+      UsePseudoProbes = true;
     }
   }
 
   if (ShowPseudoProbe)
     ProbeDecoder.printGUID2FuncDescMap(outs());
 }
 
 bool ProfiledBinary::dissassembleSymbol(std::size_t SI, ArrayRef<uint8_t> Bytes,

llvm/tools/llvm-profgen/PseudoProbe.h

public:
   // Decode pseudo_probe section to build address to probes map.
   void buildAddress2ProbeMap(const uint8_t *Start, std::size_t Size);
 
   // Print pseudo_probe_desc section info
   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;
+
+  // 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]
+  void
+  getInlineContextForProbe(const PseudoProbe *Probe,
+                           SmallVector<std::string, 16> &InlineContextStack,
+                           bool IncludeLeaf) const;
 };
 
 } // end namespace sampleprof
 } // end namespace llvm
 
 #endif

llvm/tools/llvm-profgen/PseudoProbe.cpp

void PseudoProbeDecoder::printProbeForAddress(raw_ostream &OS,
   if (It != Address2ProbesMap.end()) {
     for (auto &Probe : It->second) {
       OS << " [Probe]:\t";
       Probe.print(OS, GUID2FuncDescMap, true);
     }
   }
 }
 
+const PseudoProbe *
+PseudoProbeDecoder::getCallProbeForAddr(uint64_t Address) const {
+  auto It = Address2ProbesMap.find(Address);
+  if (It == Address2ProbesMap.end())
+    return nullptr;
+  const std::vector<PseudoProbe> &Probes = It->second;
+
+  const PseudoProbe *CallProbe = nullptr;
+  for (const auto &Probe : Probes) {
+    if (Probe.isCall()) {
+      assert(!CallProbe &&
+             "There should be only one call probe corresponding to address "
+             "which is a callsite.");
+      CallProbe = &Probe;
+    }
+  }
+  return CallProbe;
+}
+
+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);
+}
+
 } // end namespace sampleprof
 } // end namespace llvm