[BOLT] Attach pseudo probes to blocks in YAML profile

[aaupov]

Read pseudo probes in regular and BAT YAML profile generation, and
attach them to YAML profile basic blocks. This exposes GUID, probe id,
and probe type in profile for future use in stale profile matching.

Test Plan: updated pseudoprobe-decoding-inline.test

[llvmbot]

@llvm/pr-subscribers-bolt

Author: Amir Ayupov (aaupov)

Changes

Read pseudo probes in regular and BAT YAML profile generation, and
attach them to YAML profile basic blocks. This exposes GUID, probe id,
and probe type in profile for future use in stale profile matching.

Test Plan: updated pseudoprobe-decoding-inline.test

---

Full diff: https://github.com/llvm/llvm-project/pull/99554.diff

4 Files Affected:

• (modified) bolt/include/bolt/Profile/ProfileYAMLMapping.h (+28)
• (modified) bolt/lib/Profile/DataAggregator.cpp (+27)
• (modified) bolt/lib/Profile/YAMLProfileWriter.cpp (+15)
• (modified) bolt/test/X86/pseudoprobe-decoding-inline.test (+8)

diff --git a/bolt/include/bolt/Profile/ProfileYAMLMapping.h b/bolt/include/bolt/Profile/ProfileYAMLMapping.h
index c4e5f2b284a86..2bc6901e5f591 100644
--- a/bolt/include/bolt/Profile/ProfileYAMLMapping.h
+++ b/bolt/include/bolt/Profile/ProfileYAMLMapping.h
@@ -93,11 +93,36 @@ template <> struct MappingTraits<bolt::SuccessorInfo> {
   static const bool flow = true;
 };
 
+namespace bolt {
+struct PseudoProbeInfo {
+  llvm::yaml::Hex64 GUID;
+  uint64_t Index;
+  uint8_t Type;
+
+  bool operator==(const PseudoProbeInfo &Other) const {
+    return Index == Other.Index;
+  }
+  bool operator!=(const PseudoProbeInfo &Other) const {
+    return !(*this == Other);
+  }
+};
+} // end namespace bolt
+
+template <> struct MappingTraits<bolt::PseudoProbeInfo> {
+  static void mapping(IO &YamlIO, bolt::PseudoProbeInfo &PI) {
+    YamlIO.mapRequired("guid", PI.GUID);
+    YamlIO.mapRequired("id", PI.Index);
+    YamlIO.mapRequired("type", PI.Type);
+  }
+
+  static const bool flow = true;
+};
 } // end namespace yaml
 } // end namespace llvm
 
 LLVM_YAML_IS_FLOW_SEQUENCE_VECTOR(llvm::yaml::bolt::CallSiteInfo)
 LLVM_YAML_IS_FLOW_SEQUENCE_VECTOR(llvm::yaml::bolt::SuccessorInfo)
