[BOLT][NFCI] Migrate pseudo probes to MetadataRewriter interface

Use new MetdataRewriter interface to update pseudo probes and move
ProbeDecoder out of BinaryContext into new PseudoProbeRewriter class.

Depends on D154021

Reviewed By: Amir

Differential Revision: https://reviews.llvm.org/D154022
Differential Revision: https://reviews.llvm.org/D154023

7 files changed, 424 insertions, 380 deletions

diff --git a/bolt/include/bolt/Core/BinaryContext.h b/bolt/include/bolt/Core/BinaryContext.h
index bbe385f..40aa959 100644
--- a/bolt/include/bolt/Core/BinaryContext.h
+++ b/bolt/include/bolt/Core/BinaryContext.h
@@ -29,7 +29,6 @@
 #include "llvm/MC/MCContext.h"
 #include "llvm/MC/MCObjectFileInfo.h"
 #include "llvm/MC/MCObjectWriter.h"
-#include "llvm/MC/MCPseudoProbe.h"
 #include "llvm/MC/MCSectionELF.h"
 #include "llvm/MC/MCSectionMachO.h"
 #include "llvm/MC/MCStreamer.h"
@@ -681,9 +680,6 @@ public:
   /// and are referenced from BinaryFunction.
   std::list<std::pair<BinaryFunction *, uint64_t>> InterproceduralReferences;
 
-  /// PseudoProbe decoder
-  MCPseudoProbeDecoder ProbeDecoder;
-
   /// DWARF encoding. Available encoding types defined in BinaryFormat/Dwarf.h
   /// enum Constants, e.g. DW_EH_PE_omit.
   unsigned LSDAEncoding = dwarf::DW_EH_PE_omit;
diff --git a/bolt/include/bolt/Core/BinaryFunction.h b/bolt/include/bolt/Core/BinaryFunction.h
index d0927af..f907f9a 100644
--- a/bolt/include/bolt/Core/BinaryFunction.h
+++ b/bolt/include/bolt/Core/BinaryFunction.h
@@ -423,21 +423,6 @@ private:
     return BB->getIndex();
   }
 
-  /// Return basic block that originally contained offset \p Offset
-  /// from the function start.
-  BinaryBasicBlock *getBasicBlockContainingOffset(uint64_t Offset);
-
-  const BinaryBasicBlock *getBasicBlockContainingOffset(uint64_t Offset) const {
-    return const_cast<BinaryFunction *>(this)->getBasicBlockContainingOffset(
-        Offset);
-  }
-
-  /// Return basic block that started at offset \p Offset.
-  BinaryBasicBlock *getBasicBlockAtOffset(uint64_t Offset) {
-    BinaryBasicBlock *BB = getBasicBlockContainingOffset(Offset);
-    return BB && BB->getOffset() == Offset ? BB : nullptr;
-  }
-
   /// Release memory taken by the list.
   template <typename T> BinaryFunction &clearList(T &List) {
     T TempList;
@@ -900,6 +885,21 @@ public:
     return LabelToBB.lookup(Label);
   }
 
+  /// Return basic block that originally contained offset \p Offset
+  /// from the function start.
+  BinaryBasicBlock *getBasicBlockContainingOffset(uint64_t Offset);
+
+  const BinaryBasicBlock *getBasicBlockContainingOffset(uint64_t Offset) const {
+    return const_cast<BinaryFunction *>(this)->getBasicBlockContainingOffset(
+        Offset);
+  }
+
+  /// Return basic block that started at offset \p Offset.
+  BinaryBasicBlock *getBasicBlockAtOffset(uint64_t Offset) {
+    BinaryBasicBlock *BB = getBasicBlockContainingOffset(Offset);
+    return BB && BB->getOffset() == Offset ? BB : nullptr;
+  }
+
   /// Retrieve the landing pad BB associated with invoke instruction \p Invoke
   /// that is in \p BB. Return nullptr if none exists
   BinaryBasicBlock *getLandingPadBBFor(const BinaryBasicBlock &BB,
diff --git a/bolt/include/bolt/Rewrite/MetadataRewriters.h b/bolt/include/bolt/Rewrite/MetadataRewriters.h
index 56de1fa..c80c9fc 100644
--- a/bolt/include/bolt/Rewrite/MetadataRewriters.h
+++ b/bolt/include/bolt/Rewrite/MetadataRewriters.h
@@ -19,6 +19,8 @@ class BinaryContext;
 
 // The list of rewriter build functions.
 
+std::unique_ptr<MetadataRewriter> createPseudoProbeRewriter(BinaryContext &);
+
 std::unique_ptr<MetadataRewriter> createSDTRewriter(BinaryContext &);
 
 } // namespace bolt
diff --git a/bolt/include/bolt/Rewrite/RewriteInstance.h b/bolt/include/bolt/Rewrite/RewriteInstance.h
index 1fe5333..9cabef8 100644
--- a/bolt/include/bolt/Rewrite/RewriteInstance.h
+++ b/bolt/include/bolt/Rewrite/RewriteInstance.h
@@ -201,12 +201,6 @@ private:
   /// Update LKMarkers' locations for the output binary.
   void updateLKMarkers();
 
-  /// Update address of MCDecodedPseudoProbe.
-  void updatePseudoProbes();
-
-  /// Encode MCDecodedPseudoProbe.
-  void encodePseudoProbes();
-
   /// Return the list of code sections in the output order.
   std::vector<BinarySection *> getCodeSections();
 
@@ -395,10 +389,6 @@ private:
   /// of appending contents to it.
   bool willOverwriteSection(StringRef SectionName);
 
