[BOLT] Match blocks with pseudo probes

bolt/include/bolt/Profile/ProfileYAMLMapping.h

@@ -174,6 +174,9 @@ struct InlineTreeNode {
   uint32_t CallSiteProbe;
   // Index in PseudoProbeDesc.GUID, UINT32_MAX for same as previous (omitted)
   uint32_t GUIDIndex;
+  // Decoded contents, ParentIndexDelta becomes absolute value.
+  uint64_t GUID;
+  uint64_t Hash;
   bool operator==(const InlineTreeNode &) const { return false; }
 };
 } // end namespace bolt
@@ -270,12 +273,12 @@ template <> struct MappingTraits<bolt::BinaryProfileHeader> {
 };
 
 namespace bolt {
-struct PseudoProbeDesc {
+struct ProfilePseudoProbeDesc {
   std::vector<Hex64> GUID;
   std::vector<Hex64> Hash;
   std::vector<uint32_t> GUIDHashIdx; // Index of hash for that GUID in Hash
 
-  bool operator==(const PseudoProbeDesc &Other) const {
+  bool operator==(const ProfilePseudoProbeDesc &Other) const {
     // Only treat empty Desc as equal
     return GUID.empty() && Other.GUID.empty() && Hash.empty() &&
            Other.Hash.empty() && GUIDHashIdx.empty() &&
@@ -284,8 +287,8 @@ struct PseudoProbeDesc {
 };
 } // end namespace bolt
 
-template <> struct MappingTraits<bolt::PseudoProbeDesc> {
-  static void mapping(IO &YamlIO, bolt::PseudoProbeDesc &PD) {
+template <> struct MappingTraits<bolt::ProfilePseudoProbeDesc> {
+  static void mapping(IO &YamlIO, bolt::ProfilePseudoProbeDesc &PD) {
     YamlIO.mapRequired("gs", PD.GUID);
     YamlIO.mapRequired("gh", PD.GUIDHashIdx);
     YamlIO.mapRequired("hs", PD.Hash);
@@ -295,7 +298,7 @@ template <> struct MappingTraits<bolt::PseudoProbeDesc> {
 } // end namespace llvm
 
 LLVM_YAML_IS_SEQUENCE_VECTOR(llvm::yaml::bolt::BinaryFunctionProfile)
-LLVM_YAML_IS_SEQUENCE_VECTOR(llvm::yaml::bolt::PseudoProbeDesc)
+LLVM_YAML_IS_SEQUENCE_VECTOR(llvm::yaml::bolt::ProfilePseudoProbeDesc)
 
 namespace llvm {
 namespace yaml {
@@ -304,7 +307,7 @@ namespace bolt {
 struct BinaryProfile {
   BinaryProfileHeader Header;
   std::vector<BinaryFunctionProfile> Functions;
-  PseudoProbeDesc PseudoProbeDesc;
+  ProfilePseudoProbeDesc PseudoProbeDesc;
 };
 } // namespace bolt
 
@@ -313,7 +316,7 @@ template <> struct MappingTraits<bolt::BinaryProfile> {
     YamlIO.mapRequired("header", BP.Header);
     YamlIO.mapRequired("functions", BP.Functions);
     YamlIO.mapOptional("pseudo_probe_desc", BP.PseudoProbeDesc,
-                       bolt::PseudoProbeDesc());
+                       bolt::ProfilePseudoProbeDesc());
   }
 };
 

bolt/include/bolt/Profile/YAMLProfileReader.h

@@ -102,12 +102,28 @@ class YAMLProfileReader : public ProfileReaderBase {
   };
 
   // A class for matching inline tree nodes between profile and binary.
+  // Provides the mapping from profile inline tree node id to a
+  // corresponding binary MCDecodedPseudoProbeInlineTree node.
+  //
+  // The whole mapping process is the following:
+  //
+  //     (profile)                             (binary)
+  //      | blocks                                ^
+  //      v                                       |
+  // yaml::bolt::BinaryBasicBlockProfile ~= FlowBlock
+  //     ||| probes                               ^  (majority vote)
+  //      v                                      ||| BBPseudoProbeToBlock
+  // yaml::bolt::PseudoProbeInfo            MCDecodedPseudoProbe
+  //      | InlineTreeIndex                       ^
+  //      v                                       | probe id
+  // [ profile node id (uint32_t)    ->     MCDecodedPseudoProbeInlineTree *]
+  //                     InlineTreeNodeMapTy
   class InlineTreeNodeMapTy {
     DenseMap<uint32_t, const MCDecodedPseudoProbeInlineTree *> Map;
 
-    void mapInlineTreeNode(uint32_t ProfileNode,
+    void mapInlineTreeNode(uint32_t ProfileNodeIdx,
                            const MCDecodedPseudoProbeInlineTree *BinaryNode) {
-      auto Res = Map.try_emplace(ProfileNode, BinaryNode);
+      auto Res = Map.try_emplace(ProfileNodeIdx, BinaryNode);
       assert(Res.second &&
              "Duplicate mapping from profile node index to binary inline tree");
       (void)Res;
@@ -123,20 +139,23 @@ class YAMLProfileReader : public ProfileReaderBase {
       return It->second;
     }
 
-    // Match up YAML inline tree with binary inline tree.
-    // \p GetRootCallback is invoked for matching up the first YAML inline tree
-    // node and has the following signature:
-    // const MCDecodedPseudoProbeInlineTree *GetRootCallback(uint64_t RootGUID)
-    void matchInlineTrees(
+    // Match up \p YamlInlineTree with binary inline tree rooted at \p Root.
+    // Return the number of matched nodes.
+    //
+    // This function populates the mapping from profile inline tree node id to a
+    // corresponding binary MCDecodedPseudoProbeInlineTree node.
+    size_t matchInlineTrees(
         const MCPseudoProbeDecoder &Decoder,
-        const yaml::bolt::PseudoProbeDesc &YamlPD,
         const std::vector<yaml::bolt::InlineTreeNode> &YamlInlineTree,
-        llvm::function_ref<const MCDecodedPseudoProbeInlineTree *(uint64_t)>
-            GetRootCallback);
-
-    size_t size() const { return Map.size(); }
+        const MCDecodedPseudoProbeInlineTree *Root);
   };
 
+  // Partial probe matching specification: matched inline tree and corresponding
+  // BinaryFunctionProfile
+  using ProbeMatchSpec =
+      std::pair<InlineTreeNodeMapTy,
+                std::reference_wrapper<yaml::bolt::BinaryFunctionProfile>>;
+
 private:
   /// Adjustments for basic samples profiles (without LBR).
   bool NormalizeByInsnCount{false};
@@ -146,7 +165,7 @@ class YAMLProfileReader : public ProfileReaderBase {
   yaml::bolt::BinaryProfile YamlBP;
 
   /// Map a function ID from a YAML profile to a BinaryFunction object.
-  std::vector<BinaryFunction *> YamlProfileToFunction;
+  DenseMap<uint32_t, BinaryFunction *> YamlProfileToFunction;
 
   using FunctionSet = std::unordered_set<const BinaryFunction *>;
   /// To keep track of functions that have a matched profile before the profile
@@ -173,6 +192,10 @@ class YAMLProfileReader : public ProfileReaderBase {
   // Pseudo probe function GUID to inline tree node
   GUIDInlineTreeMap TopLevelGUIDToInlineTree;
 
+  // Mapping from a binary function to its partial match specification
+  // (YAML profile and its inline tree mapping to binary).
+  DenseMap<BinaryFunction *, std::vector<ProbeMatchSpec>> BFToProbeMatchSpecs;
+
   /// Populate \p Function profile with the one supplied in YAML format.
   bool parseFunctionProfile(BinaryFunction &Function,
                             const yaml::bolt::BinaryFunctionProfile &YamlBF);
@@ -183,7 +206,8 @@ class YAMLProfileReader : public ProfileReaderBase {
 
   /// Infer function profile from stale data (collected on older binaries).
   bool inferStaleProfile(BinaryFunction &Function,
-                         const yaml::bolt::BinaryFunctionProfile &YamlBF);
+                         const yaml::bolt::BinaryFunctionProfile &YamlBF,
+                         const ArrayRef<ProbeMatchSpec> ProbeMatchSpecs);
 
   /// Initialize maps for profile matching.
   void buildNameMaps(BinaryContext &BC);
@@ -200,14 +224,16 @@ class YAMLProfileReader : public ProfileReaderBase {
   /// Matches functions using the call graph.
   size_t matchWithCallGraph(BinaryContext &BC);
 
+  /// Matches functions using the call graph.
+  /// Populates BF->partial probe match spec map.
+  size_t matchWithPseudoProbes(BinaryContext &BC);
+
   /// Matches functions with similarly named profiled functions.
   size_t matchWithNameSimilarity(BinaryContext &BC);
 
   /// Update matched YAML -> BinaryFunction pair.
   void matchProfileToFunction(yaml::bolt::BinaryFunctionProfile &YamlBF,
                               BinaryFunction &BF) {
-    if (YamlBF.Id >= YamlProfileToFunction.size())
-      YamlProfileToFunction.resize(YamlBF.Id + 1);
     YamlProfileToFunction[YamlBF.Id] = &BF;
     YamlBF.Used = true;
 

bolt/include/bolt/Profile/YAMLProfileWriter.h

@@ -47,7 +47,7 @@ class YAMLProfileWriter {
   convertBFInlineTree(const MCPseudoProbeDecoder &Decoder,
                       const InlineTreeDesc &InlineTree, uint64_t GUID);
 
-  static std::tuple<yaml::bolt::PseudoProbeDesc, InlineTreeDesc>
+  static std::tuple<yaml::bolt::ProfilePseudoProbeDesc, InlineTreeDesc>
   convertPseudoProbeDesc(const MCPseudoProbeDecoder &PseudoProbeDecoder);
 
   static yaml::bolt::BinaryFunctionProfile