Index: include/llvm/Analysis/BranchProbabilityInfo.h =================================================================== --- include/llvm/Analysis/BranchProbabilityInfo.h +++ include/llvm/Analysis/BranchProbabilityInfo.h @@ -112,6 +112,11 @@ return IsLikely ? (1u << 20) - 1 : 1; } + /// Normalize a list of weights by scaling them down so that the sum of them + /// doesn't exceed UINT32_MAX. Return the scale. + template + static uint32_t normalizeEdgeWeights(WeightVec &Weights); + void calculate(Function &F, const LoopInfo& LI); private: @@ -151,6 +156,24 @@ bool calcInvokeHeuristics(BasicBlock *BB); }; +template +uint32_t BranchProbabilityInfo::normalizeEdgeWeights(WeightVec &Weights) { + // First we compute the sum with 64-bits of precision. + uint64_t Sum = std::accumulate(Weights.begin(), Weights.end(), uint64_t(0)); + + // If the computed sum fits in 32-bits, we're done. + if (Sum <= UINT32_MAX) + return 1; + + // Otherwise, compute the scale necessary to cause the weights to fit, and + // re-sum with that scale applied. + assert((Sum / UINT32_MAX) < UINT32_MAX); + uint32_t Scale = (Sum / UINT32_MAX) + 1; + for (auto &W : Weights) + W /= Scale; + return Scale; +} + /// \brief Legacy analysis pass which computes \c BranchProbabilityInfo. class BranchProbabilityInfoWrapperPass : public FunctionPass { BranchProbabilityInfo BPI; Index: lib/CodeGen/MachineBranchProbabilityInfo.cpp =================================================================== --- lib/CodeGen/MachineBranchProbabilityInfo.cpp +++ lib/CodeGen/MachineBranchProbabilityInfo.cpp @@ -11,6 +11,7 @@ // //===----------------------------------------------------------------------===// +#include "llvm/Analysis/BranchProbabilityInfo.h" #include "llvm/CodeGen/MachineBranchProbabilityInfo.h" #include "llvm/CodeGen/MachineBasicBlock.h" #include "llvm/IR/Instructions.h" @@ -30,34 +31,14 @@ uint32_t MachineBranchProbabilityInfo:: getSumForBlock(const MachineBasicBlock *MBB, uint32_t &Scale) const { - // First we compute the sum with 64-bits of precision, ensuring that cannot - // overflow by bounding the number of weights considered. Hopefully no one - // actually needs 2^32 successors. assert(MBB->succ_size() < UINT32_MAX); - uint64_t Sum = 0; - Scale = 1; + SmallVector Weights; for (MachineBasicBlock::const_succ_iterator I = MBB->succ_begin(), - E = MBB->succ_end(); I != E; ++I) { - uint32_t Weight = getEdgeWeight(MBB, I); - Sum += Weight; - } - - // If the computed sum fits in 32-bits, we're done. - if (Sum <= UINT32_MAX) - return Sum; + E = MBB->succ_end(); I != E; ++I) + Weights.push_back(getEdgeWeight(MBB, I)); - // Otherwise, compute the scale necessary to cause the weights to fit, and - // re-sum with that scale applied. - assert((Sum / UINT32_MAX) < UINT32_MAX); - Scale = (Sum / UINT32_MAX) + 1; - Sum = 0; - for (MachineBasicBlock::const_succ_iterator I = MBB->succ_begin(), - E = MBB->succ_end(); I != E; ++I) { - uint32_t Weight = getEdgeWeight(MBB, I); - Sum += Weight / Scale; - } - assert(Sum <= UINT32_MAX); - return Sum; + Scale = BranchProbabilityInfo::normalizeEdgeWeights(Weights); + return std::accumulate(Weights.begin(), Weights.end(), 0u); } uint32_t MachineBranchProbabilityInfo::