[BOLT] Add support for reading YAML profile v2 (WIP)

[Amir]

Summary

Provide support for reading YAML profile with sub-profiles.
Correctly update the function profile using sub-profiles.

bolt/include/bolt/Profile/YAMLProfileReader.h

public:
   static bool isYAML(StringRef Filename);
 
 private:
   /// Adjustments for basic samples profiles (without LBR).
   bool NormalizeByInsnCount{false};
   bool NormalizeByCalls{false};
 
   /// Binary profile in YAML format.
-  yaml::bolt::BinaryProfile YamlBP;
+  std::vector<yaml::bolt::BinaryProfile> YamlBPs;
 
   /// Map a function ID from a YAML profile to a BinaryFunction object.
-  std::vector<BinaryFunction *> YamlProfileToFunction;
+  std::vector<std::vector<BinaryFunction *>> YamlProfileToFunction;
 
   /// To keep track of functions that have a matched profile before the profile
   /// is attributed.
-  std::unordered_set<const BinaryFunction *> ProfiledFunctions;
+  std::vector<std::unordered_set<const BinaryFunction *>> ProfiledFunctions;
 
   /// For LTO symbol resolution.
   /// Map a common LTO prefix to a list of YAML profiles matching the prefix.
-  StringMap<std::vector<yaml::bolt::BinaryFunctionProfile *>> LTOCommonNameMap;
+  std::vector<StringMap<std::vector<yaml::bolt::BinaryFunctionProfile *>>>
+      LTOCommonNameMap;
 
   /// Map a common LTO prefix to a set of binary functions.
-  StringMap<std::unordered_set<const BinaryFunction *>>
+  std::vector<StringMap<std::unordered_set<const BinaryFunction *>>>
       LTOCommonNameFunctionMap;
 
   /// Strict matching of a name in a profile to its contents.
-  StringMap<yaml::bolt::BinaryFunctionProfile *> ProfileNameToProfile;
+  std::vector<StringMap<yaml::bolt::BinaryFunctionProfile *>>
+      ProfileNameToProfile;
 
   /// Populate \p Function profile with the one supplied in YAML format.
   bool parseFunctionProfile(BinaryFunction &Function,
-                            const yaml::bolt::BinaryFunctionProfile &YamlBF);
+                            const yaml::bolt::BinaryFunctionProfile &YamlBF,
+                            const yaml::bolt::BinaryProfile &YamlBP,
+                            unsigned Epoch);
 
   /// Initialize maps for profile matching.
-  void buildNameMaps(std::map<uint64_t, BinaryFunction> &Functions);
+  void buildNameMaps(std::map<uint64_t, BinaryFunction> &Functions,
+                     unsigned Epoch);
 
   /// 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;
+                              BinaryFunction &BF, unsigned Epoch) {
+    if (YamlBF.Id >= YamlProfileToFunction[Epoch].size())
+      YamlProfileToFunction[Epoch].resize(YamlBF.Id + 1);
+    YamlProfileToFunction[Epoch][YamlBF.Id] = &BF;
     YamlBF.Used = true;
 
-    assert(!ProfiledFunctions.count(&BF) &&
+    assert(!ProfiledFunctions[Epoch].count(&BF) &&
            "function already has an assigned profile");
-    ProfiledFunctions.emplace(&BF);
+    ProfiledFunctions[Epoch].emplace(&BF);
   }
 
   /// Check if the profile uses an event with a given \p Name.
-  bool usesEvent(StringRef Name) const;
+  bool usesEvent(StringRef Name, const yaml::bolt::BinaryProfile &YamlBP) const;
 };
 
 } // namespace bolt
 } // namespace llvm
 
 #endif

bolt/lib/Profile/YAMLProfileReader.cpp

