[BOLT] Expose pseudo probe function checksum and GUID (#99389)

Add a BinaryFunction field for pseudo probe function GUID.
Populate it during pseudo probe section parsing, and emit it in YAML
profile (both regular and BAT), along with function checksum.

To be used for stale function matching.

Test Plan: update pseudoprobe-decoding-inline.test

Author: aaupov

bolt/include/bolt/Core/BinaryContext.h

@@ -32,6 +32,7 @@
 #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"
@@ -246,6 +247,9 @@ class BinaryContext {
   /// DWP Context.
   std::shared_ptr<DWARFContext> DWPContext;
 
+  /// Decoded pseudo probes.
+  std::shared_ptr<MCPseudoProbeDecoder> PseudoProbeDecoder;
+
   /// A map of DWO Ids to CUs.
   using DWOIdToCUMapType = std::unordered_map<uint64_t, DWARFUnit *>;
   DWOIdToCUMapType DWOCUs;
@@ -377,6 +381,15 @@ class BinaryContext {
     RtLibrary = std::move(Lib);
   }
 
+  const MCPseudoProbeDecoder *getPseudoProbeDecoder() const {
+    return PseudoProbeDecoder.get();
+  }
+
+  void setPseudoProbeDecoder(std::shared_ptr<MCPseudoProbeDecoder> Decoder) {
+    assert(!PseudoProbeDecoder && "Cannot set pseudo probe decoder twice.");
+    PseudoProbeDecoder = Decoder;
+  }
+
   /// Return BinaryFunction containing a given \p Address or nullptr if
   /// no registered function contains the \p Address.
   ///

bolt/include/bolt/Core/BinaryFunction.h

@@ -416,6 +416,9 @@ class BinaryFunction {
   /// different parameters by every pass.
   mutable uint64_t Hash{0};
 
+  /// Function GUID assigned externally.
+  uint64_t GUID{0};
+
   /// For PLT functions it contains a symbol associated with a function
   /// reference. It is nullptr for non-PLT functions.
   const MCSymbol *PLTSymbol{nullptr};
@@ -2256,6 +2259,11 @@ class BinaryFunction {
   /// Returns the last computed hash value of the function.
   size_t getHash() const { return Hash; }
 
+  /// Returns the function GUID.
+  uint64_t getGUID() const { return GUID; }
+
+  void setGUID(uint64_t Id) { GUID = Id; }
+
   using OperandHashFuncTy =
       function_ref<typename std::string(const MCOperand &)>;
 

bolt/include/bolt/Profile/ProfileYAMLMapping.h

@@ -151,6 +151,8 @@ struct BinaryFunctionProfile {
   llvm::yaml::Hex64 Hash{0};
   uint64_t ExecCount{0};
   std::vector<BinaryBasicBlockProfile> Blocks;
+  llvm::yaml::Hex64 GUID{0};
+  llvm::yaml::Hex64 PseudoProbeDescHash{0};
   bool Used{false};
 };
 } // end namespace bolt
@@ -164,6 +166,9 @@ template <> struct MappingTraits<bolt::BinaryFunctionProfile> {
     YamlIO.mapRequired("nblocks", BFP.NumBasicBlocks);
     YamlIO.mapOptional("blocks", BFP.Blocks,
                        std::vector<bolt::BinaryBasicBlockProfile>());
+    YamlIO.mapOptional("guid", BFP.GUID, (uint64_t)0);
+    YamlIO.mapOptional("pseudo_probe_desc_hash", BFP.PseudoProbeDescHash,
+                       (uint64_t)0);
   }
 };
 

bolt/lib/Profile/DataAggregator.cpp

@@ -2298,6 +2298,8 @@ std::error_code DataAggregator::writeBATYAML(BinaryContext &BC,
 
   yaml::bolt::BinaryProfile BP;
 
+  const MCPseudoProbeDecoder *PseudoProbeDecoder = BC.getPseudoProbeDecoder();
+
   // Fill out the header info.
   BP.Header.Version = 1;
   BP.Header.FileName = std::string(BC.getFilename());
@@ -2398,6 +2400,13 @@ std::error_code DataAggregator::writeBATYAML(BinaryContext &BC,
         const unsigned BlockIndex = BlockMap.getBBIndex(BI.To.Offset);
         YamlBF.Blocks[BlockIndex].ExecCount += BI.Branches;
       }
+      if (PseudoProbeDecoder) {
+        if ((YamlBF.GUID = BF->getGUID())) {
+          const MCPseudoProbeFuncDesc *FuncDesc =
+              PseudoProbeDecoder->getFuncDescForGUID(YamlBF.GUID);
+          YamlBF.PseudoProbeDescHash = FuncDesc->FuncHash;
+        }
+      }
       // Drop blocks without a hash, won't be useful for stale matching.
       llvm::erase_if(YamlBF.Blocks,
                      [](const yaml::bolt::BinaryBasicBlockProfile &YamlBB) {

bolt/lib/Profile/YAMLProfileWriter.cpp

@@ -57,6 +57,7 @@ YAMLProfileWriter::convert(const BinaryFunction &BF, bool UseDFS,
                            const BoltAddressTranslation *BAT) {
   yaml::bolt::BinaryFunctionProfile YamlBF;
   const BinaryContext &BC = BF.getBinaryContext();
+  const MCPseudoProbeDecoder *PseudoProbeDecoder = BC.getPseudoProbeDecoder();
 
   const uint16_t LBRProfile = BF.getProfileFlags() & BinaryFunction::PF_LBR;
 
@@ -69,6 +70,13 @@ YAMLProfileWriter::convert(const BinaryFunction &BF, bool UseDFS,
   YamlBF.Hash = BF.getHash();
   YamlBF.NumBasicBlocks = BF.size();
   YamlBF.ExecCount = BF.getKnownExecutionCount();
+  if (PseudoProbeDecoder) {
+    if ((YamlBF.GUID = BF.getGUID())) {
+      const MCPseudoProbeFuncDesc *FuncDesc =
+          PseudoProbeDecoder->getFuncDescForGUID(YamlBF.GUID);
+      YamlBF.PseudoProbeDescHash = FuncDesc->FuncHash;
+    }
+  }
 
   BinaryFunction::BasicBlockOrderType Order;
   llvm::copy(UseDFS ? BF.dfs() : BF.getLayout().blocks(),

bolt/lib/Rewrite/PseudoProbeRewriter.cpp

@@ -19,6 +19,7 @@
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/LEB128.h"
+#include <memory>
 
 #undef DEBUG_TYPE
 #define DEBUG_TYPE "pseudo-probe-rewriter"
@@ -72,23 +73,35 @@ class PseudoProbeRewriter final : public MetadataRewriter {
   void parsePseudoProbe();
 
   /// PseudoProbe decoder
-  MCPseudoProbeDecoder ProbeDecoder;
+  std::shared_ptr<MCPseudoProbeDecoder> ProbeDecoderPtr;
 
 public:
   PseudoProbeRewriter(BinaryContext &BC)
-      : MetadataRewriter("pseudo-probe-rewriter", BC) {}
+      : MetadataRewriter("pseudo-probe-rewriter", BC),
+        ProbeDecoderPtr(std::make_shared<MCPseudoProbeDecoder>()) {
+    BC.setPseudoProbeDecoder(ProbeDecoderPtr);
+  }
 
+  Error preCFGInitializer() override;
   Error postEmitFinalizer() override;
+
+  ~PseudoProbeRewriter() override { ProbeDecoderPtr.reset(); }
 };
 
-Error PseudoProbeRewriter::postEmitFinalizer() {
+Error PseudoProbeRewriter::preCFGInitializer() {
   parsePseudoProbe();
+
+  return Error::success();
+}
+
+Error PseudoProbeRewriter::postEmitFinalizer() {
   updatePseudoProbes();
 
   return Error::success();
 }
 
 void PseudoProbeRewriter::parsePseudoProbe() {
+  MCPseudoProbeDecoder &ProbeDecoder(*ProbeDecoderPtr);
   PseudoProbeDescSection = BC.getUniqueSectionByName(".pseudo_probe_desc");
   PseudoProbeSection = BC.getUniqueSectionByName(".pseudo_probe");
 
@@ -138,9 +151,18 @@ void PseudoProbeRewriter::parsePseudoProbe() {
     ProbeDecoder.printGUID2FuncDescMap(outs());
     ProbeDecoder.printProbesForAllAddresses(outs());
   }
+
+  for (const auto &[GUID, FuncDesc] : ProbeDecoder.getGUID2FuncDescMap()) {
+    if (!FuncStartAddrs.contains(GUID))
+      continue;
+    BinaryFunction *BF = BC.getBinaryFunctionAtAddress(FuncStartAddrs[GUID]);
+    assert(BF);
+    BF->setGUID(GUID);
+  }
 }
 
 void PseudoProbeRewriter::updatePseudoProbes() {
+  MCPseudoProbeDecoder &ProbeDecoder(*ProbeDecoderPtr);
   // check if there is pseudo probe section decoded
   if (ProbeDecoder.getAddress2ProbesMap().empty())
     return;
@@ -241,6 +263,7 @@ void PseudoProbeRewriter::updatePseudoProbes() {
 }
 
 void PseudoProbeRewriter::encodePseudoProbes() {
+  MCPseudoProbeDecoder &ProbeDecoder(*ProbeDecoderPtr);
   // Buffer for new pseudo probes section
   SmallString<8> Contents;
   MCDecodedPseudoProbe *LastProbe = nullptr;

bolt/lib/Rewrite/RewriteInstance.cpp

@@ -669,6 +669,7 @@ Error RewriteInstance::run() {
         opts::ProfileFormat == opts::ProfileFormatKind::PF_YAML) {
       selectFunctionsToProcess();
       disassembleFunctions();
+      processMetadataPreCFG();
       buildFunctionsCFG();
     }
     processProfileData();

bolt/test/X86/pseudoprobe-decoding-inline.test

@@ -1,5 +1,26 @@
 # REQUIRES: system-linux
-# RUN: llvm-bolt  %S/../../../llvm/test/tools/llvm-profgen/Inputs/inline-cs-pseudoprobe.perfbin --print-pseudo-probes=all -o %t.bolt 2>&1 | FileCheck %s
+# RUN: llvm-bolt  %S/../../../llvm/test/tools/llvm-profgen/Inputs/inline-cs-pseudoprobe.perfbin --print-pseudo-probes=all -o %t.bolt --lite=0 --enable-bat 2>&1 | FileCheck %s
+
+# PREAGG: B X:0 #foo# 1 0
+# PREAGG: B X:0 #bar# 1 0
+# PREAGG: B X:0 #main# 1 0
+## Check pseudo-probes in regular YAML profile (non-BOLTed binary)
+# RUN: link_fdata %s %S/../../../llvm/test/tools/llvm-profgen/Inputs/inline-cs-pseudoprobe.perfbin %t.preagg PREAGG
+# RUN: perf2bolt %S/../../../llvm/test/tools/llvm-profgen/Inputs/inline-cs-pseudoprobe.perfbin -p %t.preagg --pa -w %t.yaml -o %t.fdata
+# RUN: FileCheck --input-file %t.yaml %s --check-prefix CHECK-YAML
+## Check pseudo-probes in BAT YAML profile (BOLTed binary)
+# RUN: link_fdata %s %t.bolt %t.preagg2 PREAGG
+# RUN: perf2bolt %t.bolt -p %t.preagg2 --pa -w %t.yaml2 -o %t.fdata2
+# RUN: FileCheck --input-file %t.yaml2 %s --check-prefix CHECK-YAML
+# CHECK-YAML: name: bar
+# CHECK-YAML: guid: 0xE413754A191DB537
+# CHECK-YAML: pseudo_probe_desc_hash: 0x10E852DA94
+# CHECK-YAML: name: foo
+# CHECK-YAML: guid: 0x5CF8C24CDB18BDAC
+# CHECK-YAML: pseudo_probe_desc_hash: 0x200205A19C5B4
+# CHECK-YAML: name: main
+# CHECK-YAML: guid: 0xDB956436E78DD5FA
+# CHECK-YAML: pseudo_probe_desc_hash: 0x10000FFFFFFFF
 
 CHECK: Report of decoding input pseudo probe binaries