Replace weights by probabilities in BPI.

[congh]

Summary

This patch removes all weight-related interfaces from BPI and replace them by probability versions. With this patch, we won't use edge weight anymore in either IR or MC passes. Edge probabilitiy is a better representation in terms of CFG update and validation.

[davidxl]

Very nice changes -- I will review the patch in more details shortly..

[davidxl]

Very nice changes -- I will review the patch in more details shortly..

[davidxl]

lib/Analysis/BranchProbabilityInfo.cpp

143–146

Can we unify the code a little more:

BranchProbability UnreachableProb =

( ReachableEdges.empty() ?
        BranchProbability(1, UnreachableEdges.size()) :
        BranchProbability(UR_TAKEN_WEIGHT, (....) * UnreachableEdges.size()));

BranchProbability ReachableProb = (UR_NOTAKEN_WEIGHT, ReachableEdges.size() * ...);

for (...: UnreachableEdges)

set...

for (...)

set ...

265–268

Similarly, the code can be simplified without special case..

349

No need to guard.

352

No need to guard.

355

No need to guard.

363

The original code will adjust the weight if it is too low. Do we need to check Prob here too?

588–591

Is it better just call getEdgeProbability(Src, *I)

[congh]

lib/Analysis/BranchProbabilityInfo.cpp

143–146

OK. We can set UnreachableProb to 100% when ReachableEdges is empty.

265–268

Done.

349

This is needed to keep the original behavior: when BackEdges is empty, we don't want to normalize probabilities including the BP in this branch. Part of the reason is that we don't have a function to normalize all BB's successors' probabilities (although we have one for MBB).

352

Ditto.

355

Ditto.

363

I think this is unnecessary: the weight can become too small (like zero) after being divided but we don't have this issue in probabilities.

588–591

We need to take care of the case MapI == Probs.end() so we cannot call getEdgeProbability(Src, *I).

[congh]

lib/Analysis/BranchProbabilityInfo.cpp

143–146

I later found that I have to take care of the non-zero divider case here. To make the code more clear, I think the early exit solution might be better.

[congh]

Update on David'd comments.

[congh]

Remove an unnecessary std::max<uint32_t>.

[davidxl]

lib/Analysis/BranchProbabilityInfo.cpp

143–147

Instead of using division later, why not use the multiplication suggestion in my previous comment? The weight is small enough that there is no risk of overflowing.

349

Does it matter here? As you can see it is guarded later -- so there is no functional change here. Removing the guard makes code look cleaner

[congh]

lib/Analysis/BranchProbabilityInfo.cpp

143–147

OK.

349

Here is an example: suppose BackEdges is empty and the other two are not. Then the current behavior is that we will have 0, 124/128 and 4/128 in Probs, which will not change after normalization. If we remove those branches , we will have 124/128, 124/128, and 4/128 in Probs, and after normalization they will become roughly 1/2, 1/2, and a very small prob. Note that we won't normalize probabilities that are assigned to successors anymore. Those two lists are quite different and will result in different behaviors.

[congh]

Update on David's comments.

[davidxl]

lib/Analysis/BranchProbabilityInfo.cpp

349

Ok, I now see the normalization is done before prob assignment. However I think we can get rid of the normalization completely by computing the right denorm value in the first place. Say the number of back edges is B, number of In edges is I, and number of exiting edges are E, the denorm is

D = (B? TAKEN_WEIGHT: 0) + (I ? TAKEN_WEIGHT : 0) + (E ? NON_TAKEN_WEIGHT : 0)

So a backedge prob (if exists) is:

Prob(backedge) = TAKEN_WEIGHT/(D*B)

[congh]

Update the patch according to David's comment.

[congh]

lib/Analysis/BranchProbabilityInfo.cpp

349

OK. I have updated the patch accordingly.

[davidxl]

lgtm.

include/llvm/Analysis/BranchProbabilityInfo.h

public:
   /// \brief Print an edge's probability.
   ///
   /// Retrieves an edge's probability similarly to \see getEdgeProbability, but
   /// then prints that probability to the provided stream. That stream is then
   /// returned.
   raw_ostream &printEdgeProbability(raw_ostream &OS, const BasicBlock *Src,
                                     const BasicBlock *Dst) const;
 
-  /// \brief Get the raw edge weight calculated for the edge.
+  /// \brief Set the raw edge probability for the given edge.
   ///
-  /// This returns the raw edge weight. It is guaranteed to fall between 1 and
-  /// UINT32_MAX. Note that the raw edge weight is not meaningful in isolation.
-  /// This interface should be very carefully, and primarily by routines that
-  /// are updating the analysis by later calling setEdgeWeight.
-  uint32_t getEdgeWeight(const BasicBlock *Src,
-                         unsigned IndexInSuccessors) const;
-
-  /// \brief Get the raw edge weight calculated for the block pair.
-  ///
-  /// This returns the sum of all raw edge weights from Src to Dst.
-  /// It is guaranteed to fall between 1 and UINT32_MAX.
-  uint32_t getEdgeWeight(const BasicBlock *Src, const BasicBlock *Dst) const;
-
-  uint32_t getEdgeWeight(const BasicBlock *Src,
-                         succ_const_iterator Dst) const;
-
-  /// \brief Set the raw edge weight for a given edge.
-  ///
-  /// This allows a pass to explicitly set the edge weight for an edge. It can
-  /// be used when updating the CFG to update and preserve the branch
+  /// This allows a pass to explicitly set the edge probability for an edge. It
+  /// can be used when updating the CFG to update and preserve the branch
   /// probability information. Read the implementation of how these edge
