[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.

[kazu]

Added one more case.

[yrouban]

llvm/lib/Transforms/Scalar/JumpThreading.cpp

1097

should not we BPI->eraseBlock(BB) here? another patch?

1114

can we add an assertion on the return value of the ConstantFoldTerminator() call?

[kazu]

Thanks for the review!

llvm/lib/Transforms/Scalar/JumpThreading.cpp

1097

Thank you for spotting. Yes, I'll take care of this in another patch.

1114

Good point. Let me address this in another patch after testing. Also, I'd like to make it easier to revert the assertion patch just in case.

llvm/lib/Transforms/Scalar/JumpThreading.cpp

for (unsigned i = 0, e = BBTerm->getNumSuccessors(); i != e; ++i) {
       if (i == BestSucc) continue;
       BasicBlock *Succ = BBTerm->getSuccessor(i);
       Succ->removePredecessor(BB, true);
       Updates.push_back({DominatorTree::Delete, BB, Succ});
     }
 
     LLVM_DEBUG(dbgs() << "  In block '" << BB->getName()
                       << "' folding undef terminator: " << *BBTerm << '\n');
     BranchInst::Create(BBTerm->getSuccessor(BestSucc), BBTerm);
     BBTerm->eraseFromParent();
     DTU->applyUpdatesPermissive(Updates);
     if (FI)
       FI->eraseFromParent();
     return true;
   }
 
   // If the terminator of this block is branching on a constant, simplify the
   // terminator to an unconditional branch.  This can occur due to threading in
   // other blocks.
   if (getKnownConstant(Condition, Preference)) {
     LLVM_DEBUG(dbgs() << "  In block '" << BB->getName()
                       << "' folding terminator: " << *BB->getTerminator()
                       << '\n');
     ++NumFolds;
     ConstantFoldTerminator(BB, true, nullptr, DTU);
+    if (HasProfileData)
+      BPI->eraseBlock(BB);
     return true;
   }
 
   Instruction *CondInst = dyn_cast<Instruction>(Condition);
 
   // All the rest of our checks depend on the condition being an instruction.
   if (!CondInst) {
     // FIXME: Unify this with code below.
if (CondBr && CondConst) {
         else if (CondCmp->getParent() == BB) {
           auto *CI = Ret == LazyValueInfo::True ?
             ConstantInt::getTrue(CondCmp->getType()) :
             ConstantInt::getFalse(CondCmp->getType());
           replaceFoldableUses(CondCmp, CI);
         }
         DTU->applyUpdatesPermissive(
             {{DominatorTree::Delete, BB, ToRemoveSucc}});
+        if (HasProfileData)
+          BPI->eraseBlock(BB);
         return true;
       }
 
       // We did not manage to simplify this branch, try to see whether
       // CondCmp depends on a known phi-select pattern.
       if (tryToUnfoldSelect(CondCmp, BB))
         return true;
     }
while (CurrentPred && Iter++ < ImplicationSearchThreshold) {
     if (Implication) {
       BasicBlock *KeepSucc = BI->getSuccessor(*Implication ? 0 : 1);
       BasicBlock *RemoveSucc = BI->getSuccessor(*Implication ? 1 : 0);
       RemoveSucc->removePredecessor(BB);
       BranchInst *UncondBI = BranchInst::Create(KeepSucc, BI);
       UncondBI->setDebugLoc(BI->getDebugLoc());
       BI->eraseFromParent();
       DTU->applyUpdatesPermissive({{DominatorTree::Delete, BB, RemoveSucc}});
+      if (HasProfileData)
+        BPI->eraseBlock(BB);
       return true;
     }
     CurrentBB = CurrentPred;
     CurrentPred = CurrentBB->getSinglePredecessor();
   }
 
   return false;
 }

llvm/test/Transforms/JumpThreading/thread-prob-4.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 bb1 as we fold
+; the conditional branch in it.
+
+; CHECK: eraseBlock bb1
+
+define i32 @foo(i32 %arg) !prof !0 {
+; CHECK-LABEL: @foo
+  %cond1 = icmp eq i32 %arg, 42
+  br i1 %cond1, label %bb1, label %bb2
+
+bb2:
+  ret i32 2
+
+bb1:
+  %cond2 = icmp eq i32 %arg, 42
+  br i1 %cond2, label %bb3, label %bb4
+
+bb3:
+  ret i32 3
+
+bb4:
+; CHECK-NOT: bb4:
+  ret i32 4
+}
+
+!0 = !{!"function_entry_count", i64 0}

llvm/test/Transforms/JumpThreading/thread-prob-5.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 bb1 as we fold
+; the conditional branch in it.
+
+; CHECK: eraseBlock bb1
+
+define void @foo(i32 %i, i32 %len) !prof !0 {
+; CHECK-LABEL: @foo
+  %i.inc = add nuw i32 %i, 1
+  %c0 = icmp ult i32 %i.inc, %len
+  br i1 %c0, label %bb1, label %bb2
+
+bb1:
+; CHECK: bb1:
+  %c1 = icmp ult i32 %i, %len
+  br i1 %c1, label %bb2, label %bb3
+
+bb2:
+  ret void
+
+bb3:
+; CHECK-NOT: bb3:
+  ret void
+}
+
+!0 = !{!"function_entry_count", i64 0}

llvm/test/Transforms/JumpThreading/thread-prob-6.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 bb1 as we fold
+; the conditional branch in it.
+
+; CHECK: eraseBlock bb1
+
+define i32 @foo() !prof !0 {
+; CHECK-LABEL: @foo
+bb1:
+  br i1 true, label %bb2, label %bb3
+
+bb2:
+  ret i32 3
+
+bb3:
+; CHECK-NOT: bb3:
+  ret i32 7
+}
+
+!0 = !{!"function_entry_count", i64 0}