[JumpThreading] Set edge probabilities when creating basic blocks

This patch teaches the jump threading pass to set edge probabilities
whenever the pass creates new basic blocks.

Without this patch, the compiler sometimes produces non-deterministic
results.  The non-determinism comes from the jump threading pass using
stale edge probabilities in BranchProbabilityInfo.  Specifically, when
the jump threading pass creates a new basic block, we don't initialize
its outgoing edge probability.

Edge probabilities are maintained in:

  DenseMap<Edge, BranchProbability> Probs;

in class BranchProbabilityInfo, where Edge is an ordered pair of
BasicBlock * and a successor index declared as:

  using Edge = std::pair<const BasicBlock *, unsigned>;

Probs maps edges to their corresponding probabilities.

Now, we rarely remove entries from this map, so if we happen to
allocate a new basic block at the same address as a previously deleted
basic block with an edge probability assigned, the newly created basic
block appears to have an edge probability, albeit a stale one.

This patch fixes the problem by explicitly setting edge probabilities
whenever the jump threading pass creates new basic blocks.

Differential Revision: https://reviews.llvm.org/D90106

llvm/lib/Transforms/Scalar/JumpThreading.cpp

@@ -2236,6 +2236,14 @@ void JumpThreadingPass::ThreadThroughTwoBasicBlocks(BasicBlock *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);
+  }
+
   // 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.
@@ -2424,8 +2432,15 @@ BasicBlock *JumpThreadingPass::SplitBlockPreds(BasicBlock *BB,
       if (HasProfileData) // Update frequencies between Pred -> NewBB.
         NewBBFreq += FreqMap.lookup(Pred);
     }
-    if (HasProfileData) // Apply the summed frequency to NewBB.
+    if (HasProfileData) {
+      // Apply the summed frequency to NewBB.
       BFI->setBlockFreq(NewBB, NewBBFreq.getFrequency());
+
+      // NewBB has exactly one successor.
+      SmallVector<BranchProbability, 1> BBSuccProbs;
+      BBSuccProbs.push_back(BranchProbability::getOne());
+      BPI->setEdgeProbability(NewBB, BBSuccProbs);
+    }
   }
 
   DTU->applyUpdatesPermissive(Updates);
