Allow known and unknown probabilities coexist in MBB's successor list.

Previously it is not allowed for each MBB to have successors with both known and
unknown probabilities. However, this may be too strict as at this stage we could
not always guarantee that. It is better to remove this restriction now, and I
will work on validating MBB's successors' probabilities first (for example,
check if the sum is approximate one).

llvm-svn: 254402

1 files changed, 17 insertions, 8 deletions

diff --git a/llvm/lib/CodeGen/MachineBasicBlock.cpp b/llvm/lib/CodeGen/MachineBasicBlock.cpp
index c9c6a9d6246..de91f0db75a 100644
--- a/llvm/lib/CodeGen/MachineBasicBlock.cpp
+++ b/llvm/lib/CodeGen/MachineBasicBlock.cpp
@@ -510,13 +510,8 @@ void MachineBasicBlock::addSuccessor(MachineBasicBlock *Succ,
                                      BranchProbability Prob) {
   // Probability list is either empty (if successor list isn't empty, this means
   // disabled optimization) or has the same size as successor list.
-  if (!(Probs.empty() && !Successors.empty())) {
-    assert((Probs.empty() || (Prob.isUnknown() && Probs.back().isUnknown()) ||
-            (!Prob.isUnknown() && !Probs.back().isUnknown())) &&
-           "Successors with both known and unknwon probabilities are not "
-           "allowed.");
+  if (!(Probs.empty() && !Successors.empty()))
     Probs.push_back(Prob);
-  }
   Successors.push_back(Succ);
   Succ->addPredecessor(this);
 }
@@ -1116,10 +1111,24 @@ MachineBasicBlock::findDebugLoc(instr_iterator MBBI) {
 /// Return probability of the edge from this block to MBB.
 BranchProbability
 MachineBasicBlock::getSuccProbability(const_succ_iterator Succ) const {
-  if (Probs.empty() || Probs.back().isUnknown())
+  if (Probs.empty())
     return BranchProbability(1, succ_size());
 
-  return *getProbabilityIterator(Succ);
+  const auto &Prob = *getProbabilityIterator(Succ);
+  if (Prob.isUnknown()) {
+    // For unknown probabilities, collect the sum of all known ones, and evenly
+    // ditribute the complemental of the sum to each unknown probability.
+    unsigned KnownProbNum = 0;
+    auto Sum = BranchProbability::getZero();
+    for (auto &P : Probs) {
+      if (!P.isUnknown()) {
+        Sum += P;
+        KnownProbNum++;
+      }
+    }
+    return Sum.getCompl() / (Probs.size() - KnownProbNum);
+  } else
+    return Prob;
 }
 
 /// Set successor probability of a given iterator.