[BrachProbablityInfo] Set edge probabilities at once.

[yrouban]

Summary

Hide the method that allows setting probability for particular edge and introduce a public method that sets probabilities for all outgoing edges at once.
Setting individual edge probability is error prone. More over it is difficult to check that the total probability is 1.0 because there is no easy way to know when the user finished setting all the probabilities.

One bug is fixed in BranchProbabilityInfo::calcMetadataWeights(). Changing unreachable branch probabilities to raw(1) and distributing the rest (oldProbability - raw(1)) over the reachable branches could introduce total probability inaccuracy bigger than 1/numOfBranches.

[hjyamauchi]

llvm/include/llvm/Analysis/BranchProbabilityInfo.h

135

Can you add a short general comment for this public setEdgeProbability?

llvm/lib/Analysis/BranchProbabilityInfo.cpp

1001–1002

Looks like simply checking "TotalNumerator == BranchProbability::getDenominator()" won't work probably because of rounding errors. Can you add a brief comment?

[yrouban]

added comments

[ebrevnov]

I think the new API of setting probabilities for all edges "atomically" is a good move. There is one suggestion though. In fact, there is no need to "repack" vector of probabilities into map. We can simply copy or even move (in most cases we have rvalue) entire vector into a map. This can be done as a separate change though.

[hjyamauchi]

LGTM.

[rnk]

llvm/lib/Analysis/BranchProbabilityInfo.cpp

1008

This assert doesn't appear to hold:
https://ci.chromium.org/p/chrome/builders/ci/ToTWin/6284
Assertion failed: TotalNumerator >= BranchProbability::getDenominator() - Probs.size(), file C:\b\s\w\ir\cache\builder\src\third_party\llvm\llvm\lib\Analysis\BranchProbabilityInfo.cpp, line 995

I will gather more information.

[rnk]

llvm/lib/Analysis/BranchProbabilityInfo.cpp

1008

These are the relevant values:

TotalNumerator 2147483624
BranchProbability::getDenominator() 2147483648
Probs.size() 21

This is the IR switch:

switch i32 %0, label %sw.epilog.i [
  i32 20, label %sw.bb20.i
  i32 1, label %sw.bb1.i
  i32 2, label %sw.bb2.i
  i32 3, label %sw.bb3.i
  i32 4, label %sw.bb4.i
  i32 5, label %sw.bb5.i
  i32 6, label %sw.bb6.i
  i32 11, label %sw.bb7.i
  i32 12, label %sw.bb8.i
  i32 13, label %sw.bb9.i
  i32 7, label %sw.bb10.i
  i32 9, label %"?SkColorTypeToGrColorType@@YA?AW4GrColorType@@W4SkColorType@@@Z.exit"
  i32 8, label %sw.bb12.i
  i32 10, label %"?SkColorTypeToGrColorType@@YA?AW4GrColorType@@W4SkColorType@@@Z.exit"
  i32 14, label %sw.bb14.i
  i32 15, label %sw.bb15.i
  i32 18, label %sw.bb16.i
  i32 19, label %sw.bb17.i
  i32 16, label %sw.bb18.i
  i32 17, label %sw.bb19.i
], !dbg !77, !prof !81

I think it comes from this source construct:
https://source.chromium.org/chromium/chromium/src/+/master:third_party/skia/include/private/GrTypesPriv.h;drc=2f11470d7ad8963a9add116df64d2edd1b85d3a4;l=869?originalUrl=https:%2F%2Fcs.chromium.org%2F
Inlined into this function:
https://source.chromium.org/chromium/chromium/src/+/master:third_party/skia/src/image/SkSurface_Gpu.cpp;l=342?q=MakeRenderTarget&ss=chromium&originalUrl=https:%2F%2Fcs.chromium.org%2F

I will revert for now, since this prevents us from seeing other ToT issues on PGO builders, but feel free to reland with a fix.

[yrouban]

llvm/lib/Analysis/BranchProbabilityInfo.cpp

1008

please give me the !prof branch_weights so I can reproduce.

[rnk]

llvm/lib/Analysis/BranchProbabilityInfo.cpp

1008

