[BOLT] Set entry counts in BAT YAML profile

[aaupov]

Align with DataReader::readProfile that sets entry block counts from
FuncBranchData->EntryData.

Test Plan: updated bolt-address-translation-yaml.test

[llvmbot]

@llvm/pr-subscribers-bolt

Author: Amir Ayupov (aaupov)

Changes

Align with DataReader::readProfile that sets entry block counts from
FuncBranchData->EntryData.

Test Plan: updated bolt-address-translation-yaml.test

---

Full diff: https://github.com/llvm/llvm-project/pull/91775.diff

2 Files Affected:

• (modified) bolt/lib/Profile/DataAggregator.cpp (+11)
• (modified) bolt/test/X86/bolt-address-translation-yaml.test (+4)

diff --git a/bolt/lib/Profile/DataAggregator.cpp b/bolt/lib/Profile/DataAggregator.cpp
index afbab6ffceece..e5fdf487b6d93 100644
--- a/bolt/lib/Profile/DataAggregator.cpp
+++ b/bolt/lib/Profile/DataAggregator.cpp
@@ -2385,6 +2385,17 @@ std::error_code DataAggregator::writeBATYAML(BinaryContext &BC,
       addIndex(Branches.IntraIndex);
       addIndex(Branches.InterIndex);
 
+      // Set entry counts, similar to DataReader::readProfile.
+      for (const llvm::bolt::BranchInfo &BI : Branches.EntryData) {
+        if (!BlockMap.isInputBlock(BI.To.Offset)) {
+          if (opts::Verbosity >= 1)
+            errs() << "BOLT-WARNING: Unexpected EntryData in " << FuncName
+                   << " at 0x" << Twine::utohexstr(BI.To.Offset) << '\n';
+          continue;
+        }
+        const unsigned BlockIndex = BlockMap.getBBIndex(BI.To.Offset);
+        YamlBF.Blocks[BlockIndex].ExecCount += BI.Branches;
+      }
       // Drop blocks without a hash, won't be useful for stale matching.
       llvm::erase_if(YamlBF.Blocks,
                      [](const yaml::bolt::BinaryBasicBlockProfile &YamlBB) {
diff --git a/bolt/test/X86/bolt-address-translation-yaml.test b/bolt/test/X86/bolt-address-translation-yaml.test
index c7927f92c9dd9..e21513b7dfe59 100644
--- a/bolt/test/X86/bolt-address-translation-yaml.test
+++ b/bolt/test/X86/bolt-address-translation-yaml.test
@@ -70,6 +70,10 @@ 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-NEXT:   - bid:   0
+YAML-BAT-CHECK-NEXT:     insns: [[#]]
+YAML-BAT-CHECK-NEXT:     hash:  0x700F19D24600000
+YAML-BAT-CHECK-NEXT:     exec:  25
 YAML-BAT-CHECK:        - bid:   3
 YAML-BAT-CHECK-NEXT:     insns: [[#]]
 YAML-BAT-CHECK-NEXT:     hash:  0xDDA1DC5F69F900AC