[llvm-branch-commits] [llvm] [BOLT] Match blocks with pseudo probes (PR #99891)

[Lei Wang via llvm-branch-commits]

================
@@ -482,11 +596,68 @@ matchWeightsByHashes(BinaryContext &BC,
                       << Twine::utohexstr(BB->getHash()) << "\n");
   }
   StaleMatcher Matcher;
+  // Collects function pseudo probes for use in the StaleMatcher.
+  if (opts::StaleMatchingWithPseudoProbes) {
+    const MCPseudoProbeDecoder *Decoder = BC.getPseudoProbeDecoder();
+    assert(Decoder &&
+           "If pseudo probes are in use, pseudo probe decoder should exist");
+    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()]);
+  }
   Matcher.init(Blocks, BlendedHashes, CallHashes);
 
-  // Index in yaml profile => corresponding (matched) block
-  DenseMap<uint64_t, const FlowBlock *> MatchedBlocks;
-  // Match blocks from the profile to the blocks in CFG
+  using FlowBlockTy =
+      std::pair<const FlowBlock *, const yaml::bolt::BinaryBasicBlockProfile *>;
+  using ProfileBlockMatchMap = DenseMap<uint32_t, FlowBlockTy>;
+  // Binary profile => block index => matched block + its block profile
+  DenseMap<const yaml::bolt::BinaryFunctionProfile *, ProfileBlockMatchMap>
+      MatchedBlocks;
+
+  // Map of FlowBlock and matching method.
+  DenseMap<const FlowBlock *, StaleMatcher::MatchMethod> MatchedFlowBlocks;
+
+  // Match blocks from the profile to the blocks in CFG by strict hash.
+  for (const yaml::bolt::BinaryBasicBlockProfile &YamlBB : YamlBF.Blocks) {
+    // Update matching stats.
+    ++BC.Stats.NumStaleBlocks;
+    BC.Stats.StaleSampleCount += YamlBB.ExecCount;
+
+    assert(YamlBB.Hash != 0 && "empty hash of BinaryBasicBlockProfile");
+    BlendedBlockHash YamlHash(YamlBB.Hash);
+    const FlowBlock *MatchedBlock = nullptr;
+    StaleMatcher::MatchMethod Method;
+    std::tie(MatchedBlock, Method) = Matcher.matchBlockStrict(YamlHash);
+    if (!MatchedBlock)
+      continue;
+    MatchedFlowBlocks.try_emplace(MatchedBlock, Method);
+    MatchedBlocks[&YamlBF][YamlBB.Index] = {MatchedBlock, &YamlBB};
+  }
+  // Match blocks from the profile to the blocks in CFG by pseudo probes.
+  for (const YAMLProfileReader::ProbeMatchSpec &PS : ProbeMatchSpecs) {
+    const YAMLProfileReader::InlineTreeNodeMapTy &InlineTreeNodeMap = PS.first;
+    const yaml::bolt::BinaryFunctionProfile &YamlBP = PS.second;
+    for (const yaml::bolt::BinaryBasicBlockProfile &YamlBB : YamlBP.Blocks) {
+      if (YamlBB.PseudoProbes.empty())
+        continue;
+      const FlowBlock *MatchedBlock = nullptr;
+      StaleMatcher::MatchMethod Method;
+      std::tie(MatchedBlock, Method) =
+          Matcher.matchBlockProbe(YamlBB.PseudoProbes, InlineTreeNodeMap);
+      if (!MatchedBlock)
+        continue;
+      // Don't override earlier matches
+      if (MatchedFlowBlocks.contains(MatchedBlock))
+        continue;
+      MatchedFlowBlocks.try_emplace(MatchedBlock, Method);
+      MatchedBlocks[&YamlBP][YamlBB.Index] = {MatchedBlock, &YamlBB};
----------------
wlei-llvm wrote:

Those lines are the same for the three matchings, probably can be factored out into a helper function.

https://github.com/llvm/llvm-project/pull/99891