[BOLT] Fix doTrace in BAT mode

[aaupov]

When processing BOLTed binaries with BAT section, we used to
indiscriminately use BAT->getFallthroughsInTrace to record
fall-throughs, even if the function is not covered by BAT.

Fix that by using non-BAT CFG-based getFallthroughsInTrace if the
function is not in BAT.

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

[llvmbot]

@llvm/pr-subscribers-bolt

Author: Amir Ayupov (aaupov)

Changes

When processing BOLTed binaries with BAT section, we used to
indiscriminately use BAT->getFallthroughsInTrace to record
fall-throughs, even if the function is not covered by BAT.

Fix that by using non-BAT CFG-based getFallthroughsInTrace if the
function is not in BAT.

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

---

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

2 Files Affected:

• (modified) bolt/lib/Profile/DataAggregator.cpp (+3-2)
• (modified) bolt/test/X86/bolt-address-translation-yaml.test (+5)

diff --git a/bolt/lib/Profile/DataAggregator.cpp b/bolt/lib/Profile/DataAggregator.cpp
index a859f27569385..d20626bd5062f 100644
--- a/bolt/lib/Profile/DataAggregator.cpp
+++ b/bolt/lib/Profile/DataAggregator.cpp
@@ -831,9 +831,10 @@ bool DataAggregator::doTrace(const LBREntry &First, const LBREntry &Second,
     ParentFunc = FromFunc;
   ParentFunc->SampleCountInBytes += Count * (Second.From - First.To);
 
+  const uint64_t FuncAddress = FromFunc->getAddress();
   std::optional<BoltAddressTranslation::FallthroughListTy> FTs =
-      BAT ? BAT->getFallthroughsInTrace(FromFunc->getAddress(), First.To,
-                                        Second.From)
+      BAT && BAT->isBATFunction(FuncAddress)
+          ? BAT->getFallthroughsInTrace(FuncAddress, First.To, Second.From)
           : getFallthroughsInTrace(*FromFunc, First, Second, Count);
   if (!FTs) {
     LLVM_DEBUG(
diff --git a/bolt/test/X86/bolt-address-translation-yaml.test b/bolt/test/X86/bolt-address-translation-yaml.test
index 3778891c8d916..a6a212d9c1b38 100644
--- a/bolt/test/X86/bolt-address-translation-yaml.test
+++ b/bolt/test/X86/bolt-address-translation-yaml.test
@@ -61,6 +61,11 @@ 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 } ]
+# Check fallthroughs in non-BAT function
+YAML-BAT-CHECK-NEXT:   - bid:   27
+YAML-BAT-CHECK-NEXT:     insns: 3
+YAML-BAT-CHECK-NEXT:     hash:  0x30A1EBA77A903F0
+YAML-BAT-CHECK-NEXT:     succ:  [ { bid: 28, cnt: 1 } ]
 # Calls from no-BAT to BAT function
 YAML-BAT-CHECK:        - bid:   28
 YAML-BAT-CHECK-NEXT:     insns: 13

[dcci]

LG