[CSSPGO] Migrate and refactor the decoder of Pseudo Probe

[• luoj1]

Summary

Migrate pseudo probe decoding logic in llvm-profgen to MC, so other LLVM-base program could reuse existing codes. Redesign object layout of encoded and decoded pseudo probes.

[hoy]

P.S. lgtm in general. I have done an around of code review internally before sending out the diff.

llvm/tools/llvm-profgen/ProfiledBinary.cpp

184

Nit: can just use .pseudo_probe_desc instead of ProbeDecoder.getSectionName(). Similarly below.

[wlei]

llvm/lib/MC/MCPseudoProbe.cpp

339

Same here

406

As we use return type to handle error, SectionName here is no longer used, right? If so, this can be removed.

[• luoj1]

remove sectionName

[hoy]

LGTM.

llvm/include/llvm/MC/MCPseudoProbe.h

 //          FUNCTION BODY
 //            A FUNCTION BODY entry describing the inlined function.
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_MC_MCPSEUDOPROBE_H
 #define LLVM_MC_MCPSEUDOPROBE_H
 
 #include "llvm/ADT/MapVector.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/ADT/Twine.h"
+#include "llvm/IR/PseudoProbe.h"
 #include "llvm/MC/MCSection.h"
+#include "llvm/Support/Errc.h"
+#include "llvm/Support/Error.h"
+#include "llvm/Support/ErrorOr.h"
+#include "llvm/Support/WithColor.h"
+#include "llvm/Support/raw_ostream.h"
+#include <algorithm>
 #include <functional>
+#include <list>
 #include <map>