+LLVM_YAML_IS_FLOW_SEQUENCE_VECTOR(llvm::yaml::bolt::PseudoProbeInfo)
 
 namespace llvm {
 namespace yaml {
@@ -111,6 +136,7 @@ struct BinaryBasicBlockProfile {
   uint64_t EventCount{0};
   std::vector<CallSiteInfo> CallSites;
   std::vector<SuccessorInfo> Successors;
+  std::vector<PseudoProbeInfo> PseudoProbes;
 
   bool operator==(const BinaryBasicBlockProfile &Other) const {
     return Index == Other.Index;
@@ -132,6 +158,8 @@ template <> struct MappingTraits<bolt::BinaryBasicBlockProfile> {
                        std::vector<bolt::CallSiteInfo>());
     YamlIO.mapOptional("succ", BBP.Successors,
                        std::vector<bolt::SuccessorInfo>());
+    YamlIO.mapOptional("pseudo_probes", BBP.PseudoProbes,
+                       std::vector<bolt::PseudoProbeInfo>());
   }
 };
 
diff --git a/bolt/lib/Profile/DataAggregator.cpp b/bolt/lib/Profile/DataAggregator.cpp
index beb03a952d658..85e583c75728f 100644
--- a/bolt/lib/Profile/DataAggregator.cpp
+++ b/bolt/lib/Profile/DataAggregator.cpp
@@ -2406,6 +2406,33 @@ std::error_code DataAggregator::writeBATYAML(BinaryContext &BC,
               PseudoProbeDecoder->getFuncDescForGUID(YamlBF.GUID);
           YamlBF.PseudoProbeDescHash = FuncDesc->FuncHash;
         }
+        // Fetch probes belonging to all fragments
+        const AddressProbesMap &ProbeMap =
+            PseudoProbeDecoder->getAddress2ProbesMap();
+        BinaryFunction::FragmentsSetTy Fragments(BF->Fragments);
+        Fragments.insert(BF);
+        for (const BinaryFunction *F : Fragments) {
+          const uint64_t FuncAddr = F->getAddress();
+          const auto &FragmentProbes =
+              llvm::make_range(ProbeMap.lower_bound(FuncAddr),
+                               ProbeMap.lower_bound(FuncAddr + F->getSize()));
+          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);
+            for (const MCDecodedPseudoProbe &Probe : Probes)
+              YamlBF.Blocks[BlockIndex].PseudoProbes.emplace_back(
+                  yaml::bolt::PseudoProbeInfo{Probe.getGuid(), Probe.getIndex(),
+                                              Probe.getType()});
+          }
+        }
       }
       // Drop blocks without a hash, won't be useful for stale matching.
       llvm::erase_if(YamlBF.Blocks,
diff --git a/bolt/lib/Profile/YAMLProfileWriter.cpp b/bolt/lib/Profile/YAMLProfileWriter.cpp
index 38345da73df7c..cfad75bd5a404 100644
--- a/bolt/lib/Profile/YAMLProfileWriter.cpp
+++ b/bolt/lib/Profile/YAMLProfileWriter.cpp
@@ -185,6 +185,21 @@ YAMLProfileWriter::convert(const BinaryFunction &BF, bool UseDFS,
       ++BranchInfo;
     }
 
+    if (PseudoProbeDecoder) {
+      const AddressProbesMap &ProbeMap =
+          PseudoProbeDecoder->getAddress2ProbesMap();
+      const uint64_t FuncAddr = BF.getAddress();
+      const std::pair<uint64_t, uint64_t> &BlockRange =
+          BB->getInputAddressRange();
+      const auto &BlockProbes =
+          llvm::make_range(ProbeMap.lower_bound(FuncAddr + BlockRange.first),
+                           ProbeMap.lower_bound(FuncAddr + BlockRange.second));
+      for (const auto &[_, Probes] : BlockProbes)
+        for (const MCDecodedPseudoProbe &Probe : Probes)
+          YamlBB.PseudoProbes.emplace_back(yaml::bolt::PseudoProbeInfo{
+              Probe.getGuid(), Probe.getIndex(), Probe.getType()});
+    }
+
     YamlBF.Blocks.emplace_back(YamlBB);
   }
   return YamlBF;
diff --git a/bolt/test/X86/pseudoprobe-decoding-inline.test b/bolt/test/X86/pseudoprobe-decoding-inline.test
index c32660c41a09c..924331e34c76a 100644
--- a/bolt/test/X86/pseudoprobe-decoding-inline.test
+++ b/bolt/test/X86/pseudoprobe-decoding-inline.test
@@ -13,12 +13,20 @@
 # 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: - bid: 0
+# CHECK-YAML:   pseudo_probes: [ { guid: 0xE413754A191DB537, id: 1, type: 0 }, { guid: 0xE413754A191DB537, id: 4, type: 0 } ]
 # CHECK-YAML: guid: 0xE413754A191DB537
 # CHECK-YAML: pseudo_probe_desc_hash: 0x10E852DA94
+#
 # CHECK-YAML: name: foo
+# CHECK-YAML: - bid: 0
+# CHECK-YAML:   pseudo_probes: [ { guid: 0x5CF8C24CDB18BDAC, id: 1, type: 0 }, { guid: 0x5CF8C24CDB18BDAC, id: 2, type: 0 } ]
 # CHECK-YAML: guid: 0x5CF8C24CDB18BDAC
 # CHECK-YAML: pseudo_probe_desc_hash: 0x200205A19C5B4
+#
 # CHECK-YAML: name: main
+# CHECK-YAML: - bid: 0
+# CHECK-YAML:   pseudo_probes: [ { guid: 0xDB956436E78DD5FA, id: 1, type: 0 }, { guid: 0x5CF8C24CDB18BDAC, id: 1, type: 0 }, { guid: 0x5CF8C24CDB18BDAC, id: 2, type: 0 } ]
 # CHECK-YAML: guid: 0xDB956436E78DD5FA
 # CHECK-YAML: pseudo_probe_desc_hash: 0x10000FFFFFFFF
 

[rafaelauler]

LGTM, though not an expert in pseudo probes