Handle external origin LBR (non-BAT mode)

Created using spr 1.3.4

Author: aaupov

bolt/lib/Profile/DataAggregator.cpp

@@ -775,46 +775,86 @@ 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 (IsFrom) {
-        auto checkReturn = [&](auto MaybeInst) {
-          IsReturn = MaybeInst && BC->MIB->isReturn(*MaybeInst);
-        };
-        if (Func->hasInstructions())
-          checkReturn(Func->getInstructionAtOffset(Addr));
-        else
-          checkReturn(Func->disassembleInstructionAtOffset(Addr));
-      }
+  // Returns whether \p Offset in \p Func contains a return instruction.
+  auto checkReturn = [&](const BinaryFunction &Func, const uint64_t Offset) {
+    auto isReturn = [&](auto MI) { return MI && BC->MIB->isReturn(*MI); };
+    return Func.hasInstructions()
+               ? isReturn(Func.getInstructionAtOffset(Offset))
+               : isReturn(Func.disassembleInstructionAtOffset(Offset));
+  };
 
-      if (BAT)
-        Addr = BAT->translate(Func->getAddress(), Addr, IsFrom);
+  // Returns whether \p Offset in \p Func corresponds to a call continuation
+  // fallthrough block.
+  auto checkCallCont = [&](BinaryFunction &Func, const uint64_t Offset) {
+    // Note the use of MCInstrAnalysis: no call continuation for a tail call.
+    auto isCall = [&](auto MI) { return MI && BC->MIA->isCall(*MI); };
+
+    // No call continuation at a function start.
+    if (!Offset)
+      return false;
+
+    // FIXME: support BAT case where the function might be in empty state
+    // (split fragments declared non-simple).
+    if (!Func.hasCFG())
+      return false;
+
+    // The offset should not be an entry point or a landing pad.
+    const BinaryBasicBlock *ContBB = Func.getBasicBlockAtOffset(Offset);
+    if (!ContBB || ContBB->isEntryPoint() || ContBB->isLandingPad())
+      return false;
+
+    // Check that preceding instruction is a call.
+    const BinaryBasicBlock *CallBB =
+        Func.getBasicBlockContainingOffset(Offset - 1);
+    if (!CallBB || CallBB == ContBB)
+      return false;
+    return isCall(CallBB->getLastNonPseudoInstr());
+  };
 
-      if (BinaryFunction *ParentFunc = getBATParentFunction(*Func)) {
-        Func = ParentFunc;
-        if (IsFrom)
-          NumColdSamples += Count;
-      }
+  // Mutates \p Addr to an offset into the containing function, performing BAT
+  // offset translation and parent lookup.
+  //
+  // Returns the containing function (or BAT parent) and whether the address
+  // corresponds to a return (if \p IsFrom) or a call continuation (otherwise).
+  auto handleAddress = [&](uint64_t &Addr, bool IsFrom) {
+    BinaryFunction *Func = getBinaryFunctionContainingAddress(Addr);
+    if (!Func)
+      return std::pair{Func, false};
 
-      return Func;
-    }
-    return nullptr;
-  };
+    Addr -= Func->getAddress();
 
-  BinaryFunction *FromFunc = handleAddress(From, /*IsFrom=*/true);
-  // Record returns as call->call continuation fall-through.
-  if (IsReturn) {
-    LBREntry First{To - 1, To - 1, false};
-    LBREntry Second{To, To, false};
-    return doTrace(First, Second, Count);
-  }
+    bool IsRetOrCallCont =
+        IsFrom ? checkReturn(*Func, Addr) : checkCallCont(*Func, Addr);
+
+    if (BAT)
+      Addr = BAT->translate(Func->getAddress(), Addr, IsFrom);
+
+    BinaryFunction *ParentFunc = getBATParentFunction(*Func);
+    if (!ParentFunc)
+      return std::pair{Func, IsRetOrCallCont};
 
-  BinaryFunction *ToFunc = handleAddress(To, /*IsFrom=*/false);
+    if (IsFrom)
+      NumColdSamples += Count;
+
+    return std::pair{ParentFunc, IsRetOrCallCont};
+  };
+
+  uint64_t ToOrig = To;
+  auto [FromFunc, IsReturn] = handleAddress(From, /*IsFrom=*/true);
+  auto [ToFunc, IsCallCont] = handleAddress(To, /*IsFrom=*/false);
   if (!FromFunc && !ToFunc)
     return false;
 
+  // Record call to continuation trace.
+  if (IsCallCont && FromFunc != ToFunc) {
+    LBREntry First{ToOrig - 1, ToOrig - 1, false};
+    LBREntry Second{ToOrig, ToOrig, false};
+    return doTrace(First, Second, Count);
+  }
+  // Ignore returns.
+  if (IsReturn)
+    return true;
+
   // Treat recursive control transfers as inter-branches.
   if (FromFunc == ToFunc && To != 0) {
     recordBranch(*FromFunc, From, To, Count, Mispreds);
@@ -916,24 +956,6 @@ DataAggregator::getFallthroughsInTrace(BinaryFunction &BF,
   if (!FromBB || !ToBB)
     return std::nullopt;
 
-  // Adjust FromBB if the first LBR is a return from the last instruction in
-  // the previous block (that instruction should be a call).
-  if (From == FromBB->getOffset() && !BF.containsAddress(FirstLBR.From) &&
-      !FromBB->isEntryPoint() && !FromBB->isLandingPad()) {
-    const BinaryBasicBlock *PrevBB =
-        BF.getLayout().getBlock(FromBB->getIndex() - 1);
-    if (PrevBB->getSuccessor(FromBB->getLabel())) {
-      const MCInst *Instr = PrevBB->getLastNonPseudoInstr();
-      if (Instr && BC.MIB->isCall(*Instr))
-        FromBB = PrevBB;
-      else
-        LLVM_DEBUG(dbgs() << "invalid incoming LBR (no call): " << FirstLBR
-                          << '\n');
-    } else {
-      LLVM_DEBUG(dbgs() << "invalid incoming LBR: " << FirstLBR << '\n');
-    }
-  }
-
   // Fill out information for fall-through edges. The From and To could be
   // within the same basic block, e.g. when two call instructions are in the
   // same block. In this case we skip the processing.