[llvm-branch-commits] [BOLT] Map branch source address to the containing basic block in BAT YAML (PR #91273)

[via llvm-branch-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-bolt

Author: Amir Ayupov (aaupov)

<details>
<summary>Changes</summary>

Fix an issue where the profile for all branches that have a BRANCHENTRY
is dropped. If the branch has an entry in BAT, it will be translated to
its input offset. We used to only permit the basic block offset as a
branch source. Perform a lookup of containing basic block instead.

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


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


2 Files Affected:

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


``````````diff
diff --git a/bolt/lib/Profile/DataAggregator.cpp b/bolt/lib/Profile/DataAggregator.cpp
index e4a7324c38175c..656b025a5d79d7 100644
--- a/bolt/lib/Profile/DataAggregator.cpp
+++ b/bolt/lib/Profile/DataAggregator.cpp
@@ -2378,10 +2378,16 @@ std::error_code DataAggregator::writeBATYAML(BinaryContext &BC,
         return CSI;
       };
 
+      // Lookup containing basic block offset and index
+      auto getBlock = [&BlockMap](uint32_t Offset) {
+        auto BlockIt = BlockMap.upper_bound(Offset);
+        assert(BlockIt != BlockMap.begin());
+        --BlockIt;
+        return std::pair(BlockIt->first, BlockIt->second.getBBIndex());
+      };
+
       for (const auto &[FromOffset, SuccKV] : Branches.IntraIndex) {
-        if (!BlockMap.isInputBlock(FromOffset))
-          continue;
-        const unsigned Index = BlockMap.getBBIndex(FromOffset);
+        const auto &[_, Index] = getBlock(FromOffset);
         yaml::bolt::BinaryBasicBlockProfile &YamlBB = YamlBF.Blocks[Index];
         for (const auto &[SuccOffset, SuccDataIdx] : SuccKV)
           if (BlockMap.isInputBlock(SuccOffset))
@@ -2389,10 +2395,7 @@ std::error_code DataAggregator::writeBATYAML(BinaryContext &BC,
                 getSuccessorInfo(SuccOffset, SuccDataIdx));
       }
       for (const auto &[FromOffset, CallTo] : Branches.InterIndex) {
-        auto BlockIt = BlockMap.upper_bound(FromOffset);
-        --BlockIt;
-        const unsigned BlockOffset = BlockIt->first;
-        const unsigned BlockIndex = BlockIt->second.getBBIndex();
+        const auto &[BlockOffset, BlockIndex] = getBlock(FromOffset);
         yaml::bolt::BinaryBasicBlockProfile &YamlBB = YamlBF.Blocks[BlockIndex];
         const uint32_t Offset = FromOffset - BlockOffset;
         for (const auto &[CallToLoc, CallToIdx] : CallTo)
diff --git a/bolt/test/X86/bolt-address-translation-yaml.test b/bolt/test/X86/bolt-address-translation-yaml.test
index b3d8a88394503c..f67cc6361c9ef8 100644
--- a/bolt/test/X86/bolt-address-translation-yaml.test
+++ b/bolt/test/X86/bolt-address-translation-yaml.test
@@ -5,6 +5,17 @@ RUN: llvm-bolt %t.exe -o %t.out --pa -p %p/Inputs/blarge_new.preagg.txt \
 RUN:   --reorder-blocks=ext-tsp --split-functions --split-strategy=cdsplit \
 RUN:   --reorder-functions=cdsort --enable-bat --dyno-stats --skip-funcs=main \
 RUN:   2>&1 | FileCheck --check-prefix WRITE-BAT-CHECK %s
+# Check that branch with entry in BAT is accounted for.
+RUN: perf2bolt %t.out --pa -p %p/Inputs/blarge_new_bat_branchentry.preagg.txt \
+RUN:   -w %t.yaml -o %t.fdata
+RUN: llvm-bolt %t.exe -data %t.fdata -w %t.yaml-fdata -o %t.null
+RUN: FileCheck --input-file %t.yaml --check-prefix BRANCHENTRY-YAML-CHECK %s
+RUN: FileCheck --input-file %t.yaml-fdata --check-prefix BRANCHENTRY-YAML-CHECK %s
+BRANCHENTRY-YAML-CHECK:    - name: SolveCubic
+BRANCHENTRY-YAML-CHECK:      bid: 0
+BRANCHENTRY-YAML-CHECK:      hash: 0x700F19D24600000
+BRANCHENTRY-YAML-CHECK-NEXT: succ: [ { bid: 7, cnt: 1 }
+# Large profile test
 RUN: perf2bolt %t.out --pa -p %p/Inputs/blarge_new_bat.preagg.txt -w %t.yaml -o %t.fdata \
 RUN:   2>&1 | FileCheck --check-prefix READ-BAT-CHECK %s
 RUN: FileCheck --input-file %t.yaml --check-prefix YAML-BAT-CHECK %s

``````````

</details>


https://github.com/llvm/llvm-project/pull/91273