Provide an interface normalizeSuccWeights in MachineBasicBlock to normalize its successors' weights and use it in other places.

[congh]

Summary

The implementation of the new function is borrowed from MachineBranchProbabilityInfo::getSumForBlock. The type of weight list is a template parameter so that we can pass weights in either uint32_t or uint64_t.

This function can be useful when we need to normalize a list of edge weights.

[davidxl]

include/llvm/Analysis/BranchProbabilityInfo.h

165 ↗

(On Diff #30436)

Why not doing normalization unconditionally to avoid precision loss introduced by incremental update (for small weights)?

[congh]

include/llvm/Analysis/BranchProbabilityInfo.h

165 ↗

(On Diff #30436)

The incremental update itself can deal with precision loss, then this function can be called to guarantee that the sum of weights won't exceed UINT32_MAX.

If we decide to normalize all edge weights with scaling up, we should call this normalization function every time we set or update the edge weight. Have we decided this?

[davidxl]

include/llvm/Analysis/BranchProbabilityInfo.h

165 ↗

(On Diff #30436)

You will be needing this for precision preservation purpose, so this interface needs to take an extra argument to decide if up-scaling needs to be done -- by default it can be off.

[congh]

include/llvm/Analysis/BranchProbabilityInfo.h

165 ↗

(On Diff #30436)

OK. But I am wondering how to tell if the returned scale is for scaling up or down? This is the difficult part...

[congh]

Several updates:

1. Move normalizeEdgeWeights() from BranchProbabilityInfo to MachineBranchProbabilityInfo.
2. Provide an interface in MachineBasicBlock to normalize its successors' weights.
3. Add a flag int MachineBasicBlock that tracks whether its successors' weights are normalized.
4. Provide an overload of getSumForBlock that accepts a non-const pointer to a MBB so that it can force normalizing this MBB's successors' weights.
5. Update several uses of getSumForBlock() by eliminating the once needed weight scale.

[davidxl]

include/llvm/CodeGen/MachineBasicBlock.h

88

Put this field after int Number field so that the padding space can be used (for 64bit).

[congh]

include/llvm/CodeGen/MachineBasicBlock.h

88

OK. Updated.

[congh]

Reduce the size of MachineBasicBlock by relocating SuccWeightsNormalized which is a boolean.

[davidxl]

In general, I find doing normalization like this is error prone. Once the weights gets normalized, the new weight added in later may become meaningless and there is no good way to check against it.

I suggest using this opportunity to use edge weight to represent fixed point representation of the branch probability (aka scaled probability). We can pick a large scale factor such as 100000. It has the nice property that the sum of all 'weights' will guarantee to be always '100000'. Another nice side effect is that the computation of getEdgeProbability is now super fast -- there is no need to keep the getSumWeight interface any more.

The addSuccessor interface will also need to be modified such that instead of passing raw 'weight', it needs to pass in BranchProbablity(N,M).

In the longer term, we should remove uses of getEdgeWeight interface completely and replace it with getEdgeProbabity instead (as it is already efficient). Weight will simply be an implementation detail that does not need to be exposed to users.

lib/CodeGen/IfConversion.cpp

1245

why not move this before getEdgeWeight so that the explicit call to normalize weight is not needed?

[congh]

In D11442#215298, @davidxl wrote:

In general, I find doing normalization like this is error prone. Once the weights gets normalized, the new weight added in later may become meaningless and there is no good way to check against it.

Yes, but this depends on when the normalization is applied. If the user calls it explicitly, he/she should be aware that the edge weights are already changed so the newly added one needs to be adjusted. Otherwise, the user also needs to check the current normalized weights to determine the new successor's weight (otherwise how to assign weight to a new successor?). But this design can still bring potential errors.

I suggest using this opportunity to use edge weight to represent fixed point representation of the branch probability (aka scaled probability). We can pick a large scale factor such as 100000. It has the nice property that the sum of all 'weights' will guarantee to be always '100000'. Another nice side effect is that the computation of getEdgeProbability is now super fast -- there is no need to keep the getSumWeight interface any more.

I agree this design: the only use of weights is to compute edge probabilities. We should adjust the edge "weight" in terms of probabilities instead of an absolute weight. This is a more flexible representation.

The addSuccessor interface will also need to be modified such that instead of passing raw 'weight', it needs to pass in BranchProbablity(N,M).

So the probabilities of old successors can be adjusted accordingly.

In the longer term, we should remove uses of getEdgeWeight interface completely and replace it with getEdgeProbabity instead (as it is already efficient). Weight will simply be an implementation detail that does not need to be exposed to users.

[congh]

lib/CodeGen/IfConversion.cpp

1245

This is to explicitly stating that we should normalize the edge weights then get them. I am afraid if we just simply put getSumForBlock() before getEdgeWeight(), some people who are not aware that getSumForBlock() does weights normalization may switch them later by mistake.

[congh]

Update the patch by turning zero edge weights into one in BPI so that we won't get zero weights in BPI anymore.

[davidxl]

test/Transforms/SampleProfile/branch.ll

41 ↗

(On Diff #31882)

Since BFI always ignores zero weight which leads to inconsistency (as BPI did not ignore zero weight before this fix, only BFI does it).

If the test case really intend to test a highly biased branch here, I think the test case should be modified to not use zero weight after this fix (as it now will be completely ignored). Changing the the results to 50% may violate what is intended by the test.

[congh]

test/Transforms/SampleProfile/branch.ll

41 ↗

(On Diff #31882)

I did an investigation on this. The weight 0 and 1 are not indicated in this test cased but are read from a sample profile file, in which most lines have zero weights except two which are large weights. Where does 1 come from? When a zero weight is read from sample profile, it is converted to 1 (see include/llvm/ProfileData/SampleProf.h:175). In the comment, it says that the weight is set to 0 if the profile data is missing. If the profile data exist and the number of samples is zero, it will be converted to 1 to differentiate those two different cases. However, this will lead to a possible result that two edges from the same block have weight 0 and 1 (0 for missing profile and 1 for 0 sample): now how to interpret them? I think here it is a good idea to convert 0 to 1 in BPI for this case. For the original test case the probability 0 and 100% are quite wrong, but 50% and 50% as in the updated test case is more reasonable.

[davidxl]

test/Transforms/SampleProfile/branch.ll

41 ↗

(On Diff #31882)

In AutoFDO, lines with instructions generated but without any samples are annotated with 0 counts -- they actually should not be converted to 1 as 0 counts are strong signals of coldness.

One the other hand, lines that do not have any instructions associated (due to CSE, PRE, IVOPT etc) do not have any counts either, but they should not be marked with 0 counts and should be specially handled. It seems current implementation tries to overload zero weight again to represent this state which is wrong because zero weight does not do what it thinks it does.

At this point, I think this test does not matters that much (Dehao and Diego are working to overhaul SampleFDO support anyway) -- so I am ok with this change.

Duncan, WDYT?

include/llvm/CodeGen/MachineBasicBlock.h

private:
   void createNode(const MachineInstr &);
 };
 
 class MachineBasicBlock : public ilist_node<MachineBasicBlock> {
   typedef ilist<MachineInstr> Instructions;
   Instructions Insts;
   const BasicBlock *BB;
   int Number;
+
+  /// SuccWeightsNormalized - A flag tracking whether the weights of all
+  /// successors are normalized.
+  bool SuccWeightsNormalized;
+
   MachineFunction *xParent;
 
   /// Predecessors/Successors - Keep track of the predecessor / successor
   /// basicblocks.
   std::vector<MachineBasicBlock *> Predecessors;
   std::vector<MachineBasicBlock *> Successors;
 
   /// Weights - Keep track of the weights to the successors. This vector
   /// has the same order as Successors, or it is empty if we don't use it
   /// (disable optimization).
   std::vector<uint32_t> Weights;
   typedef std::vector<uint32_t>::iterator weight_iterator;
   typedef std::vector<uint32_t>::const_iterator const_weight_iterator;
 
   /// LiveIns - Keep track of the physical registers that are livein of
   /// the basicblock.
   std::vector<unsigned> LiveIns;
 
   /// Alignment - Alignment of the basic block. Zero if the basic block does
   /// not need to be aligned.
   /// The alignment is specified as log2(bytes).
   unsigned Alignment;
 
   /// IsLandingPad - Indicate that this basic block is entered via an
   /// exception handler.
public:
   /// is the target of an indirect branch.
   void setHasAddressTaken() { AddressTaken = true; }
 
   /// getParent - Return the MachineFunction containing this basic block.
   ///
   const MachineFunction *getParent() const { return xParent; }
   MachineFunction *getParent() { return xParent; }
 
+  /// succWeightsNormalized - Return whether all weights of successors are
+  /// normalized.
+  bool succWeightsNormalized() const { return SuccWeightsNormalized; }
+
 
   /// bundle_iterator - MachineBasicBlock iterator that automatically skips over
   /// MIs that are inside bundles (i.e. walk top level MIs only).
   template<typename Ty, typename IterTy>
   class bundle_iterator
     : public std::iterator<std::bidirectional_iterator_tag, Ty, ptrdiff_t> {
     IterTy MII;
 
public:
   ///
   /// Note that duplicate Machine CFG edges are not allowed.
   ///
   void addSuccessor(MachineBasicBlock *succ, uint32_t weight = 0);
 
   /// Set successor weight of a given iterator.
   void setSuccWeight(succ_iterator I, uint32_t weight);
 
+  /// Normalize all succesor weights so that the sum of them does not exceed
+  /// UINT32_MAX. Return true if the weights are modified and false otherwise.
+  /// Note that weights that are modified after calling this function are not
+  /// guaranteed to be normalized.
+  bool normalizeSuccWeights();
+
   /// removeSuccessor - Remove successor from the successors list of this
   /// MachineBasicBlock. The Predecessors list of succ is automatically updated.
   ///
   void removeSuccessor(MachineBasicBlock *succ);
 
   /// removeSuccessor - Remove specified successor from the successors list of
   /// this MachineBasicBlock. The Predecessors list of succ is automatically
   /// updated.  Return the iterator to the element after the one removed.

include/llvm/CodeGen/MachineBranchProbabilityInfo.h

public:
   uint32_t getEdgeWeight(const MachineBasicBlock *Src,
                          MachineBasicBlock::const_succ_iterator Dst) const;
 
   // Get sum of the block successors' weights, potentially scaling them to fit
   // within 32-bits. If scaling is required, sets Scale based on the necessary
   // adjustment. Any edge weights used with the sum should be divided by Scale.
   uint32_t getSumForBlock(const MachineBasicBlock *MBB, uint32_t &Scale) const;
 
+  // Get sum of the block successors' weights, and force normalizing the
+  // successors' weights of MBB so that their sum fit within 32-bits.
+  uint32_t getSumForBlock(MachineBasicBlock *MBB) const;
+
   // A 'Hot' edge is an edge which probability is >= 80%.
   bool isEdgeHot(const MachineBasicBlock *Src,
                  const MachineBasicBlock *Dst) const;
 
   // Return a hot successor for the block BB or null if there isn't one.
   // NB: This routine's complexity is linear on the number of successors.
   MachineBasicBlock *getHotSucc(MachineBasicBlock *MBB) const;
 
   // Return a probability as a fraction between 0 (0% probability) and
   // 1 (100% probability), however the value is never equal to 0, and can be 1
   // only iff SRC block has only one successor.
   // NB: This routine's complexity is linear on the number of successors of
   // Src. Querying sequentially for each successor's probability is a quadratic
   // query pattern.
   BranchProbability getEdgeProbability(const MachineBasicBlock *Src,
                                        const MachineBasicBlock *Dst) const;
 
   // Print value between 0 (0% probability) and 1 (100% probability),
   // however the value is never equal to 0, and can be 1 only iff SRC block
   // has only one successor.
   raw_ostream &printEdgeProbability(raw_ostream &OS,
                                     const MachineBasicBlock *Src,
                                     const MachineBasicBlock *Dst) const;
+
+  // Normalize a list of weights by scaling them down so that the sum of them
+  // doesn't exceed UINT32_MAX. Return the scale.
+  template <class WeightList>
+  static uint32_t normalizeEdgeWeights(WeightList &Weights);
 };
 
+template <class WeightList>
+uint32_t
+MachineBranchProbabilityInfo::normalizeEdgeWeights(WeightList &Weights) {
+  assert(Weights.size() < UINT32_MAX && "Too many weights in the list!");
+  // 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 &&
+         "The sum of weights exceeds UINT32_MAX^2!");
+  uint32_t Scale = (Sum / UINT32_MAX) + 1;
+  for (auto &W : Weights)
+    W /= Scale;
+  return Scale;
+}
+
 }
 
 
 #endif

lib/CodeGen/IfConversion.cpp

bool IfConverter::IfConvertTriangle(BBInfo &BBI, IfcvtKind Kind) {
   Redefs.init(TRI);
   Redefs.addLiveIns(CvtBBI->BB);
   Redefs.addLiveIns(NextBBI->BB);
 
   DontKill.clear();
 
   bool HasEarlyExit = CvtBBI->FalseBB != nullptr;
   uint64_t CvtNext = 0, CvtFalse = 0, BBNext = 0, BBCvt = 0, SumWeight = 0;
-  uint32_t WeightScale = 0;
 
   if (HasEarlyExit) {
     // Get weights before modifying CvtBBI->BB and BBI.BB.
+    // Explictly normalize the weights of all edges from CvtBBI->BB so that we
+    // are aware that the edge weights obtained below are normalized.
+    CvtBBI->BB->normalizeSuccWeights();
     CvtNext = MBPI->getEdgeWeight(CvtBBI->BB, NextBBI->BB);
     CvtFalse = MBPI->getEdgeWeight(CvtBBI->BB, CvtBBI->FalseBB);
     BBNext = MBPI->getEdgeWeight(BBI.BB, NextBBI->BB);
     BBCvt = MBPI->getEdgeWeight(BBI.BB, CvtBBI->BB);
-    SumWeight = MBPI->getSumForBlock(CvtBBI->BB, WeightScale);
+    SumWeight = MBPI->getSumForBlock(CvtBBI->BB);
   }
 
   if (CvtBBI->BB->pred_size() > 1) {
     BBI.NonPredSize -= TII->RemoveBranch(*BBI.BB);
     // Copy instructions in the true block, predicate them, and add them to
     // the entry block.
     CopyAndPredicateBlock(BBI, *CvtBBI, Cond, true);
 
if (HasEarlyExit) {
     BBI.BB->addSuccessor(CvtBBI->FalseBB);
     // Update the edge weight for both CvtBBI->FalseBB and NextBBI.
     // New_Weight(BBI.BB, NextBBI->BB) =
     //   Weight(BBI.BB, NextBBI->BB) * getSumForBlock(CvtBBI->BB) +
     //   Weight(BBI.BB, CvtBBI->BB) * Weight(CvtBBI->BB, NextBBI->BB)
     // New_Weight(BBI.BB, CvtBBI->FalseBB) =
     //   Weight(BBI.BB, CvtBBI->BB) * Weight(CvtBBI->BB, CvtBBI->FalseBB)
 
-    uint64_t NewNext = BBNext * SumWeight + (BBCvt * CvtNext) / WeightScale;
-    uint64_t NewFalse = (BBCvt * CvtFalse) / WeightScale;
+    uint64_t NewNext = BBNext * SumWeight + BBCvt * CvtNext;
+    uint64_t NewFalse = BBCvt * CvtFalse;
     // We need to scale down all weights of BBI.BB to fit uint32_t.
     // Here BBI.BB is connected to CvtBBI->FalseBB and will fall through to
     // the next block.
     ScaleWeights(NewNext, NewFalse, BBI.BB, getNextBlock(BBI.BB),
                  CvtBBI->FalseBB, MBPI);
   }
 
   // Merge in the 'false' block if the 'false' block has no other

lib/CodeGen/MachineBasicBlock.cpp

 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/CodeGen/MachineBasicBlock.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/CodeGen/LiveIntervalAnalysis.h"
 #include "llvm/CodeGen/LiveVariables.h"
+#include "llvm/CodeGen/MachineBranchProbabilityInfo.h"
 #include "llvm/CodeGen/MachineDominators.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineInstrBuilder.h"
 #include "llvm/CodeGen/MachineLoopInfo.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/SlotIndexes.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/ModuleSlotTracker.h"
 #include "llvm/MC/MCAsmInfo.h"
 #include "llvm/MC/MCContext.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Target/TargetInstrInfo.h"
 #include "llvm/Target/TargetMachine.h"
 #include "llvm/Target/TargetRegisterInfo.h"
 #include "llvm/Target/TargetSubtargetInfo.h"
 #include <algorithm>
 using namespace llvm;
 
 #define DEBUG_TYPE "codegen"
 
 MachineBasicBlock::MachineBasicBlock(MachineFunction &mf, const BasicBlock *bb)
-  : BB(bb), Number(-1), xParent(&mf), Alignment(0), IsLandingPad(false),
-    AddressTaken(false), CachedMCSymbol(nullptr) {
+    : BB(bb), Number(-1), SuccWeightsNormalized(false), xParent(&mf),
+      Alignment(0), IsLandingPad(false), AddressTaken(false),
+      CachedMCSymbol(nullptr) {
   Insts.Parent = this;
 }
 
 MachineBasicBlock::~MachineBasicBlock() {
 }
 
 /// getSymbol - Return the MCSymbol for this basic block.
 ///

 
 void MachineBasicBlock::addSuccessor(MachineBasicBlock *succ, uint32_t weight) {
 
   // If we see non-zero value for the first time it means we actually use Weight
   // list, so we fill all Weights with 0's.
   if (weight != 0 && Weights.empty())
     Weights.resize(Successors.size());
 
-  if (weight != 0 || !Weights.empty())
+  if (weight != 0 || !Weights.empty()) {
     Weights.push_back(weight);
+    SuccWeightsNormalized = false;
+  }
 
    Successors.push_back(succ);
    succ->addPredecessor(this);
  }
 
 void MachineBasicBlock::removeSuccessor(MachineBasicBlock *succ) {
   succ->removePredecessor(this);
   succ_iterator I = std::find(Successors.begin(), Successors.end(), succ);
uint32_t MachineBasicBlock::getSuccWeight(const_succ_iterator Succ) const {
 
   return *getWeightIterator(Succ);
 }
 
 /// Set successor weight of a given iterator.
 void MachineBasicBlock::setSuccWeight(succ_iterator I, uint32_t weight) {
   if (Weights.empty())
     return;
-  *getWeightIterator(I) = weight;
+  auto WeightIter = getWeightIterator(I);
+  uint32_t OldWeight = *WeightIter;
+  *WeightIter = weight;
+  if (weight > OldWeight)
+    SuccWeightsNormalized = false;
+}
+
+/// Normalize all succesor weights so that the sum of them does not exceed
+/// UINT32_MAX. Return true if the weights are modified and false otherwise.
+/// Note that weights that are modified after calling this function are not
+/// guaranteed to be normalized.
+bool MachineBasicBlock::normalizeSuccWeights() {
+  if (!SuccWeightsNormalized) {
+    uint32_t Scale =
+        MachineBranchProbabilityInfo::normalizeEdgeWeights(Weights);
+    SuccWeightsNormalized = true;
+    return Scale != 1;
+  }
+  return false;
 }
 
 /// getWeightIterator - Return wight iterator corresonding to the I successor
 /// iterator
 MachineBasicBlock::weight_iterator MachineBasicBlock::
 getWeightIterator(MachineBasicBlock::succ_iterator I) {
   assert(Weights.size() == Successors.size() && "Async weight list!");
   size_t index = std::distance(Successors.begin(), I);

lib/CodeGen/MachineBlockPlacement.cpp

MachineBlockPlacement::selectBestSuccessor(MachineBasicBlock *BB,
   MachineBasicBlock *BestSucc = nullptr;
   // FIXME: Due to the performance of the probability and weight routines in
   // the MBPI analysis, we manually compute probabilities using the edge
   // weights. This is suboptimal as it means that the somewhat subtle
   // definition of edge weight semantics is encoded here as well. We should
   // improve the MBPI interface to efficiently support query patterns such as
   // this.
   uint32_t BestWeight = 0;
-  uint32_t WeightScale = 0;
-  uint32_t SumWeight = MBPI->getSumForBlock(BB, WeightScale);
+  uint32_t SumWeight = MBPI->getSumForBlock(BB);
   DEBUG(dbgs() << "Attempting merge from: " << getBlockName(BB) << "\n");
   for (MachineBasicBlock *Succ : BB->successors()) {
     if (BlockFilter && !BlockFilter->count(Succ))
       continue;
     BlockChain &SuccChain = *BlockToChain[Succ];
     if (&SuccChain == &Chain) {
       DEBUG(dbgs() << "    " << getBlockName(Succ) << " -> Already merged!\n");
       continue;
     }
     if (Succ != *SuccChain.begin()) {
       DEBUG(dbgs() << "    " << getBlockName(Succ) << " -> Mid chain!\n");
       continue;
     }
 
     uint32_t SuccWeight = MBPI->getEdgeWeight(BB, Succ);
-    BranchProbability SuccProb(SuccWeight / WeightScale, SumWeight);
+    BranchProbability SuccProb(SuccWeight, SumWeight);
 
     // If we outline optional branches, look whether Succ is unavoidable, i.e.
     // dominates all terminators of the MachineFunction. If it does, other
     // successors must be optional. Don't do this for cold branches.
     if (OutlineOptionalBranches && SuccProb > HotProb.getCompl() &&
         UnavoidableBlocks.count(Succ) > 0) {
       auto HasShortOptionalBranch = [&]() {
         for (MachineBasicBlock *Pred : Succ->predecessors()) {
for (MachineBasicBlock *MBB : L.getBlocks()) {
     // We store the old exiting state and restore it if a viable looping
     // successor isn't found.
     MachineBasicBlock *OldExitingBB = ExitingBB;
     BlockFrequency OldBestExitEdgeFreq = BestExitEdgeFreq;
     bool HasLoopingSucc = false;
     // FIXME: Due to the performance of the probability and weight routines in
     // the MBPI analysis, we use the internal weights and manually compute the
     // probabilities to avoid quadratic behavior.
-    uint32_t WeightScale = 0;
-    uint32_t SumWeight = MBPI->getSumForBlock(MBB, WeightScale);
+    uint32_t SumWeight = MBPI->getSumForBlock(MBB);
     for (MachineBasicBlock *Succ : MBB->successors()) {
       if (Succ->isLandingPad())
         continue;
       if (Succ == MBB)
         continue;
       BlockChain &SuccChain = *BlockToChain[Succ];
       // Don't split chains, either this chain or the successor's chain.
       if (&Chain == &SuccChain) {
for (MachineBasicBlock *Succ : MBB->successors()) {
 
       unsigned SuccLoopDepth = 0;
       if (MachineLoop *ExitLoop = MLI->getLoopFor(Succ)) {
         SuccLoopDepth = ExitLoop->getLoopDepth();
         if (ExitLoop->contains(&L))
           BlocksExitingToOuterLoop.insert(MBB);
       }
 
-      BranchProbability SuccProb(SuccWeight / WeightScale, SumWeight);
+      BranchProbability SuccProb(SuccWeight, SumWeight);
       BlockFrequency ExitEdgeFreq = MBFI->getBlockFreq(MBB) * SuccProb;
       DEBUG(dbgs() << "    exiting: " << getBlockName(MBB) << " -> "
                    << getBlockName(Succ) << " [L:" << SuccLoopDepth << "] (";
             MBFI->printBlockFreq(dbgs(), ExitEdgeFreq) << ")\n");
       // Note that we bias this toward an existing layout successor to retain
       // incoming order in the absence of better information. The exit must have
       // a frequency higher than the current exit before we consider breaking
       // the layout.

lib/CodeGen/MachineBranchProbabilityInfo.cpp

INITIALIZE_PASS_BEGIN(MachineBranchProbabilityInfo, "machine-branch-prob",
                       "Machine Branch Probability Analysis", false, true)
 INITIALIZE_PASS_END(MachineBranchProbabilityInfo, "machine-branch-prob",
                     "Machine Branch Probability Analysis", false, true)
 
 char MachineBranchProbabilityInfo::ID = 0;
 
 void MachineBranchProbabilityInfo::anchor() { }
 
-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;
+uint32_t
+MachineBranchProbabilityInfo::getSumForBlock(MachineBasicBlock *MBB) const {
+  // Normalize the weights of MBB's all successors so that the sum is guaranteed
+  // to be no greater than UINT32_MAX.
+  MBB->normalizeSuccWeights();
+
+  SmallVector<uint32_t, 8> 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;
-  }
+                                              E = MBB->succ_end();
+       I != E; ++I)
+    Weights.push_back(getEdgeWeight(MBB, I));
+
+  return std::accumulate(Weights.begin(), Weights.end(), 0u);
 
-  // If the computed sum fits in 32-bits, we're done.
-  if (Sum <= UINT32_MAX)
-    return Sum;
-
-  // 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;
+uint32_t
+MachineBranchProbabilityInfo::getSumForBlock(const MachineBasicBlock *MBB,
+                                             uint32_t &Scale) const {
+  SmallVector<uint32_t, 8> 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 / Scale;
-  }
-  assert(Sum <= UINT32_MAX);
-  return Sum;
+                                              E = MBB->succ_end();
+       I != E; ++I)
+    Weights.push_back(getEdgeWeight(MBB, I));
+
+  if (MBB->succWeightsNormalized())
+    Scale = 1;
+  else
+    Scale = MachineBranchProbabilityInfo::normalizeEdgeWeights(Weights);
+  return std::accumulate(Weights.begin(), Weights.end(), 0u);
 }
 
 uint32_t MachineBranchProbabilityInfo::
 getEdgeWeight(const MachineBasicBlock *Src,
               MachineBasicBlock::const_succ_iterator Dst) const {
   uint32_t Weight = Src->getSuccWeight(Dst);
   if (!Weight)
     return DEFAULT_WEIGHT;