[BOLT] Lookup top-level inline tree node in YAMLProfileWriter

[aaupov]

Top-level (binary) functions don't have a unique GUID mapping, with different
causes namely coroutine fragments sharing the same parent source function GUID.

Replace the top-level inline tree node GUID lookup with probe lookup coupled
with walk up the inline tree.

Test Plan: added test-coro-probes.yaml

[llvmbot]

@llvm/pr-subscribers-llvm-mc

@llvm/pr-subscribers-bolt

Author: Amir Ayupov (aaupov)

Changes

Top-level (binary) functions don't have a unique GUID mapping, with different
causes namely coroutine fragments sharing the same parent source function GUID.

Replace the top-level inline tree node GUID lookup with probe lookup coupled
with walk up the inline tree.

Test Plan: TBD

---

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

3 Files Affected:

• (modified) bolt/include/bolt/Profile/YAMLProfileWriter.h (+2-3)
• (modified) bolt/lib/Profile/DataAggregator.cpp (+1-1)
• (modified) bolt/lib/Profile/YAMLProfileWriter.cpp (+10-10)

diff --git a/bolt/include/bolt/Profile/YAMLProfileWriter.h b/bolt/include/bolt/Profile/YAMLProfileWriter.h
index 50ee78d342df8..9a7d5f6a1b5b8 100644
--- a/bolt/include/bolt/Profile/YAMLProfileWriter.h
+++ b/bolt/include/bolt/Profile/YAMLProfileWriter.h
@@ -36,16 +36,15 @@ class YAMLProfileWriter {
       DenseMap<const MCDecodedPseudoProbeInlineTree *, uint32_t>;
   struct InlineTreeDesc {
     template <typename T> using GUIDMapTy = std::unordered_map<uint64_t, T>;
-    using GUIDNodeMap = GUIDMapTy<const MCDecodedPseudoProbeInlineTree *>;
     using GUIDNumMap = GUIDMapTy<uint32_t>;
-    GUIDNodeMap TopLevelGUIDToInlineTree;
     GUIDNumMap GUIDIdxMap;
     GUIDNumMap HashIdxMap;
   };
 
   static std::tuple<std::vector<yaml::bolt::InlineTreeNode>, InlineTreeMapTy>
   convertBFInlineTree(const MCPseudoProbeDecoder &Decoder,
-                      const InlineTreeDesc &InlineTree, uint64_t GUID);
+                      const InlineTreeDesc &InlineTree,
+                      const BinaryFunction &BF);
 
   static std::tuple<yaml::bolt::ProfilePseudoProbeDesc, InlineTreeDesc>
   convertPseudoProbeDesc(const MCPseudoProbeDecoder &PseudoProbeDecoder);
diff --git a/bolt/lib/Profile/DataAggregator.cpp b/bolt/lib/Profile/DataAggregator.cpp
index 6a12efdcb96c0..5cd19803789bf 100644
--- a/bolt/lib/Profile/DataAggregator.cpp
+++ b/bolt/lib/Profile/DataAggregator.cpp
@@ -2376,7 +2376,7 @@ std::error_code DataAggregator::writeBATYAML(BinaryContext &BC,
         if (BF->getGUID()) {
           std::tie(YamlBF.InlineTree, InlineTreeNodeId) =
               YAMLProfileWriter::convertBFInlineTree(*PseudoProbeDecoder,
-                                                     InlineTree, BF->getGUID());
+                                                     InlineTree, *BF);
         }
         // Fetch probes belonging to all fragments
         const AddressProbesMap &ProbeMap =
diff --git a/bolt/lib/Profile/YAMLProfileWriter.cpp b/bolt/lib/Profile/YAMLProfileWriter.cpp
index 82f2114a7a3d2..7c49da5beffb5 100644
--- a/bolt/lib/Profile/YAMLProfileWriter.cpp
+++ b/bolt/lib/Profile/YAMLProfileWriter.cpp
@@ -87,10 +87,6 @@ YAMLProfileWriter::convertPseudoProbeDesc(const MCPseudoProbeDecoder &Decoder) {
   yaml::bolt::ProfilePseudoProbeDesc Desc;
   InlineTreeDesc InlineTree;
 
-  for (const MCDecodedPseudoProbeInlineTree &TopLev :
-       Decoder.getDummyInlineRoot().getChildren())
-    InlineTree.TopLevelGUIDToInlineTree[TopLev.Guid] = &TopLev;
-
   for (const auto &FuncDesc : Decoder.getGUID2FuncDescMap())
     ++InlineTree.HashIdxMap[FuncDesc.FuncHash];
 
@@ -191,14 +187,18 @@ std::tuple<std::vector<yaml::bolt::InlineTreeNode>,
            YAMLProfileWriter::InlineTreeMapTy>
 YAMLProfileWriter::convertBFInlineTree(const MCPseudoProbeDecoder &Decoder,
                                        const InlineTreeDesc &InlineTree,
-                                       uint64_t GUID) {
+                                       const BinaryFunction &BF) {
   DenseMap<const MCDecodedPseudoProbeInlineTree *, uint32_t> InlineTreeNodeId;
   std::vector<yaml::bolt::InlineTreeNode> YamlInlineTree;
-  auto It = InlineTree.TopLevelGUIDToInlineTree.find(GUID);
-  if (It == InlineTree.TopLevelGUIDToInlineTree.end())
+  uint64_t Addr = BF.getAddress();
+  uint64_t Size = BF.getSize();
+  auto Probes = Decoder.getAddress2ProbesMap().find(Addr, Addr + Size);
+  if (Probes.empty())
     return {YamlInlineTree, InlineTreeNodeId};
-  const MCDecodedPseudoProbeInlineTree *Root = It->second;
-  assert(Root && "Malformed TopLevelGUIDToInlineTree");
+  const MCDecodedPseudoProbe &Probe = *Probes.begin();
+  const MCDecodedPseudoProbeInlineTree *Root = Probe.getInlineTreeNode();
+  while (Root->hasInlineSite())
+    Root = (const MCDecodedPseudoProbeInlineTree *)Root->Parent;
   uint32_t Index = 0;
   uint32_t PrevParent = 0;
   uint32_t PrevGUIDIdx = 0;
@@ -242,7 +242,7 @@ YAMLProfileWriter::convert(const BinaryFunction &BF, bool UseDFS,
   DenseMap<const MCDecodedPseudoProbeInlineTree *, uint32_t> InlineTreeNodeId;
   if (PseudoProbeDecoder && BF.getGUID()) {
     std::tie(YamlBF.InlineTree, InlineTreeNodeId) =
-        convertBFInlineTree(*PseudoProbeDecoder, InlineTree, BF.getGUID());
+        convertBFInlineTree(*PseudoProbeDecoder, InlineTree, BF);
   }
 
   BinaryFunction::BasicBlockOrderType Order;

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[aaupov]

Ping @apolloww @wlei-llvm @HighW4y2H3ll