[BPI] Ignore tail while distributing the remaining probability from unreachanble

[skatkov]

Summary

This is a follow up patch for https://reviews.llvm.org/rL300440
to address a comment.

To make implementation to be consistent with other cases we just
ignore the tail after distribution of remaining probability between
reachable edges.

If we reduced the probability of some edges coming to unreachable
blocks we should distribute the remaining part across other edges
coming to reachable blocks to satisfy the condition that sum of all
probabilities should be equal to one. If this remaining part is not
divided by number of "reachable" edges then we get this tail.
This tail probability should be pretty small. Other cases just ignore
if the some of probabilities is not equal to one so we do the same.

[skatkov]

Please review.

[vsk]

Thanks for the cleanup.

One note is that it would help to explicitly explain what your patch changes, and why, in your description (the title would be a very condensed version of this). This would make your patch easier to review.

lib/Analysis/BranchProbabilityInfo.cpp

336–338

IIUC, the "tail" is supposed to be some small, non-negative epsilon value. Maybe it's worth asserting this? E.g assert(tail >= 0 && tail < 1).

[chandlerc]

Thanks for the updated patch description! Another minor request here.

include/llvm/Support/BranchProbability.h

115–121

Add a unittest? I know that much of BranchProbability isn't well covered by unittests, but we should really have good coverage for basic primitives like this and you can start covering the mutation operators.

[skatkov]

Sure, will do.

[skatkov]

Some test added.

[skatkov]

BTW, probably we should not add the tail to any edge at all... It might be that the information that two edges have exactly the same probability is mot important than the requirement that the sum of all probabilities should be 1. it seems that sum == 1 is not a strict requirement. For example, if we have three edges with the same weight then as a result the sum of probabilities will not be equal to one. What do you think?

[chandlerc]

In D32124#731921, @skatkov wrote:

Some test added.

The tests and added operators are awesome.

If you want, split that off into a separate change just to the BranchProbability stuff and go ahead and submit that?

[chandlerc]

In D32124#732170, @skatkov wrote:

BTW, probably we should not add the tail to any edge at all... It might be that the information that two edges have exactly the same probability is mot important than the requirement that the sum of all probabilities should be 1. it seems that sum == 1 is not a strict requirement. For example, if we have three edges with the same weight then as a result the sum of probabilities will not be equal to one. What do you think?

I think we already handle this kind of case in a few places in BPI and we should be consistent there. I can go digging for it but you may already know where it is...

[skatkov]

In D32124#733015, @chandlerc wrote:

In D32124#732170, @skatkov wrote:

BTW, probably we should not add the tail to any edge at all... It might be that the information that two edges have exactly the same probability is mot important than the requirement that the sum of all probabilities should be 1. it seems that sum == 1 is not a strict requirement. For example, if we have three edges with the same weight then as a result the sum of probabilities will not be equal to one. What do you think?

I think we already handle this kind of case in a few places in BPI and we should be consistent there. I can go digging for it but you may already know where it is...

Hi Chandler, I'm not sure what exactly you mean but take a look into the last test @test_unreachable_with_switch_prof4 in test/Analysis/BranchProbabilityInfo/basic.ll:
For the five unreachable edges with the same weight we generate the following probability
; CHECK: edge entry -> case_a probability is 0x1999999a / 0x80000000 = 20.00%
; CHECK: edge entry -> case_b probability is 0x1999999a / 0x80000000 = 20.00%
; CHECK: edge entry -> case_c probability is 0x1999999a / 0x80000000 = 20.00%
; CHECK: edge entry -> case_d probability is 0x1999999a / 0x80000000 = 20.00%
; CHECK: edge entry -> case_e probability is 0x1999999a / 0x80000000 = 20.00%

and 0x1999999a * 5 = 0x80000002 > 0x80000000, so the sum of probabilities is greater than 1.

[skatkov]

After reading of the code a bit more I would say that in general no one worries about sum of probabilities is equal to 1. They follow it when operating on weights. But the trouble is that after weights are transformed to probabilities there are rounding errors which breaks the restriction.

[skatkov]

Please review.

[skatkov]

Ping, could please review this change.

[skatkov]

ping

[reames]

LGTM w/comment addressed.

lib/Analysis/BranchProbabilityInfo.cpp

333

I would suggest leaving the decrement in the loop and asserting after the loop that the remainder is less than the number of edges. A comment stating that we explicitly drop the remainder to be consistent with other code seems called for as well.

[reames]

lib/Analysis/BranchProbabilityInfo.cpp

333

Actually, don't bother with this. The assert would follow trivially from the code just above it, that's fine. Thus, LGTM w/out comments. :)

include/llvm/Support/BranchProbability.h

public:
 
   BranchProbability &operator*=(BranchProbability RHS) {
     assert(N != UnknownN && RHS.N != UnknownN &&
            "Unknown probability cannot participate in arithmetics.");
     N = (static_cast<uint64_t>(N) * RHS.N + D / 2) / D;
     return *this;
   }
 
+  BranchProbability &operator*=(uint32_t RHS) {
+    assert(N != UnknownN &&
+           "Unknown probability cannot participate in arithmetics.");
+    N = (uint64_t(N) * RHS > D) ? D : N * RHS;
+    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;
   }
 
   BranchProbability operator-(BranchProbability RHS) const {
     BranchProbability Prob(*this);
     return Prob -= RHS;
   }
 
   BranchProbability operator*(BranchProbability RHS) const {
     BranchProbability Prob(*this);
     return Prob *= RHS;
   }
 
