[llvm] [BOLT] Attach pseudo probes to blocks in YAML profile (PR #99554)

Amir Ayupov via llvm-commits llvm-commits at lists.llvm.org
Thu Jul 18 21:01:24 PDT 2024


https://github.com/aaupov updated https://github.com/llvm/llvm-project/pull/99554

>From 4aff8f3539cfc193e8c5de6b7376ecc24ccbc92d Mon Sep 17 00:00:00 2001
From: Amir Ayupov <aaupov at fb.com>
Date: Thu, 18 Jul 2024 12:36:46 -0700
Subject: [PATCH 1/3] =?UTF-8?q?[=F0=9D=98=80=F0=9D=97=BD=F0=9D=97=BF]=20ch?=
 =?UTF-8?q?anges=20to=20main=20this=20commit=20is=20based=20on?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Created using spr 1.3.4

[skip ci]
---
 bolt/include/bolt/Core/BinaryContext.h        | 15 ++++++++++++
 bolt/include/bolt/Core/BinaryFunction.h       |  9 ++++++++
 .../include/bolt/Profile/ProfileYAMLMapping.h |  5 ++++
 bolt/lib/Profile/DataAggregator.cpp           |  9 ++++++++
 bolt/lib/Profile/YAMLProfileWriter.cpp        |  8 +++++++
 bolt/lib/Rewrite/PseudoProbeRewriter.cpp      | 22 +++++++++++++++---
 bolt/lib/Rewrite/RewriteInstance.cpp          |  1 +
 .../test/X86/pseudoprobe-decoding-inline.test | 23 ++++++++++++++++++-
 llvm/include/llvm/MC/MCPseudoProbe.h          |  4 ++--
 9 files changed, 90 insertions(+), 6 deletions(-)

diff --git a/bolt/include/bolt/Core/BinaryContext.h b/bolt/include/bolt/Core/BinaryContext.h
index 73932c4ca2fb3..ea924ebe1ab54 100644
--- a/bolt/include/bolt/Core/BinaryContext.h
+++ b/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::unique_ptr<MCPseudoProbeDecoder> PseudoProbeDecoder;
+
   /// A map of DWO Ids to CUs.
   using DWOIdToCUMapType = std::unordered_map<uint64_t, DWARFUnit *>;
   DWOIdToCUMapType DWOCUs;
@@ -377,6 +381,17 @@ class BinaryContext {
     RtLibrary = std::move(Lib);
   }
 
+  const MCPseudoProbeDecoder *getPseudoProbeDecoder() const {
+    return PseudoProbeDecoder.get();
+  }
+
+  MCPseudoProbeDecoder &
+  setPseudoProbeDecoder(std::unique_ptr<MCPseudoProbeDecoder> Decoder) {
+    assert(!PseudoProbeDecoder && "Cannot set pseudo probe decoder twice.");
+    PseudoProbeDecoder = std::move(Decoder);
+    return *PseudoProbeDecoder.get();
+  }
+
   /// Return BinaryFunction containing a given \p Address or nullptr if
   /// no registered function contains the \p Address.
   ///
diff --git a/bolt/include/bolt/Core/BinaryFunction.h b/bolt/include/bolt/Core/BinaryFunction.h
index 7d0afee4d1b78..4bebac33ced3f 100644
--- a/bolt/include/bolt/Core/BinaryFunction.h
+++ b/bolt/include/bolt/Core/BinaryFunction.h
@@ -416,6 +416,10 @@ class BinaryFunction {
   /// different parameters by every pass.
   mutable uint64_t Hash{0};
 
+  /// Function GUID assigned by the compiler to the source function.
+  /// If non-null, is a valid GUID in pseudo probe section.
+  uint64_t PseudoProbeGUID{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 +2260,11 @@ class BinaryFunction {
   /// Returns the last computed hash value of the function.
   size_t getHash() const { return Hash; }
 
+  /// Returns the function GUID from pseudo-probe description of the function.
+  uint64_t getPseudoProbeGUID() const { return PseudoProbeGUID; }
+
+  void setPseudoProbeGUID(uint64_t GUID) { PseudoProbeGUID = GUID; }
+
   using OperandHashFuncTy =
       function_ref<typename std::string(const MCOperand &)>;
 
diff --git a/bolt/include/bolt/Profile/ProfileYAMLMapping.h b/bolt/include/bolt/Profile/ProfileYAMLMapping.h
index 9dd3920dbf094..c4e5f2b284a86 100644
--- a/bolt/include/bolt/Profile/ProfileYAMLMapping.h
+++ b/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);
   }
 };
 
diff --git a/bolt/lib/Profile/DataAggregator.cpp b/bolt/lib/Profile/DataAggregator.cpp
index ce6ec0a04ac16..beb03a952d658 100644
--- a/bolt/lib/Profile/DataAggregator.cpp
+++ b/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->getPseudoProbeGUID())) {
+          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) {
diff --git a/bolt/lib/Profile/YAMLProfileWriter.cpp b/bolt/lib/Profile/YAMLProfileWriter.cpp
index 9adbfdc5ff089..38345da73df7c 100644
--- a/bolt/lib/Profile/YAMLProfileWriter.cpp
+++ b/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.getPseudoProbeGUID())) {
+      const MCPseudoProbeFuncDesc *FuncDesc =
+          PseudoProbeDecoder->getFuncDescForGUID(YamlBF.GUID);
+      YamlBF.PseudoProbeDescHash = FuncDesc->FuncHash;
+    }
+  }
 
   BinaryFunction::BasicBlockOrderType Order;
   llvm::copy(UseDFS ? BF.dfs() : BF.getLayout().blocks(),
diff --git a/bolt/lib/Rewrite/PseudoProbeRewriter.cpp b/bolt/lib/Rewrite/PseudoProbeRewriter.cpp
index 51038dbead330..9249861ba8e3a 100644
--- a/bolt/lib/Rewrite/PseudoProbeRewriter.cpp
+++ b/bolt/lib/Rewrite/PseudoProbeRewriter.cpp
@@ -72,17 +72,25 @@ class PseudoProbeRewriter final : public MetadataRewriter {
   void parsePseudoProbe();
 
   /// PseudoProbe decoder
-  MCPseudoProbeDecoder ProbeDecoder;
+  MCPseudoProbeDecoder &ProbeDecoder;
 
 public:
   PseudoProbeRewriter(BinaryContext &BC)
-      : MetadataRewriter("pseudo-probe-rewriter", BC) {}
+      : MetadataRewriter("pseudo-probe-rewriter", BC),
+        ProbeDecoder(BC.setPseudoProbeDecoder(
+            std::make_unique<MCPseudoProbeDecoder>())) {}
 
+  Error preCFGInitializer() override;
   Error postEmitFinalizer() override;
 };
 
-Error PseudoProbeRewriter::postEmitFinalizer() {
+Error PseudoProbeRewriter::preCFGInitializer() {
   parsePseudoProbe();
+
+  return Error::success();
+}
+
+Error PseudoProbeRewriter::postEmitFinalizer() {
   updatePseudoProbes();
 
   return Error::success();
@@ -138,6 +146,14 @@ 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->setPseudoProbeGUID(GUID);
+  }
 }
 
 void PseudoProbeRewriter::updatePseudoProbes() {
diff --git a/bolt/lib/Rewrite/RewriteInstance.cpp b/bolt/lib/Rewrite/RewriteInstance.cpp
index 4ae802dc97ccd..33ebae3b6e6de 100644
--- a/bolt/lib/Rewrite/RewriteInstance.cpp
+++ b/bolt/lib/Rewrite/RewriteInstance.cpp
@@ -669,6 +669,7 @@ Error RewriteInstance::run() {
         opts::ProfileFormat == opts::ProfileFormatKind::PF_YAML) {
       selectFunctionsToProcess();
       disassembleFunctions();
+      processMetadataPreCFG();
       buildFunctionsCFG();
     }
     processProfileData();
diff --git a/bolt/test/X86/pseudoprobe-decoding-inline.test b/bolt/test/X86/pseudoprobe-decoding-inline.test
index 15e93b1630a66..c32660c41a09c 100644
--- a/bolt/test/X86/pseudoprobe-decoding-inline.test
+++ b/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
 
diff --git a/llvm/include/llvm/MC/MCPseudoProbe.h b/llvm/include/llvm/MC/MCPseudoProbe.h
index 4d82da70f3df8..5344dea4141b3 100644
--- a/llvm/include/llvm/MC/MCPseudoProbe.h
+++ b/llvm/include/llvm/MC/MCPseudoProbe.h
@@ -61,6 +61,7 @@
 #include "llvm/IR/PseudoProbe.h"
 #include "llvm/Support/ErrorOr.h"
 #include <list>
+#include <map>
 #include <memory>
 #include <string>
 #include <tuple>
@@ -102,8 +103,7 @@ using MCPseudoProbeInlineStack = SmallVector<InlineSite, 8>;
 using GUIDProbeFunctionMap =
     std::unordered_map<uint64_t, MCPseudoProbeFuncDesc>;
 // Address to pseudo probes map.
-using AddressProbesMap =
-    std::unordered_map<uint64_t, std::list<MCDecodedPseudoProbe>>;
+using AddressProbesMap = std::map<uint64_t, std::list<MCDecodedPseudoProbe>>;
 
 class MCDecodedPseudoProbeInlineTree;
 

>From 9498f8f38cea050fd363d5d4591e8401e5de8bd5 Mon Sep 17 00:00:00 2001
From: Amir Ayupov <aaupov at fb.com>
Date: Thu, 18 Jul 2024 12:49:23 -0700
Subject: [PATCH 2/3] Fix operator==

Created using spr 1.3.4
---
 bolt/include/bolt/Profile/ProfileYAMLMapping.h | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/bolt/include/bolt/Profile/ProfileYAMLMapping.h b/bolt/include/bolt/Profile/ProfileYAMLMapping.h
index 2bc6901e5f591..2a0514d7d9304 100644
--- a/bolt/include/bolt/Profile/ProfileYAMLMapping.h
+++ b/bolt/include/bolt/Profile/ProfileYAMLMapping.h
@@ -100,7 +100,7 @@ struct PseudoProbeInfo {
   uint8_t Type;
 
   bool operator==(const PseudoProbeInfo &Other) const {
-    return Index == Other.Index;
+    return GUID == Other.GUID && Index == Other.Index;
   }
   bool operator!=(const PseudoProbeInfo &Other) const {
     return !(*this == Other);

>From 956084a3b4078f10d953ecab20720e9cf2d25554 Mon Sep 17 00:00:00 2001
From: Amir Ayupov <aaupov at fb.com>
Date: Thu, 18 Jul 2024 12:55:50 -0700
Subject: [PATCH 3/3] Fix probe mapping to block in BAT mode

Created using spr 1.3.4
---
 bolt/lib/Profile/DataAggregator.cpp | 9 +--------
 1 file changed, 1 insertion(+), 8 deletions(-)

diff --git a/bolt/lib/Profile/DataAggregator.cpp b/bolt/lib/Profile/DataAggregator.cpp
index 85e583c75728f..1da018c346d36 100644
--- a/bolt/lib/Profile/DataAggregator.cpp
+++ b/bolt/lib/Profile/DataAggregator.cpp
@@ -2419,14 +2419,7 @@ std::error_code DataAggregator::writeBATYAML(BinaryContext &BC,
           for (const auto &[OutputAddress, Probes] : FragmentProbes) {
             const uint32_t InputOffset = BAT->translate(
                 FuncAddr, OutputAddress - FuncAddr, /*IsBranchSrc=*/true);
-            if (!BlockMap.isInputBlock(InputOffset)) {
-              if (opts::Verbosity >= 1)
-                errs() << "BOLT-WARNING: Couldn't map pseudo probe at 0x"
-                       << Twine::utohexstr(InputOffset) << " to a block in "
-                       << F->getPrintName() << '\n';
-              continue;
-            }
-            const unsigned BlockIndex = BlockMap.getBBIndex(InputOffset);
+            const unsigned BlockIndex = getBlock(InputOffset).second;
             for (const MCDecodedPseudoProbe &Probe : Probes)
               YamlBF.Blocks[BlockIndex].PseudoProbes.emplace_back(
                   yaml::bolt::PseudoProbeInfo{Probe.getGuid(), Probe.getIndex(),



More information about the llvm-commits mailing list