<0xd3589a0> = !{!"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}

Dumping them produced this:
0x00000001 / 0x80000000 = 0.00%
0x004761ef / 0x80000000 = 0.22%
0x004761ef / 0x80000000 = 0.22%
0x004761ef / 0x80000000 = 0.22%
0x004761ef / 0x80000000 = 0.22%
0x77c7e9e7 / 0x80000000 = 93.58%
0x004761ef / 0x80000000 = 0.22%
0x03333332 / 0x80000000 = 2.50%
0x004761ef / 0x80000000 = 0.22%
0x004761ef / 0x80000000 = 0.22%
0x004761ef / 0x80000000 = 0.22%
0x004761ef / 0x80000000 = 0.22%
0x004761ef / 0x80000000 = 0.22%
0x004761ef / 0x80000000 = 0.22%
0x004761ef / 0x80000000 = 0.22%
0x004761ef / 0x80000000 = 0.22%
0x004761ef / 0x80000000 = 0.22%
0x004761ef / 0x80000000 = 0.22%
0x004761ef / 0x80000000 = 0.22%
0x004761ef / 0x80000000 = 0.22%
0x004761ef / 0x80000000 = 0.22%

[yrouban]

llvm/lib/Analysis/BranchProbabilityInfo.cpp

1008

Great! we have found a bug in calcMetadataWeights(). I will fix and add a test case.

[yrouban]

Fixed bug in BranchProbabilityInfo::calcMetadataWeights() to keep total probability close to 1.0. Added a new testcase for this and fixed another one.

[yrouban]

Please review again.

[yrouban]

llvm/lib/Analysis/BranchProbabilityInfo.cpp

369–378

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:

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)

[ebrevnov]

llvm/lib/Analysis/BranchProbabilityInfo.cpp

362

It's really hard to parse the comment since it heavily depends on old code which will be unavailable. Please try to rephrase.

369–378

I think it makes sense to put that comment into the code directly

[yrouban]

In BranchProbabilityInfo::calcMetadataWeights() rephrased one comment and added one TODO.

[ebrevnov]

Thanks! LGTM.

llvm/include/llvm/Analysis/BranchProbabilityInfo.h

public:
   /// 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;
 
+protected:
   /// Set the raw edge probability for the given edge.
   ///
   /// 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
   /// probabilities are calculated carefully before using!
   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<BranchProbability> &Probs);
+
   static BranchProbability getBranchProbStackProtector(bool IsLikely) {
     static const BranchProbability LikelyProb((1u << 20) - 1, 1u << 20);
     return IsLikely ? LikelyProb : LikelyProb.getCompl();
   }
 
   void calculate(const Function &F, const LoopInfo &LI,
                  const TargetLibraryInfo *TLI, PostDominatorTree *PDT);
 

llvm/lib/Analysis/BranchProbabilityInfo.cpp

if (PostDominatedByUnreachable.count(*I))
       UnreachableEdges.push_back(I.getSuccessorIndex());
     else
       ReachableEdges.push_back(I.getSuccessorIndex());
 
   // Skip probabilities if all were reachable.
   if (UnreachableEdges.empty())
     return false;
 
