[JumpThreading] Set edge probabilities when creating basic blocks

[kazu]

Summary

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.

[MaskRay]

I've debugged this two weeks ago:)

Now, we rarely remove entries from this map, so if we happen to ...

Thanks for narrowing it down to the issue! LGTM

Do you know whether we have missed some BranchProbabilityInfo::eraseBlock calls? Why is BasicBlockCallbackVH::deleted() not effective when some BasicBlock's are removed?

[wmi]

Do you know whether we have missed some BranchProbabilityInfo::eraseBlock calls? Why is BasicBlockCallbackVH::deleted() not effective when some BasicBlock's are removed?

Same question here. There shouldn't be stale value in Prob if the value handler is functioning correctly.

[kazu]

In D90106#2355030, @wmi wrote:

Do you know whether we have missed some BranchProbabilityInfo::eraseBlock calls? Why is BasicBlockCallbackVH::deleted() not effective when some BasicBlock's are removed?

Same question here. There shouldn't be stale value in Prob if the value handler is functioning correctly.

Please see the description of the follow-up patch I just posted: https://reviews.llvm.org/D90272

[dblaikie]

Probably would'ev been worth waiting to check that the two reviewers with the question found the answer addressed their concerns (& getting a clear approval) before committing. (bit awkward to commit things that haven't been formally approved - though I realize @MaskRay's comment may be close enough to be construed as approval)

[MaskRay]

In D90106#2357780, @dblaikie wrote:

Probably would'ev been worth waiting to check that the two reviewers with the question found the answer addressed their concerns (& getting a clear approval) before committing. (bit awkward to commit things that haven't been formally approved - though I realize @MaskRay's comment may be close enough to be construed as approval)

Thanks @dblaikie!

llvm/lib/Transforms/Scalar/JumpThreading.cpp

2517

With D90272, do we still need this chunk?

getEdgeProbability by default knows the probability is 1.

2734

With D90272, do we still need this chunk?

getEdgeProbability by default knows the probability is 1.

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

2

debug-only tests require REQUIRES: asserts

The 3 new tests are useful but I wonder -debug-only=branch-prob is the best way testing them. Such tests should also be used prudently as debug only codes tends to be very specific internal details.

[MaskRay]

I would have required the description of this patch to be fixed if I saw D90272 earlier.

[anna]

Hi @kazu, just wanted to point out - We have seen multiple assertion failures in benchmarks downstream with this patch (or a related one built on this patch). Trying to revert this downstream shows there are dependent patches. Is this assertion seen and already fixed in trunk?

Assertion failure is BranchProbabilityInfo.cpp:1163: void llvm::BranchProbabilityInfo::setEdgeProbability(const llvm::BasicBlock*, const llvm::SmallVectorImpl&): Assertion TotalNumerator <= BranchProbability::getDenominator() + Probs.size()'

@yrouban narrowed it down to:
"Debugging showed that the assertion detects incorrect set of probabilities (summ != 1.0) when they are copied from PredBB to NewBB in JumpThreadingPass::ThreadThroughTwoBasicBlocks(). The probabilities are corrupted in PredBB when it gets cloned instructions from its successor in JumpThreadingPass::DuplicateCondBranchOnPHIIntoPred(). So the bug is in this method which does not fix the probabilities. "
He is working on a fix for this with upstream testcase.

[kazu]

In D90106#2371631, @anna wrote:

Hi @kazu, just wanted to point out - We have seen multiple assertion failures in benchmarks downstream with this patch (or a related one built on this patch). Trying to revert this downstream shows there are dependent patches. Is this assertion seen and already fixed in trunk?

Assertion failure is BranchProbabilityInfo.cpp:1163: void llvm::BranchProbabilityInfo::setEdgeProbability(const llvm::BasicBlock*, const llvm::SmallVectorImpl&): Assertion TotalNumerator <= BranchProbability::getDenominator() + Probs.size()'

@yrouban narrowed it down to:
"Debugging showed that the assertion detects incorrect set of probabilities (summ != 1.0) when they are copied from PredBB to NewBB in JumpThreadingPass::ThreadThroughTwoBasicBlocks(). The probabilities are corrupted in PredBB when it gets cloned instructions from its successor in JumpThreadingPass::DuplicateCondBranchOnPHIIntoPred(). So the bug is in this method which does not fix the probabilities. "
He is working on a fix for this with upstream testcase.

Those follow-up patches do not touch the code to call setEdgeProbability from JumpThreadingPass::ThreadThroughTwoBasicBlocks, so I'm pretty sure that the issue isn't fixed on the trunk yet. Thank you for tracking down the bug. I'd appreciate if you put me as a reviewer on your patch.

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);
+  }
+
   // 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);
       PredPredTerm->setSuccessor(i, NewBB);
for (auto NewBB : NewBBs) {
     BlockFrequency NewBBFreq(0);
     Updates.push_back({DominatorTree::Insert, NewBB, BB});
     for (auto Pred : predecessors(NewBB)) {
       Updates.push_back({DominatorTree::Delete, Pred, BB});
       Updates.push_back({DominatorTree::Insert, Pred, NewBB});
       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);
   return NewBBs[0];
 }
 
 bool JumpThreadingPass::doesBlockHaveProfileData(BasicBlock *BB) {
   const Instruction *TI = BB->getTerminator();
else {
     // Normalize edge probabilities so that they sum up to one.
     BranchProbability::normalizeProbabilities(BBSuccProbs.begin(),
                                               BBSuccProbs.end());
   }
 
   // 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
   // estimated.  (This could occur despite the function having an entry
   // frequency in completely cold parts of the CFG.)
   //
   // In this case we don't want to suggest to subsequent passes that the
   // calculated weights are fully consistent.  Consider this graph:
void JumpThreadingPass::UnfoldSelectInstr(BasicBlock *Pred, BasicBlock *BB,
   DTU->applyUpdatesPermissive({{DominatorTree::Insert, NewBB, BB},
                                {DominatorTree::Insert, Pred, NewBB}});
 
   // Update any other PHI nodes in BB.
   for (BasicBlock::iterator BI = BB->begin();
        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) {
   PHINode *CondPHI = dyn_cast<PHINode>(SI->getCondition());
 
   if (!CondPHI || CondPHI->getParent() != BB)
     return false;
 

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

+; 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

+; 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

+; 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}