@@ -2498,6 +2513,11 @@ void JumpThreadingPass::UpdateBlockFreqAndEdgeWeight(BasicBlock *PredBB,
   // Update edge probabilities in BPI.
   BPI->setEdgeProbability(BB, BBSuccProbs);
 
+  // NewBB has exactly one successor.
+  SmallVector<BranchProbability, 1> NewBBSuccProbs;
+  NewBBSuccProbs.push_back(BranchProbability::getOne());
+  BPI->setEdgeProbability(NewBB, NewBBSuccProbs);
+
   // Update the profile metadata as well.
   //
   // Don't do this if the profile of the transformed blocks was statically
@@ -2708,6 +2728,13 @@ void JumpThreadingPass::UnfoldSelectInstr(BasicBlock *Pred, BasicBlock *BB,
        PHINode *Phi = dyn_cast<PHINode>(BI); ++BI)
     if (Phi != SIUse)
       Phi->addIncoming(Phi->getIncomingValueForBlock(Pred), NewBB);
+
+  if (HasProfileData) {
+    // NewBB has exactly one successor.
+    SmallVector<BranchProbability, 1> BBSuccProbs;
+    BBSuccProbs.push_back(BranchProbability::getOne());
+    BPI->setEdgeProbability(NewBB, BBSuccProbs);
+  }
 }
 
 bool JumpThreadingPass::TryToUnfoldSelect(SwitchInst *SI, BasicBlock *BB) {

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

@@ -0,0 +1,42 @@
+; RUN: opt -debug-only=branch-prob -jump-threading -S %s 2>&1 | FileCheck %s
+
+; Make sure that we set the branch probability for the newly created
+; basic block.
+
+define void @foo(i1 %arg1, i1 %arg2, i32 %arg3) !prof !0 !PGOFuncName !1 {
+entry:
+  call void @bar(i32 0)
+  br i1 %arg1, label %bb3, label %bb1, !prof !2
+
+bb1:
+  call void @bar(i32 1)
+  br i1 %arg2, label %bb2, label %bb3, !prof !3
+
+bb2:
+  call void @bar(i32 2)
+  br label %bb3
+
+bb3:
+; CHECK: set edge bb3.thr_comm -> 0 successor probability to 0x80000000 / 0x80000000
+%ptr = phi i32 [ 0, %bb1 ], [ 0, %entry ], [ %arg3, %bb2 ]
+  call void @bar(i32 3)
+  %bool = icmp eq i32 %ptr, 0
+  br i1 %bool, label %exit, label %bb4, !prof !4
+; CHECK: set edge bb3.thread -> 0 successor probability to 0x80000000 / 0x80000000
+
+bb4:
+  call void @bar(i32 %ptr)
+  br label %exit
+
+exit:
+  ret void
+}
+
+declare void @bar(i32)
+
+!0 = !{!"function_entry_count", i64 15985}
+!1 = !{!"foo:foo"}
+!2 = !{!"branch_weights", i32 15973, i32 36865}
+!3 = !{!"branch_weights", i32 2957, i32 5798}
+!4 = !{!"branch_weights", i32 1807, i32 35058}
+!5 = !{!"branch_weights", i32 38, i32 287958}

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

@@ -0,0 +1,32 @@
+; RUN: opt -debug-only=branch-prob -jump-threading -S %s 2>&1 | FileCheck %s
+
+; Make sure that we set the branch probability for the newly created
+; basic block.
+
+define void @foo(i32 %v0, i1 %arg2) !prof !0 !PGOFuncName !1 {
+entry:
+  %bool1 = icmp eq i32 %v0, 0
+  br i1 %bool1, label %bb2, label %bb1, !prof !2
+
+bb1:
+  %sel = select i1 %arg2, i32 %v0, i32 0, !prof !3
+  br label %bb2
+; CHECK: set edge select.unfold -> 0 successor probability to 0x80000000 / 0x80000000
+
+bb2:
+  %phi = phi i32 [ %sel, %bb1 ], [ 0, %entry ]
+  %bool2 = icmp eq i32 %phi, 0
+  br i1 %bool2, label %exit, label %bb3, !prof !4
+
+bb3:
+  br label %exit
+
+exit:
+  ret void
+}
+
+!0 = !{!"function_entry_count", i64 15985}
+!1 = !{!"foo.cpp:foo"}
+!2 = !{!"branch_weights", i32 0, i32 36865}
+!3 = !{!"branch_weights", i32 35058, i32 1807}
+!4 = !{!"branch_weights", i32 1807, i32 35058}

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

@@ -0,0 +1,50 @@
+; RUN: opt -debug-only=branch-prob -jump-threading -S %s 2>&1 | FileCheck %s
+
+; Make sure that we set the branch probability for the newly created
+; basic block.
+
+define void @foo(i1 %cond1, i1 %cond2) !prof !0 !PGOFuncName !1 {
+entry:
+  br i1 %cond1, label %bb.f1, label %bb.f2, !prof !2
+
+bb.f1:
+  call void @f1()
+  br label %bb.cond2
+
+bb.f2:
+  call void @f2()
+  br label %bb.cond2
+
+bb.cond2:
+  br i1 %cond2, label %exit, label %bb.cond1again, !prof !3
+; CHECK: set edge bb.cond2.thread -> 0 successor probability to 0x79b9d244 / 0x80000000
+; CHECK: set edge bb.cond2.thread -> 1 successor probability to 0x06462dbc / 0x80000000 = 4.90
+
+bb.cond1again:
+  br i1 %cond1, label %bb.f3, label %bb.f4, !prof !4
+
+bb.f3:
+  call void @f3()
+  br label %exit
+
+bb.f4:
+  call void @f4()
+  br label %exit
+
+exit:
+  ret void
+}
+
+declare void @f1()
+
+declare void @f2()
+
+declare void @f3()
+
+declare void @f4()
+
+!0 = !{!"function_entry_count", i64 15985}
+!1 = !{!"foo.cpp:foo"}
+!2 = !{!"branch_weights", i32 0, i32 36865}
+!3 = !{!"branch_weights", i32 35058, i32 1807}
+!4 = !{!"branch_weights", i32 1807, i32 35058}