+  SmallVector<BranchProbability, 4> 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;
   }
 
   auto UnreachableProb = UR_TAKEN_PROB;
   auto ReachableProb =
       (BranchProbability::getOne() - UR_TAKEN_PROB * UnreachableEdges.size()) /
       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;
 }
 
 // Propagate existing explicit probabilities from either profile data or
 // 'expect' intrinsic processing. Examine metadata against unreachable
 // heuristic. The probability of the edge coming to unreachable block is
 // set to min of metadata and unreachable heuristic.
 bool BranchProbabilityInfo::calcMetadataWeights(const BasicBlock *BB) {
bool BranchProbabilityInfo::calcMetadataWeights(const BasicBlock *BB) {
   // Set the probability.
   SmallVector<BranchProbability, 2> BP;
   for (unsigned i = 0, e = TI->getNumSuccessors(); i != e; ++i)
     BP.push_back({ Weights[i], static_cast<uint32_t>(WeightSum) });
 
   // 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)
+    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;
 }
 
 /// Calculate edge weights for edges leading to cold blocks.
 ///
 /// A cold block is one post-dominated by  a block with a call to a
 /// cold function.  Those edges are unlikely to be taken, so we give
if (PostDominatedByColdCall.count(*I))
       ColdEdges.push_back(I.getSuccessorIndex());
     else
       NormalEdges.push_back(I.getSuccessorIndex());
 
   // Skip probabilities if no cold edges.
   if (ColdEdges.empty())
     return false;
 
+  SmallVector<BranchProbability, 4> 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;
   }
 
   auto ColdProb = BranchProbability::getBranchProbability(
       CC_TAKEN_WEIGHT,
       (CC_TAKEN_WEIGHT + CC_NONTAKEN_WEIGHT) * uint64_t(ColdEdges.size()));
   auto NormalProb = BranchProbability::getBranchProbability(
       CC_NONTAKEN_WEIGHT,
       (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;
 }
 
 // Calculate Edge Weights using "Pointer Heuristics". Predict a comparison
 // between two pointer or pointer and NULL will fail.
 bool BranchProbabilityInfo::calcPointerHeuristics(const BasicBlock *BB) {
   const BranchInst *BI = dyn_cast<BranchInst>(BB->getTerminator());
   if (!BI || !BI->isConditional())
     return false;
 
   Value *Cond = BI->getCondition();
   ICmpInst *CI = dyn_cast<ICmpInst>(Cond);
   if (!CI || !CI->isEquality())
     return false;
 
   Value *LHS = CI->getOperand(0);
 
   if (!LHS->getType()->isPointerTy())
     return false;
 
   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<BranchProbability, 2>({TakenProb, UntakenProb}));
   return true;
 }
 
 static int getSCCNum(const BasicBlock *BB,
                      const BranchProbabilityInfo::SccInfo &SccI) {
   auto SccIt = SccI.SccNums.find(BB);
   if (SccIt == SccI.SccNums.end())
     return -1;
bool BranchProbabilityInfo::calcLoopBranchHeuristics(const BasicBlock *BB,
 
   // Collect the sum of probabilities of back-edges/in-edges/exiting-edges, and
   // normalize them so that they sum up to one.
   unsigned Denom = (BackEdges.empty() ? 0 : LBH_TAKEN_WEIGHT) +
                    (InEdges.empty() ? 0 : LBH_TAKEN_WEIGHT) +
                    (UnlikelyEdges.empty() ? 0 : LBH_UNLIKELY_WEIGHT) +
                    (ExitingEdges.empty() ? 0 : LBH_NONTAKEN_WEIGHT);
 
+  SmallVector<BranchProbability, 4> 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()) {
     BranchProbability NotTakenProb = BranchProbability(LBH_NONTAKEN_WEIGHT,
                                                        Denom);
     auto Prob = NotTakenProb / numExitingEdges;
     for (unsigned SuccIdx : ExitingEdges)
-      setEdgeProbability(BB, SuccIdx, Prob);
+      EdgeProbabilities[SuccIdx] = Prob;
   }
 
   if (uint32_t numUnlikelyEdges = UnlikelyEdges.size()) {
     BranchProbability UnlikelyProb = BranchProbability(LBH_UNLIKELY_WEIGHT,
                                                        Denom);
     auto Prob = UnlikelyProb / numUnlikelyEdges;
     for (unsigned SuccIdx : UnlikelyEdges)
-      setEdgeProbability(BB, SuccIdx, Prob);
+      EdgeProbabilities[SuccIdx] = Prob;
   }
 
+  setEdgeProbability(BB, EdgeProbabilities);
   return true;
 }
 
 bool BranchProbabilityInfo::calcZeroHeuristics(const BasicBlock *BB,
                                                const TargetLibraryInfo *TLI) {
   const BranchInst *BI = dyn_cast<BranchInst>(BB->getTerminator());
   if (!BI || !BI->isConditional())
     return false;
case CmpInst::ICMP_SGT:
       break;
     default:
       return false;
     }
   } else {
     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<BranchProbability, 2>({TakenProb, UntakenProb}));
   return true;
 }
 
 bool BranchProbabilityInfo::calcFloatingPointHeuristics(const BasicBlock *BB) {
   const BranchInst *BI = dyn_cast<BranchInst>(BB->getTerminator());
   if (!BI || !BI->isConditional())
     return false;
 
if (FCmp->isEquality()) {
     // isnan -> Unlikely
     isProb = false;
     TakenWeight = FPH_ORD_WEIGHT;
     NontakenWeight = FPH_UNO_WEIGHT;
   } else {
     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<BranchProbability, 2>({TakenProb, UntakenProb}));
   return true;
 }
 
 bool BranchProbabilityInfo::calcInvokeHeuristics(const BasicBlock *BB) {
   const InvokeInst *II = dyn_cast<InvokeInst>(BB->getTerminator());
   if (!II)
     return false;
 
   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<BranchProbability, 2>({TakenProb, TakenProb.getCompl()}));
   return true;
 }
 
 void BranchProbabilityInfo::releaseMemory() {
   Probs.clear();
   Handles.clear();
 }
 
