[BranchProbability] Remove the restriction that known and unknown probabilities canno...

[oota-llvm.git] / include / llvm / Support / BranchProbability.h

diff --git a/include/llvm/Support/BranchProbability.h b/include/llvm/Support/BranchProbability.h
index 3620d4d5d7722b96cdfd35ee5496687db2fef43a..1d6a7d84fbf47131edc20a558ed700e2531d31c9 100644 (file)
--- a/include/llvm/Support/BranchProbability.h
+++ b/include/llvm/Support/BranchProbability.h
@@ -53,6 +53,9 @@ public:
   // Create a BranchProbability object with the given numerator and 1<<31
   // as denominator.
   static BranchProbability getRaw(uint32_t N) { return BranchProbability(N); }
+  // Create a BranchProbability object from 64-bit integers.
+  static BranchProbability getBranchProbability(uint64_t Numerator,
+                                                uint64_t Denominator);
 
   // Normalize given probabilties so that the sum of them becomes approximate
   // one.
@@ -114,6 +117,14 @@ public:
     return *this;
   }
 
+  BranchProbability &operator/=(uint32_t RHS) {
+    assert(N != UnknownN &&
+           "Unknown probability cannot participate in arithmetics.");
+    assert(RHS > 0 && "The divider cannot be zero.");
+    N /= RHS;
+    return *this;
+  }
+
   BranchProbability operator+(BranchProbability RHS) const {
     BranchProbability Prob(*this);
     return Prob += RHS;
@@ -129,41 +140,75 @@ public:
     return Prob *= RHS;
   }
 
+  BranchProbability operator/(uint32_t RHS) const {
+    BranchProbability Prob(*this);
+    return Prob /= RHS;
+  }
+
   bool operator==(BranchProbability RHS) const { return N == RHS.N; }
   bool operator!=(BranchProbability RHS) const { return !(*this == RHS); }
-  bool operator<(BranchProbability RHS) const { return N < RHS.N; }
-  bool operator>(BranchProbability RHS) const { return RHS < *this; }
-  bool operator<=(BranchProbability RHS) const { return !(RHS < *this); }
-  bool operator>=(BranchProbability RHS) const { return !(*this < RHS); }
+
+  bool operator<(BranchProbability RHS) const {
+    assert(N != UnknownN && RHS.N != UnknownN &&
+           "Unknown probability cannot participate in comparisons.");
+    return N < RHS.N;
+  }
+
+  bool operator>(BranchProbability RHS) const {
+    assert(N != UnknownN && RHS.N != UnknownN &&
+           "Unknown probability cannot participate in comparisons.");
+    return RHS < *this;
+  }
+
+  bool operator<=(BranchProbability RHS) const {
+    assert(N != UnknownN && RHS.N != UnknownN &&
+           "Unknown probability cannot participate in comparisons.");
+    return !(RHS < *this);
+  }
+
+  bool operator>=(BranchProbability RHS) const {
+    assert(N != UnknownN && RHS.N != UnknownN &&
+           "Unknown probability cannot participate in comparisons.");
+    return !(*this < RHS);
+  }
 };
 
 inline raw_ostream &operator<<(raw_ostream &OS, BranchProbability Prob) {
   return Prob.print(OS);
 }
 
-inline BranchProbability operator/(BranchProbability LHS, uint32_t RHS) {
-  assert(LHS != BranchProbability::getUnknown() &&
-         "Unknown probability cannot participate in arithmetics.");
-  return BranchProbability::getRaw(LHS.getNumerator() / RHS);
-}
-
 template <class ProbabilityIter>
 void BranchProbability::normalizeProbabilities(ProbabilityIter Begin,
                                                ProbabilityIter End) {
   if (Begin == End)
     return;
 
-  auto UnknownProbCount =
-      std::count(Begin, End, BranchProbability::getUnknown());
-  assert((UnknownProbCount == 0 ||
-          UnknownProbCount == std::distance(Begin, End)) &&
-         "Cannot normalize probabilities with known and unknown ones.");
-  (void)UnknownProbCount;
-
-  uint64_t Sum = std::accumulate(
-      Begin, End, uint64_t(0),
-      [](uint64_t S, const BranchProbability &BP) { return S + BP.N; });
-
+  unsigned UnknownProbCount = 0;
+  uint64_t Sum = std::accumulate(Begin, End, uint64_t(0),
+                                 [&](uint64_t S, const BranchProbability &BP) {
+                                   if (!BP.isUnknown())
+                                     return S + BP.N;
+                                   UnknownProbCount++;
+                                   return S;
+                                 });
+
+  if (UnknownProbCount > 0) {
+    BranchProbability ProbForUnknown = BranchProbability::getZero();
+    // If the sum of all known probabilities is less than one, evenly distribute
+    // the complement of sum to unknown probabilities. Otherwise, set unknown
+    // probabilities to zeros and continue to normalize known probabilities.
+    if (Sum < BranchProbability::getDenominator())
+      ProbForUnknown = BranchProbability::getRaw(
+          (BranchProbability::getDenominator() - Sum) / UnknownProbCount);
+
+    std::replace_if(Begin, End,
+                    [](const BranchProbability &BP) { return BP.isUnknown(); },
+                    ProbForUnknown);
+
+    if (Sum <= BranchProbability::getDenominator())
+      return;
+  }
+ 
   if (Sum == 0) {
     BranchProbability BP(1, std::distance(Begin, End));
     std::fill(Begin, End, BP);