diff --git a/llvm/include/llvm/Analysis/BranchProbabilityInfo.h b/llvm/include/llvm/Analysis/BranchProbabilityInfo.h --- a/llvm/include/llvm/Analysis/BranchProbabilityInfo.h +++ b/llvm/include/llvm/Analysis/BranchProbabilityInfo.h @@ -131,6 +131,12 @@ void setEdgeProbability(const BasicBlock *Src, const SmallVectorImpl &Probs); + /// Copy outgoing edge probabilities from \p Src to \p Dst. + /// + /// This allows to keep probabilities unset for the destination if they were + /// unset for source. + void copyEdgeProbabilities(BasicBlock *Src, BasicBlock *Dst); + static BranchProbability getBranchProbStackProtector(bool IsLikely) { static const BranchProbability LikelyProb((1u << 20) - 1, 1u << 20); return IsLikely ? LikelyProb : LikelyProb.getCompl(); diff --git a/llvm/lib/Analysis/BranchProbabilityInfo.cpp b/llvm/lib/Analysis/BranchProbabilityInfo.cpp --- a/llvm/lib/Analysis/BranchProbabilityInfo.cpp +++ b/llvm/lib/Analysis/BranchProbabilityInfo.cpp @@ -1157,6 +1157,25 @@ assert(TotalNumerator >= BranchProbability::getDenominator() - Probs.size()); } +void BranchProbabilityInfo::copyEdgeProbabilities(BasicBlock *Src, + BasicBlock *Dst) { + eraseBlock(Dst); // Erase stale data if any. + unsigned NumSuccessors = Src->getTerminator()->getNumSuccessors(); + assert(NumSuccessors == Dst->getTerminator()->getNumSuccessors()); + if (NumSuccessors == 0) + return; // Nothing to set. + if (this->Probs.find(std::make_pair(Src, 0)) == this->Probs.end()) + return; // No probability is set for edges from Src. Keep the same for Dst. + + Handles.insert(BasicBlockCallbackVH(Dst, this)); + for (unsigned SuccIdx = 0; SuccIdx < NumSuccessors; ++SuccIdx) { + auto Prob = this->Probs[std::make_pair(Src, SuccIdx)]; + this->Probs[std::make_pair(Dst, SuccIdx)] = Prob; + LLVM_DEBUG(dbgs() << "set edge " << Dst->getName() << " -> " << SuccIdx + << " successor probability to " << Prob << "\n"); + } +} + raw_ostream & BranchProbabilityInfo::printEdgeProbability(raw_ostream &OS, const BasicBlock *Src,