[BOLT] Add CallSiteInfo entries in YAMLBAT (#76896)

Attach call counters to YAML profile, covering inter-function control
flow.

Depends on: #86218

Test Plan: 
Updated bolt/test/X86/bolt-address-translation-yaml.test

Author: aaupov

bolt/include/bolt/Profile/BoltAddressTranslation.h

@@ -118,6 +118,11 @@ class BoltAddressTranslation {
   /// True if a given \p Address is a function with translation table entry.
   bool isBATFunction(uint64_t Address) const { return Maps.count(Address); }
 
+  /// Returns branch offsets grouped by containing basic block in a given
+  /// function.
+  std::unordered_map<uint32_t, std::vector<uint32_t>>
+  getBFBranches(uint64_t FuncOutputAddress) const;
+
 private:
   /// Helper to update \p Map by inserting one or more BAT entries reflecting
   /// \p BB for function located at \p FuncAddress. At least one entry will be

bolt/include/bolt/Profile/DataAggregator.h

@@ -467,6 +467,9 @@ class DataAggregator : public DataReader {
   std::error_code writeBATYAML(BinaryContext &BC,
                                StringRef OutputFilename) const;
 
+  /// Fixup profile collected on BOLTed binary, namely handle split functions.
+  void fixupBATProfile(BinaryContext &BC);
+
   /// Filter out binaries based on PID
   void filterBinaryMMapInfo();
 

bolt/lib/Profile/BoltAddressTranslation.cpp

@@ -577,5 +577,25 @@ void BoltAddressTranslation::saveMetadata(BinaryContext &BC) {
   }
 }
 
+std::unordered_map<uint32_t, std::vector<uint32_t>>
+BoltAddressTranslation::getBFBranches(uint64_t OutputAddress) const {
+  std::unordered_map<uint32_t, std::vector<uint32_t>> Branches;
+  auto FuncIt = Maps.find(OutputAddress);
+  assert(FuncIt != Maps.end());
+  std::vector<uint32_t> InputOffsets;
+  for (const auto &KV : FuncIt->second)
+    InputOffsets.emplace_back(KV.second);
+  // Sort with LSB BRANCHENTRY bit.
+  llvm::sort(InputOffsets);
+  uint32_t BBOffset{0};
+  for (uint32_t InOffset : InputOffsets) {
+    if (InOffset & BRANCHENTRY)
+      Branches[BBOffset].push_back(InOffset >> 1);
+    else
+      BBOffset = InOffset >> 1;
+  }
+  return Branches;
+}
+
 } // namespace bolt
 } // namespace llvm

bolt/lib/Profile/DataAggregator.cpp

