[BOLT] Align BranchInfo and FuncBranchData in DataAggregator::recordTrace

[Amir]

Summary

DataAggregator::recordTrace serves two purposes:

• Attaching LBR fallthrough ("trace") information to CFG (getBranchInfo), which eventually gets emitted as YAML profile.
• Populating vector of offsets that gets added to FuncBranchData, which eventually gets emitted as fdata profile.

recordTrace is invoked from getFallthroughsInTrace which checks its return
status and passes on the collected vector of offsets to doTrace.

However, if a malformed trace is passed to recordTrace it might partially
attach the profile to CFG and exit with false, not propagating the vector of
offsets to doTrace. This leads to a difference between fdata and yaml profile
collected from the same binary and the same perf file.

(Skylake LBR errata might produce such malformed traces where the last entry
is duplicated, resulting in invalid fallthrough path between the last two
entries).

There are two ways to handle this mismatch: conservative (aligned with fdata),
or aggressive (aligned with yaml). Conservative approach would discard the
trace entirely, buffering the CFG updates until all fallthroughs are confirmed.
Aggressive approach would apply CFG updates and return the matching
fallthroughs in the vector even if the trace is invalid (doesn't correspond to
a valid fallthrough path). I chose to go with the former (conservative/fdata)
approach which produces more accurate profile.

We can't rely on pre-filtering such traces early (in LBR sample processing) as
DataAggregator is used for both perf samples and pre-aggregated perf information
which loses branch stack information.

Test Plan: https://github.com/rafaelauler/bolt-tests/pull/22

[rafauler]

LGTM with nits

bolt/lib/Profile/DataAggregator.cpp

902–903

Remove this check and assert at function start that Branches is never nullpr? Since now you're relying on branches being non-null..

[Amir]

Pass Branches by ref, clang-format

bolt/include/bolt/Profile/DataAggregator.h

private:
 
   BoltAddressTranslation *BAT{nullptr};
 
   /// Update function execution profile with a recorded trace.
   /// A trace is region of code executed between two LBR entries supplied in
   /// execution order.
   ///
   /// Return true if the trace is valid, false otherwise.
-  bool recordTrace(
-      BinaryFunction &BF, const LBREntry &First, const LBREntry &Second,
-      uint64_t Count = 1,
-      SmallVector<std::pair<uint64_t, uint64_t>, 16> *Branches = nullptr) const;
+  bool
+  recordTrace(BinaryFunction &BF, const LBREntry &First, const LBREntry &Second,
+              uint64_t Count,
+              SmallVector<std::pair<uint64_t, uint64_t>, 16> &Branches) const;
 
   /// Return a vector of offsets corresponding to a trace in a function
   /// (see recordTrace() above).
   std::optional<SmallVector<std::pair<uint64_t, uint64_t>, 16>>
   getFallthroughsInTrace(BinaryFunction &BF, const LBREntry &First,
                          const LBREntry &Second, uint64_t Count = 1) const;
 
   /// Record external entry into the function \p BF.

bolt/lib/Profile/DataAggregator.cpp

bool DataAggregator::doTrace(const LBREntry &First, const LBREntry &Second,
   for (const std::pair<uint64_t, uint64_t> &Pair : *FTs)
     doIntraBranch(*FromFunc, Pair.first + FromFunc->getAddress(),
                   Pair.second + FromFunc->getAddress(), Count, false);
 
   return true;
 }
 
 bool DataAggregator::recordTrace(
-    BinaryFunction &BF,
-    const LBREntry &FirstLBR,
-    const LBREntry &SecondLBR,
+    BinaryFunction &BF, const LBREntry &FirstLBR, const LBREntry &SecondLBR,
     uint64_t Count,
-    SmallVector<std::pair<uint64_t, uint64_t>, 16> *Branches) const {
+    SmallVector<std::pair<uint64_t, uint64_t>, 16> &Branches) const {
   BinaryContext &BC = BF.getBinaryContext();
 
   if (!BF.isSimple())
     return false;
 
   assert(BF.hasCFG() && "can only record traces in CFG state");
 
   // Offsets of the trace within this function.
while (BB != ToBB) {
 
     // Check for bad LBRs.
     if (!BB->getSuccessor(NextBB->getLabel())) {
       LLVM_DEBUG(dbgs() << "no fall-through for the trace:\n"
                         << "  " << FirstLBR << '\n'
                         << "  " << SecondLBR << '\n');
       return false;
     }
 
-    // Record fall-through jumps
-    BinaryBasicBlock::BinaryBranchInfo &BI = BB->getBranchInfo(*NextBB);
-    BI.Count += Count;
-
-    if (Branches) {
-      const MCInst *Instr = BB->getLastNonPseudoInstr();
-      uint64_t Offset = 0;
-      if (Instr)
-        Offset = BC.MIB->getOffsetWithDefault(*Instr, 0);
-      else
-        Offset = BB->getOffset();
 
-      Branches->emplace_back(Offset, NextBB->getOffset());
-    }
+    Branches.emplace_back(Offset, NextBB->getOffset());
 
     BB = NextBB;
   }
 
+  // Record fall-through jumps
+  for (const auto &[FromOffset, ToOffset] : Branches) {
+    BinaryBasicBlock *FromBB = BF.getBasicBlockContainingOffset(FromOffset);
+    BinaryBasicBlock *ToBB = BF.getBasicBlockAtOffset(ToOffset);
+    assert(FromBB && ToBB);
+    BinaryBasicBlock::BinaryBranchInfo &BI = FromBB->getBranchInfo(*ToBB);
+    BI.Count += Count;
+  }
+
   return true;
 }
 
 std::optional<SmallVector<std::pair<uint64_t, uint64_t>, 16>>
 DataAggregator::getFallthroughsInTrace(BinaryFunction &BF,
                                        const LBREntry &FirstLBR,
                                        const LBREntry &SecondLBR,
                                        uint64_t Count) const {
   SmallVector<std::pair<uint64_t, uint64_t>, 16> Res;
 
-  if (!recordTrace(BF, FirstLBR, SecondLBR, Count, &Res))
+  if (!recordTrace(BF, FirstLBR, SecondLBR, Count, Res))
     return std::nullopt;
 
   return Res;
 }
 
 bool DataAggregator::recordEntry(BinaryFunction &BF, uint64_t To, bool Mispred,
                                  uint64_t Count) const {
   if (To > BF.getSize())