void BranchProbabilityInfo::setEdgeProbability(const BasicBlock *Src,
                                                BranchProbability Prob) {
   Probs[std::make_pair(Src, IndexInSuccessors)] = Prob;
   Handles.insert(BasicBlockCallbackVH(Src, this));
   LLVM_DEBUG(dbgs() << "set edge " << Src->getName() << " -> "
                     << IndexInSuccessors << " successor probability to " << Prob
                     << "\n");
 }
 
+/// Set the edge probability for all edges at once.
+void BranchProbabilityInfo::setEdgeProbability(
+    const BasicBlock *Src, const SmallVectorImpl<BranchProbability> &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,
                                             const BasicBlock *Dst) const {
   const BranchProbability Prob = getEdgeProbability(Src, Dst);
   OS << "edge " << Src->getName() << " -> " << Dst->getName()
      << " probability is " << Prob
      << (isEdgeHot(Src, Dst) ? " [HOT edge]\n" : "\n");

llvm/lib/Transforms/Scalar/JumpThreading.cpp

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]);
+  BPI->setEdgeProbability(BB, BBSuccProbs);
 
   // Update the profile metadata as well.
   //
   // Don't do this if the profile of the transformed blocks was statically
   // estimated.  (This could occur despite the function having an entry
   // frequency in completely cold parts of the CFG.)
   //
   // In this case we don't want to suggest to subsequent passes that the

llvm/lib/Transforms/Utils/BreakCriticalEdges.cpp

if (ShouldUpdateAnalysis) {
            I < E; ++I)
         EdgeProbabilities.emplace_back(BPI->getEdgeProbability(Target, I));
       BPI->eraseBlock(Target);
     }
 
     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.
     // In this case, the indirectbr now comes from BodyBlock.
     if (IBRPred == Target)
       IBRPred = BodyBlock;
 
     // At this point Target only has PHIs, and BodyBlock has the rest of the

llvm/lib/Transforms/Utils/CodeExtractor.cpp

void CodeExtractor::calculateNewCallTerminatorWeights(
 
   // Update the branch weights for the exit block.
   Instruction *TI = CodeReplacer->getTerminator();
   SmallVector<unsigned, 8> BranchWeights(TI->getNumSuccessors(), 0);
 
   // Block Frequency distribution with dummy node.
   Distribution BranchDist;
 
+  SmallVector<BranchProbability, 4> 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);
     uint64_t ExitFreq = ExitWeights[TI->getSuccessor(i)].getFrequency();
     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();
 
   // Create normalized branch weights and set the metadata.
   for (unsigned I = 0, E = BranchDist.Weights.size(); I < E; ++I) {
     const auto &Weight = BranchDist.Weights[I];
 
     // 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));
 }
 
 /// Erase debug info intrinsics which refer to values in \p F but aren't in
 /// \p F.
 static void eraseDebugIntrinsicsWithNonLocalRefs(Function &F) {

llvm/test/Analysis/BranchProbabilityInfo/basic.ll

 ; CHECK: Printing analysis {{.*}} for function 'test_unreachable_with_switch_prof2'
 entry:
   switch i32 %i, label %case_a [ i32 1, label %case_b
                                  i32 2, label %case_c
                                  i32 3, label %case_d
                                  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
 
 case_b:
   unreachable
 
 case_c:

llvm/unittests/Transforms/Utils/BasicBlockUtilsTest.cpp

TEST(BasicBlockUtils, SplitIndirectBrCriticalEdge) {
   EXPECT_TRUE(Split);
 
   // Check that successors of the split block get their probability correct.
   BasicBlock *SplitBB = Block("bb1").getTerminator()->getSuccessor(0);
   EXPECT_EQ(2u, SplitBB->getTerminator()->getNumSuccessors());
   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<Module> 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));
+}