[BrachProbablityInfo] Proportional distribution of reachable probabilities

[yrouban]

Summary

When fixing probability of unreachable edges in BranchProbabilityInfo::calcMetadataWeights() proportionally distribute remainder probability over the reachable edges. The old implementation distributes the remainder probability evenly. See examples in the fixed tests.

[ebrevnov]

llvm/lib/Analysis/BranchProbabilityInfo.cpp

410

It would help to understand the code if you split this series of calculations:

uint64_t Mul = static_cast<uint64_t>(...)
uint32_t Div = static_cast<uint_32_t>l(lvm::divideNearest(...));
BP[i] = BranchProbability::getRaw(Div);

[hjyamauchi]

llvm/lib/Analysis/BranchProbabilityInfo.cpp

105

"equally" -> "proportionally"?

372

'independed on' -> 'independent of'?

385

Capitalize 'i' -> 'I'.

This loop could be combined with the above loop, but would be less clear?

392

Capitalize 'i' -> 'I'

[yrouban]

llvm/lib/Analysis/BranchProbabilityInfo.cpp

385

The old code above has the index i lowcased. I would propose to keep the style.
Yes, the separate loop is for clarity.

392

ditto, if you do not insist

[yrouban]

addressed comments.

[hjyamauchi]

llvm/lib/Analysis/BranchProbabilityInfo.cpp

385

OK, though in general I'd encourage to update to the capitalized variable name style in newly added code.

[craig.topper]

This seems to have been pushed to a branch called D80611-land2 and not to trunk.

[yrouban]

In D80611#2067956, @craig.topper wrote:

This seems to have been pushed to a branch called D80611-land2 and not to trunk.

Sorry, I used a wrong push command. I will land it soon.

llvm/lib/Analysis/BranchProbabilityInfo.cpp

 static const uint32_t LBH_NONTAKEN_WEIGHT = 4;
 // Unlikely edges within a loop are half as likely as other edges
 static const uint32_t LBH_UNLIKELY_WEIGHT = 62;
 
 /// Unreachable-terminating branch taken probability.
 ///
 /// This is the probability for a branch being taken to a block that terminates
 /// (eventually) in unreachable. These are predicted as unlikely as possible.
 /// All reachable probability will equally share the remaining part.
 static const BranchProbability UR_TAKEN_PROB = BranchProbability::getRaw(1);
 
 /// Weight for a branch taken going into a cold block.
 ///
 /// This is the weight for a branch taken toward a block marked
 /// cold.  A block is marked cold if it's postdominated by a
 /// block containing a call to a cold function.  Cold functions
 /// are those marked with attribute 'cold'.