-  /// weights are calculated carefully before using!
-  void setEdgeWeight(const BasicBlock *Src, unsigned IndexInSuccessors,
-                     uint32_t Weight);
-
-  static uint32_t getBranchWeightStackProtector(bool IsLikely) {
-    return IsLikely ? (1u << 20) - 1 : 1;
-  }
+  /// probabilities are calculated carefully before using!
+  void setEdgeProbability(const BasicBlock *Src, unsigned IndexInSuccessors,
+                          BranchProbability Prob);
 
   static BranchProbability getBranchProbStackProtector(bool IsLikely) {
     static const BranchProbability LikelyProb((1u << 20) - 1, 1u << 20);
     return IsLikely ? LikelyProb : LikelyProb.getCompl();
   }
 
   void calculate(Function &F, const LoopInfo& LI);
 
 private:
   // Since we allow duplicate edges from one basic block to another, we use
   // a pair (PredBlock and an index in the successors) to specify an edge.
   typedef std::pair<const BasicBlock *, unsigned> Edge;
 
   // Default weight value. Used when we don't have information about the edge.
   // TODO: DEFAULT_WEIGHT makes sense during static predication, when none of
   // the successors have a weight yet. But it doesn't make sense when providing
   // weight to an edge that may have siblings with non-zero weights. This can
   // be handled various ways, but it's probably fine for an edge with unknown
   // weight to just "inherit" the non-zero weight of an adjacent successor.
   static const uint32_t DEFAULT_WEIGHT = 16;
 
-  DenseMap<Edge, uint32_t> Weights;
+  DenseMap<Edge, BranchProbability> Probs;
 
   /// \brief Track the last function we run over for printing.
   Function *LastF;
 
   /// \brief Track the set of blocks directly succeeded by a returning block.
   SmallPtrSet<BasicBlock *, 16> PostDominatedByUnreachable;
 
   /// \brief Track the set of blocks that always lead to a cold call.
   SmallPtrSet<BasicBlock *, 16> PostDominatedByColdCall;
 
-  /// \brief Get sum of the block successors' weights.
-  uint32_t getSumForBlock(const BasicBlock *BB) const;
-
   bool calcUnreachableHeuristics(BasicBlock *BB);
   bool calcMetadataWeights(BasicBlock *BB);
   bool calcColdCallHeuristics(BasicBlock *BB);
   bool calcPointerHeuristics(BasicBlock *BB);
   bool calcLoopBranchHeuristics(BasicBlock *BB, const LoopInfo &LI);
   bool calcZeroHeuristics(BasicBlock *BB);
   bool calcFloatingPointHeuristics(BasicBlock *BB);
   bool calcInvokeHeuristics(BasicBlock *BB);

include/llvm/Support/BranchProbability.h

static BranchProbability getBranchProbability(uint64_t Numerator,
                                                 uint64_t Denominator);
 
   // Normalize given probabilties so that the sum of them becomes approximate
   // one.
   template <class ProbabilityIter>
   static void normalizeProbabilities(ProbabilityIter Begin,
                                      ProbabilityIter End);
 
-  // Normalize a list of weights by scaling them down so that the sum of them
-  // doesn't exceed UINT32_MAX.
-  template <class WeightListIter>
-  static void normalizeEdgeWeights(WeightListIter Begin, WeightListIter End);
-
   uint32_t getNumerator() const { return N; }
   static uint32_t getDenominator() { return D; }
 
   // Return (1 - Probability).
   BranchProbability getCompl() const { return BranchProbability(D - N); }
 
   raw_ostream &print(raw_ostream &OS) const;
 
if (Sum == 0) {
     std::fill(Begin, End, BP);
     return;
   }
 
   for (auto I = Begin; I != End; ++I)
     I->N = (I->N * uint64_t(D) + Sum / 2) / Sum;
 }
 
-template <class WeightListIter>
-void BranchProbability::normalizeEdgeWeights(WeightListIter Begin,
-                                             WeightListIter End) {
-  // First we compute the sum with 64-bits of precision.
-  uint64_t Sum = std::accumulate(Begin, End, uint64_t(0));
-
-  if (Sum > UINT32_MAX) {
-    // 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 I = Begin; I != End; ++I)
-      *I /= Scale;
-    Sum = std::accumulate(Begin, End, uint64_t(0));
-  }
-
-  // Eliminate zero weights.
-  auto ZeroWeightNum = std::count(Begin, End, 0u);
-  if (ZeroWeightNum > 0) {
-    // If all weights are zeros, replace them by 1.
-    if (Sum == 0)
-      std::fill(Begin, End, 1u);
-    else {
-      // We are converting zeros into ones, and here we need to make sure that
-      // after this the sum won't exceed UINT32_MAX.
-      if (Sum + ZeroWeightNum > UINT32_MAX) {
-        for (auto I = Begin; I != End; ++I)
-          *I /= 2;
-        ZeroWeightNum = std::count(Begin, End, 0u);
-        Sum = std::accumulate(Begin, End, uint64_t(0));
-      }
-      // Scale up non-zero weights and turn zero weights into ones.
-      uint64_t ScalingFactor = (UINT32_MAX - ZeroWeightNum) / Sum;
-      assert(ScalingFactor >= 1);
-      if (ScalingFactor > 1)
-        for (auto I = Begin; I != End; ++I)
-          *I *= ScalingFactor;
-      std::replace(Begin, End, 0u, 1u);
-    }
-  }
-}
-
 }
 
 #endif

