[BPI] Fix two potential divide-by-zero operations that are introduced in r256263.

llvm-svn: 256303

2 files changed, 20 insertions, 9 deletions

diff --git a/llvm/lib/Analysis/BranchProbabilityInfo.cpp b/llvm/lib/Analysis/BranchProbabilityInfo.cpp
index c7c2f436e8c..cf0cc8da6ef 100644
--- a/llvm/lib/Analysis/BranchProbabilityInfo.cpp
+++ b/llvm/lib/Analysis/BranchProbabilityInfo.cpp
@@ -221,8 +221,14 @@ bool BranchProbabilityInfo::calcMetadataWeights(BasicBlock *BB) {
     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)});
+
+  if (WeightSum == 0) {
+    for (unsigned i = 0, e = TI->getNumSuccessors(); i != e; ++i)
+      setEdgeProbability(BB, i, {1, e});
+  } else {
+    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");
diff --git a/llvm/lib/Transforms/Scalar/JumpThreading.cpp b/llvm/lib/Transforms/Scalar/JumpThreading.cpp
index d30c336be21..5c16b2c5de0 100644
--- a/llvm/lib/Transforms/Scalar/JumpThreading.cpp
+++ b/llvm/lib/Transforms/Scalar/JumpThreading.cpp
@@ -1647,14 +1647,19 @@ void JumpThreading::UpdateBlockFreqAndEdgeWeight(BasicBlock *PredBB,
 
   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());
+  SmallVector<BranchProbability, 4> BBSuccProbs;
+  if (MaxBBSuccFreq == 0)
+    BBSuccProbs.assign(BBSuccFreq.size(),
+                       {1, static_cast<unsigned>(BBSuccFreq.size())});
+  else {
+    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++)