[llvm-branch-commits] [llvm] [BOLT] Impute missing trace fall-through (PR #145258)

[Amir Ayupov via llvm-branch-commits]

https://github.com/aaupov updated https://github.com/llvm/llvm-project/pull/145258

>From 9aef8e0a499fa4b9e6bbde910a678a65a0ab7f92 Mon Sep 17 00:00:00 2001
From: Amir Ayupov <aaupov at fb.com>
Date: Mon, 23 Jun 2025 12:54:06 -0700
Subject: [PATCH] unified checkReturn and checkUncondJump, logging

Created using spr 1.3.4
---
 bolt/include/bolt/Profile/DataAggregator.h | 34 ++++++++++++++-
 bolt/lib/Profile/DataAggregator.cpp        | 49 ++++++----------------
 2 files changed, 46 insertions(+), 37 deletions(-)

diff --git a/bolt/include/bolt/Profile/DataAggregator.h b/bolt/include/bolt/Profile/DataAggregator.h
index 00c6f56520fdc..364dbdbed841c 100644
--- a/bolt/include/bolt/Profile/DataAggregator.h
+++ b/bolt/include/bolt/Profile/DataAggregator.h
@@ -137,7 +137,7 @@ class DataAggregator : public DataReader {
   std::vector<std::pair<Trace, TakenBranchInfo>> Traces;
   /// Pre-populated addresses of returns, coming from pre-aggregated data or
   /// disassembly. Used to disambiguate call-continuation fall-throughs.
-  std::unordered_set<uint64_t> Returns;
+  std::unordered_map<uint64_t, bool> Returns;
   std::unordered_map<uint64_t, uint64_t> BasicSamples;
   std::vector<PerfMemSample> MemSamples;
 
@@ -514,6 +514,38 @@ class DataAggregator : public DataReader {
   void printBasicSamplesDiagnostics(uint64_t OutOfRangeSamples) const;
   void printBranchStacksDiagnostics(uint64_t IgnoredSamples) const;
 
+  template <typename T>
+  std::optional<T>
+  testInstructionAt(const uint64_t Addr,
+                    std::function<T(const MCInst &)> Callback) const {
+    BinaryFunction *Func = getBinaryFunctionContainingAddress(Addr);
+    if (!Func)
+      return std::nullopt;
+    const uint64_t Offset = Addr - Func->getAddress();
+    if (Func->hasInstructions()) {
+      if (auto *MI = Func->getInstructionAtOffset(Offset))
+        return Callback(*MI);
+    } else {
+      if (auto MI = Func->disassembleInstructionAtOffset(Offset))
+        return Callback(*MI);
+    }
+    return std::nullopt;
+  }
+
+  template <typename T>
+  std::optional<T> testAndSet(const uint64_t Addr,
+                              std::function<T(const MCInst &)> Callback,
+                              std::unordered_map<uint64_t, T> &Map) {
+    auto It = Map.find(Addr);
+    if (It != Map.end())
+      return It->second;
+    if (std::optional<T> Res = testInstructionAt<T>(Addr, Callback)) {
+      Map.emplace(Addr, *Res);
+      return *Res;
+    }
+    return std::nullopt;
+  }
+
 public:
   /// If perf.data was collected without build ids, the buildid-list may contain
   /// incomplete entries. Return true if the buffer containing
diff --git a/bolt/lib/Profile/DataAggregator.cpp b/bolt/lib/Profile/DataAggregator.cpp
index 2fcc2561cc212..1323c4f30049d 100644
--- a/bolt/lib/Profile/DataAggregator.cpp
+++ b/bolt/lib/Profile/DataAggregator.cpp
@@ -546,23 +546,10 @@ void DataAggregator::imputeFallThroughs() {
   // Helper map with whether the instruction is a call/ret/unconditional branch
   std::unordered_map<uint64_t, bool> IsUncondJumpMap;
   auto checkUncondJump = [&](const uint64_t Addr) {
-    auto isUncondJump = [&](auto MI) {
-      return MI && BC->MIB->IsUnconditionalJump(*MI);
+    auto isUncondJump = [&](const MCInst &MI) -> bool {
+      return BC->MIB->IsUnconditionalJump(MI);
     };
-    auto It = IsUncondJumpMap.find(Addr);
-    if (It != IsUncondJumpMap.end())
-      return It->second;
-    BinaryFunction *Func = getBinaryFunctionContainingAddress(Addr);
-    if (!Func)
-      return false;
-    const uint64_t Offset = Addr - Func->getAddress();
-    if (Func->hasInstructions()
-            ? isUncondJump(Func->getInstructionAtOffset(Offset))
-            : isUncondJump(Func->disassembleInstructionAtOffset(Offset))) {
-      IsUncondJumpMap.emplace(Addr, true);
-      return true;
-    }
-    return false;
+    return testAndSet<bool>(Addr, isUncondJump, IsUncondJumpMap).value_or(true);
   };
 
   for (auto &[Trace, Info] : Traces) {
@@ -574,7 +561,8 @@ void DataAggregator::imputeFallThroughs() {
                                    ? AggregateFallthroughSize / AggregateCount
                                    : !checkUncondJump(Trace.From);
       Trace.To = Trace.From + InferredBytes;
-      outs() << "Inferred " << Trace << " " << InferredBytes << '\n';
+      LLVM_DEBUG(dbgs() << "imputed " << Trace << " (" << InferredBytes
+                        << " bytes)\n");
       ++InferredTraces;
     } else {
       if (CurrentBranch != PrevBranch)
@@ -585,7 +573,8 @@ void DataAggregator::imputeFallThroughs() {
     }
     PrevBranch = CurrentBranch;
   }
-  outs() << "Inferred " << InferredTraces << " traces\n";
+  if (opts::Verbosity >= 1)
+    outs() << "BOLT-INFO: imputed " << InferredTraces << " traces\n";
 }
 
 Error DataAggregator::preprocessProfile(BinaryContext &BC) {
@@ -804,22 +793,10 @@ bool DataAggregator::doInterBranch(BinaryFunction *FromFunc,
 }
 
 bool DataAggregator::checkReturn(uint64_t Addr) {
-  auto isReturn = [&](auto MI) { return MI && BC->MIB->isReturn(*MI); };
-  if (llvm::is_contained(Returns, Addr))
-    return true;
-
-  BinaryFunction *Func = getBinaryFunctionContainingAddress(Addr);
-  if (!Func)
-    return false;
-
-  const uint64_t Offset = Addr - Func->getAddress();
-  if (Func->hasInstructions()
-          ? isReturn(Func->getInstructionAtOffset(Offset))
-          : isReturn(Func->disassembleInstructionAtOffset(Offset))) {
-    Returns.emplace(Addr);
-    return true;
-  }
-  return false;
+  auto isReturn = [&](const MCInst &MI) -> bool {
+    return BC->MIB->isReturn(MI);
+  };
+  return testAndSet<bool>(Addr, isReturn, Returns).value_or(false);
 }
 
 bool DataAggregator::doBranch(uint64_t From, uint64_t To, uint64_t Count,
@@ -1409,7 +1386,7 @@ std::error_code DataAggregator::parseAggregatedLBREntry() {
     if (!Addr[0]->Offset)
       Addr[0]->Offset = Trace::FT_EXTERNAL_RETURN;
     else
-      Returns.emplace(Addr[0]->Offset);
+      Returns.emplace(Addr[0]->Offset, true);
   }
 
   /// Record a trace.
@@ -1670,7 +1647,7 @@ void DataAggregator::processBranchEvents() {
   NamedRegionTimer T("processBranch", "Processing branch events",
                      TimerGroupName, TimerGroupDesc, opts::TimeAggregator);
 
-  Returns.emplace(Trace::FT_EXTERNAL_RETURN);
+  Returns.emplace(Trace::FT_EXTERNAL_RETURN, true);
   for (const auto &[Trace, Info] : Traces) {
     bool IsReturn = checkReturn(Trace.Branch);
     // Ignore returns.