lib/Analysis/BranchProbabilityInfo.cpp

 static const uint32_t IH_TAKEN_WEIGHT = 1024 * 1024 - 1;
 
 /// \brief Invoke-terminating normal branch not-taken weight.
 ///
 /// This is the weight for branching to the unwind destination of an invoke
 /// instruction. This is essentially never taken.
 static const uint32_t IH_NONTAKEN_WEIGHT = 1;
 
-// Standard weight value. Used when none of the heuristics set weight for
-// the edge.
-static const uint32_t NORMAL_WEIGHT = 16;
-
-// Minimum weight of an edge. Please note, that weight is NEVER 0.
-static const uint32_t MIN_WEIGHT = 1;
-
 /// \brief Calculate edge weights for successors lead to unreachable.
 ///
 /// Predict that a successor which leads necessarily to an
 /// unreachable-terminated block as extremely unlikely.
 bool BranchProbabilityInfo::calcUnreachableHeuristics(BasicBlock *BB) {
   TerminatorInst *TI = BB->getTerminator();
   if (TI->getNumSuccessors() == 0) {
     if (isa<UnreachableInst>(TI))
bool BranchProbabilityInfo::calcUnreachableHeuristics(BasicBlock *BB) {
   if (UnreachableEdges.size() == TI->getNumSuccessors())
     PostDominatedByUnreachable.insert(BB);
 
   // Skip probabilities if this block has a single successor or if all were
   // reachable.
   if (TI->getNumSuccessors() == 1 || UnreachableEdges.empty())
     return false;
 
-  uint32_t UnreachableWeight =
-    std::max(UR_TAKEN_WEIGHT / (unsigned)UnreachableEdges.size(), MIN_WEIGHT);
-  for (SmallVectorImpl<unsigned>::iterator I = UnreachableEdges.begin(),
-                                           E = UnreachableEdges.end();
-       I != E; ++I)
-    setEdgeWeight(BB, *I, UnreachableWeight);
-
-  if (ReachableEdges.empty())
+  if (ReachableEdges.empty()) {
+    BranchProbability Prob(1, UnreachableEdges.size());
+    for (unsigned SuccIdx : UnreachableEdges)
+      setEdgeProbability(BB, SuccIdx, Prob);
     return true;
-  uint32_t ReachableWeight =
-    std::max(UR_NONTAKEN_WEIGHT / (unsigned)ReachableEdges.size(),
-             NORMAL_WEIGHT);
-  for (SmallVectorImpl<unsigned>::iterator I = ReachableEdges.begin(),
-                                           E = ReachableEdges.end();
-       I != E; ++I)
-    setEdgeWeight(BB, *I, ReachableWeight);
+  }
+
+  BranchProbability UnreachableProb(UR_TAKEN_WEIGHT,
+                                    (UR_TAKEN_WEIGHT + UR_NONTAKEN_WEIGHT) *
+                                        UnreachableEdges.size());
+  BranchProbability ReachableProb(UR_NONTAKEN_WEIGHT,
+                                  (UR_TAKEN_WEIGHT + UR_NONTAKEN_WEIGHT) *
+                                      ReachableEdges.size());
+
+  for (unsigned SuccIdx : UnreachableEdges)
+    setEdgeProbability(BB, SuccIdx, UnreachableProb);
+  for (unsigned SuccIdx : ReachableEdges)
+    setEdgeProbability(BB, SuccIdx, ReachableProb);
 
   return true;
 }
 
 // Propagate existing explicit probabilities from either profile data or
 // 'expect' intrinsic processing.
 bool BranchProbabilityInfo::calcMetadataWeights(BasicBlock *BB) {
   TerminatorInst *TI = BB->getTerminator();
bool BranchProbabilityInfo::calcMetadataWeights(BasicBlock *BB) {
 
   // If the sum of weights does not fit in 32 bits, scale every weight down
   // accordingly.
   uint64_t ScalingFactor =
       (WeightSum > UINT32_MAX) ? WeightSum / UINT32_MAX + 1 : 1;
 
   WeightSum = 0;
   for (unsigned i = 0, e = TI->getNumSuccessors(); i != e; ++i) {
-    uint32_t W = Weights[i] / ScalingFactor;
-    WeightSum += W;
-    setEdgeWeight(BB, i, W);
+    Weights[i] /= ScalingFactor;
+    WeightSum += Weights[i];
   }
+  for (unsigned i = 0, e = TI->getNumSuccessors(); i != e; ++i)
+    setEdgeProbability(BB, i, {Weights[i], static_cast<uint32_t>(WeightSum)});
+
   assert(WeightSum <= UINT32_MAX &&
          "Expected weights to scale down to 32 bits");
 
   return true;
 }
 
 /// \brief Calculate edge weights for edges leading to cold blocks.
 ///
for (BasicBlock::iterator I = BB->begin(), E = BB->end(); I != E; ++I)
           break;
         }
   }
 
   // Skip probabilities if this block has a single successor.
   if (TI->getNumSuccessors() == 1 || ColdEdges.empty())
     return false;
 
-  uint32_t ColdWeight =
-      std::max(CC_TAKEN_WEIGHT / (unsigned) ColdEdges.size(), MIN_WEIGHT);
-  for (SmallVectorImpl<unsigned>::iterator I = ColdEdges.begin(),
-                                           E = ColdEdges.end();
-       I != E; ++I)
-    setEdgeWeight(BB, *I, ColdWeight);
-
-  if (NormalEdges.empty())
+  if (NormalEdges.empty()) {
+    BranchProbability Prob(1, ColdEdges.size());
+    for (unsigned SuccIdx : ColdEdges)
+      setEdgeProbability(BB, SuccIdx, Prob);
     return true;
-  uint32_t NormalWeight = std::max(
-      CC_NONTAKEN_WEIGHT / (unsigned) NormalEdges.size(), NORMAL_WEIGHT);
-  for (SmallVectorImpl<unsigned>::iterator I = NormalEdges.begin(),
-                                           E = NormalEdges.end();
-       I != E; ++I)
-    setEdgeWeight(BB, *I, NormalWeight);
+  }
+
+  BranchProbability ColdProb(CC_TAKEN_WEIGHT,
+                             (CC_TAKEN_WEIGHT + CC_NONTAKEN_WEIGHT) *
+                                 ColdEdges.size());
+  BranchProbability NormalProb(CC_NONTAKEN_WEIGHT,
+                               (CC_TAKEN_WEIGHT + CC_NONTAKEN_WEIGHT) *
+                                   NormalEdges.size());
+
+  for (unsigned SuccIdx : ColdEdges)
+    setEdgeProbability(BB, SuccIdx, ColdProb);
+  for (unsigned SuccIdx : NormalEdges)
+    setEdgeProbability(BB, SuccIdx, NormalProb);
 
   return true;
 }
 
 // Calculate Edge Weights using "Pointer Heuristics". Predict a comparsion
 // between two pointer or pointer and NULL will fail.
 bool BranchProbabilityInfo::calcPointerHeuristics(BasicBlock *BB) {
   BranchInst * BI = dyn_cast<BranchInst>(BB->getTerminator());
bool BranchProbabilityInfo::calcPointerHeuristics(BasicBlock *BB) {
   // 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);
 
-  setEdgeWeight(BB, TakenIdx, PH_TAKEN_WEIGHT);
-  setEdgeWeight(BB, NonTakenIdx, PH_NONTAKEN_WEIGHT);
+  BranchProbability TakenProb(PH_TAKEN_WEIGHT,
+                              PH_TAKEN_WEIGHT + PH_NONTAKEN_WEIGHT);
+  setEdgeProbability(BB, TakenIdx, TakenProb);
+  setEdgeProbability(BB, NonTakenIdx, TakenProb.getCompl());
   return true;
 }
 
 // Calculate Edge Weights using "Loop Branch Heuristics". Predict backedges
 // as taken, exiting edges as not-taken.
 bool BranchProbabilityInfo::calcLoopBranchHeuristics(BasicBlock *BB,
                                                      const LoopInfo &LI) {
   Loop *L = LI.getLoopFor(BB);
else if (L->getHeader() == *I)
       BackEdges.push_back(I.getSuccessorIndex());
     else
       InEdges.push_back(I.getSuccessorIndex());
   }
 
   if (BackEdges.empty() && ExitingEdges.empty())
     return false;
 
+  // Collect the sum of probabilities of back-edges/in-edges/exiting-edges, and
+  // normalize them so that they sum up to one.
+  SmallVector<BranchProbability, 4> Probs(3, BranchProbability::getZero());
+  if (!BackEdges.empty())
+    Probs[0] = BranchProbability(LBH_TAKEN_WEIGHT,
+                                 LBH_TAKEN_WEIGHT + LBH_NONTAKEN_WEIGHT);
+  if (!InEdges.empty())
+    Probs[1] = BranchProbability(LBH_TAKEN_WEIGHT,
+                                 LBH_TAKEN_WEIGHT + LBH_NONTAKEN_WEIGHT);
+  if (!ExitingEdges.empty())
+    Probs[2] = BranchProbability(LBH_NONTAKEN_WEIGHT,
+                                 LBH_TAKEN_WEIGHT + LBH_NONTAKEN_WEIGHT);
+  BranchProbability::normalizeProbabilities(Probs.begin(), Probs.end());
+
   if (uint32_t numBackEdges = BackEdges.size()) {
-    uint32_t backWeight = LBH_TAKEN_WEIGHT / numBackEdges;
-    if (backWeight < NORMAL_WEIGHT)
-      backWeight = NORMAL_WEIGHT;
-
-    for (SmallVectorImpl<unsigned>::iterator EI = BackEdges.begin(),
-         EE = BackEdges.end(); EI != EE; ++EI) {
-      setEdgeWeight(BB, *EI, backWeight);
-    }
+    auto Prob = Probs[0] / numBackEdges;
+    for (unsigned SuccIdx : BackEdges)
+      setEdgeProbability(BB, SuccIdx, Prob);
   }
 
   if (uint32_t numInEdges = InEdges.size()) {
-    uint32_t inWeight = LBH_TAKEN_WEIGHT / numInEdges;
-    if (inWeight < NORMAL_WEIGHT)
-      inWeight = NORMAL_WEIGHT;
-
-    for (SmallVectorImpl<unsigned>::iterator EI = InEdges.begin(),
-         EE = InEdges.end(); EI != EE; ++EI) {
-      setEdgeWeight(BB, *EI, inWeight);
-    }
+    auto Prob = Probs[1] / numInEdges;
+    for (unsigned SuccIdx : InEdges)
+      setEdgeProbability(BB, SuccIdx, Prob);
   }
 
   if (uint32_t numExitingEdges = ExitingEdges.size()) {
-    uint32_t exitWeight = LBH_NONTAKEN_WEIGHT / numExitingEdges;
-    if (exitWeight < MIN_WEIGHT)
-      exitWeight = MIN_WEIGHT;
-
-    for (SmallVectorImpl<unsigned>::iterator EI = ExitingEdges.begin(),
-         EE = ExitingEdges.end(); EI != EE; ++EI) {
-      setEdgeWeight(BB, *EI, exitWeight);
-    }
+    auto Prob = Probs[2] / numExitingEdges;
+    for (unsigned SuccIdx : ExitingEdges)
+      setEdgeProbability(BB, SuccIdx, Prob);
   }
 
   return true;
 }
 
 bool BranchProbabilityInfo::calcZeroHeuristics(BasicBlock *BB) {
   BranchInst * BI = dyn_cast<BranchInst>(BB->getTerminator());
   if (!BI || !BI->isConditional())
if (CV->isZero()) {
     return false;
   }
 
   unsigned TakenIdx = 0, NonTakenIdx = 1;
 
   if (!isProb)
     std::swap(TakenIdx, NonTakenIdx);
 
-  setEdgeWeight(BB, TakenIdx, ZH_TAKEN_WEIGHT);
-  setEdgeWeight(BB, NonTakenIdx, ZH_NONTAKEN_WEIGHT);
-
+  BranchProbability TakenProb(ZH_TAKEN_WEIGHT,
+                              ZH_TAKEN_WEIGHT + ZH_NONTAKEN_WEIGHT);
+  setEdgeProbability(BB, TakenIdx, TakenProb);
+  setEdgeProbability(BB, NonTakenIdx, TakenProb.getCompl());
   return true;
 }
 
 bool BranchProbabilityInfo::calcFloatingPointHeuristics(BasicBlock *BB) {
   BranchInst *BI = dyn_cast<BranchInst>(BB->getTerminator());
   if (!BI || !BI->isConditional())
     return false;
 
if (FCmp->isEquality()) {
     return false;
   }
 
   unsigned TakenIdx = 0, NonTakenIdx = 1;
 
   if (!isProb)
     std::swap(TakenIdx, NonTakenIdx);
 
-  setEdgeWeight(BB, TakenIdx, FPH_TAKEN_WEIGHT);
-  setEdgeWeight(BB, NonTakenIdx, FPH_NONTAKEN_WEIGHT);
-
+  BranchProbability TakenProb(FPH_TAKEN_WEIGHT,
+                              FPH_TAKEN_WEIGHT + FPH_NONTAKEN_WEIGHT);
+  setEdgeProbability(BB, TakenIdx, TakenProb);
+  setEdgeProbability(BB, NonTakenIdx, TakenProb.getCompl());
   return true;
 }
 
 bool BranchProbabilityInfo::calcInvokeHeuristics(BasicBlock *BB) {
   InvokeInst *II = dyn_cast<InvokeInst>(BB->getTerminator());
   if (!II)
     return false;
 
-  setEdgeWeight(BB, 0/*Index for Normal*/, IH_TAKEN_WEIGHT);
-  setEdgeWeight(BB, 1/*Index for Unwind*/, IH_NONTAKEN_WEIGHT);
+  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());
   return true;
 }
 
 void BranchProbabilityInfo::releaseMemory() {
-  Weights.clear();
+  Probs.clear();
 }
 
 void BranchProbabilityInfo::print(raw_ostream &OS) const {
   OS << "---- Branch Probabilities ----\n";
   // We print the probabilities from the last function the analysis ran over,
   // or the function it is currently running over.
   assert(LastF && "Cannot print prior to running over a function");
   for (const auto &BI : *LastF) {
     for (succ_const_iterator SI = succ_begin(&BI), SE = succ_end(&BI); SI != SE;
          ++SI) {
       printEdgeProbability(OS << "  ", &BI, *SI);
     }
   }
 }
 
