@@ -778,42 +778,75 @@ bool DataAggregator::doInterBranch(BinaryFunction *FromFunc,
778778}
779779
780780bool DataAggregator::doBranch (uint64_t From, uint64_t To, uint64_t Count,
781- uint64_t Mispreds) {
782- bool IsReturn = false ;
783- auto handleAddress = [&](uint64_t &Addr, bool IsFrom) -> BinaryFunction * {
784- if (BinaryFunction *Func = getBinaryFunctionContainingAddress (Addr)) {
785- Addr -= Func->getAddress ();
786- if (IsFrom) {
787- auto checkReturn = [&](auto MaybeInst) {
788- IsReturn = MaybeInst && BC->MIB ->isReturn (*MaybeInst);
789- };
790- if (Func->hasInstructions ())
791- checkReturn (Func->getInstructionAtOffset (Addr));
792- else
793- checkReturn (Func->disassembleInstructionAtOffset (Addr));
794- }
781+ uint64_t Mispreds, bool IsPreagg) {
782+ // Returns whether \p Offset in \p Func contains a return instruction.
783+ auto checkReturn = [&](const BinaryFunction &Func, const uint64_t Offset) {
784+ auto isReturn = [&](auto MI) { return MI && BC->MIB ->isReturn (*MI); };
785+ return Func.hasInstructions ()
786+ ? isReturn (Func.getInstructionAtOffset (Offset))
787+ : isReturn (Func.disassembleInstructionAtOffset (Offset));
788+ };
795789
796- if (BAT)
797- Addr = BAT->translate (Func->getAddress (), Addr, IsFrom);
790+ // Returns whether \p Offset in \p Func may be a call continuation excluding
791+ // entry points and landing pads.
792+ auto checkCallCont = [&](const BinaryFunction &Func, const uint64_t Offset) {
793+ // No call continuation at a function start.
794+ if (!Offset)
795+ return false ;
796+
797+ // FIXME: support BAT case where the function might be in empty state
798+ // (split fragments declared non-simple).
799+ if (!Func.hasCFG ())
800+ return false ;
801+
802+ // The offset should not be an entry point or a landing pad.
803+ const BinaryBasicBlock *ContBB = Func.getBasicBlockAtOffset (Offset);
804+ return ContBB && !ContBB->isEntryPoint () && !ContBB->isLandingPad ();
805+ };
798806
799- if (BinaryFunction *ParentFunc = getBATParentFunction (*Func)) {
800- Func = ParentFunc;
801- if (IsFrom)
802- NumColdSamples += Count;
803- }
807+ // Mutates \p Addr to an offset into the containing function, performing BAT
808+ // offset translation and parent lookup.
809+ //
810+ // Returns the containing function (or BAT parent) and whether the address
811+ // corresponds to a return (if \p IsFrom) or a call continuation (otherwise).
812+ auto handleAddress = [&](uint64_t &Addr, bool IsFrom) {
813+ BinaryFunction *Func = getBinaryFunctionContainingAddress (Addr);
814+ if (!Func)
815+ return std::pair{Func, false };
804816
805- return Func;
806- }
807- return nullptr ;
817+ Addr -= Func->getAddress ();
818+
819+ bool IsRetOrCallCont =
820+ IsFrom ? checkReturn (*Func, Addr) : checkCallCont (*Func, Addr);
821+
822+ if (BAT)
823+ Addr = BAT->translate (Func->getAddress (), Addr, IsFrom);
824+
825+ BinaryFunction *ParentFunc = getBATParentFunction (*Func);
826+ if (!ParentFunc)
827+ return std::pair{Func, IsRetOrCallCont};
828+
829+ if (IsFrom)
830+ NumColdSamples += Count;
831+
832+ return std::pair{ParentFunc, IsRetOrCallCont};
808833 };
809834
810- BinaryFunction *FromFunc = handleAddress (From, /* IsFrom=*/ true );
835+ uint64_t ToOrig = To;
836+ auto [FromFunc, IsReturn] = handleAddress (From, /* IsFrom*/ true );
837+ auto [ToFunc, IsCallCont] = handleAddress (To, /* IsFrom*/ false );
838+ if (!FromFunc && !ToFunc)
839+ return false ;
840+
841+ // Record call to continuation trace.
842+ if (IsPreagg && FromFunc != ToFunc && (IsReturn || IsCallCont)) {
843+ LBREntry First{ToOrig - 1 , ToOrig - 1 , false };
844+ LBREntry Second{ToOrig, ToOrig, false };
845+ return doTrace (First, Second, Count);
846+ }
811847 // Ignore returns.
812848 if (IsReturn)
813849 return true ;
814- BinaryFunction *ToFunc = handleAddress (To, /* IsFrom=*/ false );
815- if (!FromFunc && !ToFunc)
816- return false ;
817850
818851 // Treat recursive control transfers as inter-branches.
819852 if (FromFunc == ToFunc && To != 0 ) {
@@ -830,10 +863,19 @@ bool DataAggregator::doTrace(const LBREntry &First, const LBREntry &Second,
830863 BinaryFunction *ToFunc = getBinaryFunctionContainingAddress (Second.From );
831864 if (!FromFunc || !ToFunc) {
832865 LLVM_DEBUG ({
833- dbgs () << " Out of range trace starting in " << FromFunc->getPrintName ()
834- << formatv (" @ {0:x}" , First.To - FromFunc->getAddress ())
835- << " and ending in " << ToFunc->getPrintName ()
836- << formatv (" @ {0:x}\n " , Second.From - ToFunc->getAddress ());
866+ dbgs () << " Out of range trace starting in " ;
867+ if (FromFunc)
868+ dbgs () << formatv (" {0} @ {1:x}" , *FromFunc,
869+ First.To - FromFunc->getAddress ());
870+ else
871+ dbgs () << Twine::utohexstr (First.To );
872+ dbgs () << " and ending in " ;
873+ if (ToFunc)
874+ dbgs () << formatv (" {0} @ {1:x}" , *ToFunc,
875+ Second.From - ToFunc->getAddress ());
876+ else
877+ dbgs () << Twine::utohexstr (Second.From );
878+ dbgs () << ' \n ' ;
837879 });
838880 NumLongRangeTraces += Count;
839881 return false ;
@@ -1620,7 +1662,8 @@ void DataAggregator::processBranchEvents() {
16201662 for (const auto &AggrLBR : BranchLBRs) {
16211663 const Trace &Loc = AggrLBR.first ;
16221664 const TakenBranchInfo &Info = AggrLBR.second ;
1623- doBranch (Loc.From , Loc.To , Info.TakenCount , Info.MispredCount );
1665+ doBranch (Loc.From , Loc.To , Info.TakenCount , Info.MispredCount ,
1666+ /* IsPreagg*/ false );
16241667 }
16251668}
16261669
@@ -1781,7 +1824,7 @@ void DataAggregator::processPreAggregated() {
17811824 switch (AggrEntry.EntryType ) {
17821825 case AggregatedLBREntry::BRANCH:
17831826 doBranch (AggrEntry.From .Offset , AggrEntry.To .Offset , AggrEntry.Count ,
1784- AggrEntry.Mispreds );
1827+ AggrEntry.Mispreds , /* IsPreagg */ true );
17851828 break ;
17861829 case AggregatedLBREntry::FT:
17871830 case AggregatedLBREntry::FT_EXTERNAL_ORIGIN: {
0 commit comments