+#include <set>
+#include <sstream>
+#include <string>
+#include <system_error>
+#include <unordered_map>
+#include <unordered_set>
 #include <vector>
 
 namespace llvm {
 
 class MCStreamer;
 class MCSymbol;
 class MCObjectStreamer;
 
 enum class MCPseudoProbeFlag {
   // If set, indicates that the probe is encoded as an address delta
   // instead of a real code address.
   AddressDelta = 0x1,
 };
 
+// Function descriptor decoded from .pseudo_probe_desc section
+struct MCPseudoProbeFuncDesc {
+  uint64_t FuncGUID = 0;
+  uint64_t FuncHash = 0;
+  std::string FuncName;
+
+  MCPseudoProbeFuncDesc(uint64_t GUID, uint64_t Hash, StringRef Name)
+      : FuncGUID(GUID), FuncHash(Hash), FuncName(Name){};
+
+  void print(raw_ostream &OS);
+};
+
+class MCPseudoProbe;
+class MCDecodedPseudoProbe;
+
+// An inline frame has the form <Guid, ProbeID>
+using InlineSite = std::tuple<uint64_t, uint32_t>;
+using MCPseudoProbeInlineStack = SmallVector<InlineSite, 8>;
+// GUID to PseudoProbeFuncDesc map
+using GUIDProbeFunctionMap =
+    std::unordered_map<uint64_t, MCPseudoProbeFuncDesc>;
+// Address to pseudo probes map.
+using AddressProbesMap =
+    std::unordered_map<uint64_t, std::list<MCDecodedPseudoProbe>>;
+
+class MCPseudoProbeInlineTree;
+class MCDecodedPseudoProbeInlineTree;
+
+class MCPseudoProbeBase {
+protected:
+  uint64_t Guid;
+  uint64_t Index;
+  uint8_t Attributes;
+  uint8_t Type;
+  // The value should be equal to PseudoProbeReservedId::Last + 1 which is
+  // defined in SampleProfileProbe.h. The header file is not included here to
+  // reduce the dependency from MC to IPO.
+  const static uint32_t PseudoProbeFirstId = 1;
+
+public:
+  MCPseudoProbeBase(uint64_t G, uint64_t I, uint64_t At, uint8_t T)
+      : Guid(G), Index(I), Attributes(At), Type(T) {}
+
+  bool isEntry() const { return Index == PseudoProbeFirstId; }
+
+  bool isTailCall() const {
+    return Attributes & static_cast<uint8_t>(PseudoProbeAttributes::Reserved);
+  }
+
+  uint64_t getGuid() const { return Guid; }
+
+  uint64_t getIndex() const { return Index; }
+
+  uint8_t getAttributes() const { return Attributes; }
+
+  uint8_t getType() const { return Type; }
+
+  bool isBlock() const {
+    return Type == static_cast<uint8_t>(PseudoProbeType::Block);
+  }
+
+  bool isIndirectCall() const {
+    return Type == static_cast<uint8_t>(PseudoProbeType::IndirectCall);
+  }
+
+  bool isDirectCall() const {
+    return Type == static_cast<uint8_t>(PseudoProbeType::DirectCall);
+  }
+
+  bool isCall() const { return isIndirectCall() || isDirectCall(); }
+
+  void setAttributes(uint8_t Attr) { Attributes = Attr; }
+};
+
 /// Instances of this class represent a pseudo probe instance for a pseudo probe
 /// table entry, which is created during a machine instruction is assembled and
 /// uses an address from a temporary label created at the current address in the
 /// current section.
-class MCPseudoProbe {
+class MCPseudoProbe : public MCPseudoProbeBase {
   MCSymbol *Label;
-  uint64_t Guid;
-  uint64_t Index;
-  uint8_t Type;
-  uint8_t Attributes;
 
 public:
   MCPseudoProbe(MCSymbol *Label, uint64_t Guid, uint64_t Index, uint64_t Type,
                 uint64_t Attributes)
-      : Label(Label), Guid(Guid), Index(Index), Type(Type),
-        Attributes(Attributes) {
+      : MCPseudoProbeBase(Guid, Index, Attributes, Type), Label(Label) {
     assert(Type <= 0xFF && "Probe type too big to encode, exceeding 2^8");
     assert(Attributes <= 0xFF &&
            "Probe attributes too big to encode, exceeding 2^16");
   }
 
   MCSymbol *getLabel() const { return Label; }
+  void emit(MCObjectStreamer *MCOS, const MCPseudoProbe *LastProbe) const;
+};
 
-  uint64_t getGuid() const { return Guid; }
+class MCDecodedPseudoProbe : public MCPseudoProbeBase {
+  uint64_t Address;
+  MCDecodedPseudoProbeInlineTree *InlineTree;
 
-  uint64_t getIndex() const { return Index; }
+public:
+  MCDecodedPseudoProbe(uint64_t Ad, uint64_t G, uint32_t I, PseudoProbeType K,
+                       uint8_t At, MCDecodedPseudoProbeInlineTree *Tree)
+      : MCPseudoProbeBase(G, I, At, static_cast<uint8_t>(K)), Address(Ad),
+        InlineTree(Tree){};
 
-  uint8_t getType() const { return Type; }
+  uint64_t getAddress() const { return Address; }
 
-  uint8_t getAttributes() const { return Attributes; }
+  void setAddress(uint64_t Addr) { Address = Addr; }
 
-  void emit(MCObjectStreamer *MCOS, const MCPseudoProbe *LastProbe) const;
+  MCDecodedPseudoProbeInlineTree *getInlineTreeNode() const {
+    return InlineTree;
+  }
+
+  // Get the inlined context by traversing current inline tree backwards,
+  // each tree node has its InlineSite which is taken as the context.
+  // \p ContextStack is populated in root to leaf order
+  void getInlineContext(SmallVectorImpl<std::string> &ContextStack,
+                        const GUIDProbeFunctionMap &GUID2FuncMAP,
+                        bool ShowName) const;
+
+  // Helper function to get the string from context stack
+  std::string getInlineContextStr(const GUIDProbeFunctionMap &GUID2FuncMAP,
+                                  bool ShowName) const;
+
+  // Print pseudo probe while disassembling
+  void print(raw_ostream &OS, const GUIDProbeFunctionMap &GUID2FuncMAP,
+             bool ShowName) const;
 };
 
-// An inline frame has the form <Guid, ProbeID>
-using InlineSite = std::tuple<uint64_t, uint32_t>;
-using MCPseudoProbeInlineStack = SmallVector<InlineSite, 8>;
+template <typename ProbeType, typename DerivedProbeInlineTreeType>
+class MCPseudoProbeInlineTreeBase {
+  struct InlineSiteHash {
+    uint64_t operator()(const InlineSite &Site) const {
+      return std::get<0>(Site) ^ std::get<1>(Site);
+    }
+  };
+
+protected:
+  // Track children (e.g. inlinees) of current context
+  using InlinedProbeTreeMap = std::unordered_map<
+      InlineSite, std::unique_ptr<DerivedProbeInlineTreeType>, InlineSiteHash>;
+  InlinedProbeTreeMap Children;
+  // Set of probes that come with the function.
+  std::vector<ProbeType> Probes;
+  MCPseudoProbeInlineTreeBase() {
+    static_assert(std::is_base_of<MCPseudoProbeInlineTreeBase,
+                                  DerivedProbeInlineTreeType>::value,
+                  "DerivedProbeInlineTreeType must be subclass of "
+                  "MCPseudoProbeInlineTreeBase");
+  }
+
+public:
+  uint64_t Guid = 0;
+
+  // Root node has a GUID 0.
+  bool isRoot() const { return Guid == 0; }
+  InlinedProbeTreeMap &getChildren() { return Children; }
+  const InlinedProbeTreeMap &getChildren() const { return Children; }
+  std::vector<ProbeType> &getProbes() { return Probes; }
+  void addProbes(ProbeType Probe) { Probes.push_back(Probe); }
+  // Caller node of the inline site
+  MCPseudoProbeInlineTreeBase<ProbeType, DerivedProbeInlineTreeType> *Parent;
+  DerivedProbeInlineTreeType *getOrAddNode(const InlineSite &Site) {
+    auto Ret = Children.emplace(
+        Site, std::make_unique<DerivedProbeInlineTreeType>(Site));
+    Ret.first->second->Parent = this;
+    return Ret.first->second.get();
+  };
+};
 
 // A Tri-tree based data structure to group probes by inline stack.
 // A tree is allocated for a standalone .text section. A fake
 // instance is created as the root of a tree.
 // A real instance of this class is created for each function, either an
 // unlined function that has code in .text section or an inlined function.
-class MCPseudoProbeInlineTree {
-  uint64_t Guid;
-  // Set of probes that come with the function.
-  std::vector<MCPseudoProbe> Probes;
-  // Use std::map for a deterministic output.
-  std::map<InlineSite, MCPseudoProbeInlineTree *> Inlinees;
-
-  // Root node has a GUID 0.
-  bool isRoot() { return Guid == 0; }
-  MCPseudoProbeInlineTree *getOrAddNode(InlineSite Site);
 
+class MCPseudoProbeInlineTree
+    : public MCPseudoProbeInlineTreeBase<MCPseudoProbe,
+                                         MCPseudoProbeInlineTree> {
 public:
   MCPseudoProbeInlineTree() = default;
-  MCPseudoProbeInlineTree(uint64_t Guid) : Guid(Guid) {}
-  ~MCPseudoProbeInlineTree();
+  MCPseudoProbeInlineTree(uint64_t Guid) { this->Guid = Guid; }
+  MCPseudoProbeInlineTree(const InlineSite &Site) {
+    this->Guid = std::get<0>(Site);
+  }
+
+  // MCPseudoProbeInlineTree method based on Inlinees
   void addPseudoProbe(const MCPseudoProbe &Probe,
                       const MCPseudoProbeInlineStack &InlineStack);
   void emit(MCObjectStreamer *MCOS, const MCPseudoProbe *&LastProbe);
 };
 
+// inline tree node for the decoded pseudo probe
+class MCDecodedPseudoProbeInlineTree
+    : public MCPseudoProbeInlineTreeBase<MCDecodedPseudoProbe *,
+                                         MCDecodedPseudoProbeInlineTree> {
+public:
+  InlineSite ISite;
+  // Used for decoding
+  uint32_t ChildrenToProcess = 0;
+
+  MCDecodedPseudoProbeInlineTree(){};
+  MCDecodedPseudoProbeInlineTree(const InlineSite &Site) : ISite(Site){};
+
+  // Return false if it's a dummy inline site
+  bool hasInlineSite() const { return std::get<0>(ISite) != 0; }
+};
+
 /// Instances of this class represent the pseudo probes inserted into a compile
 /// unit.
 class MCPseudoProbeSection {
 public:
   void addPseudoProbe(MCSection *Sec, const MCPseudoProbe &Probe,
                       const MCPseudoProbeInlineStack &InlineStack) {
     MCProbeDivisions[Sec].addPseudoProbe(Probe, InlineStack);
   }
public:
   static void emit(MCObjectStreamer *MCOS);
 
   MCPseudoProbeSection &getProbeSections() { return MCProbeSections; }
 
 #ifndef NDEBUG
   static int DdgPrintIndent;
 #endif
 };
+
+class MCPseudoProbeDecoder {
+  // GUID to PseudoProbeFuncDesc map.
+  GUIDProbeFunctionMap GUID2FuncDescMap;
+
+  // Address to probes map.
+  AddressProbesMap Address2ProbesMap;
+
+  // The dummy root of the inline trie, all the outlined function will directly
+  // be the children of the dummy root, all the inlined function will be the
+  // children of its inlineer. So the relation would be like:
+  // DummyRoot --> OutlinedFunc --> InlinedFunc1 --> InlinedFunc2
+  MCDecodedPseudoProbeInlineTree DummyInlineRoot;
+
+  /// Points to the current location in the buffer.
+  const uint8_t *Data = nullptr;
+
+  /// Points to the end of the buffer.
+  const uint8_t *End = nullptr;
+
+  /// SectionName used for debug
+  std::string SectionName;
+
+  // Decoding helper function
+  template <typename T> ErrorOr<T> readUnencodedNumber();
+  template <typename T> ErrorOr<T> readUnsignedNumber();
+  template <typename T> ErrorOr<T> readSignedNumber();
+  ErrorOr<StringRef> readString(uint32_t Size);
+
+public:
+  // Decode pseudo_probe_desc section to build GUID to PseudoProbeFuncDesc map.
+  bool buildGUID2FuncDescMap(const uint8_t *Start, std::size_t Size);
+
+  // Decode pseudo_probe section to build address to probes map.
+  bool 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);
+
+  // do printProbeForAddress for all addresses
+  void printProbesForAllAddresses(raw_ostream &OS);
+
+  // Look up the probe of a call for the input address
+  const MCDecodedPseudoProbe *getCallProbeForAddr(uint64_t Address) const;
+
+  const MCPseudoProbeFuncDesc *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]
+  void
+  getInlineContextForProbe(const MCDecodedPseudoProbe *Probe,
+                           SmallVectorImpl<std::string> &InlineContextStack,
+                           bool IncludeLeaf) const;
+
+  const AddressProbesMap &getAddress2ProbesMap() const {
+    return Address2ProbesMap;
+  }
+
+  AddressProbesMap &getAddress2ProbesMap() { return Address2ProbesMap; }
+
+  const GUIDProbeFunctionMap &getGUID2FuncDescMap() const {
+    return GUID2FuncDescMap;
+  }
+
+  const MCPseudoProbeFuncDesc *
+  getInlinerDescForProbe(const MCDecodedPseudoProbe *Probe) const;
+
+  std::string getSectionName() const { return SectionName; };
+
+  const MCDecodedPseudoProbeInlineTree &getDummyInlineRoot() const {
+    return DummyInlineRoot;
+  }
+};
+
 } // end namespace llvm
 
 #endif // LLVM_MC_MCPSEUDOPROBE_H

llvm/lib/MC/CMakeLists.txt

add_llvm_component_library(LLVMMC
 
   ADDITIONAL_HEADER_DIRS
   ${LLVM_MAIN_INCLUDE_DIR}/llvm/MC
 
   LINK_COMPONENTS
   Support
   BinaryFormat
   DebugInfoCodeView
+  ProfileData
   )
 
 add_subdirectory(MCParser)
 add_subdirectory(MCDisassembler)

llvm/lib/MC/MCPseudoProbe.cpp

 //===- lib/MC/MCPseudoProbe.cpp - Pseudo probe encoding support ----------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/MC/MCPseudoProbe.h"
 #include "llvm/MC/MCAsmInfo.h"
 #include "llvm/MC/MCContext.h"
 #include "llvm/MC/MCObjectFileInfo.h"
 #include "llvm/MC/MCObjectStreamer.h"
 #include "llvm/MC/MCStreamer.h"
+#include "llvm/ProfileData/SampleProf.h"
+#include "llvm/Support/Endian.h"
+#include "llvm/Support/LEB128.h"
+#include "llvm/Support/raw_ostream.h"
+#include <limits>
+#include <memory>
 
 #define DEBUG_TYPE "mcpseudoprobe"
 
 using namespace llvm;
+using namespace support;
+using namespace sampleprof;
 
 #ifndef NDEBUG
 int MCPseudoProbeTable::DdgPrintIndent = 0;
 #endif
 
 static const MCExpr *buildSymbolDiff(MCObjectStreamer *MCOS, const MCSymbol *A,
                                      const MCSymbol *B) {
   MCContext &Context = MCOS->getContext();
void MCPseudoProbe::emit(MCObjectStreamer *MCOS,
   }
 
   LLVM_DEBUG({
     dbgs().indent(MCPseudoProbeTable::DdgPrintIndent);
     dbgs() << "Probe: " << Index << "\n";
   });
 }
 
-MCPseudoProbeInlineTree::~MCPseudoProbeInlineTree() {
-  for (auto &Inlinee : Inlinees)
-    delete Inlinee.second;
-}
-
-MCPseudoProbeInlineTree *
-MCPseudoProbeInlineTree::getOrAddNode(InlineSite Site) {
-  auto Iter = Inlinees.find(Site);
-  if (Iter == Inlinees.end()) {
-    auto *Node = new MCPseudoProbeInlineTree(std::get<0>(Site));
-    Inlinees[Site] = Node;
-    return Node;
-  } else {
-    return Iter->second;
-  }
-}
-
 void MCPseudoProbeInlineTree::addPseudoProbe(
     const MCPseudoProbe &Probe, const MCPseudoProbeInlineStack &InlineStack) {
   // The function should not be called on the root.
   assert(isRoot() && "Should not be called on root");
 
   // When it comes here, the input look like:
   //    Probe: GUID of C, ...
   //    InlineStack: [88, A], [66, B]
LLVM_DEBUG({
       dbgs().indent(MCPseudoProbeTable::DdgPrintIndent);
       dbgs() << "GUID: " << Guid << "\n";
     });
     // Emit Guid
     MCOS->emitInt64(Guid);
     // Emit number of probes in this node
     MCOS->emitULEB128IntValue(Probes.size());
     // Emit number of direct inlinees
-    MCOS->emitULEB128IntValue(Inlinees.size());
+    MCOS->emitULEB128IntValue(Children.size());
     // Emit probes in this group
     for (const auto &Probe : Probes) {
       Probe.emit(MCOS, LastProbe);
       LastProbe = &Probe;
     }
   } else {
     assert(Probes.empty() && "Root should not have probes");
   }
 
-  // Emit descendent
+  // Emit sorted descendant
+  // InlineSite is unique for each pair,
+  // so there will be no ordering of Inlinee based on MCPseudoProbeInlineTree*
+  std::map<InlineSite, MCPseudoProbeInlineTree *> Inlinees;
+  for (auto Child = Children.begin(); Child != Children.end(); ++Child)
+    Inlinees[Child->first] = Child->second.get();
+
   for (const auto &Inlinee : Inlinees) {
     if (Guid) {
       // Emit probe index
       MCOS->emitULEB128IntValue(std::get<1>(Inlinee.first));
       LLVM_DEBUG({
         dbgs().indent(MCPseudoProbeTable::DdgPrintIndent);
         dbgs() << "InlineSite: " << std::get<1>(Inlinee.first) << "\n";
       });
void MCPseudoProbeTable::emit(MCObjectStreamer *MCOS) {
   if (ProbeSections.empty())
     return;
 
   LLVM_DEBUG(MCPseudoProbeTable::DdgPrintIndent = 0);
 
   // Put out the probe.
   ProbeSections.emit(MCOS);
 }
+
+static StringRef getProbeFNameForGUID(const GUIDProbeFunctionMap &GUID2FuncMAP,
+                                      uint64_t GUID) {
+  auto It = GUID2FuncMAP.find(GUID);
+  assert(It != GUID2FuncMAP.end() &&
+         "Probe function must exist for a valid GUID");
+  return It->second.FuncName;
+}
+
+void MCPseudoProbeFuncDesc::print(raw_ostream &OS) {
+  OS << "GUID: " << FuncGUID << " Name: " << FuncName << "\n";
+  OS << "Hash: " << FuncHash << "\n";
+}
+
+void MCDecodedPseudoProbe::getInlineContext(
+    SmallVectorImpl<std::string> &ContextStack,
+    const GUIDProbeFunctionMap &GUID2FuncMAP, bool ShowName) const {
+  uint32_t Begin = ContextStack.size();
+  MCDecodedPseudoProbeInlineTree *Cur = InlineTree;
+  // It will add the string of each node's inline site during iteration.
+  // Note that it won't include the probe's belonging function(leaf location)
+  while (Cur->hasInlineSite()) {
+    std::string ContextStr;
+    if (ShowName) {
+      StringRef FuncName =
+          getProbeFNameForGUID(GUID2FuncMAP, std::get<0>(Cur->ISite));
+      ContextStr += FuncName.str();
+    } else {
+      ContextStr += Twine(std::get<0>(Cur->ISite)).str();
+    }
+    ContextStr += ":";
+    ContextStr += Twine(std::get<1>(Cur->ISite)).str();
+    ContextStack.emplace_back(ContextStr);
+    Cur = static_cast<MCDecodedPseudoProbeInlineTree *>(Cur->Parent);
+  }
+  // Make the ContextStack in caller-callee order
+  std::reverse(ContextStack.begin() + Begin, ContextStack.end());
+}
+
+std::string MCDecodedPseudoProbe::getInlineContextStr(
+    const GUIDProbeFunctionMap &GUID2FuncMAP, bool ShowName) const {
+  std::ostringstream OContextStr;
+  SmallVector<std::string, 16> ContextStack;
+  getInlineContext(ContextStack, GUID2FuncMAP, ShowName);
+  for (auto &CxtStr : ContextStack) {
+    if (OContextStr.str().size())
+      OContextStr << " @ ";
+    OContextStr << CxtStr;
+  }
+  return OContextStr.str();
+}
+
+static const char *PseudoProbeTypeStr[3] = {"Block", "IndirectCall",
+                                            "DirectCall"};
+
+void MCDecodedPseudoProbe::print(raw_ostream &OS,
+                                 const GUIDProbeFunctionMap &GUID2FuncMAP,
+                                 bool ShowName) const {
+  OS << "FUNC: ";
+  if (ShowName) {
+    StringRef FuncName = getProbeFNameForGUID(GUID2FuncMAP, Guid);
+    OS << FuncName.str() << " ";
+  } else {
+    OS << Guid << " ";
+  }
+  OS << "Index: " << Index << "  ";
+  OS << "Type: " << PseudoProbeTypeStr[static_cast<uint8_t>(Type)] << "  ";
+  if (isTailCall()) {
+    OS << "TailCall  ";
+  }
+  std::string InlineContextStr = getInlineContextStr(GUID2FuncMAP, ShowName);
+  if (InlineContextStr.size()) {
+    OS << "Inlined: @ ";
+    OS << InlineContextStr;
+  }
+  OS << "\n";
+}
+
+template <typename T> ErrorOr<T> MCPseudoProbeDecoder::readUnencodedNumber() {
+  if (Data + sizeof(T) > End) {
+    return std::error_code();
+  }
+  T Val = endian::readNext<T, little, unaligned>(Data);
+  return ErrorOr<T>(Val);
+}
+
+template <typename T> ErrorOr<T> MCPseudoProbeDecoder::readUnsignedNumber() {
+  unsigned NumBytesRead = 0;
+  uint64_t Val = decodeULEB128(Data, &NumBytesRead);
+  if (Val > std::numeric_limits<T>::max() || (Data + NumBytesRead > End)) {
+    return std::error_code();
+  }
+  Data += NumBytesRead;
+  return ErrorOr<T>(static_cast<T>(Val));
+}
+
+template <typename T> ErrorOr<T> MCPseudoProbeDecoder::readSignedNumber() {
+  unsigned NumBytesRead = 0;
+  int64_t Val = decodeSLEB128(Data, &NumBytesRead);
+  if (Val > std::numeric_limits<T>::max() || (Data + NumBytesRead > End)) {
+    return std::error_code();
+  }
+  Data += NumBytesRead;
+  return ErrorOr<T>(static_cast<T>(Val));
+}
+
+ErrorOr<StringRef> MCPseudoProbeDecoder::readString(uint32_t Size) {
+  StringRef Str(reinterpret_cast<const char *>(Data), Size);
+  if (Data + Size > End) {
+    return std::error_code();
+  }
+  Data += Size;
+  return ErrorOr<StringRef>(Str);
+}
+
+bool MCPseudoProbeDecoder::buildGUID2FuncDescMap(const uint8_t *Start,
+                                                 std::size_t Size) {
+  // The pseudo_probe_desc section has a format like:
+  // .section .pseudo_probe_desc,"",@progbits
+  // .quad -5182264717993193164   // GUID
+  // .quad 4294967295             // Hash
+  // .uleb 3                      // Name size
+  // .ascii "foo"                 // Name
+  // .quad -2624081020897602054
+  // .quad 174696971957
+  // .uleb 34
+  // .ascii "main"
+#ifndef NDEBUG
+  SectionName = "pseudo_probe_desc";
+#endif
+  Data = Start;
+  End = Data + Size;
+
+  while (Data < End) {
+    auto ErrorOrGUID = readUnencodedNumber<uint64_t>();
+    if (!ErrorOrGUID)
+      return false;
+
+    auto ErrorOrHash = readUnencodedNumber<uint64_t>();
+    if (!ErrorOrHash)
+      return false;
+
+    auto ErrorOrNameSize = readUnsignedNumber<uint32_t>();
+    if (!ErrorOrNameSize)
+      return false;
+    uint32_t NameSize = std::move(*ErrorOrNameSize);
+
+    auto ErrorOrName = readString(NameSize);
+    if (!ErrorOrName)
+      return false;
+
+    uint64_t GUID = std::move(*ErrorOrGUID);
+    uint64_t Hash = std::move(*ErrorOrHash);
+    StringRef Name =
+        FunctionSamples::getCanonicalFnName(std::move(*ErrorOrName));
+
+    // Initialize PseudoProbeFuncDesc and populate it into GUID2FuncDescMap
+    GUID2FuncDescMap.emplace(GUID, MCPseudoProbeFuncDesc(GUID, Hash, Name));
+  }
+  assert(Data == End && "Have unprocessed data in pseudo_probe_desc section");
+  return true;
+}
+
+bool MCPseudoProbeDecoder::buildAddress2ProbeMap(const uint8_t *Start,
+                                                 std::size_t Size) {
+  // The pseudo_probe section encodes an inline forest and each tree has a
+  // format like:
+  //  FUNCTION BODY (one for each uninlined function present in the text
+  //  section)
+  //     GUID (uint64)
+  //         GUID of the function
+  //     NPROBES (ULEB128)
+  //         Number of probes originating from this function.
+  //     NUM_INLINED_FUNCTIONS (ULEB128)
+  //         Number of callees inlined into this function, aka number of
+  //         first-level inlinees
+  //     PROBE RECORDS
+  //         A list of NPROBES entries. Each entry contains:
+  //           INDEX (ULEB128)
+  //           TYPE (uint4)
+  //             0 - block probe, 1 - indirect call, 2 - direct call
+  //           ATTRIBUTE (uint3)
+  //             1 - tail call, 2 - dangling
+  //           ADDRESS_TYPE (uint1)
+  //             0 - code address, 1 - address delta
+  //           CODE_ADDRESS (uint64 or ULEB128)
+  //             code address or address delta, depending on Flag
+  //     INLINED FUNCTION RECORDS
+  //         A list of NUM_INLINED_FUNCTIONS entries describing each of the
+  //         inlined callees.  Each record contains:
+  //           INLINE SITE
+  //             Index of the callsite probe (ULEB128)
+  //           FUNCTION BODY
+  //             A FUNCTION BODY entry describing the inlined function.
+#ifndef NDEBUG
+  SectionName = "pseudo_probe";
+#endif
+  Data = Start;
+  End = Data + Size;
+
+  MCDecodedPseudoProbeInlineTree *Root = &DummyInlineRoot;
+  MCDecodedPseudoProbeInlineTree *Cur = &DummyInlineRoot;
+  uint64_t LastAddr = 0;
+  uint32_t Index = 0;
+  // A DFS-based decoding
+  while (Data < End) {
+    if (Root == Cur) {
+      // Use a sequential id for top level inliner.
+      Index = Root->getChildren().size();
+    } else {
+      // Read inline site for inlinees
+      auto ErrorOrIndex = readUnsignedNumber<uint32_t>();
+      if (!ErrorOrIndex)
+        return false;
+      Index = std::move(*ErrorOrIndex);
+    }
+    // Switch/add to a new tree node(inlinee)
+    Cur = Cur->getOrAddNode(std::make_tuple(Cur->Guid, Index));
+    // Read guid
+    auto ErrorOrCurGuid = readUnencodedNumber<uint64_t>();
+    if (!ErrorOrCurGuid)
+      return false;
+    Cur->Guid = std::move(*ErrorOrCurGuid);
+    // Read number of probes in the current node.
+    auto ErrorOrNodeCount = readUnsignedNumber<uint32_t>();
+    if (!ErrorOrNodeCount)
+      return false;
+    uint32_t NodeCount = std::move(*ErrorOrNodeCount);
+    // Read number of direct inlinees
+    auto ErrorOrCurChildrenToProcess = readUnsignedNumber<uint32_t>();
+    if (!ErrorOrCurChildrenToProcess)
+      return false;
+    Cur->ChildrenToProcess = std::move(*ErrorOrCurChildrenToProcess);
+    // Read all probes in this node
+    for (std::size_t I = 0; I < NodeCount; I++) {
+      // Read index
+      auto ErrorOrIndex = readUnsignedNumber<uint32_t>();
+      if (!ErrorOrIndex)
+        return false;
+      uint32_t Index = std::move(*ErrorOrIndex);
+      // Read type | flag.
+      auto ErrorOrValue = readUnencodedNumber<uint8_t>();
+      if (!ErrorOrValue)
+        return false;
+      uint8_t Value = std::move(*ErrorOrValue);
+      uint8_t Kind = Value & 0xf;
+      uint8_t Attr = (Value & 0x70) >> 4;
+      // Read address
+      uint64_t Addr = 0;
+      if (Value & 0x80) {
+        auto ErrorOrOffset = readSignedNumber<int64_t>();
+        if (!ErrorOrOffset)
+          return false;
+        int64_t Offset = std::move(*ErrorOrOffset);
+        Addr = LastAddr + Offset;
+      } else {
+        auto ErrorOrAddr = readUnencodedNumber<int64_t>();
+        if (!ErrorOrAddr)
+          return false;
+        Addr = std::move(*ErrorOrAddr);
+      }
+      // Populate Address2ProbesMap
+      auto &Probes = Address2ProbesMap[Addr];
+      Probes.emplace_back(Addr, Cur->Guid, Index, PseudoProbeType(Kind), Attr,
+                          Cur);
+      Cur->addProbes(&Probes.back());
+      LastAddr = Addr;
+    }
+
+    // Look for the parent for the next node by subtracting the current
+    // node count from tree counts along the parent chain. The first node
+    // in the chain that has a non-zero tree count is the target.
+    while (Cur != Root) {
+      if (Cur->ChildrenToProcess == 0) {
+        Cur = static_cast<MCDecodedPseudoProbeInlineTree *>(Cur->Parent);
+        if (Cur != Root) {
+          assert(Cur->ChildrenToProcess > 0 &&
+                 "Should have some unprocessed nodes");
+          Cur->ChildrenToProcess -= 1;
+        }
+      } else {
+        break;
+      }
+    }
+  }
+
+  assert(Data == End && "Have unprocessed data in pseudo_probe section");
+  assert(Cur == Root &&
+         " Cur should point to root when the forest is fully built up");
+  return true;
+}
+
+void MCPseudoProbeDecoder::printGUID2FuncDescMap(raw_ostream &OS) {
+  OS << "Pseudo Probe Desc:\n";
+  // Make the output deterministic
+  std::map<uint64_t, MCPseudoProbeFuncDesc> OrderedMap(GUID2FuncDescMap.begin(),
+                                                       GUID2FuncDescMap.end());
+  for (auto &I : OrderedMap) {
+    I.second.print(OS);
+  }
+}
+
+void MCPseudoProbeDecoder::printProbeForAddress(raw_ostream &OS,
+                                                uint64_t Address) {
+  auto It = Address2ProbesMap.find(Address);
+  if (It != Address2ProbesMap.end()) {
+    for (auto &Probe : It->second) {
+      OS << " [Probe]:\t";
+      Probe.print(OS, GUID2FuncDescMap, true);
+    }
+  }
+}
+
+void MCPseudoProbeDecoder::printProbesForAllAddresses(raw_ostream &OS) {
+  std::vector<uint64_t> Addresses;
+  for (auto Entry : Address2ProbesMap)
+    Addresses.push_back(Entry.first);
+  std::sort(Addresses.begin(), Addresses.end());
+  for (auto K : Addresses) {
+    OS << "Address:\t";
+    OS << K;
+    OS << "\n";
+    printProbeForAddress(OS, K);
+  }
+}
+
+const MCDecodedPseudoProbe *
+MCPseudoProbeDecoder::getCallProbeForAddr(uint64_t Address) const {
+  auto It = Address2ProbesMap.find(Address);
+  if (It == Address2ProbesMap.end())
+    return nullptr;
+  const auto &Probes = It->second;
+
+  const MCDecodedPseudoProbe *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;
+}
+
+const MCPseudoProbeFuncDesc *
+MCPseudoProbeDecoder::getFuncDescForGUID(uint64_t GUID) const {
+  auto It = GUID2FuncDescMap.find(GUID);
+  assert(It != GUID2FuncDescMap.end() && "Function descriptor doesn't exist");
+  return &It->second;
+}
+
+void MCPseudoProbeDecoder::getInlineContextForProbe(
+    const MCDecodedPseudoProbe *Probe,
+    SmallVectorImpl<std::string> &InlineContextStack, bool IncludeLeaf) const {
+  Probe->getInlineContext(InlineContextStack, GUID2FuncDescMap, true);
+  if (!IncludeLeaf)
+    return;
+  // Note that the context from probe doesn't include leaf frame,
+  // hence we need to retrieve and prepend leaf if requested.
+  const auto *FuncDesc = getFuncDescForGUID(Probe->getGuid());
+  InlineContextStack.emplace_back(FuncDesc->FuncName + ":" +
+                                  Twine(Probe->getIndex()).str());
+}
+
+const MCPseudoProbeFuncDesc *MCPseudoProbeDecoder::getInlinerDescForProbe(
+    const MCDecodedPseudoProbe *Probe) const {
+  MCDecodedPseudoProbeInlineTree *InlinerNode = Probe->getInlineTreeNode();
+  if (!InlinerNode->hasInlineSite())
+    return nullptr;
+  return getFuncDescForGUID(std::get<0>(InlinerNode->ISite));
+}

llvm/tools/llvm-profgen/CMakeLists.txt

set(LLVM_LINK_COMPONENTS
   )
 
 add_llvm_tool(llvm-profgen
   llvm-profgen.cpp
   PerfReader.cpp
   CSPreInliner.cpp
   ProfiledBinary.cpp
   ProfileGenerator.cpp
-  PseudoProbe.cpp
   )

llvm/tools/llvm-profgen/PerfReader.h

 };
 
 // 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;
+  SmallVector<const MCDecodedPseudoProbe *, 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);
struct FrameStack {
   void popFrame() {
     if (!Stack.empty())
       Stack.pop_back();
   }
   std::shared_ptr<StringBasedCtxKey> getContextKey();
 };
 
 struct ProbeStack {
-  SmallVector<const PseudoProbe *, 16> Stack;
+  SmallVector<const MCDecodedPseudoProbe *, 16> Stack;
   const ProfiledBinary *Binary;
   ProbeStack(const ProfiledBinary *B) : Binary(B) {}
   bool pushFrame(UnwindState::ProfiledFrame *Cur) {
-    const PseudoProbe *CallProbe = Binary->getCallProbeForAddr(Cur->Address);
+    const MCDecodedPseudoProbe *CallProbe =
+        Binary->getCallProbeForAddr(Cur->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 inliner will
     // not know how to consume a context with unknown callsites.
     if (!CallProbe) {
       if (!Cur->isLeafFrame())
         WithColor::warning()
             << "Untracked frame at " << format("%" PRIx64, Cur->Address)

llvm/tools/llvm-profgen/PerfReader.cpp

 }
 
 std::shared_ptr<ProbeBasedCtxKey> ProbeStack::getContextKey() {
   std::shared_ptr<ProbeBasedCtxKey> ProbeBasedKey =
       std::make_shared<ProbeBasedCtxKey>();
   for (auto CallProbe : Stack) {
     ProbeBasedKey->Probes.emplace_back(CallProbe);
   }
-  CSProfileGenerator::compressRecursionContext<const PseudoProbe *>(
+  CSProfileGenerator::compressRecursionContext<const MCDecodedPseudoProbe *>(
       ProbeBasedKey->Probes);
   ProbeBasedKey->genHashCode();
   return ProbeBasedKey;
 }
 
 template <typename T>
 void VirtualUnwinder::collectSamplesFromFrame(UnwindState::ProfiledFrame *Cur,
                                               T &Stack) {

llvm/tools/llvm-profgen/ProfileGenerator.h

private:
   void populateInferredFunctionSamples();
 
 public:
   // Deduplicate adjacent repeated context sequences up to a given sequence
   // length. -1 means no size limit.
   static int32_t MaxCompressionSize;
 };
 
-using ProbeCounterMap = std::unordered_map<const PseudoProbe *, uint64_t>;
+using ProbeCounterMap =
+    std::unordered_map<const MCDecodedPseudoProbe *, uint64_t>;
 
 class PseudoProbeCSProfileGenerator : public CSProfileGenerator {
 
 public:
   PseudoProbeCSProfileGenerator(const BinarySampleCounterMap &Counters)
       : CSProfileGenerator(Counters) {}
   void generateProfile() override;
 
populateBodySamplesWithProbes(const RangeSample &RangeCounter,
                                 ProfiledBinary *Binary);
   // Fill in boundary samples for a call probe
   void populateBoundarySamplesWithProbes(
       const BranchSample &BranchCounter,
       SmallVectorImpl<std::string> &ContextStrStack, ProfiledBinary *Binary);
   // Helper function to get FunctionSamples for the leaf inlined context
   FunctionSamples &
   getFunctionProfileForLeafProbe(SmallVectorImpl<std::string> &ContextStrStack,
-                                 const PseudoProbeFuncDesc *LeafFuncDesc,
+                                 const MCPseudoProbeFuncDesc *LeafFuncDesc,
                                  bool WasLeafInlined);
   // Helper function to get FunctionSamples for the leaf probe
   FunctionSamples &
   getFunctionProfileForLeafProbe(SmallVectorImpl<std::string> &ContextStrStack,
-                                 const PseudoProbe *LeafProbe,
+                                 const MCDecodedPseudoProbe *LeafProbe,
                                  ProfiledBinary *Binary);
 };
 
 } // end namespace sampleprof
 } // end namespace llvm
 
 #endif

llvm/tools/llvm-profgen/ProfileGenerator.cpp

void CSProfileGenerator::write(std::unique_ptr<SampleProfileWriter> Writer,
                                StringMap<FunctionSamples> &ProfileMap) {
   if (std::error_code EC = Writer->write(ProfileMap))
     exitWithError(std::move(EC));
 }
 
 // Helper function to extract context prefix string stack
 // Extract context stack for reusing, leaf context stack will
 // be added compressed while looking up function profile
 static void
 extractPrefixContextStack(SmallVectorImpl<std::string> &ContextStrStack,
-                          const SmallVectorImpl<const PseudoProbe *> &Probes,
-                          ProfiledBinary *Binary) {
+    const SmallVectorImpl<const MCDecodedPseudoProbe *> &Probes,
   for (const auto *P : Probes) {
     Binary->getInlineContextForProbe(P, ContextStrStack, true);
   }
 }
 
 void PseudoProbeCSProfileGenerator::generateProfile() {
   // Enable pseudo probe functionalities in SampleProf
   FunctionSamples::ProfileIsProbeBased = true;

 
 void PseudoProbeCSProfileGenerator::populateBodySamplesWithProbes(
     const RangeSample &RangeCounter,
     SmallVectorImpl<std::string> &ContextStrStack, ProfiledBinary *Binary) {
   ProbeCounterMap ProbeCounter;
   // Extract the top frame probes by looking up each address among the range in
   // the Address2ProbeMap
   extractProbesFromRange(RangeCounter, ProbeCounter, Binary);
-  std::unordered_map<PseudoProbeInlineTree *, FunctionSamples *> FrameSamples;
+  std::unordered_map<MCDecodedPseudoProbeInlineTree *, FunctionSamples *>
+      FrameSamples;
   for (auto PI : ProbeCounter) {
-    const PseudoProbe *Probe = PI.first;
+    const MCDecodedPseudoProbe *Probe = PI.first;
     uint64_t Count = PI.second;
     FunctionSamples &FunctionProfile =
         getFunctionProfileForLeafProbe(ContextStrStack, Probe, Binary);
     // 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()] = &FunctionProfile;
-    FunctionProfile.addBodySamplesForProbe(Probe->Index, Count);
+    FunctionProfile.addBodySamplesForProbe(Probe->getIndex(), Count);
     FunctionProfile.addTotalSamples(Count);
     if (Probe->isEntry()) {
       FunctionProfile.addHeadSamples(Count);
       // Look up for the caller's function profile
       const auto *InlinerDesc = Binary->getInlinerDescForProbe(Probe);
       if (InlinerDesc != nullptr) {
         // Since the context id will be compressed, we have to use callee's
         // context id to infer caller's context id to ensure they share the
for (auto PI : ProbeCounter) {
     }
 
     // Assign zero count for remaining probes without sample hits to
     // differentiate from probes optimized away, of which the counts are unknown
     // and will be inferred by the compiler.
     for (auto &I : FrameSamples) {
       auto *FunctionProfile = I.second;
       for (auto *Probe : I.first->getProbes()) {
-          FunctionProfile->addBodySamplesForProbe(Probe->Index, 0);
+        FunctionProfile->addBodySamplesForProbe(Probe->getIndex(), 0);
       }
     }
   }
 }
 
 void PseudoProbeCSProfileGenerator::populateBoundarySamplesWithProbes(
     const BranchSample &BranchCounter,
     SmallVectorImpl<std::string> &ContextStrStack, 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);
+    const MCDecodedPseudoProbe *CallProbe =
+        Binary->getCallProbeForAddr(SourceAddress);
     if (CallProbe == nullptr)
       continue;
     FunctionSamples &FunctionProfile =
         getFunctionProfileForLeafProbe(ContextStrStack, CallProbe, Binary);
-    FunctionProfile.addBodySamples(CallProbe->Index, 0, Count);
+    FunctionProfile.addBodySamples(CallProbe->getIndex(), 0, Count);
     FunctionProfile.addTotalSamples(Count);
     StringRef CalleeName = FunctionSamples::getCanonicalFnName(
         Binary->getFuncFromStartOffset(TargetOffset));
     if (CalleeName.size() == 0)
       continue;
-    FunctionProfile.addCalledTargetSamples(CallProbe->Index, 0, CalleeName,
+    FunctionProfile.addCalledTargetSamples(CallProbe->getIndex(), 0, CalleeName,
                                            Count);
   }
 }
 
 FunctionSamples &PseudoProbeCSProfileGenerator::getFunctionProfileForLeafProbe(
     SmallVectorImpl<std::string> &ContextStrStack,
-    const PseudoProbeFuncDesc *LeafFuncDesc, bool WasLeafInlined) {
+    const MCPseudoProbeFuncDesc *LeafFuncDesc, bool WasLeafInlined) {
   assert(ContextStrStack.size() && "Profile context must have the leaf frame");
   // Compress the context string except for the leaf frame
   std::string LeafFrame = ContextStrStack.back();
   ContextStrStack.pop_back();
   CSProfileGenerator::compressRecursionContext(ContextStrStack);
 
   std::ostringstream OContextStr;
   for (uint32_t I = 0; I < ContextStrStack.size(); I++) {
FunctionSamples &PseudoProbeCSProfileGenerator::getFunctionProfileForLeafProbe(
 
   FunctionSamples &FunctionProile =
       getFunctionProfileForContext(OContextStr.str(), WasLeafInlined);
   FunctionProile.setFunctionHash(LeafFuncDesc->FuncHash);
   return FunctionProile;
 }
 
 FunctionSamples &PseudoProbeCSProfileGenerator::getFunctionProfileForLeafProbe(
-    SmallVectorImpl<std::string> &ContextStrStack, const PseudoProbe *LeafProbe,
-    ProfiledBinary *Binary) {
+    SmallVectorImpl<std::string> &ContextStrStack,
+    const MCDecodedPseudoProbe *LeafProbe, ProfiledBinary *Binary) {
+
   // Explicitly copy the context for appending the leaf context
   SmallVector<std::string, 16> ContextStrStackCopy(ContextStrStack.begin(),
                                                    ContextStrStack.end());
   Binary->getInlineContextForProbe(LeafProbe, ContextStrStackCopy, true);
-  const auto *FuncDesc = Binary->getFuncDescForGUID(LeafProbe->GUID);
-  bool WasLeafInlined = LeafProbe->InlineTree->hasInlineSite();
+  const auto *FuncDesc = Binary->getFuncDescForGUID(LeafProbe->getGuid());
+  bool WasLeafInlined = LeafProbe->getInlineTreeNode()->hasInlineSite();
   return getFunctionProfileForLeafProbe(ContextStrStackCopy, FuncDesc,
                                         WasLeafInlined);
 }
 
 } // end namespace sampleprof
 } // end namespace llvm

llvm/tools/llvm-profgen/ProfiledBinary.h

 //===-- ProfiledBinary.h - Binary decoder -----------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_TOOLS_LLVM_PROFGEN_PROFILEDBINARY_H
 #define LLVM_TOOLS_LLVM_PROFGEN_PROFILEDBINARY_H
 
 #include "CallContext.h"
-#include "PseudoProbe.h"
 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/DebugInfo/Symbolize/Symbolize.h"
 #include "llvm/MC/MCAsmInfo.h"
 #include "llvm/MC/MCContext.h"
 #include "llvm/MC/MCDisassembler/MCDisassembler.h"
 #include "llvm/MC/MCInst.h"
 #include "llvm/MC/MCInstPrinter.h"
 #include "llvm/MC/MCInstrAnalysis.h"
 #include "llvm/MC/MCInstrInfo.h"
 #include "llvm/MC/MCObjectFileInfo.h"
+#include "llvm/MC/MCPseudoProbe.h"
 #include "llvm/MC/MCRegisterInfo.h"
 #include "llvm/MC/MCSubtargetInfo.h"
 #include "llvm/MC/MCTargetOptions.h"
 #include "llvm/Object/ELFObjectFile.h"
 #include "llvm/ProfileData/SampleProf.h"
 #include "llvm/Support/Path.h"
 #include <list>
 #include <set>
class ProfiledBinary {
   std::unordered_set<uint64_t> RetAddrs;
 
   PrologEpilogTracker ProEpilogTracker;
 
   // The symbolizer used to get inline context for an instruction.
   std::unique_ptr<symbolize::LLVMSymbolizer> Symbolizer;
 
   // Pseudo probe decoder
-  PseudoProbeDecoder ProbeDecoder;
+  MCPseudoProbeDecoder ProbeDecoder;
 
   bool UsePseudoProbes = false;
 
   void setPreferredTextSegmentAddresses(const ELFObjectFileBase *O);
 
   template <class ELFT>
   void setPreferredTextSegmentAddresses(const ELFFile<ELFT> &Obj, StringRef FileName);
 
public:
   }
   uint64_t offsetToVirtualAddr(uint64_t Offset) const {
     return Offset + BaseAddress;
   }
   StringRef getPath() const { return Path; }
   StringRef getName() const { return llvm::sys::path::filename(Path); }
   uint64_t getBaseAddress() const { return BaseAddress; }
   void setBaseAddress(uint64_t Address) { BaseAddress = Address; }
-  
   // Return the preferred load address for the first executable segment.
   uint64_t getPreferredBaseAddress() const { return PreferredTextSegmentAddresses[0]; }
   // Return the file offset for the first executable segment.
   uint64_t getTextSegmentOffset() const { return TextSegmentOffsets[0]; }
   const std::vector<uint64_t> &getPreferredTextSegmentAddresses() const {
     return PreferredTextSegmentAddresses;
   }
   const std::vector<uint64_t> &getTextSegmentOffsets() const {
public:
   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 SmallVectorImpl<uint64_t> &Stack,
                                     bool &WasLeafInlined) const;
 
-  const PseudoProbe *getCallProbeForAddr(uint64_t Address) const {
+  const MCDecodedPseudoProbe *getCallProbeForAddr(uint64_t Address) const {
     return ProbeDecoder.getCallProbeForAddr(Address);
   }
+
   void
-  getInlineContextForProbe(const PseudoProbe *Probe,
+  getInlineContextForProbe(const MCDecodedPseudoProbe *Probe,
                            SmallVectorImpl<std::string> &InlineContextStack,
                            bool IncludeLeaf = false) const {
     return ProbeDecoder.getInlineContextForProbe(Probe, InlineContextStack,
                                                  IncludeLeaf);
   }
   const AddressProbesMap &getAddress2ProbesMap() const {
     return ProbeDecoder.getAddress2ProbesMap();
   }
-  const PseudoProbeFuncDesc *getFuncDescForGUID(uint64_t GUID) {
+  const MCPseudoProbeFuncDesc *getFuncDescForGUID(uint64_t GUID) {
     return ProbeDecoder.getFuncDescForGUID(GUID);
   }
-  const PseudoProbeFuncDesc *getInlinerDescForProbe(const PseudoProbe *Probe) {
+
+  const MCPseudoProbeFuncDesc *
+  getInlinerDescForProbe(const MCDecodedPseudoProbe *Probe) {
     return ProbeDecoder.getInlinerDescForProbe(Probe);
   }
 };
 
 } // end namespace sampleprof
 } // end namespace llvm
 
 #endif

llvm/tools/llvm-profgen/ProfiledBinary.cpp

void ProfiledBinary::decodePseudoProbe(const ELFObjectFileBase *Obj) {
   StringRef FileName = Obj->getFileName();
   for (section_iterator SI = Obj->section_begin(), SE = Obj->section_end();
        SI != SE; ++SI) {
     const SectionRef &Section = *SI;
     StringRef SectionName = unwrapOrError(Section.getName(), FileName);
 
     if (SectionName == ".pseudo_probe_desc") {
       StringRef Contents = unwrapOrError(Section.getContents(), FileName);
-      ProbeDecoder.buildGUID2FuncDescMap(
-          reinterpret_cast<const uint8_t *>(Contents.data()), Contents.size());
+      if (!ProbeDecoder.buildGUID2FuncDescMap(
+              reinterpret_cast<const uint8_t *>(Contents.data()),
+              Contents.size()))
+        exitWithError("Pseudo Probe decoder fail in " +
+                      ProbeDecoder.getSectionName());
     } else if (SectionName == ".pseudo_probe") {
       StringRef Contents = unwrapOrError(Section.getContents(), FileName);
-      ProbeDecoder.buildAddress2ProbeMap(
-          reinterpret_cast<const uint8_t *>(Contents.data()), Contents.size());
+      if (!ProbeDecoder.buildAddress2ProbeMap(
+              reinterpret_cast<const uint8_t *>(Contents.data()),
+              Contents.size()))
+        exitWithError("Pseudo Probe decoder fail in " +
+                      ProbeDecoder.getSectionName());
       // set UsePseudoProbes flag, used for PerfReader
       UsePseudoProbes = true;
     }
   }
 
   if (ShowPseudoProbe)
     ProbeDecoder.printGUID2FuncDescMap(outs());
 }

llvm/tools/llvm-profgen/PseudoProbe.h

-//===--- PseudoProbe.h - Pseudo probe decoding utilities ---------*- C++-*-===//
-//
-// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
-// See https://llvm.org/LICENSE.txt for license information.
-// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
-//
-//===----------------------------------------------------------------------===//
-#ifndef LLVM_TOOLS_LLVM_PROFGEN_PSEUDOPROBE_H
-#define LLVM_TOOLS_LLVM_PROFGEN_PSEUDOPROBE_H
-
-#include "llvm/ADT/StringRef.h"
-#include "llvm/ADT/Twine.h"
-#include "llvm/IR/PseudoProbe.h"
-#include "llvm/Support/raw_ostream.h"
-#include "llvm/Transforms/IPO/SampleProfileProbe.h"
-#include <algorithm>
-#include <set>
-#include <sstream>
-#include <string>
-#include <unordered_map>
-#include <unordered_set>
-#include <vector>
-
-namespace llvm {
-namespace sampleprof {
-
-enum PseudoProbeAttributes { RESERVED = 1 };
-
-// Use func GUID and index as the location info of the inline site
-using InlineSite = std::tuple<uint64_t, uint32_t>;
-
-struct PseudoProbe;
-
-// Tree node to represent the inline relation and its inline site, we use a
-// dummy root in the PseudoProbeDecoder to lead the tree, the outlined
-// function will directly be the children of the dummy root. For the inlined
-// function, all the inlinee will be connected to its inlineer, then further to
-// its outlined function. Pseudo probes originating from the function stores the
-// tree's leaf node which we can process backwards to get its inline context
-class PseudoProbeInlineTree {
-  std::vector<PseudoProbe *> ProbeVector;
-
-  struct InlineSiteHash {
-    uint64_t operator()(const InlineSite &Site) const {
-      return std::get<0>(Site) ^ std::get<1>(Site);
-    }
-  };
-  using InlinedProbeTreeMap =
-      std::unordered_map<InlineSite, std::unique_ptr<PseudoProbeInlineTree>,
-                         InlineSiteHash>;
-  InlinedProbeTreeMap Children;
-
-public:
-  // Inlinee function GUID
-  uint64_t GUID = 0;
-  // Inline site to indicate the location in its inliner. As the node could also
-  // be an outlined function, it will use a dummy InlineSite whose GUID and
-  // Index is 0 connected to the dummy root
-  InlineSite ISite;
-  // Used for decoding
-  uint32_t ChildrenToProcess = 0;
-  // Caller node of the inline site
-  PseudoProbeInlineTree *Parent;
-
-  PseudoProbeInlineTree(){};
-  PseudoProbeInlineTree(const InlineSite &Site) : ISite(Site){};
-
-  PseudoProbeInlineTree *getOrAddNode(const InlineSite &Site) {
-    auto Ret =
-        Children.emplace(Site, std::make_unique<PseudoProbeInlineTree>(Site));
-    Ret.first->second->Parent = this;
-    return Ret.first->second.get();
-  }
-
-  InlinedProbeTreeMap &getChildren() { return Children; }
-  std::vector<PseudoProbe *> &getProbes() { return ProbeVector; }
-  void addProbes(PseudoProbe *Probe) { ProbeVector.push_back(Probe); }
-  // Return false if it's a dummy inline site
-  bool hasInlineSite() const { return std::get<0>(ISite) != 0; }
-};
-
-// Function descriptor decoded from .pseudo_probe_desc section
-struct PseudoProbeFuncDesc {
-  uint64_t FuncGUID = 0;
-  uint64_t FuncHash = 0;
-  std::string FuncName;
-
-  PseudoProbeFuncDesc(uint64_t GUID, uint64_t Hash, StringRef Name)
-      : FuncGUID(GUID), FuncHash(Hash), FuncName(Name){};
-
-  void print(raw_ostream &OS);
-};
-
-// GUID to PseudoProbeFuncDesc map
-using GUIDProbeFunctionMap = std::unordered_map<uint64_t, PseudoProbeFuncDesc>;
-// Address to pseudo probes map.
-using AddressProbesMap = std::unordered_map<uint64_t, std::list<PseudoProbe>>;
-
-/*
-A pseudo probe has the format like below:
-  INDEX (ULEB128)
-  TYPE (uint4)
-    0 - block probe, 1 - indirect call, 2 - direct call
-  ATTRIBUTE (uint3)
-    1 - reserved
-  ADDRESS_TYPE (uint1)
-    0 - code address, 1 - address delta
-  CODE_ADDRESS (uint64 or ULEB128)
-  code address or address delta, depending on Flag
-*/
-struct PseudoProbe {
-  uint64_t Address;
-  uint64_t GUID;
-  uint32_t Index;
-  PseudoProbeType Type;
-  uint8_t Attribute;
-  PseudoProbeInlineTree *InlineTree;
-  const static uint32_t PseudoProbeFirstId =
-      static_cast<uint32_t>(PseudoProbeReservedId::Last) + 1;
-
-  PseudoProbe(uint64_t Ad, uint64_t G, uint32_t I, PseudoProbeType K,
-              uint8_t At, PseudoProbeInlineTree *Tree)
-      : Address(Ad), GUID(G), Index(I), Type(K), Attribute(At),
-        InlineTree(Tree){};
-
-  bool isEntry() const { return Index == PseudoProbeFirstId; }
-  bool isBlock() const { return Type == PseudoProbeType::Block; }
-  bool isIndirectCall() const { return Type == PseudoProbeType::IndirectCall; }
-  bool isDirectCall() const { return Type == PseudoProbeType::DirectCall; }
-  bool isCall() const { return isIndirectCall() || isDirectCall(); }
-
-  PseudoProbeInlineTree *getInlineTreeNode() const { return InlineTree; }
-
-  // Get the inlined context by traversing current inline tree backwards,
-  // each tree node has its InlineSite which is taken as the context.
-  // \p ContextStack is populated in root to leaf order
-  void getInlineContext(SmallVectorImpl<std::string> &ContextStack,
-                        const GUIDProbeFunctionMap &GUID2FuncMAP,
-                        bool ShowName) const;
-  // Helper function to get the string from context stack
-  std::string getInlineContextStr(const GUIDProbeFunctionMap &GUID2FuncMAP,
-                                  bool ShowName) const;
-  // Print pseudo probe while disassembling
-  void print(raw_ostream &OS, const GUIDProbeFunctionMap &GUID2FuncMAP,
-             bool ShowName);
-};
-
-/*
-Decode pseudo probe info from ELF section, used along with ELF reader
-Two sections are decoded here:
-  1) \fn buildGUID2FunctionMap is responsible for .pseudo_probe_desc
-  section which encodes all function descriptors.
-  2) \fn buildAddress2ProbeMap is responsible for .pseudoprobe section
-    which encodes an inline function forest and each tree includes its
-    inlined function and all pseudo probes inside the function.
-see \file MCPseudoProbe.h for the details of the section encoding format.
-*/
-class PseudoProbeDecoder {
-  // GUID to PseudoProbeFuncDesc map.
-  GUIDProbeFunctionMap GUID2FuncDescMap;
-
-  // Address to probes map.
-  AddressProbesMap Address2ProbesMap;
-
-  // The dummy root of the inline trie, all the outlined function will directly
-  // be the children of the dummy root, all the inlined function will be the
-  // children of its inlineer. So the relation would be like:
-  // DummyRoot --> OutlinedFunc --> InlinedFunc1 --> InlinedFunc2
-  PseudoProbeInlineTree DummyInlineRoot;
-
-  /// Points to the current location in the buffer.
-  const uint8_t *Data = nullptr;
-
-  /// Points to the end of the buffer.
-  const uint8_t *End = nullptr;
-
-  /// SectionName used for debug
-  std::string SectionName;
-
-  // Decoding helper function
-  template <typename T> T readUnencodedNumber();
-  template <typename T> T readUnsignedNumber();
-  template <typename T> T readSignedNumber();
-  StringRef readString(uint32_t Size);
-
-public:
-  // Decode pseudo_probe_desc section to build GUID to PseudoProbeFuncDesc map.
-  void buildGUID2FuncDescMap(const uint8_t *Start, std::size_t Size);
-
-  // 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;
-
-  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]
-  void
-  getInlineContextForProbe(const PseudoProbe *Probe,
-                           SmallVectorImpl<std::string> &InlineContextStack,
-                           bool IncludeLeaf) const;
-
-  const AddressProbesMap &getAddress2ProbesMap() const {
-    return Address2ProbesMap;
-  }
-
-  const PseudoProbeFuncDesc *
-  getInlinerDescForProbe(const PseudoProbe *Probe) const;
-};
-
-} // end namespace sampleprof
-} // end namespace llvm
-
-#endif

llvm/tools/llvm-profgen/PseudoProbe.cpp

-//===--- PseudoProbe.cpp - Pseudo probe decoding utilities  ------*- C++-*-===//
-//
-// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
-// See https://llvm.org/LICENSE.txt for license information.
-// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
-//
-//===----------------------------------------------------------------------===//
-
-#include "PseudoProbe.h"
-#include "ErrorHandling.h"
-#include "llvm/Support/Endian.h"
-#include "llvm/Support/LEB128.h"
-#include "llvm/Support/raw_ostream.h"
-#include <limits>
-#include <memory>
-
-using namespace llvm;
-using namespace sampleprof;
-using namespace support;
-
-namespace llvm {
-namespace sampleprof {
-
-static StringRef getProbeFNameForGUID(const GUIDProbeFunctionMap &GUID2FuncMAP,
-                                      uint64_t GUID) {
-  auto It = GUID2FuncMAP.find(GUID);
-  assert(It != GUID2FuncMAP.end() &&
-         "Probe function must exist for a valid GUID");
-  return It->second.FuncName;
-}
-
-void PseudoProbeFuncDesc::print(raw_ostream &OS) {
-  OS << "GUID: " << FuncGUID << " Name: " << FuncName << "\n";
-  OS << "Hash: " << FuncHash << "\n";
-}
-
-void PseudoProbe::getInlineContext(SmallVectorImpl<std::string> &ContextStack,
-                                   const GUIDProbeFunctionMap &GUID2FuncMAP,
-                                   bool ShowName) const {
-  uint32_t Begin = ContextStack.size();
-  PseudoProbeInlineTree *Cur = InlineTree;
-  // It will add the string of each node's inline site during iteration.
-  // Note that it won't include the probe's belonging function(leaf location)
-  while (Cur->hasInlineSite()) {
-    std::string ContextStr;
-    if (ShowName) {
-      StringRef FuncName =
-          getProbeFNameForGUID(GUID2FuncMAP, std::get<0>(Cur->ISite));
-      ContextStr += FuncName.str();
-    } else {
-      ContextStr += Twine(std::get<0>(Cur->ISite)).str();
-    }
-    ContextStr += ":";
-    ContextStr += Twine(std::get<1>(Cur->ISite)).str();
-    ContextStack.emplace_back(ContextStr);
-    Cur = Cur->Parent;
-  }
-  // Make the ContextStack in caller-callee order
-  std::reverse(ContextStack.begin() + Begin, ContextStack.end());
-}
-
-std::string
-PseudoProbe::getInlineContextStr(const GUIDProbeFunctionMap &GUID2FuncMAP,
-                                 bool ShowName) const {
-  std::ostringstream OContextStr;
-  SmallVector<std::string, 16> ContextStack;
-  getInlineContext(ContextStack, GUID2FuncMAP, ShowName);
-  for (auto &CxtStr : ContextStack) {
-    if (OContextStr.str().size())
-      OContextStr << " @ ";
-    OContextStr << CxtStr;
-  }
-  return OContextStr.str();
-}
-
-static const char *PseudoProbeTypeStr[3] = {"Block", "IndirectCall",
-                                            "DirectCall"};
-
-void PseudoProbe::print(raw_ostream &OS,
-                        const GUIDProbeFunctionMap &GUID2FuncMAP,
-                        bool ShowName) {
-  OS << "FUNC: ";
-  if (ShowName) {
-    StringRef FuncName = getProbeFNameForGUID(GUID2FuncMAP, GUID);
-    OS << FuncName.str() << " ";
-  } else {
-    OS << GUID << " ";
-  }
-  OS << "Index: " << Index << "  ";
-  OS << "Type: " << PseudoProbeTypeStr[static_cast<uint8_t>(Type)] << "  ";
-
-  std::string InlineContextStr = getInlineContextStr(GUID2FuncMAP, ShowName);
-  if (InlineContextStr.size()) {
-    OS << "Inlined: @ ";
-    OS << InlineContextStr;
-  }
-  OS << "\n";
-}
-
-template <typename T> T PseudoProbeDecoder::readUnencodedNumber() {
-  if (Data + sizeof(T) > End) {
-    exitWithError("Decode unencoded number error in " + SectionName +
-                  " section");
-  }
-  T Val = endian::readNext<T, little, unaligned>(Data);
-  return Val;
-}
-
-template <typename T> T PseudoProbeDecoder::readUnsignedNumber() {
-  unsigned NumBytesRead = 0;
-  uint64_t Val = decodeULEB128(Data, &NumBytesRead);
-  if (Val > std::numeric_limits<T>::max() || (Data + NumBytesRead > End)) {
-    exitWithError("Decode number error in " + SectionName + " section");
-  }
-  Data += NumBytesRead;
-  return static_cast<T>(Val);
-}
-
-template <typename T> T PseudoProbeDecoder::readSignedNumber() {
-  unsigned NumBytesRead = 0;
-  int64_t Val = decodeSLEB128(Data, &NumBytesRead);
-  if (Val > std::numeric_limits<T>::max() || (Data + NumBytesRead > End)) {
-    exitWithError("Decode number error in " + SectionName + " section");
-  }
-  Data += NumBytesRead;
-  return static_cast<T>(Val);
-}
-
-StringRef PseudoProbeDecoder::readString(uint32_t Size) {
-  StringRef Str(reinterpret_cast<const char *>(Data), Size);
-  if (Data + Size > End) {
-    exitWithError("Decode string error in " + SectionName + " section");
-  }
-  Data += Size;
-  return Str;
-}
-
-void PseudoProbeDecoder::buildGUID2FuncDescMap(const uint8_t *Start,
-                                               std::size_t Size) {
-  // The pseudo_probe_desc section has a format like:
-  // .section .pseudo_probe_desc,"",@progbits
-  // .quad -5182264717993193164   // GUID
-  // .quad 4294967295             // Hash
-  // .uleb 3                      // Name size
-  // .ascii "foo"                 // Name
-  // .quad -2624081020897602054
-  // .quad 174696971957
-  // .uleb 34
-  // .ascii "main"
-#ifndef NDEBUG
-  SectionName = "pseudo_probe_desc";
-#endif
-  Data = Start;
-  End = Data + Size;
-
-  while (Data < End) {
-    uint64_t GUID = readUnencodedNumber<uint64_t>();
-    uint64_t Hash = readUnencodedNumber<uint64_t>();
-    uint32_t NameSize = readUnsignedNumber<uint32_t>();
-    StringRef Name = FunctionSamples::getCanonicalFnName(readString(NameSize));
-
-    // Initialize PseudoProbeFuncDesc and populate it into GUID2FuncDescMap
-    GUID2FuncDescMap.emplace(GUID, PseudoProbeFuncDesc(GUID, Hash, Name));
-  }
-  assert(Data == End && "Have unprocessed data in pseudo_probe_desc section");
-}
-
-void PseudoProbeDecoder::buildAddress2ProbeMap(const uint8_t *Start,
-                                               std::size_t Size) {
-  // The pseudo_probe section encodes an inline forest and each tree has a
-  // format like:
-  //  FUNCTION BODY (one for each uninlined function present in the text
-  //  section)
-  //     GUID (uint64)
-  //         GUID of the function
-  //     NPROBES (ULEB128)
-  //         Number of probes originating from this function.
-  //     NUM_INLINED_FUNCTIONS (ULEB128)
-  //         Number of callees inlined into this function, aka number of
-  //         first-level inlinees
-  //     PROBE RECORDS
-  //         A list of NPROBES entries. Each entry contains:
-  //           INDEX (ULEB128)
-  //           TYPE (uint4)
-  //             0 - block probe, 1 - indirect call, 2 - direct call
-  //           ATTRIBUTE (uint3)
-  //             1 - reserved
-  //           ADDRESS_TYPE (uint1)
-  //             0 - code address, 1 - address delta
-  //           CODE_ADDRESS (uint64 or ULEB128)
-  //             code address or address delta, depending on Flag
-  //     INLINED FUNCTION RECORDS
-  //         A list of NUM_INLINED_FUNCTIONS entries describing each of the
-  //         inlined callees.  Each record contains:
-  //           INLINE SITE
-  //             Index of the callsite probe (ULEB128)
-  //           FUNCTION BODY
-  //             A FUNCTION BODY entry describing the inlined function.
-#ifndef NDEBUG
-  SectionName = "pseudo_probe";
-#endif
-  Data = Start;
-  End = Data + Size;
-
-  PseudoProbeInlineTree *Root = &DummyInlineRoot;
-  PseudoProbeInlineTree *Cur = &DummyInlineRoot;
-  uint64_t LastAddr = 0;
-  uint32_t Index = 0;
-  // A DFS-based decoding
-  while (Data < End) {
-    if (Root == Cur) {
-      // Use a sequential id for top level inliner.
-      Index = Root->getChildren().size();
-    } else {
-      // Read inline site for inlinees
-      Index = readUnsignedNumber<uint32_t>();
-    }
-    // Switch/add to a new tree node(inlinee)
-    Cur = Cur->getOrAddNode(std::make_tuple(Cur->GUID, Index));
-    // Read guid
-    Cur->GUID = readUnencodedNumber<uint64_t>();
-    // Read number of probes in the current node.
-    uint32_t NodeCount = readUnsignedNumber<uint32_t>();
-    // Read number of direct inlinees
-    Cur->ChildrenToProcess = readUnsignedNumber<uint32_t>();
-    // Read all probes in this node
-    for (std::size_t I = 0; I < NodeCount; I++) {
-      // Read index
-      uint32_t Index = readUnsignedNumber<uint32_t>();
-      // Read type | flag.
-      uint8_t Value = readUnencodedNumber<uint8_t>();
-      uint8_t Kind = Value & 0xf;
-      uint8_t Attr = (Value & 0x70) >> 4;
-      // Read address
-      uint64_t Addr = 0;
-      if (Value & 0x80) {
-        int64_t Offset = readSignedNumber<int64_t>();
-        Addr = LastAddr + Offset;
-      } else {
-        Addr = readUnencodedNumber<int64_t>();
-      }
-      // Populate Address2ProbesMap
-      auto &Probes = Address2ProbesMap[Addr];
-      Probes.emplace_back(Addr, Cur->GUID, Index, PseudoProbeType(Kind), Attr,
-                          Cur);
-      Cur->addProbes(&Probes.back());
-      LastAddr = Addr;
-    }
-
-    // Look for the parent for the next node by subtracting the current
-    // node count from tree counts along the parent chain. The first node
-    // in the chain that has a non-zero tree count is the target.
-    while (Cur != Root) {
-      if (Cur->ChildrenToProcess == 0) {
-        Cur = Cur->Parent;
-        if (Cur != Root) {
-          assert(Cur->ChildrenToProcess > 0 &&
-                 "Should have some unprocessed nodes");
-          Cur->ChildrenToProcess -= 1;
-        }
-      } else {
-        break;
-      }
-    }
-  }
-
-  assert(Data == End && "Have unprocessed data in pseudo_probe section");
-  assert(Cur == Root &&
-         " Cur should point to root when the forest is fully built up");
-}
-
-void PseudoProbeDecoder::printGUID2FuncDescMap(raw_ostream &OS) {
-  OS << "Pseudo Probe Desc:\n";
-  // Make the output deterministic
-  std::map<uint64_t, PseudoProbeFuncDesc> OrderedMap(GUID2FuncDescMap.begin(),
-                                                     GUID2FuncDescMap.end());
-  for (auto &I : OrderedMap) {
-    I.second.print(OS);
-  }
-}
-
-void PseudoProbeDecoder::printProbeForAddress(raw_ostream &OS,
-                                              uint64_t Address) {
-  auto It = Address2ProbesMap.find(Address);
-  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 auto &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;
-}
-
-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, SmallVectorImpl<std::string> &InlineContextStack,
-    bool IncludeLeaf) const {
-  Probe->getInlineContext(InlineContextStack, GUID2FuncDescMap, true);
-  if (!IncludeLeaf)
-    return;
-  // Note that the context from probe doesn't include leaf frame,
-  // hence we need to retrieve and prepend leaf if requested.
-  const auto *FuncDesc = getFuncDescForGUID(Probe->GUID);
-  InlineContextStack.emplace_back(FuncDesc->FuncName + ":" +
-                                  Twine(Probe->Index).str());
-}
-
-const PseudoProbeFuncDesc *
-PseudoProbeDecoder::getInlinerDescForProbe(const PseudoProbe *Probe) const {
-  PseudoProbeInlineTree *InlinerNode = Probe->InlineTree;
-  if (!InlinerNode->hasInlineSite())
-    return nullptr;
-  return getFuncDescForGUID(std::get<0>(InlinerNode->ISite));
-}
-
-} // end namespace sampleprof
-} // end namespace llvm