Index: llvm/include/llvm/Analysis/BranchProbabilityInfo.h =================================================================== --- llvm/include/llvm/Analysis/BranchProbabilityInfo.h +++ llvm/include/llvm/Analysis/BranchProbabilityInfo.h @@ -121,6 +121,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 @@ -130,6 +131,15 @@ void setEdgeProbability(const BasicBlock *Src, unsigned IndexInSuccessors, BranchProbability Prob); +public: + /// Set the raw probabilities for all edges from the given block. + /// + /// This allows a pass to explicitly set edge probabilities for a block. It + /// can be used when updating the CFG to update the branch probability + /// information. + 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 @@ -251,10 +251,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; } @@ -264,10 +267,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; } @@ -346,25 +350,37 @@ // Examine the metadata against unreachable heuristic. // If the unreachable heuristic is more strong then we use it for this edge. if (UnreachableIdxs.size() > 0 && ReachableIdxs.size() > 0) { - auto ToDistribute = BranchProbability::getZero(); auto UnreachableProb = UR_TAKEN_PROB; for (auto i : UnreachableIdxs) if (UnreachableProb < BP[i]) { - ToDistribute += BP[i] - UnreachableProb; BP[i] = UnreachableProb; } + // Because of possible rounding errors and the above fix up for + // the unreachable heuristic the sum of probabilities of all edges may be + // less than 1.0. Distribute the remaining probability (calculated as + // 1.0 - (sum of BP[i])) evenly among all the reachable edges. + auto ToDistribute = BranchProbability::getOne(); + for (auto &P : BP) + ToDistribute -= P; + // If we modified the probability of some edges then we must distribute // the difference between reachable blocks. - if (ToDistribute > BranchProbability::getZero()) { - BranchProbability PerEdge = ToDistribute / ReachableIdxs.size(); + // TODO: This spreads ToDistribute evenly upon the reachable edges. A better + // distribution would be proportional. So the relation between weights of + // the reachable edges would be kept unchanged. That is for any reachable + // edges i and j: + // newBP[i] / newBP[j] == oldBP[i] / oldBP[j] + // newBP[i] / oldBP[i] == newBP[j] / oldBP[j] == + // == Denominator / (Denominator - ToDistribute) + // newBP[i] = oldBP[i] * Denominator / (Denominator - ToDistribute) + BranchProbability PerEdge = ToDistribute / ReachableIdxs.size(); + if (PerEdge > BranchProbability::getZero()) for (auto i : ReachableIdxs) BP[i] += PerEdge; - } } - for (unsigned i = 0, e = TI->getNumSuccessors(); i != e; ++i) - setEdgeProbability(BB, i, BP[i]); + setEdgeProbability(BB, BP); return true; } @@ -397,10 +413,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; } @@ -412,10 +431,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; } @@ -438,19 +458,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; } @@ -647,18 +669,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()) { @@ -666,7 +690,7 @@ Denom); auto Prob = NotTakenProb / numExitingEdges; for (unsigned SuccIdx : ExitingEdges) - setEdgeProbability(BB, SuccIdx, Prob); + EdgeProbabilities[SuccIdx] = Prob; } if (uint32_t numUnlikelyEdges = UnlikelyEdges.size()) { @@ -674,9 +698,10 @@ Denom); auto Prob = UnlikelyProb / numUnlikelyEdges; for (unsigned SuccIdx : UnlikelyEdges) - setEdgeProbability(BB, SuccIdx, Prob); + EdgeProbabilities[SuccIdx] = Prob; } + setEdgeProbability(BB, EdgeProbabilities); return true; } @@ -787,15 +812,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; } @@ -830,14 +855,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; } @@ -848,8 +872,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; } @@ -962,6 +986,28 @@ << "\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(); + } + + // Because of rounding errors the total probability cannot be checked to be + // 1.0 exactly. That is TotalNumerator == BranchProbability::getDenominator. + // Instead, every single probability in Probs must be as accurate as possible. + // This results in error 1/denominator at most, thus the total absolute error + // should be within Probs.size / BranchProbability::getDenominator. + 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 @@ -2469,8 +2469,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 @@ -1364,6 +1364,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); @@ -1371,12 +1374,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(); @@ -1388,8 +1393,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)); Index: llvm/test/Analysis/BranchProbabilityInfo/basic.ll =================================================================== --- llvm/test/Analysis/BranchProbabilityInfo/basic.ll +++ llvm/test/Analysis/BranchProbabilityInfo/basic.ll @@ -513,9 +513,9 @@ i32 4, label %case_e ], !prof !9 ; CHECK: edge entry -> case_a probability is 0x00000001 / 0x80000000 = 0.00% ; CHECK: edge entry -> case_b probability is 0x00000001 / 0x80000000 = 0.00% -; CHECK: edge entry -> case_c probability is 0x6aaaaaa9 / 0x80000000 = 83.33% [HOT edge] -; CHECK: edge entry -> case_d probability is 0x0aaaaaa9 / 0x80000000 = 8.33% -; CHECK: edge entry -> case_e probability is 0x0aaaaaa9 / 0x80000000 = 8.33% +; CHECK: edge entry -> case_c probability is 0x6aaaaaaa / 0x80000000 = 83.33% [HOT edge] +; CHECK: edge entry -> case_d probability is 0x0aaaaaaa / 0x80000000 = 8.33% +; CHECK: edge entry -> case_e probability is 0x0aaaaaaa / 0x80000000 = 8.33% case_a: unreachable Index: llvm/unittests/Transforms/Utils/BasicBlockUtilsTest.cpp =================================================================== --- llvm/unittests/Transforms/Utils/BasicBlockUtilsTest.cpp +++ llvm/unittests/Transforms/Utils/BasicBlockUtilsTest.cpp @@ -181,3 +181,96 @@ EXPECT_EQ(BranchProbability(1, 2), BPI.getEdgeProbability(SplitBB, 0u)); EXPECT_EQ(BranchProbability(1, 2), BPI.getEdgeProbability(SplitBB, 1u)); } + +TEST(BasicBlockUtils, SetEdgeProbability) { + LLVMContext C; + + std::unique_ptr M = parseIR( + C, "define void @edge_probability(i32 %0) {\n" + "entry:\n" + "switch i32 %0, label %LD [\n" + " i32 700, label %L0\n" + " i32 701, label %L1\n" + " i32 702, label %L2\n" + " i32 703, label %L3\n" + " i32 704, label %L4\n" + " i32 705, label %L5\n" + " i32 706, label %L6\n" + " i32 707, label %L7\n" + " i32 708, label %L8\n" + " i32 709, label %L9\n" + " i32 710, label %L10\n" + " i32 711, label %L11\n" + " i32 712, label %L12\n" + " i32 713, label %L13\n" + " i32 714, label %L14\n" + " i32 715, label %L15\n" + " i32 716, label %L16\n" + " i32 717, label %L17\n" + " i32 718, label %L18\n" + " i32 719, label %L19\n" + "], !prof !{!\"branch_weights\", i32 1, i32 1, i32 1, i32 1, i32 1, " + "i32 451, i32 1, i32 12, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, " + "i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1}\n" + "LD:\n" + " unreachable\n" + "L0:\n" + " ret void\n" + "L1:\n" + " ret void\n" + "L2:\n" + " ret void\n" + "L3:\n" + " ret void\n" + "L4:\n" + " ret void\n" + "L5:\n" + " ret void\n" + "L6:\n" + " ret void\n" + "L7:\n" + " ret void\n" + "L8:\n" + " ret void\n" + "L9:\n" + " ret void\n" + "L10:\n" + " ret void\n" + "L11:\n" + " ret void\n" + "L12:\n" + " ret void\n" + "L13:\n" + " ret void\n" + "L14:\n" + " ret void\n" + "L15:\n" + " ret void\n" + "L16:\n" + " ret void\n" + "L17:\n" + " ret void\n" + "L18:\n" + " ret void\n" + "L19:\n" + " ret void\n" + "}\n"); + + auto *F = M->getFunction("edge_probability"); + DominatorTree DT(*F); + LoopInfo LI(DT); + BranchProbabilityInfo BPI(*F, LI); + + auto Block = [&F](StringRef BBName) -> const BasicBlock & { + for (auto &BB : *F) + if (BB.getName() == BBName) + return BB; + llvm_unreachable("Block not found"); + }; + + // Check that the unreachable block has the minimal probability. + const BasicBlock &EntryBB = Block("entry"); + const BasicBlock &UnreachableBB = Block("LD"); + EXPECT_EQ(BranchProbability::getRaw(1), + BPI.getEdgeProbability(&EntryBB, &UnreachableBB)); +}