[BOLT] Set entry counts in BAT YAML profile (#91775)

aaupov · web-flow · commit 4f127667ca57 · 2024-05-10T22:23:45.000-07:00
Align with DataReader::readProfile that sets entry block counts from
FuncBranchData-&gt;EntryData.

Test Plan: updated bolt-address-translation-yaml.test
diff --git a/bolt/lib/Profile/DataAggregator.cpp b/bolt/lib/Profile/DataAggregator.cpp
@@ -2403,6 +2403,17 @@ std::error_code DataAggregator::writeBATYAML(BinaryContext &BC,
           return A.Offset < B.Offset;
         });
       }
+      // 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
@@ -48,6 +48,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
