Remove the restriction that known and unknown probabilities cannot coexist when being normalized.

[congh]

Summary

The current BranchProbability::normalizeProbabilities() forbids known and unknown probabilities to coexist in the list. This was once used to help capture probability exceptions but has caused some reported build failures (https://llvm.org/bugs/show_bug.cgi?id=25838).

This patch removes this restriction by evenly distributing the complement of the sum of all known probabilities to unknown ones. We could still treat this as an abnormal behavior, but it is better to emit warnings in our future profile validator.

[davidxl]

While it makes sense to do this -- can we start to teach producers not to pass unknown Probabilities in the first place?

[congh]

In D15548#312148, @davidxl wrote:

While it makes sense to do this -- can we start to teach producers not to pass unknown Probabilities in the first place?

If you search addSuccessor in LLVM, you will find about 200 users for different platforms without probability passed explicitly. Some are conditionally used such as:

MBB->addSuccessor(Succ1);
if (...)

MBB->addSuccessor(Succ2);

It seems that a default probability is more convenient but I don't object to add probabilities explicitly. There are just some tradeoffs.

[davidxl]

include/llvm/Support/BranchProbability.h

194

why guarded with Sum > 0?

200

may be name it 'ProbForUnknown'

[congh]

include/llvm/Support/BranchProbability.h

194

After reviewing the patch, I think this guard should be removed. Thanks for pointing it out!

[congh]

Update the patch on David' comment.

[davidxl]

lgtm

include/llvm/Support/BranchProbability.h

 }
 
 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);
     return;
   }
 
   for (auto I = Begin; I != End; ++I)
     I->N = (I->N * uint64_t(D) + Sum / 2) / Sum;

unittests/Support/BranchProbabilityTest.cpp

TEST(BranchProbabilityTest, scaleBruteForce) {
   };
 
   for (const auto &T : Tests) {
     EXPECT_EQ(T.Result, BP(T.Prob[0], T.Prob[1]).scale(T.Num));
   }
 }
 
 TEST(BranchProbabilityTest, NormalizeProbabilities) {
+  const auto UnknownProb = BranchProbability::getUnknown();
   {
     SmallVector<BranchProbability, 2> Probs{{0, 1}, {0, 1}};
     BranchProbability::normalizeProbabilities(Probs.begin(), Probs.end());
     EXPECT_EQ(BranchProbability::getDenominator() / 2, Probs[0].getNumerator());
     EXPECT_EQ(BranchProbability::getDenominator() / 2, Probs[1].getNumerator());
   }
   {
     SmallVector<BranchProbability, 2> Probs{{0, 1}, {1, 1}};
const auto UnknownProb = BranchProbability::getUnknown();
     BranchProbability::normalizeProbabilities(Probs.begin(), Probs.end());
     EXPECT_EQ(BranchProbability::getDenominator() / 3 + 1,
               Probs[0].getNumerator());
     EXPECT_EQ(BranchProbability::getDenominator() / 3 + 1,
               Probs[1].getNumerator());
     EXPECT_EQ(BranchProbability::getDenominator() / 3 + 1,
               Probs[2].getNumerator());
   }
+  {
+    SmallVector<BranchProbability, 2> Probs{{0, 1}, UnknownProb};
+    BranchProbability::normalizeProbabilities(Probs.begin(), Probs.end());
+    EXPECT_EQ(0, Probs[0].getNumerator());
+    EXPECT_EQ(BranchProbability::getDenominator(), Probs[1].getNumerator());
+  }
+  {
+    SmallVector<BranchProbability, 2> Probs{{1, 1}, UnknownProb};
+    BranchProbability::normalizeProbabilities(Probs.begin(), Probs.end());
+    EXPECT_EQ(BranchProbability::getDenominator(), Probs[0].getNumerator());
+    EXPECT_EQ(0, Probs[1].getNumerator());
+  }
+  {
+    SmallVector<BranchProbability, 2> Probs{{1, 2}, UnknownProb};
+    BranchProbability::normalizeProbabilities(Probs.begin(), Probs.end());
+    EXPECT_EQ(BranchProbability::getDenominator() / 2, Probs[0].getNumerator());
+    EXPECT_EQ(BranchProbability::getDenominator() / 2, Probs[1].getNumerator());
+  }
+  {
+    SmallVector<BranchProbability, 4> Probs{
+        {1, 2}, {1, 2}, {1, 2}, UnknownProb};
+    BranchProbability::normalizeProbabilities(Probs.begin(), Probs.end());
+    EXPECT_EQ(BranchProbability::getDenominator() / 3 + 1,
+              Probs[0].getNumerator());
+    EXPECT_EQ(BranchProbability::getDenominator() / 3 + 1,
+              Probs[1].getNumerator());
+    EXPECT_EQ(BranchProbability::getDenominator() / 3 + 1,
+              Probs[2].getNumerator());
+    EXPECT_EQ(0, Probs[3].getNumerator());
+  }
 }
 
 }