@@ -604,6 +604,8 @@ Error DataAggregator::readProfile(BinaryContext &BC) {
     // BAT YAML is handled by DataAggregator since normal YAML output requires
     // CFG which is not available in BAT mode.
     if (usesBAT()) {
+      // Postprocess split function profile for BAT
+      fixupBATProfile(BC);
       if (opts::ProfileFormat == opts::ProfileFormatKind::PF_YAML)
         if (std::error_code EC = writeBATYAML(BC, opts::OutputFilename))
           report_error("cannot create output data file", EC);
@@ -2271,6 +2273,29 @@ DataAggregator::writeAggregatedFile(StringRef OutputFilename) const {
   return std::error_code();
 }
 
+void DataAggregator::fixupBATProfile(BinaryContext &BC) {
+  for (auto &[FuncName, Branches] : NamesToBranches) {
+    BinaryData *BD = BC.getBinaryDataByName(FuncName);
+    assert(BD);
+    uint64_t FuncAddress = BD->getAddress();
+    if (!BAT->isBATFunction(FuncAddress))
+      continue;
+    // Filter out cold fragments
+    if (!BD->getSectionName().equals(BC.getMainCodeSectionName()))
+      continue;
+    // Convert inter-branches between hot and cold fragments into
+    // intra-branches.
+    for (auto &[OffsetFrom, CallToMap] : Branches.InterIndex) {
+      for (auto &[CallToLoc, CallToIdx] : CallToMap) {
+        if (CallToLoc.Name != FuncName)
+          continue;
+        Branches.IntraIndex[OffsetFrom][CallToLoc.Offset] = CallToIdx;
+        Branches.InterIndex[OffsetFrom].erase(CallToLoc);
+      }
+    }
+  }
+}
+
 std::error_code DataAggregator::writeBATYAML(BinaryContext &BC,
                                              StringRef OutputFilename) const {
   std::error_code EC;
@@ -2343,6 +2368,63 @@ std::error_code DataAggregator::writeBATYAML(BinaryContext &BC,
         YamlBB.Successors.emplace_back(SI);
       };
 
+      std::unordered_map<uint32_t, std::vector<uint32_t>> BFBranches =
+          BAT->getBFBranches(FuncAddress);
+
+      auto addCallsProfile = [&](yaml::bolt::BinaryBasicBlockProfile &YamlBB,
+                                 uint64_t Offset) {
+        // Iterate over BRANCHENTRY records in the current block
+        for (uint32_t BranchOffset : BFBranches[Offset]) {
+          if (!Branches.InterIndex.contains(BranchOffset))
+            continue;
+          for (const auto &[CallToLoc, CallToIdx] :
+               Branches.InterIndex.at(BranchOffset)) {
+            const llvm::bolt::BranchInfo &BI = Branches.Data.at(CallToIdx);
+            yaml::bolt::CallSiteInfo YamlCSI;
+            YamlCSI.DestId = 0; // designated for unknown functions
+            YamlCSI.EntryDiscriminator = 0;
+            YamlCSI.Count = BI.Branches;
+            YamlCSI.Mispreds = BI.Mispreds;
+            YamlCSI.Offset = BranchOffset - Offset;
+            BinaryData *CallTargetBD = BC.getBinaryDataByName(CallToLoc.Name);
+            if (!CallTargetBD) {
+              YamlBB.CallSites.emplace_back(YamlCSI);
+              continue;
+            }
+            uint64_t CallTargetAddress = CallTargetBD->getAddress();
+            BinaryFunction *CallTargetBF =
+                BC.getBinaryFunctionAtAddress(CallTargetAddress);
+            if (!CallTargetBF) {
+              YamlBB.CallSites.emplace_back(YamlCSI);
+              continue;
+            }
+            // Calls between hot and cold fragments must be handled in
+            // fixupBATProfile.
+            assert(CallTargetBF != BF && "invalid CallTargetBF");
+            YamlCSI.DestId = CallTargetBF->getFunctionNumber();
+            if (CallToLoc.Offset) {
+              if (BAT->isBATFunction(CallTargetAddress)) {
+                LLVM_DEBUG(dbgs() << "BOLT-DEBUG: Unsupported secondary "
+                                     "entry point in BAT function "
+                                  << CallToLoc.Name << '\n');
+              } else if (const BinaryBasicBlock *CallTargetBB =
+                             CallTargetBF->getBasicBlockAtOffset(
+                                 CallToLoc.Offset)) {
+                // Only record true call information, ignoring returns (normally
+                // won't have a target basic block) and jumps to the landing
+                // pads (not an entry point).
+                if (CallTargetBB->isEntryPoint()) {
+                  YamlCSI.EntryDiscriminator =
+                      CallTargetBF->getEntryIDForSymbol(
+                          CallTargetBB->getLabel());
+                }
+              }
+            }
+            YamlBB.CallSites.emplace_back(YamlCSI);
+          }
+        }
+      };
+
       for (const auto &[FromOffset, SuccKV] : Branches.IntraIndex) {
         yaml::bolt::BinaryBasicBlockProfile YamlBB;
         if (!BlockMap.isInputBlock(FromOffset))
@@ -2351,7 +2433,9 @@ std::error_code DataAggregator::writeBATYAML(BinaryContext &BC,
         YamlBB.Hash = BlockMap.getBBHash(FromOffset);
         for (const auto &[SuccOffset, SuccDataIdx] : SuccKV)
           addSuccProfile(YamlBB, SuccOffset, SuccDataIdx);
-        if (YamlBB.ExecCount || !YamlBB.Successors.empty())
+        addCallsProfile(YamlBB, FromOffset);
+        if (YamlBB.ExecCount || !YamlBB.Successors.empty() ||
+            !YamlBB.CallSites.empty())
           YamlBF.Blocks.emplace_back(YamlBB);
       }
       BP.Functions.emplace_back(YamlBF);

bolt/test/X86/bolt-address-translation-yaml.test

@@ -36,6 +36,18 @@ YAML-BAT-CHECK-NEXT:   - bid:   0
 YAML-BAT-CHECK-NEXT:     insns: 26
 YAML-BAT-CHECK-NEXT:     hash:  0xA900AE79CFD40000
 YAML-BAT-CHECK-NEXT:     succ:  [ { bid: 3, cnt: 0 }, { bid: 1, cnt: 0 } ]
+# Function covered by BAT with calls
+YAML-BAT-CHECK:      - name:    SolveCubic
+YAML-BAT-CHECK-NEXT:   fid:     [[#]]
+YAML-BAT-CHECK-NEXT:   hash:    0x6AF7E61EA3966722
+YAML-BAT-CHECK-NEXT:   exec:    25
+YAML-BAT-CHECK-NEXT:   nblocks: 15
+YAML-BAT-CHECK-NEXT:   blocks:
+YAML-BAT-CHECK:        - bid:   3
+YAML-BAT-CHECK-NEXT:     insns: [[#]]
+YAML-BAT-CHECK-NEXT:     hash:  0xDDA1DC5F69F900AC
+YAML-BAT-CHECK-NEXT:     calls: [ { off: 0x26, fid: [[#]], cnt: [[#]] } ]
+YAML-BAT-CHECK-NEXT:     succ:  [ { bid: 5, cnt: [[#]] }
 # Function covered by BAT - doesn't have insns in basic block
 YAML-BAT-CHECK:      - name:    usqrt
 YAML-BAT-CHECK-NEXT:   fid:     [[#]]