[JumpThreading] Call eraseBlock when folding a conditional branch

[kazu]

Summary

This patch teaches the jump threading pass to call BPI->eraseBlock
when it folds a conditional branch.

Without this patch, BranchProbabilityInfo could end up with stale edge
probabilities for the basic block containing the conditional branch --
one edge probability with less than 1.0 and the other for a removed
edge.

This patch is one of the steps before we can safely re-apply D91017.

[yrouban]

The patch looks good, but I have some concerns.

This patch is one of the steps before we can safely re-apply D91017.

Will there be any other patches? Please add links if any and dependencies to D91017.
Could you please describe the problem that prevents D91017?

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

13

Would it be possible to add checks that describe the CFG change like:
; CHECK-LABEL: @foo
; CHECK: bb.entry:
; CHECK-NEXT: br i1 %cond, label %bb.31, label %bb.12
; CHECK-NOT: bb.cond:

[kazu]

In D91511#2396561, @yrouban wrote:

The patch looks good, but I have some concerns.

This patch is one of the steps before we can safely re-apply D91017.

Will there be any other patches? Please add links if any and dependencies to D91017.
Could you please describe the problem that prevents D91017?

We need more calls to BPI->eraseBlock() if we want to require all basic blocks, including dead ones yet to be removed, to satisfy:

SrcProbs.size() == Src->getTerminator()->getNumSuccessors()

at all times during the jump threading pass except when it is actively modifying the CFG.

If we are OK with the consistency just for those basic blocks that are live, then the call to BPI->eraseBlock() being introduced in this patch may be enough. (I haven't checked.)

The problem caused by D91017 is due to the jump threading pass not updating edge probabilities (as addressed in this patch).

Once this patch lands, I'm going to build large pieces of software with PGO enabled to see what else is broken.

[kazu]

Updated the testcase.

[yrouban]

In D91511#2396576, @kazu wrote:

The problem caused by D91017 is due to the jump threading pass not updating edge probabilities (as addressed in this patch).

D91017 is NFC, why does it reveal the bug in JumpThreading while it has not been seen without D91017? Can we create a test case, which passes without D91017 but fails with it?

[kazu]

In D91511#2396581, @yrouban wrote:

In D91511#2396576, @kazu wrote:

The problem caused by D91017 is due to the jump threading pass not updating edge probabilities (as addressed in this patch).

D91017 is NFC, why does it reveal the bug in JumpThreading while it has not been seen without D91017? Can we create a test case, which passes without D91017 but fails with it?

D91017 introduces an additional assert in getEdgeProbability(const BasicBlock *Src, const BasicBlock *Dst), and that's causing trouble. The testcase in this patch should serve the testcase you are describing.

Once I do enough testing with profile guided optimizations to exercise BPI related code in the jump threading pass, I'm going to re-introduce the assert above by itself without D91017, wait for a bit, and then re-apply D91017. Thoughts?

[yrouban]

agree. See my comments in D91017.

llvm/lib/Analysis/BranchProbabilityInfo.cpp

BranchProbabilityInfo::printEdgeProbability(raw_ostream &OS,
   OS << "edge " << Src->getName() << " -> " << Dst->getName()
      << " probability is " << Prob
      << (isEdgeHot(Src, Dst) ? " [HOT edge]\n" : "\n");
 
   return OS;
 }
 
 void BranchProbabilityInfo::eraseBlock(const BasicBlock *BB) {
+  LLVM_DEBUG(dbgs() << "eraseBlock " << BB->getName() << "\n");
+
   // Note that we cannot use successors of BB because the terminator of BB may
   // have changed when eraseBlock is called as a BasicBlockCallbackVH callback.
   // Instead we remove prob data for the block by iterating successors by their
   // indices from 0 till the last which exists. There could not be prob data for
   // a pair (BB, N) if there is no data for (BB, N-1) because the data is always
   // set for all successors from 0 to M at once by the method
   // setEdgeProbability().
   Handles.erase(BasicBlockCallbackVH(BB, this));

llvm/lib/Transforms/Scalar/JumpThreading.cpp

if (BB->hasNPredecessors(PredToDestList.size())) {
         }
       }
 
       // Finally update the terminator.
       Instruction *Term = BB->getTerminator();
       BranchInst::Create(OnlyDest, Term);
       Term->eraseFromParent();
       DTU->applyUpdatesPermissive(Updates);
+      if (HasProfileData)
+        BPI->eraseBlock(BB);
 
       // If the condition is now dead due to the removal of the old terminator,
       // erase it.
       if (auto *CondInst = dyn_cast<Instruction>(Cond)) {
         if (CondInst->use_empty() && !CondInst->mayHaveSideEffects())
           CondInst->eraseFromParent();
         // We can safely replace *some* uses of the CondInst if it has
         // exactly one value as returned by LVI. RAUW is incorrect in the

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

+; RUN: opt -debug-only=branch-prob -jump-threading -S %s 2>&1 | FileCheck %s
+; REQUIRES: asserts
+
+; Make sure that we clear edge probabilities for bb.cond as we fold
+; the conditional branch in it.
+
+define i32 @foo(i1 %cond) !prof !0 {
+bb.entry:
+  br i1 %cond, label %bb.31, label %bb.cond
+
+bb.cond:
+  br i1 %cond, label %bb.31, label %bb.12
+; CHECK: eraseBlock bb.cond
+
+bb.31:
+  ret i32 31
+
+bb.12:
+  ret i32 12
+}
+
+!0 = !{!"function_entry_count", i64 0}