if (std::error_code EC = MB.getError())
     report_error(Filename, EC);
   StringRef Buffer = MB.get()->getBuffer();
   if (Buffer.startswith("---\n"))
     return true;
   return false;
 }
 
 void YAMLProfileReader::buildNameMaps(
-    std::map<uint64_t, BinaryFunction> &Functions) {
+    std::map<uint64_t, BinaryFunction> &Functions, unsigned Epoch) {
+  yaml::bolt::BinaryProfile &YamlBP = YamlBPs[Epoch];
   for (yaml::bolt::BinaryFunctionProfile &YamlBF : YamlBP.Functions) {
     StringRef Name = YamlBF.Name;
     const size_t Pos = Name.find("(*");
     if (Pos != StringRef::npos)
       Name = Name.substr(0, Pos);
-    ProfileNameToProfile[Name] = &YamlBF;
+    ProfileNameToProfile[Epoch][Name] = &YamlBF;
     if (const std::optional<StringRef> CommonName = getLTOCommonName(Name))
-      LTOCommonNameMap[*CommonName].push_back(&YamlBF);
+      LTOCommonNameMap[Epoch][*CommonName].push_back(&YamlBF);
   }
   for (auto &BFI : Functions) {
     const BinaryFunction &Function = BFI.second;
     for (StringRef Name : Function.getNames())
       if (const std::optional<StringRef> CommonName = getLTOCommonName(Name))
-        LTOCommonNameFunctionMap[*CommonName].insert(&Function);
+        LTOCommonNameFunctionMap[Epoch][*CommonName].insert(&Function);
   }
 }
 
 bool YAMLProfileReader::hasLocalsWithFileName() const {
-  for (const StringMapEntry<yaml::bolt::BinaryFunctionProfile *> &KV :
-       ProfileNameToProfile) {
-    const StringRef &FuncName = KV.getKey();
-    if (FuncName.count('/') == 2 && FuncName[0] != '/')
-      return true;
-  }
-  return false;
+  auto CheckFuncName = [](const StringRef &FuncName) {
+    return FuncName.count('/') == 2 && FuncName[0] != '/';
+  };
+  auto CheckMapKeys =
+      [&](const StringMap<yaml::bolt::BinaryFunctionProfile *> &Map) {
+        return llvm::any_of(Map.keys(), CheckFuncName);
+      };
+  return llvm::any_of(ProfileNameToProfile, CheckMapKeys);
 }
 
 bool YAMLProfileReader::parseFunctionProfile(
-    BinaryFunction &BF, const yaml::bolt::BinaryFunctionProfile &YamlBF) {
+    BinaryFunction &BF, const yaml::bolt::BinaryFunctionProfile &YamlBF,
+    const yaml::bolt::BinaryProfile &YamlBP, unsigned Epoch) {
   BinaryContext &BC = BF.getBinaryContext();
 
   bool ProfileMatched = true;
   uint64_t MismatchedBlocks = 0;
   uint64_t MismatchedCalls = 0;
   uint64_t MismatchedEdges = 0;
 
   uint64_t FunctionExecutionCount = 0;
 
-  BF.setExecutionCount(YamlBF.ExecCount);
+  BF.setExecutionCount(BF.getExecutionCount() + YamlBF.ExecCount);
 
   if (!opts::IgnoreHash && YamlBF.Hash != BF.computeHash(/*UseDFS=*/true)) {
     if (opts::Verbosity >= 1)
       errs() << "BOLT-WARNING: function hash mismatch\n";
     ProfileMatched = false;
   }
 
   if (YamlBF.NumBasicBlocks != BF.size()) {
if (YamlBB.Index >= DFSOrder.size()) {
       continue;
     }
 
     BinaryBasicBlock &BB = *DFSOrder[YamlBB.Index];
 
     // Basic samples profile (without LBR) does not have branches information
     // and needs a special processing.
     if (YamlBP.Header.Flags & BinaryFunction::PF_SAMPLE) {
-      if (!YamlBB.EventCount) {
-        BB.setExecutionCount(0);
+      if (!YamlBB.EventCount)
         continue;
-      }
       uint64_t NumSamples = YamlBB.EventCount * 1000;
       if (NormalizeByInsnCount && BB.getNumNonPseudos())
         NumSamples /= BB.getNumNonPseudos();
       else if (NormalizeByCalls)
         NumSamples /= BB.getNumCalls() + 1;
 
-      BB.setExecutionCount(NumSamples);
+      BB.setExecutionCount(BB.getExecutionCount() + NumSamples);
       if (BB.isEntryPoint())
         FunctionExecutionCount += NumSamples;
       continue;
     }
 
-    BB.setExecutionCount(YamlBB.ExecCount);
+    BB.setExecutionCount(BB.getExecutionCount() + YamlBB.ExecCount);
 
     for (const yaml::bolt::CallSiteInfo &YamlCSI : YamlBB.CallSites) {
-      BinaryFunction *Callee = YamlCSI.DestId < YamlProfileToFunction.size()
-                                   ? YamlProfileToFunction[YamlCSI.DestId]
-                                   : nullptr;
+      BinaryFunction *Callee =
+          YamlCSI.DestId < YamlProfileToFunction[Epoch].size()
+              ? YamlProfileToFunction[Epoch][YamlCSI.DestId]
       bool IsFunction = Callee ? true : false;
       MCSymbol *CalleeSymbol = nullptr;
       if (IsFunction)
         CalleeSymbol = Callee->getSymbolForEntryID(YamlCSI.EntryDiscriminator);
 
       BF.getAllCallSites().emplace_back(CalleeSymbol, YamlCSI.Count,
                                         YamlCSI.Mispreds, YamlCSI.Offset);
 
bool YAMLProfileReader::parseFunctionProfile(
   }
 
   // If basic block profile wasn't read it should be 0.
   for (BinaryBasicBlock &BB : BF)
     if (BB.getExecutionCount() == BinaryBasicBlock::COUNT_NO_PROFILE)
       BB.setExecutionCount(0);
 
   if (YamlBP.Header.Flags & BinaryFunction::PF_SAMPLE) {
-    BF.setExecutionCount(FunctionExecutionCount);
+    BF.setExecutionCount(BF.getExecutionCount() + FunctionExecutionCount);
     estimateEdgeCounts(BF);
   }
 
   ProfileMatched &= !MismatchedBlocks && !MismatchedCalls && !MismatchedEdges;
 
   if (ProfileMatched)
     BF.markProfiled(YamlBP.Header.Flags);
 
ErrorOr<std::unique_ptr<MemoryBuffer>> MB =
       MemoryBuffer::getFileOrSTDIN(Filename);
   if (std::error_code EC = MB.getError()) {
     errs() << "ERROR: cannot open " << Filename << ": " << EC.message() << "\n";
     return errorCodeToError(EC);
   }
   yaml::Input YamlInput(MB.get()->getBuffer());
 
   // Consume YAML file.
-  YamlInput >> YamlBP;
+  YamlInput >> YamlBPs;
   if (YamlInput.error()) {
     errs() << "BOLT-ERROR: syntax error parsing profile in " << Filename
            << " : " << YamlInput.error().message() << '\n';
     return errorCodeToError(YamlInput.error());
   }
 
+  ProfileNameToProfile.resize(YamlBPs.size());
+  LTOCommonNameMap.resize(YamlBPs.size());
+  LTOCommonNameFunctionMap.resize(YamlBPs.size());
+  for (unsigned Epoch = 0, End = YamlBPs.size(); Epoch != End; ++Epoch) {
+    yaml::bolt::BinaryProfile YamlBP = YamlBPs[Epoch];
     // Sanity check.
-  if (YamlBP.Header.Version != 1)
+    if (YamlBP.Header.Version != 2)
       return make_error<StringError>(
           Twine("cannot read profile : unsupported version"),
           inconvertibleErrorCode());
 
     if (YamlBP.Header.EventNames.find(',') != StringRef::npos)
       return make_error<StringError>(
           Twine("multiple events in profile are not supported"),
           inconvertibleErrorCode());
 
     // Match profile to function based on a function name.
-  buildNameMaps(BC.getBinaryFunctions());
+    buildNameMaps(BC.getBinaryFunctions(), Epoch);
 
     // Preliminary assign function execution count.
     for (auto &KV : BC.getBinaryFunctions()) {
       BinaryFunction &BF = KV.second;
       for (StringRef Name : BF.getNames()) {
-      auto PI = ProfileNameToProfile.find(Name);
-      if (PI != ProfileNameToProfile.end()) {
+        auto PI = ProfileNameToProfile[Epoch].find(Name);
+        if (PI != ProfileNameToProfile[Epoch].end()) {
           yaml::bolt::BinaryFunctionProfile &YamlBF = *PI->getValue();
-        BF.setExecutionCount(YamlBF.ExecCount);
+          BF.setExecutionCount(BF.getExecutionCount() + YamlBF.ExecCount);
           break;
         }
       }
     }
+  }
 
   return Error::success();
 }
 
 bool YAMLProfileReader::mayHaveProfileData(const BinaryFunction &BF) {
   for (StringRef Name : BF.getNames()) {
-    if (ProfileNameToProfile.find(Name) != ProfileNameToProfile.end())
+    for (unsigned Epoch = 0, End = ProfileNameToProfile.size(); Epoch != End;
+         ++Epoch) {
+      if (ProfileNameToProfile[Epoch].find(Name) !=
+          ProfileNameToProfile[Epoch].end())
         return true;
       if (const std::optional<StringRef> CommonName = getLTOCommonName(Name)) {
-      if (LTOCommonNameMap.find(*CommonName) != LTOCommonNameMap.end())
+        if (LTOCommonNameMap[Epoch].find(*CommonName) !=
+            LTOCommonNameMap[Epoch].end())
           return true;
       }
     }
+  }
 
   return false;
 }
 
 Error YAMLProfileReader::readProfile(BinaryContext &BC) {
-  YamlProfileToFunction.resize(YamlBP.Functions.size() + 1);
+  uint64_t NumUnused = 0;
+  YamlProfileToFunction.resize(YamlBPs.size());
+  ProfiledFunctions.resize(YamlBPs.size());
+  for (unsigned Epoch = 0, End = YamlBPs.size(); Epoch != End; ++Epoch) {
+    auto YamlBP = YamlBPs[Epoch];
+    YamlProfileToFunction[Epoch].resize(YamlBP.Functions.size() + 1);
 
     auto profileMatches = [](const yaml::bolt::BinaryFunctionProfile &Profile,
                              BinaryFunction &BF) {
       if (opts::IgnoreHash && Profile.NumBasicBlocks == BF.size())
         return true;
       if (!opts::IgnoreHash &&
           Profile.Hash == static_cast<uint64_t>(BF.getHash()))
         return true;
       return false;
     };
 
     // We have to do 2 passes since LTO introduces an ambiguity in function
     // names. The first pass assigns profiles that match 100% by name and
     // by hash. The second pass allows name ambiguity for LTO private functions.
     for (auto &BFI : BC.getBinaryFunctions()) {
       BinaryFunction &Function = BFI.second;
 
       // Clear function call count that may have been set while pre-processing
       // the profile.
       Function.setExecutionCount(BinaryFunction::COUNT_NO_PROFILE);
 
       // Recompute hash once per function.
       if (!opts::IgnoreHash)
         Function.computeHash(/*UseDFS=*/true);
 
       for (StringRef FunctionName : Function.getNames()) {
-      auto PI = ProfileNameToProfile.find(FunctionName);
-      if (PI == ProfileNameToProfile.end())
+        auto PI = ProfileNameToProfile[Epoch].find(FunctionName);
+        if (PI == ProfileNameToProfile[Epoch].end())
           continue;
 
         yaml::bolt::BinaryFunctionProfile &YamlBF = *PI->getValue();
         if (profileMatches(YamlBF, Function))
-        matchProfileToFunction(YamlBF, Function);
+          matchProfileToFunction(YamlBF, Function, Epoch);
       }
     }
 
     for (auto &BFI : BC.getBinaryFunctions()) {
       BinaryFunction &Function = BFI.second;
 
-    if (ProfiledFunctions.count(&Function))
+      if (ProfiledFunctions[Epoch].count(&Function))
         continue;
 
       for (StringRef FunctionName : Function.getNames()) {
         const std::optional<StringRef> CommonName =
             getLTOCommonName(FunctionName);
         if (CommonName) {
-        auto I = LTOCommonNameMap.find(*CommonName);
-        if (I == LTOCommonNameMap.end())
+          auto I = LTOCommonNameMap[Epoch].find(*CommonName);
+          if (I == LTOCommonNameMap[Epoch].end())
             continue;
 
           bool ProfileMatched = false;
           std::vector<yaml::bolt::BinaryFunctionProfile *> &LTOProfiles =
               I->getValue();
           for (yaml::bolt::BinaryFunctionProfile *YamlBF : LTOProfiles) {
             if (YamlBF->Used)
               continue;
             if ((ProfileMatched = profileMatches(*YamlBF, Function))) {
-            matchProfileToFunction(*YamlBF, Function);
+              matchProfileToFunction(*YamlBF, Function, Epoch);
               break;
             }
           }
           if (ProfileMatched)
             break;
 
           // If there's only one function with a given name, try to
           // match it partially.
           if (LTOProfiles.size() == 1 &&
-            LTOCommonNameFunctionMap[*CommonName].size() == 1 &&
+              LTOCommonNameFunctionMap[Epoch][*CommonName].size() == 1 &&
               !LTOProfiles.front()->Used) {
-          matchProfileToFunction(*LTOProfiles.front(), Function);
+            matchProfileToFunction(*LTOProfiles.front(), Function, Epoch);
             break;
           }
         } else {
-        auto PI = ProfileNameToProfile.find(FunctionName);
-        if (PI == ProfileNameToProfile.end())
+          auto PI = ProfileNameToProfile[Epoch].find(FunctionName);
+          if (PI == ProfileNameToProfile[Epoch].end())
             continue;
 
           yaml::bolt::BinaryFunctionProfile &YamlBF = *PI->getValue();
           if (!YamlBF.Used) {
-          matchProfileToFunction(YamlBF, Function);
+            matchProfileToFunction(YamlBF, Function, Epoch);
             break;
           }
         }
       }
     }
 
     for (yaml::bolt::BinaryFunctionProfile &YamlBF : YamlBP.Functions)
       if (!YamlBF.Used && opts::Verbosity >= 1)
         errs() << "BOLT-WARNING: profile ignored for function " << YamlBF.Name
                << '\n';
 
     // Set for parseFunctionProfile().
-  NormalizeByInsnCount = usesEvent("cycles") || usesEvent("instructions");
-  NormalizeByCalls = usesEvent("branches");
+    NormalizeByInsnCount =
+        usesEvent("cycles", YamlBP) || usesEvent("instructions", YamlBP);
+    NormalizeByCalls = usesEvent("branches", YamlBP);
 
-  uint64_t NumUnused = 0;
     for (yaml::bolt::BinaryFunctionProfile &YamlBF : YamlBP.Functions) {
-    if (YamlBF.Id >= YamlProfileToFunction.size()) {
+      if (YamlBF.Id >= YamlProfileToFunction[Epoch].size()) {
         // Such profile was ignored.
         ++NumUnused;
         continue;
       }
-    if (BinaryFunction *BF = YamlProfileToFunction[YamlBF.Id])
-      parseFunctionProfile(*BF, YamlBF);
+      if (BinaryFunction *BF = YamlProfileToFunction[Epoch][YamlBF.Id])
+        parseFunctionProfile(*BF, YamlBF, YamlBP, Epoch);
       else
         ++NumUnused;
     }
+  }
 
   BC.setNumUnusedProfiledObjects(NumUnused);
 
   return Error::success();
 }
 
-bool YAMLProfileReader::usesEvent(StringRef Name) const {
+bool YAMLProfileReader::usesEvent(
+    StringRef Name, const yaml::bolt::BinaryProfile &YamlBP) const {
   return YamlBP.Header.EventNames.find(std::string(Name)) != StringRef::npos;
 }
 
 } // end namespace bolt
 } // end namespace llvm