[BOLT] Emit intra-function control flow in YAMLBAT

Attach branch counters to YAML profile, covering intra-function control
flow.

Depends on: #86353

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

Reviewers: rafaelauler, dcci, ayermolo, maksfb

Reviewed By: rafaelauler

Pull Request: #76911

Author: aaupov

bolt/include/bolt/Profile/BoltAddressTranslation.h

@@ -257,6 +257,11 @@ class BoltAddressTranslation {
         std::as_const(*this).getBBHashMap(FuncOutputAddress));
   }
 
+  /// Returns the number of basic blocks in a function.
+  size_t getNumBasicBlocks(uint64_t OutputAddress) const {
+    return NumBasicBlocksMap.at(OutputAddress);
+  }
+
 private:
   FuncHashesTy FuncHashes;
 };

bolt/lib/Profile/DataAggregator.cpp

@@ -2310,6 +2310,52 @@ std::error_code DataAggregator::writeBATYAML(BinaryContext &BC,
       BP.Functions.emplace_back(
           YAMLProfileWriter::convert(Function, /*UseDFS=*/false));
     }
+
+    for (const auto &KV : NamesToBranches) {
+      const StringRef FuncName = KV.first;
+      const FuncBranchData &Branches = KV.second;
+      yaml::bolt::BinaryFunctionProfile YamlBF;
+      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;
+      BinaryFunction *BF = BC.getBinaryFunctionAtAddress(FuncAddress);
+      assert(BF);
+      YamlBF.Name = FuncName.str();
+      YamlBF.Id = BF->getFunctionNumber();
+      YamlBF.Hash = BAT->getBFHash(FuncAddress);
+      YamlBF.ExecCount = BF->getKnownExecutionCount();
+      YamlBF.NumBasicBlocks = BAT->getNumBasicBlocks(FuncAddress);
+      const BoltAddressTranslation::BBHashMapTy &BlockMap =
+          BAT->getBBHashMap(FuncAddress);
+
+      auto addSuccProfile = [&](yaml::bolt::BinaryBasicBlockProfile &YamlBB,
+                                uint64_t SuccOffset, unsigned SuccDataIdx) {
+        const llvm::bolt::BranchInfo &BI = Branches.Data.at(SuccDataIdx);
+        yaml::bolt::SuccessorInfo SI;
+        SI.Index = BlockMap.getBBIndex(SuccOffset);
+        SI.Count = BI.Branches;
+        SI.Mispreds = BI.Mispreds;
+        YamlBB.Successors.emplace_back(SI);
+      };
+
+      for (const auto &[FromOffset, SuccKV] : Branches.IntraIndex) {
+        yaml::bolt::BinaryBasicBlockProfile YamlBB;
+        if (!BlockMap.isInputBlock(FromOffset))
+          continue;
+        YamlBB.Index = BlockMap.getBBIndex(FromOffset);
+        YamlBB.Hash = BlockMap.getBBHash(FromOffset);
+        for (const auto &[SuccOffset, SuccDataIdx] : SuccKV)
+          addSuccProfile(YamlBB, SuccOffset, SuccDataIdx);
+        if (YamlBB.ExecCount || !YamlBB.Successors.empty())
+          YamlBF.Blocks.emplace_back(YamlBB);
+      }
+      BP.Functions.emplace_back(YamlBF);
+    }
   }
 
   // Write the profile.

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

@@ -25,6 +25,7 @@ READ-BAT-CHECK: BOLT-INFO: Parsed 5 BAT entries
 READ-BAT-CHECK: PERF2BOLT: read 79 aggregated LBR entries
 
 YAML-BAT-CHECK:      functions:
+# Function not covered by BAT - has insns in basic block
 YAML-BAT-CHECK:      - name:    main
 YAML-BAT-CHECK-NEXT:   fid:     2
 YAML-BAT-CHECK-NEXT:   hash:    0x9895746D48B2C876
@@ -35,6 +36,17 @@ 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 - doesn't have insns in basic block
+YAML-BAT-CHECK:      - name:    usqrt
+YAML-BAT-CHECK-NEXT:   fid:     [[#]]
+YAML-BAT-CHECK-NEXT:   hash:    0x99E67ED32A203023
+YAML-BAT-CHECK-NEXT:   exec:    21
+YAML-BAT-CHECK-NEXT:   nblocks: 5
+YAML-BAT-CHECK-NEXT:   blocks:
+YAML-BAT-CHECK:        - bid:   1
+YAML-BAT-CHECK-NEXT:       insns: [[#]]
+YAML-BAT-CHECK-NEXT:       hash:  0xD70DC695320E0010
+YAML-BAT-CHECK-NEXT:       succ:  {{.*}} { bid: 2, cnt: [[#]] }
 
 CHECK-BOLT-YAML:      pre-processing profile using YAML profile reader
-CHECK-BOLT-YAML-NEXT: 1 out of 16 functions in the binary (6.2%) have non-empty execution profile
+CHECK-BOLT-YAML-NEXT: 5 out of 16 functions in the binary (31.2%) have non-empty execution profile