[BOLT] Ignore returns in DataAggregator

[aaupov]

Returns are ignored in perf/pre-aggregated/fdata profile reader (see
DataReader::convertBranchData). They are also omitted in
YAMLProfileWriter by virtue of not having the profile attached to them
in the reader, and YAMLProfileWriter converting the profile attached to
BinaryFunctions. Thus, return profile is universally ignored across all
profile types except BAT YAML.

To make returns ignored for YAML produced in BAT mode, we can:

1. ignore them in YAMLProfileReader,
2. omit them from YAML profile in profile conversion/writing.

The first option is prone to profile staleness issue, where the profiled
binary doesn't match the one to be optimized, and thus returns in the
profile can no longer be reliably detected (as we don't distinguish them
from calls in the profile).

The second option is robust to staleness but requires disassembling the
branch source instruction.

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

[llvmbot]

@llvm/pr-subscribers-bolt

Author: Amir Ayupov (aaupov)

Changes

Returns are ignored in perf/pre-aggregated/fdata profile reader (see
DataReader::convertBranchData). They are also omitted in
YAMLProfileWriter by virtue of not having the profile attached to them
in the reader, and YAMLProfileWriter converting the profile attached to
BinaryFunctions. Thus, return profile is universally ignored across all
profile types except BAT YAML.

To make returns ignored for YAML produced in BAT mode, we can:

1. ignore them in YAMLProfileReader,
2. omit them from YAML profile in profile conversion/writing.

The first option is prone to profile staleness issue, where the profiled
binary doesn't match the one to be optimized, and thus returns in the
profile can no longer be reliably detected (as we don't distinguish them
from calls in the profile).

The second option is more robust but requires disassembling the
functions. We already do that in BAT YAML profile generation, primarily
for non-BAT functions.

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

---

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

2 Files Affected:

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

diff --git a/bolt/lib/Profile/DataAggregator.cpp b/bolt/lib/Profile/DataAggregator.cpp
index 5108392c824c10..0fd6579517e8f6 100644
--- a/bolt/lib/Profile/DataAggregator.cpp
+++ b/bolt/lib/Profile/DataAggregator.cpp
@@ -773,9 +773,13 @@ bool DataAggregator::doInterBranch(BinaryFunction *FromFunc,
 
 bool DataAggregator::doBranch(uint64_t From, uint64_t To, uint64_t Count,
                               uint64_t Mispreds) {
+  bool IsReturn = false;
   auto handleAddress = [&](uint64_t &Addr, bool IsFrom) -> BinaryFunction * {
     if (BinaryFunction *Func = getBinaryFunctionContainingAddress(Addr)) {
       Addr -= Func->getAddress();
+      if (Func->hasInstructions())
+        if (const MCInst *Inst = Func->getInstructionAtOffset(Addr))
+          IsReturn = BC->MIB->isReturn(*Inst);
 
       if (BAT)
         Addr = BAT->translate(Func->getAddress(), Addr, IsFrom);
@@ -792,6 +796,9 @@ bool DataAggregator::doBranch(uint64_t From, uint64_t To, uint64_t Count,
   };
 
   BinaryFunction *FromFunc = handleAddress(From, /*IsFrom=*/true);
+  // Ignore returns.
+  if (IsReturn)
+    return true;
   BinaryFunction *ToFunc = handleAddress(To, /*IsFrom=*/false);
   if (!FromFunc && !ToFunc)
     return false;
diff --git a/bolt/test/X86/bolt-address-translation-yaml.test b/bolt/test/X86/bolt-address-translation-yaml.test
index af24c3d84a0f15..b3d8a88394503c 100644
--- a/bolt/test/X86/bolt-address-translation-yaml.test
+++ b/bolt/test/X86/bolt-address-translation-yaml.test
@@ -13,7 +13,7 @@ RUN: llvm-bolt %t.exe -data %t.fdata -w %t.yaml-fdata -o /dev/null
 RUN: FileCheck --input-file %t.yaml-fdata --check-prefix YAML-BAT-CHECK %s
 
 # Test resulting YAML profile with the original binary (no-stale mode)
-RUN: llvm-bolt %t.exe -data %t.yaml -o %t.null -dyno-stats \
+RUN: llvm-bolt %t.exe -data %t.yaml -o %t.null -dyno-stats 2>&1 \
 RUN:   | FileCheck --check-prefix CHECK-BOLT-YAML %s
 
 WRITE-BAT-CHECK: BOLT-INFO: Wrote 5 BAT maps
@@ -63,7 +63,8 @@ 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: [[#]] }
+YAML-BAT-CHECK-NEXT:       succ:  {{.*}} { bid: 2, cnt: [[#]]
 
 CHECK-BOLT-YAML:      pre-processing profile using YAML profile reader
 CHECK-BOLT-YAML-NEXT: 5 out of 16 functions in the binary (31.2%) have non-empty execution profile
+CHECK-BOLT-YAML-NOT: invalid (possibly stale) profile

[maksfb]

Please address the nits. Otherwise - good to go.

[maksfb]

nit: add message.

[maksfb]

nit: capitalize the message.

[maksfb]

nit: ditto.