bool BranchProbabilityInfo::calcMetadataWeights(const BasicBlock *BB) {
 
   // Set the probability.
   SmallVector<BranchProbability, 2> BP;
   for (unsigned i = 0, e = TI->getNumSuccessors(); i != e; ++i)
     BP.push_back({ Weights[i], static_cast<uint32_t>(WeightSum) });
 
   // Examine the metadata against unreachable heuristic.
   // If the unreachable heuristic is more strong then we use it for this edge.
-  if (UnreachableIdxs.size() > 0 && ReachableIdxs.size() > 0) {
-    auto UnreachableProb = UR_TAKEN_PROB;
-    for (auto i : UnreachableIdxs)
-      if (UnreachableProb < BP[i]) {
-        BP[i] = UnreachableProb;
-      }
+  if (UnreachableIdxs.size() == 0 || ReachableIdxs.size() == 0) {
+    setEdgeProbability(BB, BP);
+    return true;
+  }
+
 
-    // Because of possible rounding errors and the above fix up for
-    // the unreachable heuristic the sum of probabilities of all edges may be
-    // less than 1.0. Distribute the remaining probability (calculated as
-    // 1.0 - (sum of BP[i])) evenly among all the reachable edges.
-    auto ToDistribute = BranchProbability::getOne();
-    for (auto &P : BP)
-      ToDistribute -= P;
-
-    // If we modified the probability of some edges then we must distribute
-    // the difference between reachable blocks.
-    // TODO: This spreads ToDistribute evenly upon the reachable edges. A better
-    // distribution would be proportional. So the relation between weights of
-    // the reachable edges would be kept unchanged. That is for any reachable
-    // edges i and j:
-    // newBP[i] / newBP[j] == oldBP[i] / oldBP[j]
-    // newBP[i] / oldBP[i] == newBP[j] / oldBP[j] ==
-    //                     == Denominator / (Denominator - ToDistribute)
-    // newBP[i] = oldBP[i] * Denominator / (Denominator - ToDistribute)
-    BranchProbability PerEdge = ToDistribute / ReachableIdxs.size();
-    if (PerEdge > BranchProbability::getZero())
+  // Sum of all edge probabilities must be 1.0. If we modified the probability
+  // of some edges then we must distribute the introduced difference over the
+  // reachable blocks.
+  //
+  // Proportional distribution: the relation between probabilities of the
+  // reachable edges is kept unchanged. That is for any reachable edges i and j:
+  //   newBP[i] / newBP[j] == oldBP[i] / oldBP[j] =>
+  //   newBP[i] / oldBP[i] == newBP[j] / oldBP[j] == K
+  // Where K is independed on i,j.
+  //   newBP[i] == oldBP[i] * K
+  // We need to find K.
+  // Make sum of all reachables of the left and right parts:
+  //   sum_of_reachable(newBP) == K * sum_of_reachable(oldBP)
+  // Sum of newBP must be equal to 1.0:
+  //   sum_of_reachable(newBP) + sum_of_unreachable(newBP) == 1.0 =>
+  //   sum_of_reachable(newBP) = 1.0 - sum_of_unreachable(newBP)
+  // Where sum_of_unreachable(newBP) is what has been just changed.
+  // Finally:
+  //   K == sum_of_reachable(newBP) / sum_of_reachable(oldBP) =>
+  //   K == (1.0 - sum_of_unreachable(newBP)) / sum_of_reachable(oldBP)
+  BranchProbability NewUnreachableSum = BranchProbability::getZero();
+  for (auto i : UnreachableIdxs)
+    NewUnreachableSum += BP[i];
+
+  BranchProbability NewReachableSum =
+      BranchProbability::getOne() - NewUnreachableSum;
+
+  BranchProbability OldReachableSum = BranchProbability::getZero();
+  for (auto i : ReachableIdxs)
+    OldReachableSum += BP[i];
+
+  if (OldReachableSum != NewReachableSum) // Anything to dsitribute?
+    if (OldReachableSum.isZero()) {
+      // If all oldBP[i] are zeroes then the proportional distribution results
+      // in all zero probabilities and the error stays big. In this case we
+      // evenly spread NewReachableSum over the reachable edges.
+      BranchProbability PerEdge = NewReachableSum / ReachableIdxs.size();
+      for (auto i : ReachableIdxs)
+        BP[i] = PerEdge;
+    } else {
       for (auto i : ReachableIdxs)
-        BP[i] += PerEdge;
+        // We use uint64_t to avoid double rounding error of the following
+        // calculation: BP[i] = BP[i] * NewReachableSum / OldReachableSum
+        // The formula is taken from the private constructor
+        // BranchProbability(uint32_t Numerator, uint32_t Denominator)
+        BP[i] = BranchProbability::getRaw(
+            (BP[i].getNumerator() *
+                 static_cast<uint64_t>(NewReachableSum.getNumerator()) +
+             OldReachableSum.getNumerator() / 2) /
+            OldReachableSum.getNumerator());
     }
 
   setEdgeProbability(BB, BP);
 
   return true;
 }
 
 /// Calculate edge weights for edges leading to cold blocks.
 ///

llvm/test/Analysis/BranchProbabilityInfo/basic.ll

 
 define i32 @test_unreachable_with_switch_prof1(i32 %i, i32 %a, i32 %b, i32 %c, i32 %d, i32 %e) {
 ; CHECK: Printing analysis {{.*}} for function 'test_unreachable_with_switch_prof1'
 entry:
   switch i32 %i, label %case_a [ i32 1, label %case_b
                                  i32 2, label %case_c
                                  i32 3, label %case_d
                                  i32 4, label %case_e ], !prof !8
+; Reachable probabilities keep their relation: 4/64/4/4 = 5.26% / 84.21% / 5.26% / 5.26%.
 ; CHECK: edge entry -> case_a probability is 0x00000001 / 0x80000000 = 0.00%
-; CHECK: edge entry -> case_b probability is 0x07ffffff / 0x80000000 = 6.25%
-; CHECK: edge entry -> case_c probability is 0x67ffffff / 0x80000000 = 81.25% [HOT edge]
-; CHECK: edge entry -> case_d probability is 0x07ffffff / 0x80000000 = 6.25%
-; CHECK: edge entry -> case_e probability is 0x07ffffff / 0x80000000 = 6.25%
+; CHECK: edge entry -> case_b probability is 0x06bca1af / 0x80000000 = 5.26%
+; CHECK: edge entry -> case_c probability is 0x6bca1af3 / 0x80000000 = 84.21% [HOT edge]
+; CHECK: edge entry -> case_d probability is 0x06bca1af / 0x80000000 = 5.26%
+; CHECK: edge entry -> case_e probability is 0x06bca1af / 0x80000000 = 5.26%
 
 case_a:
   unreachable
 
 case_b:
   br label %exit
 ; CHECK: edge case_b -> exit probability is 0x80000000 / 0x80000000 = 100.00% [HOT edge]
 

 
 define i32 @test_unreachable_with_switch_prof2(i32 %i, i32 %a, i32 %b, i32 %c, i32 %d, i32 %e) {
 ; CHECK: Printing analysis {{.*}} for function 'test_unreachable_with_switch_prof2'
 entry:
   switch i32 %i, label %case_a [ i32 1, label %case_b
                                  i32 2, label %case_c
                                  i32 3, label %case_d
                                  i32 4, label %case_e ], !prof !9
+; Reachable probabilities keep their relation: 64/4/4 = 88.89% / 5.56% / 5.56%.
 ; CHECK: edge entry -> case_a probability is 0x00000001 / 0x80000000 = 0.00%
 ; CHECK: edge entry -> case_b probability is 0x00000001 / 0x80000000 = 0.00%
-; CHECK: edge entry -> case_c probability is 0x6aaaaaaa / 0x80000000 = 83.33% [HOT edge]
-; CHECK: edge entry -> case_d probability is 0x0aaaaaaa / 0x80000000 = 8.33%
-; CHECK: edge entry -> case_e probability is 0x0aaaaaaa / 0x80000000 = 8.33%
+; CHECK: edge entry -> case_c probability is 0x71c71c71 / 0x80000000 = 88.89% [HOT edge]
+; CHECK: edge entry -> case_d probability is 0x071c71c7 / 0x80000000 = 5.56%
+; CHECK: edge entry -> case_e probability is 0x071c71c7 / 0x80000000 = 5.56%
+
 
 case_a:
   unreachable
 
 case_b:
   unreachable
 
 case_c:

 
 define i32 @test_unreachable_with_switch_prof3(i32 %i, i32 %a, i32 %b, i32 %c, i32 %d, i32 %e) {
 ; CHECK: Printing analysis {{.*}} for function 'test_unreachable_with_switch_prof3'
 entry:
   switch i32 %i, label %case_a [ i32 1, label %case_b
                                  i32 2, label %case_c
                                  i32 3, label %case_d
                                  i32 4, label %case_e ], !prof !10
+; Reachable probabilities keep their relation: 64/4/4 = 88.89% / 5.56% / 5.56%.
 ; CHECK: edge entry -> case_a probability is 0x00000000 / 0x80000000 = 0.00%
 ; CHECK: edge entry -> case_b probability is 0x00000001 / 0x80000000 = 0.00%
-; CHECK: edge entry -> case_c probability is 0x6e08fb82 / 0x80000000 = 85.96% [HOT edge]
-; CHECK: edge entry -> case_d probability is 0x08fb823e / 0x80000000 = 7.02%
-; CHECK: edge entry -> case_e probability is 0x08fb823e / 0x80000000 = 7.02%
+; CHECK: edge entry -> case_c probability is 0x71c71c71 / 0x80000000 = 88.89% [HOT edge]
+; CHECK: edge entry -> case_d probability is 0x071c71c7 / 0x80000000 = 5.56%
+; CHECK: edge entry -> case_e probability is 0x071c71c7 / 0x80000000 = 5.56%
 
 case_a:
   unreachable
 
 case_b:
   unreachable
 
 case_c: