Memoize top-level GUID->InlineTree mapping, cuts inference time by ~30%

Created using spr 1.3.4

Author: aaupov

bolt/include/bolt/Profile/ProfileYAMLMapping.h

@@ -275,6 +275,12 @@ struct PseudoProbeDesc {
   std::vector<Hex64> GUID;
   std::vector<Hex64> Hash;
   std::vector<uint32_t> GUIDHash; // Index of hash for that GUID in Hash
+
+  bool operator==(const PseudoProbeDesc &Other) const {
+    // Only treat empty Desc as equal
+    return GUID.empty() && Other.GUID.empty() && Hash.empty() &&
+           Other.Hash.empty() && GUIDHash.empty() && Other.GUIDHash.empty();
+  }
 };
 } // end namespace bolt
 
@@ -306,7 +312,8 @@ template <> struct MappingTraits<bolt::BinaryProfile> {
   static void mapping(IO &YamlIO, bolt::BinaryProfile &BP) {
     YamlIO.mapRequired("header", BP.Header);
     YamlIO.mapRequired("functions", BP.Functions);
-    YamlIO.mapOptional("pseudo_probe_desc", BP.PseudoProbeDesc);
+    YamlIO.mapOptional("pseudo_probe_desc", BP.PseudoProbeDesc,
+                       bolt::PseudoProbeDesc());
   }
 };
 

bolt/include/bolt/Profile/YAMLProfileReader.h