-  /// Parse .pseudo_probe_desc section and .pseudo_probe section
-  /// Setup Pseudo probe decoder
-  void parsePseudoProbe();
-
 public:
   /// Standard ELF sections we overwrite.
   static constexpr const char *SectionsToOverwrite[] = {
@@ -580,15 +570,6 @@ private:
   /// .note.gnu.build-id section.
   ErrorOr<BinarySection &> BuildIDSection{std::errc::bad_address};
 
-  /// .pseudo_probe_desc section.
-  /// Contains information about pseudo probe description, like its related
-  /// function
-  ErrorOr<BinarySection &> PseudoProbeDescSection{std::errc::bad_address};
-
-  /// .pseudo_probe section.
-  /// Contains information about pseudo probe details, like its address
-  ErrorOr<BinarySection &> PseudoProbeSection{std::errc::bad_address};
-
   /// Helper for accessing sections by name.
   BinarySection *getSection(const Twine &Name) {
     ErrorOr<BinarySection &> ErrOrSection = BC->getUniqueSectionByName(Name);
diff --git a/bolt/lib/Rewrite/CMakeLists.txt b/bolt/lib/Rewrite/CMakeLists.txt
index 7d633c4..29cc5fc 100644
--- a/bolt/lib/Rewrite/CMakeLists.txt
+++ b/bolt/lib/Rewrite/CMakeLists.txt
@@ -16,6 +16,7 @@ add_llvm_library(LLVMBOLTRewrite
   JITLinkLinker.cpp
   MachORewriteInstance.cpp
   MetadataManager.cpp
+  PseudoProbeRewriter.cpp
   RewriteInstance.cpp
   SDTRewriter.cpp
 
diff --git a/bolt/lib/Rewrite/PseudoProbeRewriter.cpp b/bolt/lib/Rewrite/PseudoProbeRewriter.cpp
new file mode 100644
index 0000000..64b8a8b
--- /dev/null
+++ b/bolt/lib/Rewrite/PseudoProbeRewriter.cpp
@@ -0,0 +1,402 @@
+//===- bolt/Rewrite/PseudoProbeRewriter.cpp -------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Implement support for pseudo probes.
+//
+//===----------------------------------------------------------------------===//
+
+#include "bolt/Core/BinaryFunction.h"
+#include "bolt/Rewrite/MetadataRewriter.h"
+#include "bolt/Rewrite/MetadataRewriters.h"
+#include "bolt/Utils/CommandLineOpts.h"
+#include "llvm/IR/Function.h"
+#include "llvm/MC/MCPseudoProbe.h"
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/Debug.h"
+#include "llvm/Support/LEB128.h"
+
+#undef DEBUG_TYPE
+#define DEBUG_TYPE "pseudo-probe-rewriter"
+
+using namespace llvm;
+using namespace bolt;
+
+namespace opts {
+
+enum PrintPseudoProbesOptions {
+  PPP_None = 0,
+  PPP_Probes_Section_Decode = 0x1,
+  PPP_Probes_Address_Conversion = 0x2,
+  PPP_Encoded_Probes = 0x3,
+  PPP_All = 0xf
+};
+
+static cl::opt<PrintPseudoProbesOptions> PrintPseudoProbes(
+    "print-pseudo-probes", cl::desc("print pseudo probe info"),
+    cl::init(PPP_None),
+    cl::values(clEnumValN(PPP_Probes_Section_Decode, "decode",
+                          "decode probes section from binary"),
+               clEnumValN(PPP_Probes_Address_Conversion, "address_conversion",
+                          "update address2ProbesMap with output block address"),
+               clEnumValN(PPP_Encoded_Probes, "encoded_probes",
+                          "display the encoded probes in binary section"),
+               clEnumValN(PPP_All, "all", "enable all debugging printout")),
+    cl::Hidden, cl::cat(BoltCategory));
+
+} // namespace opts
+
+namespace {
+class PseudoProbeRewriter final : public MetadataRewriter {
+  /// .pseudo_probe_desc section.
+  /// Contains information about pseudo probe description, like its related
+  /// function
+  ErrorOr<BinarySection &> PseudoProbeDescSection{std::errc::bad_address};
+
+  /// .pseudo_probe section.
+  /// Contains information about pseudo probe details, like its address
+  ErrorOr<BinarySection &> PseudoProbeSection{std::errc::bad_address};
+
+  /// Update address of MCDecodedPseudoProbe.
+  void updatePseudoProbes();
+
+  /// Encode MCDecodedPseudoProbe.
+  void encodePseudoProbes();
+
+  /// Parse .pseudo_probe_desc section and .pseudo_probe section
+  /// Setup Pseudo probe decoder
+  void parsePseudoProbe();
+
+  /// PseudoProbe decoder
+  MCPseudoProbeDecoder ProbeDecoder;
+
+public:
+  PseudoProbeRewriter(BinaryContext &BC)
+      : MetadataRewriter("pseudo-probe-rewriter", BC) {}
+
+  Error postEmitFinalizer() override;
+};
+
+Error PseudoProbeRewriter::postEmitFinalizer() {
+  parsePseudoProbe();
+  updatePseudoProbes();
+
+  return Error::success();
+}
+
+void PseudoProbeRewriter::parsePseudoProbe() {
+  PseudoProbeDescSection = BC.getUniqueSectionByName(".pseudo_probe_desc");
+  PseudoProbeSection = BC.getUniqueSectionByName(".pseudo_probe");
+
+  if (!PseudoProbeDescSection && !PseudoProbeSection) {
+    // pesudo probe is not added to binary. It is normal and no warning needed.
+    return;
+  }
+
+  // If only one section is found, it might mean the ELF is corrupted.
+  if (!PseudoProbeDescSection) {
+    errs() << "BOLT-WARNING: fail in reading .pseudo_probe_desc binary\n";
+    return;
+  } else if (!PseudoProbeSection) {
+    errs() << "BOLT-WARNING: fail in reading .pseudo_probe binary\n";
+    return;
+  }
+
+  StringRef Contents = PseudoProbeDescSection->getContents();
+  if (!ProbeDecoder.buildGUID2FuncDescMap(
+          reinterpret_cast<const uint8_t *>(Contents.data()),
+          Contents.size())) {
+    errs() << "BOLT-WARNING: fail in building GUID2FuncDescMap\n";
+    return;
+  }
+
+  MCPseudoProbeDecoder::Uint64Set GuidFilter;
+  MCPseudoProbeDecoder::Uint64Map FuncStartAddrs;
+  for (const BinaryFunction *F : BC.getAllBinaryFunctions()) {
+    for (const MCSymbol *Sym : F->getSymbols()) {
+      FuncStartAddrs[Function::getGUID(NameResolver::restore(Sym->getName()))] =
+          F->getAddress();
+    }
+  }
+  Contents = PseudoProbeSection->getContents();
+  if (!ProbeDecoder.buildAddress2ProbeMap(
+          reinterpret_cast<const uint8_t *>(Contents.data()), Contents.size(),
+          GuidFilter, FuncStartAddrs)) {
+    ProbeDecoder.getAddress2ProbesMap().clear();
+    errs() << "BOLT-WARNING: fail in building Address2ProbeMap\n";
+    return;
+  }
+
+  if (opts::PrintPseudoProbes == opts::PrintPseudoProbesOptions::PPP_All ||
+      opts::PrintPseudoProbes ==
+          opts::PrintPseudoProbesOptions::PPP_Probes_Section_Decode) {
+    outs() << "Report of decoding input pseudo probe binaries \n";
+    ProbeDecoder.printGUID2FuncDescMap(outs());
+    ProbeDecoder.printProbesForAllAddresses(outs());
+  }
+}
+
+void PseudoProbeRewriter::updatePseudoProbes() {
+  // check if there is pseudo probe section decoded
+  if (ProbeDecoder.getAddress2ProbesMap().empty())
+    return;
+  // input address converted to output
+  AddressProbesMap &Address2ProbesMap = ProbeDecoder.getAddress2ProbesMap();
+  const GUIDProbeFunctionMap &GUID2Func = ProbeDecoder.getGUID2FuncDescMap();
+
+  for (auto &AP : Address2ProbesMap) {
+    BinaryFunction *F = BC.getBinaryFunctionContainingAddress(AP.first);
+    // If F is removed, eliminate all probes inside it from inline tree
+    // Setting probes' addresses as INT64_MAX means elimination
+    if (!F) {
+      for (MCDecodedPseudoProbe &Probe : AP.second)
+        Probe.setAddress(INT64_MAX);
+      continue;
+    }
+    // If F is not emitted, the function will remain in the same address as its
+    // input
+    if (!F->isEmitted())
+      continue;
+
+    uint64_t Offset = AP.first - F->getAddress();
+    const BinaryBasicBlock *BB = F->getBasicBlockContainingOffset(Offset);
+    uint64_t BlkOutputAddress = BB->getOutputAddressRange().first;
+    // Check if block output address is defined.
+    // If not, such block is removed from binary. Then remove the probes from
+    // inline tree
+    if (BlkOutputAddress == 0) {
+      for (MCDecodedPseudoProbe &Probe : AP.second)
+        Probe.setAddress(INT64_MAX);
+      continue;
+    }
+
+    unsigned ProbeTrack = AP.second.size();
+    std::list<MCDecodedPseudoProbe>::iterator Probe = AP.second.begin();
+    while (ProbeTrack != 0) {
+      if (Probe->isBlock()) {
+        Probe->setAddress(BlkOutputAddress);
+      } else if (Probe->isCall()) {
+        // A call probe may be duplicated due to ICP
+        // Go through output of InputOffsetToAddressMap to collect all related
+        // probes
+        const InputOffsetToAddressMapTy &Offset2Addr =
+            F->getInputOffsetToAddressMap();
+        auto CallOutputAddresses = Offset2Addr.equal_range(Offset);
+        auto CallOutputAddress = CallOutputAddresses.first;
+        if (CallOutputAddress == CallOutputAddresses.second) {
+          Probe->setAddress(INT64_MAX);
+        } else {
+          Probe->setAddress(CallOutputAddress->second);
+          CallOutputAddress = std::next(CallOutputAddress);
+        }
+
+        while (CallOutputAddress != CallOutputAddresses.second) {
+          AP.second.push_back(*Probe);
+          AP.second.back().setAddress(CallOutputAddress->second);
+          Probe->getInlineTreeNode()->addProbes(&(AP.second.back()));
+          CallOutputAddress = std::next(CallOutputAddress);
+        }
+      }
+      Probe = std::next(Probe);
+      ProbeTrack--;
+    }
+  }
+
+  if (opts::PrintPseudoProbes == opts::PrintPseudoProbesOptions::PPP_All ||
+      opts::PrintPseudoProbes ==
+          opts::PrintPseudoProbesOptions::PPP_Probes_Address_Conversion) {
+    outs() << "Pseudo Probe Address Conversion results:\n";
+    // table that correlates address to block
+    std::unordered_map<uint64_t, StringRef> Addr2BlockNames;
+    for (auto &F : BC.getBinaryFunctions())
+      for (BinaryBasicBlock &BinaryBlock : F.second)
+        Addr2BlockNames[BinaryBlock.getOutputAddressRange().first] =
+            BinaryBlock.getName();
+
+    // scan all addresses -> correlate probe to block when print out
+    std::vector<uint64_t> Addresses;
+    for (auto &Entry : Address2ProbesMap)
+      Addresses.push_back(Entry.first);
+    llvm::sort(Addresses);
+    for (uint64_t Key : Addresses) {
+      for (MCDecodedPseudoProbe &Probe : Address2ProbesMap[Key]) {
+        if (Probe.getAddress() == INT64_MAX)
+          outs() << "Deleted Probe: ";
+        else
+          outs() << "Address: " << format_hex(Probe.getAddress(), 8) << " ";
+        Probe.print(outs(), GUID2Func, true);
+        // print block name only if the probe is block type and undeleted.
+        if (Probe.isBlock() && Probe.getAddress() != INT64_MAX)
+          outs() << format_hex(Probe.getAddress(), 8) << " Probe is in "
+                 << Addr2BlockNames[Probe.getAddress()] << "\n";
+      }
+    }
+    outs() << "=======================================\n";
+  }
+
+  // encode pseudo probes with updated addresses
+  encodePseudoProbes();
+}
+
+void PseudoProbeRewriter::encodePseudoProbes() {
+  // Buffer for new pseudo probes section
+  SmallString<8> Contents;
+  MCDecodedPseudoProbe *LastProbe = nullptr;
+
+  auto EmitInt = [&](uint64_t Value, uint32_t Size) {
+    const bool IsLittleEndian = BC.AsmInfo->isLittleEndian();
+    uint64_t Swapped = support::endian::byte_swap(
+        Value, IsLittleEndian ? support::little : support::big);
+    unsigned Index = IsLittleEndian ? 0 : 8 - Size;
+    auto Entry = StringRef(reinterpret_cast<char *>(&Swapped) + Index, Size);
+    Contents.append(Entry.begin(), Entry.end());
+  };
+
+  auto EmitULEB128IntValue = [&](uint64_t Value) {
+    SmallString<128> Tmp;
+    raw_svector_ostream OSE(Tmp);
+    encodeULEB128(Value, OSE, 0);
+    Contents.append(OSE.str().begin(), OSE.str().end());
+  };
+
+  auto EmitSLEB128IntValue = [&](int64_t Value) {
+    SmallString<128> Tmp;
+    raw_svector_ostream OSE(Tmp);
+    encodeSLEB128(Value, OSE);
+    Contents.append(OSE.str().begin(), OSE.str().end());
+  };
+
+  // Emit indiviual pseudo probes in a inline tree node
+  // Probe index, type, attribute, address type and address are encoded
+  // Address of the first probe is absolute.
+  // Other probes' address are represented by delta
+  auto EmitDecodedPseudoProbe = [&](MCDecodedPseudoProbe *&CurProbe) {
+    assert(!isSentinelProbe(CurProbe->getAttributes()) &&
+           "Sentinel probes should not be emitted");
+    EmitULEB128IntValue(CurProbe->getIndex());
+    uint8_t PackedType = CurProbe->getType() | (CurProbe->getAttributes() << 4);
+    uint8_t Flag =
+        LastProbe ? ((int8_t)MCPseudoProbeFlag::AddressDelta << 7) : 0;
+    EmitInt(Flag | PackedType, 1);
+    if (LastProbe) {
+      // Emit the delta between the address label and LastProbe.
+      int64_t Delta = CurProbe->getAddress() - LastProbe->getAddress();
+      EmitSLEB128IntValue(Delta);
+    } else {
+      // Emit absolute address for encoding the first pseudo probe.
+      uint32_t AddrSize = BC.AsmInfo->getCodePointerSize();
+      EmitInt(CurProbe->getAddress(), AddrSize);
+    }
+  };
+
+  std::map<InlineSite, MCDecodedPseudoProbeInlineTree *,
+           std::greater<InlineSite>>
+      Inlinees;
+
+  // DFS of inline tree to emit pseudo probes in all tree node
+  // Inline site index of a probe is emitted first.
+  // Then tree node Guid, size of pseudo probes and children nodes, and detail
+  // of contained probes are emitted Deleted probes are skipped Root node is not
+  // encoded to binaries. It's a "wrapper" of inline trees of each function.
+  std::list<std::pair<uint64_t, MCDecodedPseudoProbeInlineTree *>> NextNodes;
+  const MCDecodedPseudoProbeInlineTree &Root =
+      ProbeDecoder.getDummyInlineRoot();
+  for (auto Child = Root.getChildren().begin();
+       Child != Root.getChildren().end(); ++Child)
+    Inlinees[Child->first] = Child->second.get();
+
+  for (auto Inlinee : Inlinees)
+    // INT64_MAX is "placeholder" of unused callsite index field in the pair
+    NextNodes.push_back({INT64_MAX, Inlinee.second});
+
+  Inlinees.clear();
+
+  while (!NextNodes.empty()) {
+    uint64_t ProbeIndex = NextNodes.back().first;
+    MCDecodedPseudoProbeInlineTree *Cur = NextNodes.back().second;
+    NextNodes.pop_back();
+
+    if (Cur->Parent && !Cur->Parent->isRoot())
+      // Emit probe inline site
+      EmitULEB128IntValue(ProbeIndex);
+
+    // Emit probes grouped by GUID.
+    LLVM_DEBUG({
+      dbgs().indent(MCPseudoProbeTable::DdgPrintIndent);
+      dbgs() << "GUID: " << Cur->Guid << "\n";
+    });
+    // Emit Guid
+    EmitInt(Cur->Guid, 8);
+    // Emit number of probes in this node
+    uint64_t Deleted = 0;
+    for (MCDecodedPseudoProbe *&Probe : Cur->getProbes())
+      if (Probe->getAddress() == INT64_MAX)
+        Deleted++;
+    LLVM_DEBUG(dbgs() << "Deleted Probes:" << Deleted << "\n");
+    uint64_t ProbesSize = Cur->getProbes().size() - Deleted;
+    EmitULEB128IntValue(ProbesSize);
+    // Emit number of direct inlinees
+    EmitULEB128IntValue(Cur->getChildren().size());
+    // Emit probes in this group
+    for (MCDecodedPseudoProbe *&Probe : Cur->getProbes()) {
+      if (Probe->getAddress() == INT64_MAX)
+        continue;
+      EmitDecodedPseudoProbe(Probe);
+      LastProbe = Probe;
+    }
+
+    for (auto Child = Cur->getChildren().begin();
+         Child != Cur->getChildren().end(); ++Child)
+      Inlinees[Child->first] = Child->second.get();
+    for (const auto &Inlinee : Inlinees) {
+      assert(Cur->Guid != 0 && "non root tree node must have nonzero Guid");
+      NextNodes.push_back({std::get<1>(Inlinee.first), Inlinee.second});
+      LLVM_DEBUG({
+        dbgs().indent(MCPseudoProbeTable::DdgPrintIndent);
+        dbgs() << "InlineSite: " << std::get<1>(Inlinee.first) << "\n";
+      });
+    }
+    Inlinees.clear();
+  }
+
+  // Create buffer for new contents for the section
+  // Freed when parent section is destroyed
+  uint8_t *Output = new uint8_t[Contents.str().size()];
+  memcpy(Output, Contents.str().data(), Contents.str().size());
+  BC.registerOrUpdateSection(".pseudo_probe", PseudoProbeSection->getELFType(),
+                             PseudoProbeSection->getELFFlags(), Output,
+                             Contents.str().size(), 1);
+  if (opts::PrintPseudoProbes == opts::PrintPseudoProbesOptions::PPP_All ||
+      opts::PrintPseudoProbes ==
+          opts::PrintPseudoProbesOptions::PPP_Encoded_Probes) {
+    // create a dummy decoder;
+    MCPseudoProbeDecoder DummyDecoder;
+    StringRef DescContents = PseudoProbeDescSection->getContents();
+    DummyDecoder.buildGUID2FuncDescMap(
+        reinterpret_cast<const uint8_t *>(DescContents.data()),
+        DescContents.size());
+    StringRef ProbeContents = PseudoProbeSection->getOutputContents();
+    MCPseudoProbeDecoder::Uint64Set GuidFilter;
+    MCPseudoProbeDecoder::Uint64Map FuncStartAddrs;
+    for (const BinaryFunction *F : BC.getAllBinaryFunctions()) {
+      const uint64_t Addr =
+          F->isEmitted() ? F->getOutputAddress() : F->getAddress();
+      FuncStartAddrs[Function::getGUID(
+          NameResolver::restore(F->getOneName()))] = Addr;
+    }
+    DummyDecoder.buildAddress2ProbeMap(
+        reinterpret_cast<const uint8_t *>(ProbeContents.data()),
+        ProbeContents.size(), GuidFilter, FuncStartAddrs);
+    DummyDecoder.printProbesForAllAddresses(outs());
+  }
+}
+} // namespace
+
+std::unique_ptr<MetadataRewriter>
+llvm::bolt::createPseudoProbeRewriter(BinaryContext &BC) {
+  return std::make_unique<PseudoProbeRewriter>(BC);
+}
diff --git a/bolt/lib/Rewrite/RewriteInstance.cpp b/bolt/lib/Rewrite/RewriteInstance.cpp
index 6f4faa5..a69fcaa 100644
--- a/bolt/lib/Rewrite/RewriteInstance.cpp
+++ b/bolt/lib/Rewrite/RewriteInstance.cpp
@@ -36,7 +36,6 @@
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/DebugInfo/DWARF/DWARFContext.h"
 #include "llvm/DebugInfo/DWARF/DWARFDebugFrame.h"
-#include "llvm/IR/Function.h"
 #include "llvm/MC/MCAsmBackend.h"
 #include "llvm/MC/MCAsmInfo.h"
 #include "llvm/MC/MCAsmLayout.h"
@@ -53,7 +52,6 @@
 #include "llvm/Support/Errc.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/FileSystem.h"
-#include "llvm/Support/LEB128.h"
 #include "llvm/Support/ManagedStatic.h"
 #include "llvm/Support/Regex.h"
 #include "llvm/Support/Timer.h"
@@ -196,26 +194,6 @@ static cl::opt<bool> PrintLoopInfo("print-loops",
                                    cl::desc("print loop related information"),
                                    cl::Hidden, cl::cat(BoltCategory));
 
-enum PrintPseudoProbesOptions {
-  PPP_None = 0,
-  PPP_Probes_Section_Decode = 0x1,
-  PPP_Probes_Address_Conversion = 0x2,
-  PPP_Encoded_Probes = 0x3,
-  PPP_All = 0xf
-};
-
-cl::opt<PrintPseudoProbesOptions> PrintPseudoProbes(
-    "print-pseudo-probes", cl::desc("print pseudo probe info"),
-    cl::init(PPP_None),
-    cl::values(clEnumValN(PPP_Probes_Section_Decode, "decode",
-                          "decode probes section from binary"),
-               clEnumValN(PPP_Probes_Address_Conversion, "address_conversion",
-                          "update address2ProbesMap with output block address"),
-               clEnumValN(PPP_Encoded_Probes, "encoded_probes",
-                          "display the encoded probes in binary section"),
-               clEnumValN(PPP_All, "all", "enable all debugging printout")),
-    cl::ZeroOrMore, cl::Hidden, cl::cat(BoltCategory));
-
 static cl::opt<cl::boolOrDefault> RelocationMode(
     "relocs", cl::desc("use relocations in the binary (default=autodetect)"),
     cl::cat(BoltCategory));
@@ -274,11 +252,12 @@ static cl::opt<bool> WriteBoltInfoSection(
 
 } // namespace opts
 
+// FIXME: implement a better way to mark sections for replacement.
 constexpr const char *RewriteInstance::SectionsToOverwrite[];
 std::vector<std::string> RewriteInstance::DebugSectionsToOverwrite = {
     ".debug_abbrev", ".debug_aranges",  ".debug_line",   ".debug_line_str",
     ".debug_loc",    ".debug_loclists", ".debug_ranges", ".debug_rnglists",
-    ".gdb_index",    ".debug_addr"};
+    ".gdb_index",    ".debug_addr",     ".pseudo_probe"};
 
 const char RewriteInstance::TimerGroupName[] = "rewrite";
 const char RewriteInstance::TimerGroupDesc[] = "Rewrite passes";
@@ -630,55 +609,6 @@ Error RewriteInstance::discoverStorage() {
   return Error::success();
 }
 
-void RewriteInstance::parsePseudoProbe() {
-  if (!PseudoProbeDescSection && !PseudoProbeSection) {
-    // pesudo probe is not added to binary. It is normal and no warning needed.
-    return;
-  }
-
-  // If only one section is found, it might mean the ELF is corrupted.
-  if (!PseudoProbeDescSection) {
-    errs() << "BOLT-WARNING: fail in reading .pseudo_probe_desc binary\n";
-    return;
-  } else if (!PseudoProbeSection) {
-    errs() << "BOLT-WARNING: fail in reading .pseudo_probe binary\n";
-    return;
-  }
-
-  StringRef Contents = PseudoProbeDescSection->getContents();
-  if (!BC->ProbeDecoder.buildGUID2FuncDescMap(
-          reinterpret_cast<const uint8_t *>(Contents.data()),
-          Contents.size())) {
-    errs() << "BOLT-WARNING: fail in building GUID2FuncDescMap\n";
-    return;
-  }
-
-  MCPseudoProbeDecoder::Uint64Set GuidFilter;
-  MCPseudoProbeDecoder::Uint64Map FuncStartAddrs;
-  for (const BinaryFunction *F : BC->getAllBinaryFunctions()) {
-    for (const MCSymbol *Sym : F->getSymbols()) {
-      FuncStartAddrs[Function::getGUID(NameResolver::restore(Sym->getName()))] =
-          F->getAddress();
-    }
-  }
-  Contents = PseudoProbeSection->getContents();
-  if (!BC->ProbeDecoder.buildAddress2ProbeMap(
-          reinterpret_cast<const uint8_t *>(Contents.data()), Contents.size(),
-          GuidFilter, FuncStartAddrs)) {
-    BC->ProbeDecoder.getAddress2ProbesMap().clear();
-    errs() << "BOLT-WARNING: fail in building Address2ProbeMap\n";
-    return;
-  }
-
-  if (opts::PrintPseudoProbes == opts::PrintPseudoProbesOptions::PPP_All ||
-      opts::PrintPseudoProbes ==
-          opts::PrintPseudoProbesOptions::PPP_Probes_Section_Decode) {
-    outs() << "Report of decoding input pseudo probe binaries \n";
-    BC->ProbeDecoder.printGUID2FuncDescMap(outs());
-    BC->ProbeDecoder.printProbesForAllAddresses(outs());
-  }
-}
-
 void RewriteInstance::parseBuildID() {
   if (!BuildIDSection)
     return;
@@ -1791,8 +1721,6 @@ Error RewriteInstance::readSpecialSections() {
   LSDASection = BC->getUniqueSectionByName(".gcc_except_table");
   EHFrameSection = BC->getUniqueSectionByName(".eh_frame");
   BuildIDSection = BC->getUniqueSectionByName(".note.gnu.build-id");
-  PseudoProbeDescSection = BC->getUniqueSectionByName(".pseudo_probe_desc");
-  PseudoProbeSection = BC->getUniqueSectionByName(".pseudo_probe");
 
   if (ErrorOr<BinarySection &> BATSec =
           BC->getUniqueSectionByName(BoltAddressTranslation::SECTION_NAME)) {
@@ -3180,6 +3108,8 @@ void RewriteInstance::preprocessProfileData() {
 }
 
 void RewriteInstance::initializeMetadataManager() {
+  MetadataManager.registerRewriter(createPseudoProbeRewriter(*BC));
+
   MetadataManager.registerRewriter(createSDTRewriter(*BC));
 }
 
@@ -3557,8 +3487,6 @@ void RewriteInstance::updateMetadata() {
 
   // TODO: use MetadataManager for updates.
   updateLKMarkers();
-  parsePseudoProbe();
-  updatePseudoProbes();
 
   if (opts::UpdateDebugSections) {
     NamedRegionTimer T("updateDebugInfo", "update debug info", TimerGroupName,
@@ -3570,272 +3498,6 @@ void RewriteInstance::updateMetadata() {
     addBoltInfoSection();
 }
 
-void RewriteInstance::updatePseudoProbes() {
-  // check if there is pseudo probe section decoded
-  if (BC->ProbeDecoder.getAddress2ProbesMap().empty())
-    return;
-  // input address converted to output
-  AddressProbesMap &Address2ProbesMap = BC->ProbeDecoder.getAddress2ProbesMap();
-  const GUIDProbeFunctionMap &GUID2Func =
-      BC->ProbeDecoder.getGUID2FuncDescMap();
-
-  for (auto &AP : Address2ProbesMap) {
-    BinaryFunction *F = BC->getBinaryFunctionContainingAddress(AP.first);
-    // If F is removed, eliminate all probes inside it from inline tree
-    // Setting probes' addresses as INT64_MAX means elimination
-    if (!F) {
-      for (MCDecodedPseudoProbe &Probe : AP.second)
-        Probe.setAddress(INT64_MAX);
-      continue;
-    }
-    // If F is not emitted, the function will remain in the same address as its
-    // input
-    if (!F->isEmitted())
-      continue;
-
-    uint64_t Offset = AP.first - F->getAddress();
-    const BinaryBasicBlock *BB = F->getBasicBlockContainingOffset(Offset);
-    uint64_t BlkOutputAddress = BB->getOutputAddressRange().first;
-    // Check if block output address is defined.
-    // If not, such block is removed from binary. Then remove the probes from
-    // inline tree
-    if (BlkOutputAddress == 0) {
-      for (MCDecodedPseudoProbe &Probe : AP.second)
-        Probe.setAddress(INT64_MAX);
-      continue;
-    }
-
-    unsigned ProbeTrack = AP.second.size();
-    std::list<MCDecodedPseudoProbe>::iterator Probe = AP.second.begin();
-    while (ProbeTrack != 0) {
-      if (Probe->isBlock()) {
-        Probe->setAddress(BlkOutputAddress);
-      } else if (Probe->isCall()) {
-        // A call probe may be duplicated due to ICP
-        // Go through output of InputOffsetToAddressMap to collect all related
-        // probes
-        const InputOffsetToAddressMapTy &Offset2Addr =
-            F->getInputOffsetToAddressMap();
-        auto CallOutputAddresses = Offset2Addr.equal_range(Offset);
-        auto CallOutputAddress = CallOutputAddresses.first;
-        if (CallOutputAddress == CallOutputAddresses.second) {
-          Probe->setAddress(INT64_MAX);
-        } else {
-          Probe->setAddress(CallOutputAddress->second);
-          CallOutputAddress = std::next(CallOutputAddress);
-        }
-
-        while (CallOutputAddress != CallOutputAddresses.second) {
-          AP.second.push_back(*Probe);
-          AP.second.back().setAddress(CallOutputAddress->second);
-          Probe->getInlineTreeNode()->addProbes(&(AP.second.back()));
-          CallOutputAddress = std::next(CallOutputAddress);
-        }
-      }
-      Probe = std::next(Probe);
-      ProbeTrack--;
-    }
-  }
-
-  if (opts::PrintPseudoProbes == opts::PrintPseudoProbesOptions::PPP_All ||
-      opts::PrintPseudoProbes ==
-          opts::PrintPseudoProbesOptions::PPP_Probes_Address_Conversion) {
-    outs() << "Pseudo Probe Address Conversion results:\n";
-    // table that correlates address to block
-    std::unordered_map<uint64_t, StringRef> Addr2BlockNames;
-    for (auto &F : BC->getBinaryFunctions())
-      for (BinaryBasicBlock &BinaryBlock : F.second)
-        Addr2BlockNames[BinaryBlock.getOutputAddressRange().first] =
-            BinaryBlock.getName();
-
-    // scan all addresses -> correlate probe to block when print out
-    std::vector<uint64_t> Addresses;
-    for (auto &Entry : Address2ProbesMap)
-      Addresses.push_back(Entry.first);
-    llvm::sort(Addresses);
-    for (uint64_t Key : Addresses) {
-      for (MCDecodedPseudoProbe &Probe : Address2ProbesMap[Key]) {
-        if (Probe.getAddress() == INT64_MAX)
-          outs() << "Deleted Probe: ";
-        else
-          outs() << "Address: " << format_hex(Probe.getAddress(), 8) << " ";
-        Probe.print(outs(), GUID2Func, true);
-        // print block name only if the probe is block type and undeleted.
-        if (Probe.isBlock() && Probe.getAddress() != INT64_MAX)
-          outs() << format_hex(Probe.getAddress(), 8) << " Probe is in "
-                 << Addr2BlockNames[Probe.getAddress()] << "\n";
-      }
-    }
-    outs() << "=======================================\n";
-  }
-
-  // encode pseudo probes with updated addresses
-  encodePseudoProbes();
-}
-
-template <typename F>
-static void emitLEB128IntValue(F encode, uint64_t Value,
-                               SmallString<8> &Contents) {
-  SmallString<128> Tmp;
-  raw_svector_ostream OSE(Tmp);
-  encode(Value, OSE);
-  Contents.append(OSE.str().begin(), OSE.str().end());
-}
-
-void RewriteInstance::encodePseudoProbes() {
-  // Buffer for new pseudo probes section
-  SmallString<8> Contents;
-  MCDecodedPseudoProbe *LastProbe = nullptr;
-
-  auto EmitInt = [&](uint64_t Value, uint32_t Size) {
-    const bool IsLittleEndian = BC->AsmInfo->isLittleEndian();
-    uint64_t Swapped = support::endian::byte_swap(
-        Value, IsLittleEndian ? support::little : support::big);
-    unsigned Index = IsLittleEndian ? 0 : 8 - Size;
-    auto Entry = StringRef(reinterpret_cast<char *>(&Swapped) + Index, Size);
-    Contents.append(Entry.begin(), Entry.end());
-  };
-
-  auto EmitULEB128IntValue = [&](uint64_t Value) {
-    SmallString<128> Tmp;
-    raw_svector_ostream OSE(Tmp);
-    encodeULEB128(Value, OSE, 0);
-    Contents.append(OSE.str().begin(), OSE.str().end());
-  };
-
-  auto EmitSLEB128IntValue = [&](int64_t Value) {
-    SmallString<128> Tmp;
-    raw_svector_ostream OSE(Tmp);
-    encodeSLEB128(Value, OSE);
-    Contents.append(OSE.str().begin(), OSE.str().end());
-  };
-
-  // Emit indiviual pseudo probes in a inline tree node
-  // Probe index, type, attribute, address type and address are encoded
-  // Address of the first probe is absolute.
-  // Other probes' address are represented by delta
-  auto EmitDecodedPseudoProbe = [&](MCDecodedPseudoProbe *&CurProbe) {
-    assert(!isSentinelProbe(CurProbe->getAttributes()) &&
-           "Sentinel probes should not be emitted");
-    EmitULEB128IntValue(CurProbe->getIndex());
-    uint8_t PackedType = CurProbe->getType() | (CurProbe->getAttributes() << 4);
-    uint8_t Flag =
-        LastProbe ? ((int8_t)MCPseudoProbeFlag::AddressDelta << 7) : 0;
-    EmitInt(Flag | PackedType, 1);
-    if (LastProbe) {
-      // Emit the delta between the address label and LastProbe.
-      int64_t Delta = CurProbe->getAddress() - LastProbe->getAddress();
-      EmitSLEB128IntValue(Delta);
-    } else {
-      // Emit absolute address for encoding the first pseudo probe.
-      uint32_t AddrSize = BC->AsmInfo->getCodePointerSize();
-      EmitInt(CurProbe->getAddress(), AddrSize);
-    }
-  };
-
-  std::map<InlineSite, MCDecodedPseudoProbeInlineTree *,
-           std::greater<InlineSite>>
-      Inlinees;
-
-  // DFS of inline tree to emit pseudo probes in all tree node
-  // Inline site index of a probe is emitted first.
-  // Then tree node Guid, size of pseudo probes and children nodes, and detail
-  // of contained probes are emitted Deleted probes are skipped Root node is not
-  // encoded to binaries. It's a "wrapper" of inline trees of each function.
-  std::list<std::pair<uint64_t, MCDecodedPseudoProbeInlineTree *>> NextNodes;
-  const MCDecodedPseudoProbeInlineTree &Root =
-      BC->ProbeDecoder.getDummyInlineRoot();
-  for (auto Child = Root.getChildren().begin();
-       Child != Root.getChildren().end(); ++Child)
-    Inlinees[Child->first] = Child->second.get();
-
-  for (auto Inlinee : Inlinees)
-    // INT64_MAX is "placeholder" of unused callsite index field in the pair
-    NextNodes.push_back({INT64_MAX, Inlinee.second});
-
-  Inlinees.clear();
-
-  while (!NextNodes.empty()) {
-    uint64_t ProbeIndex = NextNodes.back().first;
-    MCDecodedPseudoProbeInlineTree *Cur = NextNodes.back().second;
-    NextNodes.pop_back();
-
-    if (Cur->Parent && !Cur->Parent->isRoot())
-      // Emit probe inline site
-      EmitULEB128IntValue(ProbeIndex);
-
-    // Emit probes grouped by GUID.
-    LLVM_DEBUG({
-      dbgs().indent(MCPseudoProbeTable::DdgPrintIndent);
-      dbgs() << "GUID: " << Cur->Guid << "\n";
-    });
-    // Emit Guid
-    EmitInt(Cur->Guid, 8);
-    // Emit number of probes in this node
-    uint64_t Deleted = 0;
-    for (MCDecodedPseudoProbe *&Probe : Cur->getProbes())
-      if (Probe->getAddress() == INT64_MAX)
-        Deleted++;
-    LLVM_DEBUG(dbgs() << "Deleted Probes:" << Deleted << "\n");
-    uint64_t ProbesSize = Cur->getProbes().size() - Deleted;
-    EmitULEB128IntValue(ProbesSize);
-    // Emit number of direct inlinees
-    EmitULEB128IntValue(Cur->getChildren().size());
-    // Emit probes in this group
-    for (MCDecodedPseudoProbe *&Probe : Cur->getProbes()) {
-      if (Probe->getAddress() == INT64_MAX)
-        continue;
-      EmitDecodedPseudoProbe(Probe);
-      LastProbe = Probe;
-    }
-
-    for (auto Child = Cur->getChildren().begin();
-         Child != Cur->getChildren().end(); ++Child)
-      Inlinees[Child->first] = Child->second.get();
-    for (const auto &Inlinee : Inlinees) {
-      assert(Cur->Guid != 0 && "non root tree node must have nonzero Guid");
-      NextNodes.push_back({std::get<1>(Inlinee.first), Inlinee.second});
-      LLVM_DEBUG({
-        dbgs().indent(MCPseudoProbeTable::DdgPrintIndent);
-        dbgs() << "InlineSite: " << std::get<1>(Inlinee.first) << "\n";
-      });
-    }
-    Inlinees.clear();
-  }
-
-  // Create buffer for new contents for the section
-  // Freed when parent section is destroyed
-  uint8_t *Output = new uint8_t[Contents.str().size()];
-  memcpy(Output, Contents.str().data(), Contents.str().size());
-  addToDebugSectionsToOverwrite(".pseudo_probe");
-  BC->registerOrUpdateSection(".pseudo_probe", PseudoProbeSection->getELFType(),
-                              PseudoProbeSection->getELFFlags(), Output,
-                              Contents.str().size(), 1);
-  if (opts::PrintPseudoProbes == opts::PrintPseudoProbesOptions::PPP_All ||
-      opts::PrintPseudoProbes ==
-          opts::PrintPseudoProbesOptions::PPP_Encoded_Probes) {
-    // create a dummy decoder;
-    MCPseudoProbeDecoder DummyDecoder;
-    StringRef DescContents = PseudoProbeDescSection->getContents();
-    DummyDecoder.buildGUID2FuncDescMap(
-        reinterpret_cast<const uint8_t *>(DescContents.data()),
-        DescContents.size());
-    StringRef ProbeContents = PseudoProbeSection->getOutputContents();
-    MCPseudoProbeDecoder::Uint64Set GuidFilter;
-    MCPseudoProbeDecoder::Uint64Map FuncStartAddrs;
-    for (const BinaryFunction *F : BC->getAllBinaryFunctions()) {
-      const uint64_t Addr =
-          F->isEmitted() ? F->getOutputAddress() : F->getAddress();
-      FuncStartAddrs[Function::getGUID(
-          NameResolver::restore(F->getOneName()))] = Addr;
-    }
-    DummyDecoder.buildAddress2ProbeMap(
-        reinterpret_cast<const uint8_t *>(ProbeContents.data()),
-        ProbeContents.size(), GuidFilter, FuncStartAddrs);
-    DummyDecoder.printProbesForAllAddresses(outs());
-  }
-}
-
 void RewriteInstance::updateLKMarkers() {
   if (BC->LKMarkers.size() == 0)
     return;