[JumpThreading] Fix branch probabilities in DuplicateCondBranchOnPHIIntoPred()

[yrouban]

Summary

When instructions are cloned from block BB to PredBB in the method DuplicateCondBranchOnPHIIntoPred() number of successors of PredBB changes from 1 to number of successors of BB. So we have to copy branch probabilities from BB to PredBB. I found this place when the single successor of PredBB got 1.0 probability which was left unchanged after PredBB got the instructions from its successor along with a terminator with several successors. So the first edge from ProbBB had probability 1.0 and the second resulted in 0.5 as it was unset explicitly. I cannot create an LL test since D90284. Here is a short test that failed before:

; CMD: opt -passes=jump-threading -debug-only=branch-prob -debug-only=jump-threading -S test.ll
define void @foo(i1 %f0, i1 %f1, i1 %f2) !prof !{!"function_entry_count", i64 2} {
bb1:
  br i1 %f0, label %bb3, label %bb2

bb2:
  br label %bb3

bb3:
  %ph = phi i1 [ %f1, %bb1 ], [ %f2, %bb2 ]
  br i1 %ph, label %exit1, label %unreach

exit1:
  ret void

unreach:
  unreachable
}

Here is the log with additional printing of probabilities at the end of DuplicateCondBranchOnPHIIntoPred() (compare the probs of bb3 and bb2, they must be the same):

---- Branch Probability Info : foo ----

Computing probabilities for exit1
Computing probabilities for unreach
Computing probabilities for bb3
set edge bb3 -> 0 successor probability to 0x7fffffff / 0x80000000 = 100.00%
set edge bb3 -> 1 successor probability to 0x00000001 / 0x80000000 = 0.00%
Computing probabilities for bb2
Computing probabilities for bb1
Jump threading on function 'foo'
  Duplicating block 'bb3' into end of 'bb2' to eliminate branch on phi.  Cost: 0 block is:
bb3:                                              ; preds = %bb2, %bb1
  %ph = phi i1 [ %f1, %bb1 ], [ %f2, %bb2 ]
  br i1 %ph, label %exit1, label %unreach

Probability of bb3 -> (0) = 0x7fffffff / 0x80000000 = 100.00%
Probability of bb3 -> (1) = 0x00000001 / 0x80000000 = 0.00%
Probability of bb2 -> (0) = 0x40000000 / 0x80000000 = 50.00%
Probability of bb2 -> (1) = 0x40000000 / 0x80000000 = 50.00%

------------- Result
define void @foo(i1 %f0, i1 %f1, i1 %f2) !prof !0 {
bb1:
  br i1 %f0, label %bb3, label %bb2

bb2:                                              ; preds = %bb1
  br i1 %f2, label %exit1, label %unreach

bb3:                                              ; preds = %bb1
  %ph = phi i1 [ %f1, %bb1 ]
  br i1 %ph, label %exit1, label %unreach

exit1:                                            ; preds = %bb2, %bb3
  ret void

unreach:                                          ; preds = %bb2, %bb3
  unreachable
}

!0 = !{!"function_entry_count", i64 2}
--------------------------------------------------

Copying can be done with the newly introduced method in D90839.
The change in the method ThreadThroughTwoBasicBlocks() as a small related NFC.

[kazu]

Your patch looks reasonable. Is there any way you could come up with a testcase to make sure that the pass indeed copies the probabilities in your testcase? (Note that I'm not asking for a testcase that exposed this bug before D90284.)

[yrouban]

In D90841#2377256, @kazu wrote:

Your patch looks reasonable. Is there any way you could come up with a testcase to make sure that the pass indeed copies the probabilities in your testcase? (Note that I'm not asking for a testcase that exposed this bug before D90284.)

I could not create a simple test. One more option to create one is to change the usage of BPI/BFI to request analysis from AM. Then the test can print BPI to show the difference. I'm not sure this big change is worth this small patch. It is not clear why JumpThreading uses its own local version of BPI/BFI.

[yrouban]

added a test

[kazu]

Thank you for the testcase!

llvm/lib/Transforms/Scalar/JumpThreading.cpp

void JumpThreadingPass::threadThroughTwoBasicBlocks(BasicBlock *PredPredBB,
 
   // We are going to have to map operands from the original BB block to the new
   // copy of the block 'NewBB'.  If there are PHI nodes in PredBB, evaluate them
   // to account for entry from PredPredBB.
   DenseMap<Instruction *, Value *> ValueMapping =
       cloneInstructions(PredBB->begin(), PredBB->end(), NewBB, PredPredBB);
 
   // Copy the edge probabilities from PredBB to NewBB.
-  if (HasProfileData) {
-    SmallVector<BranchProbability, 4> Probs;
-    for (BasicBlock *Succ : successors(PredBB))
-      Probs.push_back(BPI->getEdgeProbability(PredBB, Succ));
-    BPI->setEdgeProbability(NewBB, Probs);
-  }
+  if (HasProfileData)
+    BPI->copyEdgeProbabilities(PredBB, NewBB);
 
   // Update the terminator of PredPredBB to jump to NewBB instead of PredBB.
   // This eliminates predecessors from PredPredBB, which requires us to simplify
   // any PHI nodes in PredBB.
   Instruction *PredPredTerm = PredPredBB->getTerminator();
   for (unsigned i = 0, e = PredPredTerm->getNumSuccessors(); i != e; ++i)
     if (PredPredTerm->getSuccessor(i) == PredBB) {
       PredBB->removePredecessor(PredPredBB, true);
bool JumpThreadingPass::duplicateCondBranchOnPHIIntoPred(
   updateSSA(BB, PredBB, ValueMapping);
 
   // PredBB no longer jumps to BB, remove entries in the PHI node for the edge
   // that we nuked.
   BB->removePredecessor(PredBB, true);
 
   // Remove the unconditional branch at the end of the PredBB block.
   OldPredBranch->eraseFromParent();
+  if (HasProfileData)
+    BPI->copyEdgeProbabilities(BB, PredBB);
   DTU->applyUpdatesPermissive(Updates);
 
   ++NumDupes;
   return true;
 }
 
 // Pred is a predecessor of BB with an unconditional branch to BB. SI is
 // a Select instruction in Pred. BB has other predecessors and SI is used in

llvm/test/Transforms/JumpThreading/thread-prob-3.ll

+; RUN: opt -debug-only=branch-prob -jump-threading -S %s 2>&1 | FileCheck %s
+; REQUIRES: asserts
+
+; Make sure that we set edge probabilities for bb2 as we
+; call DuplicateCondBranchOnPHIIntoPred(bb3, {bb2}).
+;
+; CHECK-LABEL: ---- Branch Probability Info : foo
+; CHECK:      set edge bb2 -> 0 successor probability to 0x7fffffff / 0x80000000 = 100.00%
+; CHECK-NEXT: set edge bb2 -> 1 successor probability to 0x00000001 / 0x80000000 = 0.00%
+define void @foo(i1 %f0, i1 %f1, i1 %f2) !prof !{!"function_entry_count", i64 0} {
+; CHECK-LABEL: @foo(
+bb1:
+  br i1 %f0, label %bb3, label %bb2
+
+bb2:
+; CHECK:      bb2:
+; CHECK-NEXT:   br i1 %f2, label %exit1, label %unreach
+  br label %bb3
+
+bb3:
+  %ph = phi i1 [ %f1, %bb1 ], [ %f2, %bb2 ]
+  br i1 %ph, label %exit1, label %unreach
+
+exit1:
+  ret void
+
+unreach:
+  unreachable
+}