+  BranchProbability operator*(uint32_t RHS) const {
+    BranchProbability Prob(*this);
+    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); }
 

lib/Analysis/BranchProbabilityInfo.cpp

bool BranchProbabilityInfo::calcMetadataWeights(const BasicBlock *BB) {
 
   if (ScalingFactor > 1) {
     WeightSum = 0;
     for (unsigned i = 0, e = TI->getNumSuccessors(); i != e; ++i) {
       Weights[i] /= ScalingFactor;
       WeightSum += Weights[i];
     }
   }
+  assert(WeightSum <= UINT32_MAX &&
+         "Expected weights to scale down to 32 bits");
 
   if (WeightSum == 0 || ReachableIdxs.size() == 0) {
     for (unsigned i = 0, e = TI->getNumSuccessors(); i != e; ++i)
       Weights[i] = 1;
     WeightSum = TI->getNumSuccessors();
   }
 
   // Set the probability.
for (auto i : UnreachableIdxs)
         ToDistribute += BP[i] - UnreachableProb;
         BP[i] = UnreachableProb;
       }
 
     // 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();
-      for (auto i : ReachableIdxs) {
+      for (auto i : ReachableIdxs)
         BP[i] += PerEdge;
-        ToDistribute -= PerEdge;
-      }
       // Tail goes to the first reachable edge.
-      BP[ReachableIdxs[0]] += ToDistribute;
+      BranchProbability Tail = ToDistribute - (PerEdge * ReachableIdxs.size());
+      assert(Tail.getNumerator() < ReachableIdxs.size() && "Tail is too big!");
+      BP[ReachableIdxs[0]] += Tail;
     }
   }
 
   for (unsigned i = 0, e = TI->getNumSuccessors(); i != e; ++i)
     setEdgeProbability(BB, i, BP[i]);
 
-  assert(WeightSum <= UINT32_MAX &&
-         "Expected weights to scale down to 32 bits");
-
   return true;
 }
 
 /// \brief 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
 /// them relatively low weight.

unittests/Support/BranchProbabilityTest.cpp

TEST(BranchProbabilityTest, MoreOperators) {
   EXPECT_FALSE(BigZero == BigOne);
   EXPECT_TRUE(BigZero != BigOne);
   EXPECT_TRUE(BigZero < BigOne);
   EXPECT_FALSE(BigZero > BigOne);
   EXPECT_TRUE(BigZero <= BigOne);
   EXPECT_FALSE(BigZero >= BigOne);
 }
 
+TEST(BranchProbabilityTest, ArithmeticOperators) {
+  BP Z(0, 1);
+  BP O(1, 1);
+  BP H(1, 2);
+  BP Q(1, 4);
+  BP Q3(3, 4);
+
+  EXPECT_EQ(Z + O, O);
+  EXPECT_EQ(H + Z, H);
+  EXPECT_EQ(H + H, O);
+  EXPECT_EQ(Q + H, Q3);
+  EXPECT_EQ(Q + Q3, O);
+  EXPECT_EQ(H + Q3, O);
+  EXPECT_EQ(Q3 + Q3, O);
+
+  EXPECT_EQ(Z - O, Z);
+  EXPECT_EQ(O - Z, O);
+  EXPECT_EQ(O - H, H);
+  EXPECT_EQ(O - Q, Q3);
+  EXPECT_EQ(Q3 - H, Q);
+  EXPECT_EQ(Q - H, Z);
+  EXPECT_EQ(Q - Q3, Z);
+
+  EXPECT_EQ(Z * O, Z);
+  EXPECT_EQ(H * H, Q);
+  EXPECT_EQ(Q * O, Q);
+  EXPECT_EQ(O * O, O);
+  EXPECT_EQ(Z * Z, Z);
+
+  EXPECT_EQ(Z * 3, Z);
+  EXPECT_EQ(Q * 3, Q3);
+  EXPECT_EQ(H * 3, O);
+  EXPECT_EQ(Q3 * 2, O);
+  EXPECT_EQ(O * UINT32_MAX, O);
+
+  EXPECT_EQ(Z / 4, Z);
+  EXPECT_EQ(O / 4, Q);
+  EXPECT_EQ(Q3 / 3, Q);
+  EXPECT_EQ(H / 2, Q);
+  EXPECT_EQ(O / 2, H);
+  EXPECT_EQ(H / UINT32_MAX, Z);
+
+  BP MIN(1, 1u << 31);
+
+  EXPECT_EQ(O / UINT32_MAX, Z);
+  EXPECT_EQ(MIN * UINT32_MAX, O);
+}
+
 TEST(BranchProbabilityTest, getCompl) {
   EXPECT_EQ(BP(5, 7), BP(2, 7).getCompl());
   EXPECT_EQ(BP(2, 7), BP(5, 7).getCompl());
   EXPECT_EQ(BP::getZero(), BP(7, 7).getCompl());
   EXPECT_EQ(BP::getOne(), BP(0, 7).getCompl());
 }
 
 TEST(BranchProbabilityTest, scale) {