@@ -14,6 +14,8 @@
 #include <unordered_set>
 
 namespace llvm {
+class MCDecodedPseudoProbeInlineTree;
+
 namespace bolt {
 
 class YAMLProfileReader : public ProfileReaderBase {
@@ -43,6 +45,9 @@ class YAMLProfileReader : public ProfileReaderBase {
   using ProfileLookupMap =
       DenseMap<uint32_t, yaml::bolt::BinaryFunctionProfile *>;
 
+  using GUIDInlineTreeMap =
+      std::unordered_map<uint64_t, const MCDecodedPseudoProbeInlineTree *>;
+
   /// A class for matching binary functions in functions in the YAML profile.
   /// First, a call graph is constructed for both profiled and binary functions.
   /// Then functions are hashed based on the names of their callee/caller
@@ -129,6 +134,9 @@ class YAMLProfileReader : public ProfileReaderBase {
   /// BinaryFunction pointers indexed by YamlBP functions.
   std::vector<BinaryFunction *> ProfileBFs;
 
+  // Pseudo probe function GUID to inline tree node
+  GUIDInlineTreeMap TopLevelGUIDToInlineTree;
+
   /// Populate \p Function profile with the one supplied in YAML format.
   bool parseFunctionProfile(BinaryFunction &Function,
                             const yaml::bolt::BinaryFunctionProfile &YamlBF);

bolt/lib/Profile/StaleProfileMatching.cpp

@@ -117,8 +117,8 @@ cl::opt<unsigned> StaleMatchingCostJumpUnknownFTInc(
         "The cost of increasing an unknown fall-through jump count by one."),
     cl::init(3), cl::ReallyHidden, cl::cat(BoltOptCategory));
 
-cl::opt<bool> StaleMatchingWithBlockPseudoProbes(
-    "stale-matching-with-block-pseudo-probes",
+cl::opt<bool> StaleMatchingWithPseudoProbes(
+    "stale-matching-with-pseudo-probes",
     cl::desc("Turns on stale matching with block pseudo probes."),
     cl::init(false), cl::ReallyHidden, cl::cat(BoltOptCategory));
 
@@ -328,7 +328,7 @@ class StaleMatcher {
   std::pair<const FlowBlock *, bool> matchWithPseudoProbes(
       const ArrayRef<yaml::bolt::PseudoProbeInfo> BlockPseudoProbes,
       const ArrayRef<yaml::bolt::InlineTreeInfo> InlineTree) const {
-    if (!opts::StaleMatchingWithBlockPseudoProbes)
+    if (!opts::StaleMatchingWithPseudoProbes)
       return {nullptr, false};
 
     DenseMap<const FlowBlock *, uint32_t> FlowBlockMatchCount;
@@ -574,7 +574,8 @@ size_t matchWeightsByHashes(
     BinaryContext &BC, const BinaryFunction::BasicBlockOrderType &BlockOrder,
     const yaml::bolt::BinaryFunctionProfile &YamlBF, FlowFunction &Func,
     HashFunction HashFunction, YAMLProfileReader::ProfileLookupMap &IdToYamlBF,
-    const BinaryFunction &BF, const yaml::bolt::PseudoProbeDesc &YamlPD) {
+    const BinaryFunction &BF, const yaml::bolt::PseudoProbeDesc &YamlPD,
+    const YAMLProfileReader::GUIDInlineTreeMap &TopLevelGUIDToInlineTree) {
 
   assert(Func.Blocks.size() == BlockOrder.size() + 2);
 
@@ -605,21 +606,19 @@ size_t matchWeightsByHashes(
   }
   StaleMatcher Matcher;
   // Collects function pseudo probes for use in the StaleMatcher.
-  if (opts::StaleMatchingWithBlockPseudoProbes) {
-    const MCPseudoProbeDecoder *PseudoProbeDecoder = BC.getPseudoProbeDecoder();
-    assert(PseudoProbeDecoder &&
+  if (opts::StaleMatchingWithPseudoProbes) {
+    const MCPseudoProbeDecoder *Decoder = BC.getPseudoProbeDecoder();
+    assert(Decoder &&
            "If pseudo probes are in use, pseudo probe decoder should exist");
-    const AddressProbesMap &ProbeMap =
-        PseudoProbeDecoder->getAddress2ProbesMap();
+    const AddressProbesMap &ProbeMap = Decoder->getAddress2ProbesMap();
     const uint64_t FuncAddr = BF.getAddress();
     for (const MCDecodedPseudoProbe &Probe :
          ProbeMap.find(FuncAddr, FuncAddr + BF.getSize()))
       if (const BinaryBasicBlock *BB =
               BF.getBasicBlockContainingOffset(Probe.getAddress() - FuncAddr))
         Matcher.mapProbeToBB(&Probe, Blocks[BB->getIndex()]);
+
     // Match inline tree nodes by GUID, checksum, parent, and call site.
-    const MCDecodedPseudoProbeInlineTree *DummyInlineRoot =
-        &PseudoProbeDecoder->getDummyInlineRoot();
     uint32_t ParentId = 0;
     uint32_t PrevGUIDIdx = 0;
     uint32_t Index = 0;
@@ -638,23 +637,24 @@ size_t matchWeightsByHashes(
       uint32_t InlineTreeNodeId = Index++;
       ParentId += InlineTreeNode.ParentIndexDelta;
       uint32_t CallSiteProbe = InlineTreeNode.CallSiteProbe;
-      const MCDecodedPseudoProbeInlineTree *ParentNode =
-          InlineTreeNodeId ? Matcher.getInlineTreeNode(ParentId)
-                           : DummyInlineRoot;
-      if (!ParentNode)
-        continue;
-      for (const MCDecodedPseudoProbeInlineTree &Child :
-           ParentNode->getChildren()) {
-        if (Child.Guid != GUID ||
-            PseudoProbeDecoder->getFuncDescForGUID(GUID)->FuncHash != Hash)
-          continue;
-        // Check inline site for non-toplev inline tree nodes.
-        if (ParentNode != DummyInlineRoot &&
-            std::get<1>(Child.getInlineSite()) != CallSiteProbe)
-          continue;
-        Matcher.mapInlineTreeNode(InlineTreeNodeId, &Child);
-        break;
+      const MCDecodedPseudoProbeInlineTree *Cur = nullptr;
+      if (!InlineTreeNodeId) {
+        auto It = TopLevelGUIDToInlineTree.find(GUID);
+        if (It != TopLevelGUIDToInlineTree.end())
+          Cur = It->second;
+      } else if (const MCDecodedPseudoProbeInlineTree *Parent =
+                     Matcher.getInlineTreeNode(ParentId)) {
+        for (const MCDecodedPseudoProbeInlineTree &Child :
+             Parent->getChildren()) {
+          if (Child.Guid == GUID) {
+            if (std::get<1>(Child.getInlineSite()) == CallSiteProbe)
+              Cur = &Child;
+            break;
+          }
+        }
       }
+      if (Cur && Decoder->getFuncDescForGUID(GUID)->FuncHash == Hash)
+        Matcher.mapInlineTreeNode(InlineTreeNodeId, Cur);
     }
   }
   Matcher.init(Blocks, BlendedHashes, CallHashes);
@@ -1028,9 +1028,10 @@ bool YAMLProfileReader::inferStaleProfile(
   FlowFunction Func = createFlowFunction(BlockOrder);
 
   // Match as many block/jump counts from the stale profile as possible
-  size_t MatchedBlocks = matchWeightsByHashes(
-      BF.getBinaryContext(), BlockOrder, YamlBF, Func,
-      YamlBP.Header.HashFunction, IdToYamLBF, BF, YamlBP.PseudoProbeDesc);
+  size_t MatchedBlocks =
+      matchWeightsByHashes(BF.getBinaryContext(), BlockOrder, YamlBF, Func,
+                           YamlBP.Header.HashFunction, IdToYamLBF, BF,
+                           YamlBP.PseudoProbeDesc, TopLevelGUIDToInlineTree);
 
   // Adjust the flow function by marking unreachable blocks Unlikely so that
   // they don't get any counts assigned.

bolt/lib/Profile/YAMLProfileReader.cpp

@@ -16,6 +16,7 @@
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/edit_distance.h"
 #include "llvm/Demangle/Demangle.h"
+#include "llvm/MC/MCPseudoProbe.h"
 #include "llvm/Support/CommandLine.h"
 
 using namespace llvm;
@@ -49,6 +50,8 @@ llvm::cl::opt<bool>
 llvm::cl::opt<bool> ProfileUseDFS("profile-use-dfs",
                                   cl::desc("use DFS order for YAML profile"),
                                   cl::Hidden, cl::cat(BoltOptCategory));
+
+extern llvm::cl::opt<bool> StaleMatchingWithPseudoProbes;
 } // namespace opts
 
 namespace llvm {
@@ -722,6 +725,15 @@ Error YAMLProfileReader::readProfile(BinaryContext &BC) {
     }
   }
 
+  if (opts::StaleMatchingWithPseudoProbes) {
+    const MCPseudoProbeDecoder *Decoder = BC.getPseudoProbeDecoder();
+    assert(Decoder &&
+           "If pseudo probes are in use, pseudo probe decoder should exist");
+    for (const MCDecodedPseudoProbeInlineTree &TopLev :
+         Decoder->getDummyInlineRoot().getChildren())
+      TopLevelGUIDToInlineTree[TopLev.Guid] = &TopLev;
+  }
+
   // Map profiled function ids to names.
   for (yaml::bolt::BinaryFunctionProfile &YamlBF : YamlBP.Functions)
     IdToYamLBF[YamlBF.Id] = &YamlBF;

bolt/lib/Rewrite/PseudoProbeRewriter.cpp

@@ -51,7 +51,7 @@ static cl::opt<PrintPseudoProbesOptions> PrintPseudoProbes(
     cl::Hidden, cl::cat(BoltCategory));
 
 extern cl::opt<bool> ProfileWritePseudoProbes;
-extern cl::opt<bool> StaleMatchingWithBlockPseudoProbes;
+extern cl::opt<bool> StaleMatchingWithPseudoProbes;
 } // namespace opts
 
 namespace {
@@ -94,14 +94,14 @@ class PseudoProbeRewriter final : public MetadataRewriter {
 
 Error PseudoProbeRewriter::preCFGInitializer() {
   if (opts::ProfileWritePseudoProbes ||
-      opts::StaleMatchingWithBlockPseudoProbes)
+      opts::StaleMatchingWithPseudoProbes)
     parsePseudoProbe(opts::ProfileWritePseudoProbes);
 
   return Error::success();
 }
 
 Error PseudoProbeRewriter::postEmitFinalizer() {
-  if (!opts::StaleMatchingWithBlockPseudoProbes)
+  if (!opts::StaleMatchingWithPseudoProbes)
     parsePseudoProbe();
   updatePseudoProbes();
 

bolt/test/X86/match-blocks-with-pseudo-probes.test

@@ -6,9 +6,9 @@
 # RUN: %clang %cflags %t.o -o %t.exe -Wl,-q -nostdlib
 # RUN: llvm-bolt %t.exe -o %t.out --data %t/yaml -v=2 \
 # RUN:   --print-cfg --funcs=main --infer-stale-profile \
-# RUN:   --stale-matching-with-block-pseudo-probes 2>&1 | FileCheck %s
+# RUN:   --stale-matching-with-pseudo-probes 2>&1 | FileCheck %s
 
-# CHECK: BOLT-INFO: inference found an exact pseudo probe match for 100.00% of basic blocks (1 out of 1 stale) responsible for -nan% samples (0 out of 0 stale)
+# CHECK: BOLT-INFO: inference found an exact pseudo probe match for 100.00% of basic blocks (1 out of 1 stale)
 
 #--- main.s
  .text

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

@@ -36,8 +36,9 @@
 ## generated
 # 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-NO-OPT
-# CHECK-NO-OPT-NOT: pseudo_probes
-# CHECK-NO-OPT-NOT: inline_tree
+# CHECK-NO-OPT-NOT: probes:
+# CHECK-NO-OPT-NOT: inline_tree:
+# CHECK-NO-OPT-NOT: pseudo_probe_desc:
 
 CHECK: Report of decoding input pseudo probe binaries