-uint32_t BranchProbabilityInfo::getSumForBlock(const BasicBlock *BB) const {
-  uint32_t Sum = 0;
-
-  for (succ_const_iterator I = succ_begin(BB), E = succ_end(BB); I != E; ++I) {
-    uint32_t Weight = getEdgeWeight(BB, I.getSuccessorIndex());
-    uint32_t PrevSum = Sum;
-
-    Sum += Weight;
-    assert(Sum >= PrevSum); (void) PrevSum;
-  }
-
-  return Sum;
-}
-
 bool BranchProbabilityInfo::
 isEdgeHot(const BasicBlock *Src, const BasicBlock *Dst) const {
   // Hot probability is at least 4/5 = 80%
   // FIXME: Compare against a static "hot" BranchProbability.
   return getEdgeProbability(Src, Dst) > BranchProbability(4, 5);
 }
 
 BasicBlock *BranchProbabilityInfo::getHotSucc(BasicBlock *BB) const {
-  uint32_t Sum = 0;
-  uint32_t MaxWeight = 0;
+  auto MaxProb = BranchProbability::getZero();
   BasicBlock *MaxSucc = nullptr;
 
   for (succ_iterator I = succ_begin(BB), E = succ_end(BB); I != E; ++I) {
     BasicBlock *Succ = *I;
-    uint32_t Weight = getEdgeWeight(BB, Succ);
-    uint32_t PrevSum = Sum;
-
-    Sum += Weight;
-    assert(Sum > PrevSum); (void) PrevSum;
-
-    if (Weight > MaxWeight) {
-      MaxWeight = Weight;
+    auto Prob = getEdgeProbability(BB, Succ);
+    if (Prob > MaxProb) {
+      MaxProb = Prob;
       MaxSucc = Succ;
     }
   }
 
   // Hot probability is at least 4/5 = 80%
-  if (BranchProbability(MaxWeight, Sum) > BranchProbability(4, 5))
+  if (MaxProb > BranchProbability(4, 5))
     return MaxSucc;
 
   return nullptr;
 }
 
-/// Get the raw edge weight for the edge. If can't find it, return
-/// DEFAULT_WEIGHT value. Here an edge is specified using PredBlock and an index
-/// to the successors.
-uint32_t BranchProbabilityInfo::
-getEdgeWeight(const BasicBlock *Src, unsigned IndexInSuccessors) const {
-  DenseMap<Edge, uint32_t>::const_iterator I =
-      Weights.find(std::make_pair(Src, IndexInSuccessors));
+/// Get the raw edge probability for the edge. If can't find it, return a
+/// default probability 1/N where N is the number of successors. Here an edge is
+/// specified using PredBlock and an
+/// index to the successors.
+BranchProbability
+BranchProbabilityInfo::getEdgeProbability(const BasicBlock *Src,
+                                          unsigned IndexInSuccessors) const {
+  auto I = Probs.find(std::make_pair(Src, IndexInSuccessors));
 
-  if (I != Weights.end())
+  if (I != Probs.end())
     return I->second;
 
-  return DEFAULT_WEIGHT;
+  return {1,
+          static_cast<uint32_t>(std::distance(succ_begin(Src), succ_end(Src)))};
 }
 
-uint32_t BranchProbabilityInfo::getEdgeWeight(const BasicBlock *Src,
-                                              succ_const_iterator Dst) const {
-  return getEdgeWeight(Src, Dst.getSuccessorIndex());
-}
-
-/// Get the raw edge weight calculated for the block pair. This returns the sum
-/// of all raw edge weights from Src to Dst.
-uint32_t BranchProbabilityInfo::
-getEdgeWeight(const BasicBlock *Src, const BasicBlock *Dst) const {
-  uint32_t Weight = 0;
-  bool FoundWeight = false;
-  DenseMap<Edge, uint32_t>::const_iterator MapI;
-  for (succ_const_iterator I = succ_begin(Src), E = succ_end(Src); I != E; ++I)
-    if (*I == Dst) {
-      MapI = Weights.find(std::make_pair(Src, I.getSuccessorIndex()));
-      if (MapI != Weights.end()) {
-        FoundWeight = true;
-        Weight += MapI->second;
-      }
-    }
-  return (!FoundWeight) ? DEFAULT_WEIGHT : Weight;
-}
-
-/// Set the edge weight for a given edge specified by PredBlock and an index
-/// to the successors.
-void BranchProbabilityInfo::
-setEdgeWeight(const BasicBlock *Src, unsigned IndexInSuccessors,
-              uint32_t Weight) {
-  Weights[std::make_pair(Src, IndexInSuccessors)] = Weight;
-  DEBUG(dbgs() << "set edge " << Src->getName() << " -> "
-               << IndexInSuccessors << " successor weight to "
-               << Weight << "\n");
-}
-
-/// Get an edge's probability, relative to other out-edges from Src.
-BranchProbability BranchProbabilityInfo::
-getEdgeProbability(const BasicBlock *Src, unsigned IndexInSuccessors) const {
-  uint32_t N = getEdgeWeight(Src, IndexInSuccessors);
-  uint32_t D = getSumForBlock(Src);
-
-  // It is possible that the edge weight on the only successor edge of Src is
-  // zero, in which case we return 100%.
-  if (N == 0 && D == 0)
-    return BranchProbability::getOne();
-
-  return BranchProbability(N, D);
-}
-
-/// Get the probability of going from Src to Dst. It returns the sum of all
-/// probabilities for edges from Src to Dst.
-BranchProbability BranchProbabilityInfo::
-getEdgeProbability(const BasicBlock *Src, const BasicBlock *Dst) const {
-
-  uint32_t N = getEdgeWeight(Src, Dst);
-  uint32_t D = getSumForBlock(Src);
-
-  // It is possible that the edge weight on the only successor edge of Src is
-  // zero, in which case we return 100%.
-  if (N == 0 && D == 0)
-    return BranchProbability::getOne();
-
-  return BranchProbability(N, D);
+BranchProbability
+BranchProbabilityInfo::getEdgeProbability(const BasicBlock *Src,
+  return getEdgeProbability(Src, Dst.getSuccessorIndex());
 }
 
+/// Get the raw edge probability calculated for the block pair. This returns the
+/// sum of all raw edge probabilities from Src to Dst.
 BranchProbability
 BranchProbabilityInfo::getEdgeProbability(const BasicBlock *Src,
-                                          succ_const_iterator Dst) const {
-  return getEdgeProbability(Src, Dst.getSuccessorIndex());
+                                          const BasicBlock *Dst) const {
+  auto Prob = BranchProbability::getZero();
+  bool FoundProb = false;
+  for (succ_const_iterator I = succ_begin(Src), E = succ_end(Src); I != E; ++I)
+    if (*I == Dst) {
+      auto MapI = Probs.find(std::make_pair(Src, I.getSuccessorIndex()));
+      if (MapI != Probs.end()) {
+        FoundProb = true;
+        Prob += MapI->second;
+      }
+    }
+  uint32_t succ_num = std::distance(succ_begin(Src), succ_end(Src));
+  return FoundProb ? Prob : BranchProbability(1, succ_num);
+}
+
+/// Set the edge probability for a given edge specified by PredBlock and an
+/// index to the successors.
+void BranchProbabilityInfo::setEdgeProbability(const BasicBlock *Src,
+                                               unsigned IndexInSuccessors,
+                                               BranchProbability Prob) {
+  Probs[std::make_pair(Src, IndexInSuccessors)] = Prob;
+  DEBUG(dbgs() << "set edge " << Src->getName() << " -> " << IndexInSuccessors
+               << " successor probability to " << Prob << "\n");
 }
 
 raw_ostream &
 BranchProbabilityInfo::printEdgeProbability(raw_ostream &OS,
                                             const BasicBlock *Src,
                                             const BasicBlock *Dst) const {
 
   const BranchProbability Prob = getEdgeProbability(Src, Dst);

lib/CodeGen/StackProtector.cpp

if (SupportsSelectionDAGSP) {
       // basic block so that it's in the "fall through" position.
       NewBB->moveAfter(BB);
 
       // Generate the stack protector instructions in the old basic block.
       IRBuilder<> B(BB);
       LoadInst *LI1 = B.CreateLoad(StackGuardVar);
       LoadInst *LI2 = B.CreateLoad(AI);
       Value *Cmp = B.CreateICmpEQ(LI1, LI2);
-      unsigned SuccessWeight =
-          BranchProbabilityInfo::getBranchWeightStackProtector(true);
-      unsigned FailureWeight =
-          BranchProbabilityInfo::getBranchWeightStackProtector(false);
+      auto SuccessProb =
+          BranchProbabilityInfo::getBranchProbStackProtector(true);
+      auto FailureProb =
+          BranchProbabilityInfo::getBranchProbStackProtector(false);
       MDNode *Weights = MDBuilder(F->getContext())
-                            .createBranchWeights(SuccessWeight, FailureWeight);
+                            .createBranchWeights(SuccessProb.getNumerator(),
+                                                 FailureProb.getNumerator());
       B.CreateCondBr(Cmp, NewBB, FailBB, Weights);
     }
   }
 
   // Return if we didn't modify any basic blocks. i.e., there are no return
   // statements in the function.
   return HasPrologue;
 }

lib/Target/PowerPC/PPCISelDAGToDAG.cpp

static unsigned getBranchHint(unsigned PCC, FunctionLoweringInfo *FuncInfo,
   const BasicBlock *BB = FuncInfo->MBB->getBasicBlock();
   const TerminatorInst *BBTerm = BB->getTerminator();
 
   if (BBTerm->getNumSuccessors() != 2) return PPC::BR_NO_HINT;
 
   const BasicBlock *TBB = BBTerm->getSuccessor(0);
   const BasicBlock *FBB = BBTerm->getSuccessor(1);
 
-  uint32_t TWeight = FuncInfo->BPI->getEdgeWeight(BB, TBB);
-  uint32_t FWeight = FuncInfo->BPI->getEdgeWeight(BB, FBB);
+  auto TProb = FuncInfo->BPI->getEdgeProbability(BB, TBB);
+  auto FProb = FuncInfo->BPI->getEdgeProbability(BB, FBB);
 
   // We only want to handle cases which are easy to predict at static time, e.g.
   // C++ throw statement, that is very likely not taken, or calling never
   // returned function, e.g. stdlib exit(). So we set Threshold to filter
   // unwanted cases.
   //
   // Below is LLVM branch weight table, we only want to handle case 1, 2
   //
   // Case                  Taken:Nontaken  Example
   // 1. Unreachable        1048575:1       C++ throw, stdlib exit(),
   // 2. Invoke-terminating 1:1048575
   // 3. Coldblock          4:64            __builtin_expect
   // 4. Loop Branch        124:4           For loop
   // 5. PH/ZH/FPH          20:12
   const uint32_t Threshold = 10000;
 
-  // Minimal weight should be at least 1
-  if (std::max(TWeight, FWeight) /
-      std::max(1u, std::min(TWeight, FWeight)) < Threshold)
+  if (std::max(TProb, FProb) / Threshold < std::min(TProb, FProb))
     return PPC::BR_NO_HINT;
 
   DEBUG(dbgs() << "Use branch hint for '" << FuncInfo->Fn->getName() << "::"
                << BB->getName() << "'\n"
-               << " -> " << TBB->getName() << ": " << TWeight << "\n"
-               << " -> " << FBB->getName() << ": " << FWeight << "\n");
+               << " -> " << TBB->getName() << ": " << TProb << "\n"
+               << " -> " << FBB->getName() << ": " << FProb << "\n");
 
   const BasicBlockSDNode *BBDN = cast<BasicBlockSDNode>(DestMBB);
 
-  // If Dest BasicBlock is False-BasicBlock (FBB), swap branch weight,
-  // because we want 'TWeight' stands for 'branch weight' to Dest BasicBlock
+  // If Dest BasicBlock is False-BasicBlock (FBB), swap branch probabilities,
+  // because we want 'TProb' stands for 'branch probability' to Dest BasicBlock
   if (BBDN->getBasicBlock()->getBasicBlock() != TBB)
-    std::swap(TWeight, FWeight);
+    std::swap(TProb, FProb);
 
-  return (TWeight > FWeight) ? PPC::BR_TAKEN_HINT : PPC::BR_NONTAKEN_HINT;
+  return (TProb > FProb) ? PPC::BR_TAKEN_HINT : PPC::BR_NONTAKEN_HINT;
 }
 
 // isOpcWithIntImmediate - This method tests to see if the node is a specific
 // opcode and that it has a immediate integer right operand.
 // If so Imm will receive the 32 bit value.
 static bool isOpcWithIntImmediate(SDNode *N, unsigned Opc, unsigned& Imm) {
   return N->getOpcode() == Opc
          && isInt32Immediate(N->getOperand(1).getNode(), Imm);

lib/Transforms/Scalar/JumpThreading.cpp

void JumpThreading::UpdateBlockFreqAndEdgeWeight(BasicBlock *PredBB,
   // frequency of BB.
   auto BBOrigFreq = BFI->getBlockFreq(BB);
   auto NewBBFreq = BFI->getBlockFreq(NewBB);
   auto BB2SuccBBFreq = BBOrigFreq * BPI->getEdgeProbability(BB, SuccBB);
   auto BBNewFreq = BBOrigFreq - NewBBFreq;
   BFI->setBlockFreq(BB, BBNewFreq.getFrequency());
 
   // Collect updated outgoing edges' frequencies from BB and use them to update
-  // edge weights.
+  // edge probabilities.
   SmallVector<uint64_t, 4> BBSuccFreq;
   for (auto I = succ_begin(BB), E = succ_end(BB); I != E; ++I) {
     auto SuccFreq = (*I == SuccBB)
                         ? BB2SuccBBFreq - NewBBFreq
                         : BBOrigFreq * BPI->getEdgeProbability(BB, *I);
     BBSuccFreq.push_back(SuccFreq.getFrequency());
   }
 
-  // Normalize edge weights in Weights64 so that the sum of them can fit in
-  BranchProbability::normalizeEdgeWeights(BBSuccFreq.begin(), BBSuccFreq.end());
+  uint64_t MaxBBSuccFreq =
+      *std::max_element(BBSuccFreq.begin(), BBSuccFreq.end());
+  SmallVector<BranchProbability, 4> BBSuccProbs;
+  for (uint64_t Freq : BBSuccFreq)
+    BBSuccProbs.push_back(
+        BranchProbability::getBranchProbability(Freq, MaxBBSuccFreq));
+
+  // Normalize edge probabilities so that they sum up to one.
+  BranchProbability::normalizeProbabilities(BBSuccProbs.begin(),
+                                            BBSuccProbs.end());
+
+  // Update edge probabilities in BPI.
+  for (int I = 0, E = BBSuccProbs.size(); I < E; I++)
+    BPI->setEdgeProbability(BB, I, BBSuccProbs[I]);
 
+  if (BBSuccProbs.size() >= 2) {
     SmallVector<uint32_t, 4> Weights;
-  for (auto Freq : BBSuccFreq)
-    Weights.push_back(static_cast<uint32_t>(Freq));
+    for (auto Prob : BBSuccProbs)
+      Weights.push_back(Prob.getNumerator());
 
-  // Update edge weights in BPI.
-  for (int I = 0, E = Weights.size(); I < E; I++)
-    BPI->setEdgeWeight(BB, I, Weights[I]);
-
-  if (Weights.size() >= 2) {
     auto TI = BB->getTerminator();
     TI->setMetadata(
         LLVMContext::MD_prof,
         MDBuilder(TI->getParent()->getContext()).createBranchWeights(Weights));
   }
 }
 
 /// DuplicateCondBranchOnPHIIntoPred - PredBB contains an unconditional branch

test/Analysis/BranchProbabilityInfo/noreturn.ll

 
 define i32 @test2(i32 %a, i32 %b) {
 ; CHECK: Printing analysis {{.*}} for function 'test2'
 entry:
   switch i32 %a, label %exit [i32 1, label %case_a
                               i32 2, label %case_b
                               i32 3, label %case_c
                               i32 4, label %case_d]
-; CHECK: edge entry -> exit probability is 0x7fffe000 / 0x80000000 = 100.00% [HOT edge]
-; CHECK: edge entry -> case_a probability is 0x00000800 / 0x80000000 = 0.00%
-; CHECK: edge entry -> case_b probability is 0x00000800 / 0x80000000 = 0.00%
-; CHECK: edge entry -> case_c probability is 0x00000800 / 0x80000000 = 0.00%
-; CHECK: edge entry -> case_d probability is 0x00000800 / 0x80000000 = 0.00%
+; CHECK: edge entry -> exit probability is 0x7ffff800 / 0x80000000 = 100.00% [HOT edge]
+; CHECK: edge entry -> case_a probability is 0x00000200 / 0x80000000 = 0.00%
+; CHECK: edge entry -> case_b probability is 0x00000200 / 0x80000000 = 0.00%
+; CHECK: edge entry -> case_c probability is 0x00000200 / 0x80000000 = 0.00%
+; CHECK: edge entry -> case_d probability is 0x00000200 / 0x80000000 = 0.00%
 
 case_a:
   br label %case_b
 
 case_b:
   br label %case_c
 
 case_c:

test/Transforms/JumpThreading/update-edge-weight.ll

 ; RUN: opt -S -jump-threading %s | FileCheck %s
 
 ; Test if edge weights are properly updated after jump threading.
 
-; CHECK: !2 = !{!"branch_weights", i32 22, i32 7}
+; CHECK: !2 = !{!"branch_weights", i32 1629125526, i32 518358122}
 
 define void @foo(i32 %n) !prof !0 {
 entry:
   %cmp = icmp sgt i32 %n, 10
   br i1 %cmp, label %if.then.1, label %if.else.1, !prof !1
 
 if.then.1:
   tail call void @a()