Index: llvm/include/llvm/Analysis/BranchProbabilityInfo.h =================================================================== --- llvm/include/llvm/Analysis/BranchProbabilityInfo.h +++ llvm/include/llvm/Analysis/BranchProbabilityInfo.h @@ -120,6 +120,7 @@ raw_ostream &printEdgeProbability(raw_ostream &OS, const BasicBlock *Src, const BasicBlock *Dst) const; +protected: /// Set the raw edge probability for the given edge. /// /// This allows a pass to explicitly set the edge probability for an edge. It @@ -129,6 +130,10 @@ void setEdgeProbability(const BasicBlock *Src, unsigned IndexInSuccessors, BranchProbability Prob); +public: + void setEdgeProbability(const BasicBlock *Src, + const SmallVectorImpl &Probs); + static BranchProbability getBranchProbStackProtector(bool IsLikely) { static const BranchProbability LikelyProb((1u << 20) - 1, 1u << 20); return IsLikely ? LikelyProb : LikelyProb.getCompl(); Index: llvm/lib/Analysis/BranchProbabilityInfo.cpp =================================================================== --- llvm/lib/Analysis/BranchProbabilityInfo.cpp +++ llvm/lib/Analysis/BranchProbabilityInfo.cpp @@ -250,10 +250,13 @@ if (UnreachableEdges.empty()) return false; + SmallVector EdgeProbabilities( + BB->getTerminator()->getNumSuccessors(), BranchProbability::getUnknown()); if (ReachableEdges.empty()) { BranchProbability Prob(1, UnreachableEdges.size()); for (unsigned SuccIdx : UnreachableEdges) - setEdgeProbability(BB, SuccIdx, Prob); + EdgeProbabilities[SuccIdx] = Prob; + setEdgeProbability(BB, EdgeProbabilities); return true; } @@ -263,10 +266,11 @@ ReachableEdges.size(); for (unsigned SuccIdx : UnreachableEdges) - setEdgeProbability(BB, SuccIdx, UnreachableProb); + EdgeProbabilities[SuccIdx] = UnreachableProb; for (unsigned SuccIdx : ReachableEdges) - setEdgeProbability(BB, SuccIdx, ReachableProb); + EdgeProbabilities[SuccIdx] = ReachableProb; + setEdgeProbability(BB, EdgeProbabilities); return true; } @@ -362,8 +366,7 @@ } } - for (unsigned i = 0, e = TI->getNumSuccessors(); i != e; ++i) - setEdgeProbability(BB, i, BP[i]); + setEdgeProbability(BB, BP); return true; } @@ -396,10 +399,13 @@ if (ColdEdges.empty()) return false; + SmallVector EdgeProbabilities( + BB->getTerminator()->getNumSuccessors(), BranchProbability::getUnknown()); if (NormalEdges.empty()) { BranchProbability Prob(1, ColdEdges.size()); for (unsigned SuccIdx : ColdEdges) - setEdgeProbability(BB, SuccIdx, Prob); + EdgeProbabilities[SuccIdx] = Prob; + setEdgeProbability(BB, EdgeProbabilities); return true; } @@ -411,10 +417,11 @@ (CC_TAKEN_WEIGHT + CC_NONTAKEN_WEIGHT) * uint64_t(NormalEdges.size())); for (unsigned SuccIdx : ColdEdges) - setEdgeProbability(BB, SuccIdx, ColdProb); + EdgeProbabilities[SuccIdx] = ColdProb; for (unsigned SuccIdx : NormalEdges) - setEdgeProbability(BB, SuccIdx, NormalProb); + EdgeProbabilities[SuccIdx] = NormalProb; + setEdgeProbability(BB, EdgeProbabilities); return true; } @@ -437,19 +444,21 @@ assert(CI->getOperand(1)->getType()->isPointerTy()); + BranchProbability TakenProb(PH_TAKEN_WEIGHT, + PH_TAKEN_WEIGHT + PH_NONTAKEN_WEIGHT); + BranchProbability UntakenProb(PH_NONTAKEN_WEIGHT, + PH_TAKEN_WEIGHT + PH_NONTAKEN_WEIGHT); + // p != 0 -> isProb = true // p == 0 -> isProb = false // p != q -> isProb = true // p == q -> isProb = false; - unsigned TakenIdx = 0, NonTakenIdx = 1; bool isProb = CI->getPredicate() == ICmpInst::ICMP_NE; if (!isProb) - std::swap(TakenIdx, NonTakenIdx); + std::swap(TakenProb, UntakenProb); - BranchProbability TakenProb(PH_TAKEN_WEIGHT, - PH_TAKEN_WEIGHT + PH_NONTAKEN_WEIGHT); - setEdgeProbability(BB, TakenIdx, TakenProb); - setEdgeProbability(BB, NonTakenIdx, TakenProb.getCompl()); + setEdgeProbability( + BB, SmallVector({TakenProb, UntakenProb})); return true; } @@ -646,18 +655,20 @@ (UnlikelyEdges.empty() ? 0 : LBH_UNLIKELY_WEIGHT) + (ExitingEdges.empty() ? 0 : LBH_NONTAKEN_WEIGHT); + SmallVector EdgeProbabilities( + BB->getTerminator()->getNumSuccessors(), BranchProbability::getUnknown()); if (uint32_t numBackEdges = BackEdges.size()) { BranchProbability TakenProb = BranchProbability(LBH_TAKEN_WEIGHT, Denom); auto Prob = TakenProb / numBackEdges; for (unsigned SuccIdx : BackEdges) - setEdgeProbability(BB, SuccIdx, Prob); + EdgeProbabilities[SuccIdx] = Prob; } if (uint32_t numInEdges = InEdges.size()) { BranchProbability TakenProb = BranchProbability(LBH_TAKEN_WEIGHT, Denom); auto Prob = TakenProb / numInEdges; for (unsigned SuccIdx : InEdges) - setEdgeProbability(BB, SuccIdx, Prob); + EdgeProbabilities[SuccIdx] = Prob; } if (uint32_t numExitingEdges = ExitingEdges.size()) { @@ -665,7 +676,7 @@ Denom); auto Prob = NotTakenProb / numExitingEdges; for (unsigned SuccIdx : ExitingEdges) - setEdgeProbability(BB, SuccIdx, Prob); + EdgeProbabilities[SuccIdx] = Prob; } if (uint32_t numUnlikelyEdges = UnlikelyEdges.size()) { @@ -673,9 +684,10 @@ Denom); auto Prob = UnlikelyProb / numUnlikelyEdges; for (unsigned SuccIdx : UnlikelyEdges) - setEdgeProbability(BB, SuccIdx, Prob); + EdgeProbabilities[SuccIdx] = Prob; } + setEdgeProbability(BB, EdgeProbabilities); return true; } @@ -786,15 +798,15 @@ return false; } - unsigned TakenIdx = 0, NonTakenIdx = 1; - - if (!isProb) - std::swap(TakenIdx, NonTakenIdx); - BranchProbability TakenProb(ZH_TAKEN_WEIGHT, ZH_TAKEN_WEIGHT + ZH_NONTAKEN_WEIGHT); - setEdgeProbability(BB, TakenIdx, TakenProb); - setEdgeProbability(BB, NonTakenIdx, TakenProb.getCompl()); + BranchProbability UntakenProb(ZH_NONTAKEN_WEIGHT, + ZH_TAKEN_WEIGHT + ZH_NONTAKEN_WEIGHT); + if (!isProb) + std::swap(TakenProb, UntakenProb); + + setEdgeProbability( + BB, SmallVector({TakenProb, UntakenProb})); return true; } @@ -829,14 +841,13 @@ return false; } - unsigned TakenIdx = 0, NonTakenIdx = 1; - + BranchProbability TakenProb(TakenWeight, TakenWeight + NontakenWeight); + BranchProbability UntakenProb(NontakenWeight, TakenWeight + NontakenWeight); if (!isProb) - std::swap(TakenIdx, NonTakenIdx); + std::swap(TakenProb, UntakenProb); - BranchProbability TakenProb(TakenWeight, TakenWeight + NontakenWeight); - setEdgeProbability(BB, TakenIdx, TakenProb); - setEdgeProbability(BB, NonTakenIdx, TakenProb.getCompl()); + setEdgeProbability( + BB, SmallVector({TakenProb, UntakenProb})); return true; } @@ -847,8 +858,8 @@ BranchProbability TakenProb(IH_TAKEN_WEIGHT, IH_TAKEN_WEIGHT + IH_NONTAKEN_WEIGHT); - setEdgeProbability(BB, 0 /*Index for Normal*/, TakenProb); - setEdgeProbability(BB, 1 /*Index for Unwind*/, TakenProb.getCompl()); + setEdgeProbability( + BB, SmallVector({TakenProb, TakenProb.getCompl()})); return true; } @@ -959,6 +970,22 @@ << "\n"); } +/// Set the edge probability for all edges at once. +void BranchProbabilityInfo::setEdgeProbability( + const BasicBlock *Src, const SmallVectorImpl &Probs) { + assert(Src->getTerminator()->getNumSuccessors() == Probs.size()); + if (Probs.size() == 0) + return; // Nothing to set. + + uint64_t TotalNumerator = 0; + for (unsigned SuccIdx = 0; SuccIdx < Probs.size(); ++SuccIdx) { + setEdgeProbability(Src, SuccIdx, Probs[SuccIdx]); + TotalNumerator += Probs[SuccIdx].getNumerator(); + } + assert(TotalNumerator <= BranchProbability::getDenominator() + Probs.size()); + assert(TotalNumerator >= BranchProbability::getDenominator() - Probs.size()); +} + raw_ostream & BranchProbabilityInfo::printEdgeProbability(raw_ostream &OS, const BasicBlock *Src, Index: llvm/lib/Transforms/Scalar/JumpThreading.cpp =================================================================== --- llvm/lib/Transforms/Scalar/JumpThreading.cpp +++ llvm/lib/Transforms/Scalar/JumpThreading.cpp @@ -2513,8 +2513,7 @@ } // Update edge probabilities in BPI. - for (int I = 0, E = BBSuccProbs.size(); I < E; I++) - BPI->setEdgeProbability(BB, I, BBSuccProbs[I]); + BPI->setEdgeProbability(BB, BBSuccProbs); // Update the profile metadata as well. // Index: llvm/lib/Transforms/Utils/BreakCriticalEdges.cpp =================================================================== --- llvm/lib/Transforms/Utils/BreakCriticalEdges.cpp +++ llvm/lib/Transforms/Utils/BreakCriticalEdges.cpp @@ -401,9 +401,7 @@ BasicBlock *BodyBlock = Target->splitBasicBlock(FirstNonPHI, ".split"); if (ShouldUpdateAnalysis) { // Copy the BFI/BPI from Target to BodyBlock. - for (unsigned I = 0, E = BodyBlock->getTerminator()->getNumSuccessors(); - I < E; ++I) - BPI->setEdgeProbability(BodyBlock, I, EdgeProbabilities[I]); + BPI->setEdgeProbability(BodyBlock, EdgeProbabilities); BFI->setBlockFreq(BodyBlock, BFI->getBlockFreq(Target).getFrequency()); } // It's possible Target was its own successor through an indirectbr. Index: llvm/lib/Transforms/Utils/CodeExtractor.cpp =================================================================== --- llvm/lib/Transforms/Utils/CodeExtractor.cpp +++ llvm/lib/Transforms/Utils/CodeExtractor.cpp @@ -1356,6 +1356,9 @@ // Block Frequency distribution with dummy node. Distribution BranchDist; + SmallVector EdgeProbabilities( + TI->getNumSuccessors(), BranchProbability::getUnknown()); + // Add each of the frequencies of the successors. for (unsigned i = 0, e = TI->getNumSuccessors(); i < e; ++i) { BlockNode ExitNode(i); @@ -1363,12 +1366,14 @@ if (ExitFreq != 0) BranchDist.addExit(ExitNode, ExitFreq); else - BPI->setEdgeProbability(CodeReplacer, i, BranchProbability::getZero()); + EdgeProbabilities[i] = BranchProbability::getZero(); } // Check for no total weight. - if (BranchDist.Total == 0) + if (BranchDist.Total == 0) { + BPI->setEdgeProbability(CodeReplacer, EdgeProbabilities); return; + } // Normalize the distribution so that they can fit in unsigned. BranchDist.normalize(); @@ -1380,8 +1385,9 @@ // Get the weight and update the current BFI. BranchWeights[Weight.TargetNode.Index] = Weight.Amount; BranchProbability BP(Weight.Amount, BranchDist.Total); - BPI->setEdgeProbability(CodeReplacer, Weight.TargetNode.Index, BP); + EdgeProbabilities[Weight.TargetNode.Index] = BP; } + BPI->setEdgeProbability(CodeReplacer, EdgeProbabilities); TI->setMetadata( LLVMContext::MD_prof, MDBuilder(TI->getContext()).createBranchWeights(BranchWeights));