[JumpThreading] [PR40992] Fix miscompile when folding a successor of an indirectbr

[brzycki]

Summary

The continue statement in JumpThreadingPass::ProcessThreadableEdges() correctly identifies the case when a predecessor is an indirectbr or callbr and must not be threaded. Unfortunately the bookkeeping for the remaining edges was incorrect and there are cases where a conditional branch in BB may be incorrectly folded based on the flawed analysis above. This patch correctly prevents folding in these cases while not fully pessimizing other threading opportunities to the same BB.

Thanks to Matthias Liedtke for creating a simplified test-case and concrete failure example.

See PR40992 for debug analysis.

[brzycki]

Passes make check and test-suite runs.

[craig.topper]

llvm/lib/Transforms/Scalar/JumpThreading.cpp

1662

Can we just use PredToDestList.size() here? And switch to BB->hasNPredecessors(PredToDestList.size()). pred_size is linear in the number of predecessors, but hasNPredecessors will only count to N and return false if that's not the end without continuing to the end.

[brzycki]

llvm/lib/Transforms/Scalar/JumpThreading.cpp

1662

That's fine by me. I only wrote it this way to make minimal changes. I'll try your suggested change and test for correctness.

[brzycki]

Address review comments.

[brzycki]

Added missing check to testcase ensuring %bb1 threads to %bb7 when %c is false.

[sebpop]

Looks good to me.

llvm/lib/Transforms/Scalar/JumpThreading.cpp

bool JumpThreadingPass::ProcessThreadableEdges(Value *Cond, BasicBlock *BB,
   SmallPtrSet<BasicBlock*, 16> SeenPreds;
   SmallVector<std::pair<BasicBlock*, BasicBlock*>, 16> PredToDestList;
 
   BasicBlock *OnlyDest = nullptr;
   BasicBlock *MultipleDestSentinel = (BasicBlock*)(intptr_t)~0ULL;
   Constant *OnlyVal = nullptr;
   Constant *MultipleVal = (Constant *)(intptr_t)~0ULL;
 
-  unsigned PredWithKnownDest = 0;
   for (const auto &PredValue : PredValues) {
     BasicBlock *Pred = PredValue.second;
     if (!SeenPreds.insert(Pred).second)
       continue;  // Duplicate predecessor entry.
 
     Constant *Val = PredValue.first;
 
     BasicBlock *DestBB;
if (PredToDestList.empty()) {
       if (OnlyDest != DestBB)
         OnlyDest = MultipleDestSentinel;
       // It possible we have same destination, but different value, e.g. default
       // case in switchinst.
       if (Val != OnlyVal)
         OnlyVal = MultipleVal;
     }
 
-    // We know where this predecessor is going.
-    ++PredWithKnownDest;
-
     // If the predecessor ends with an indirect goto, we can't change its
     // destination. Same for CallBr.
     if (isa<IndirectBrInst>(Pred->getTerminator()) ||
         isa<CallBrInst>(Pred->getTerminator()))
       continue;
 
     PredToDestList.push_back(std::make_pair(Pred, DestBB));
   }
 
   // If all edges were unthreadable, we fail.
   if (PredToDestList.empty())
     return false;
 
   // If all the predecessors go to a single known successor, we want to fold,
   // not thread. By doing so, we do not need to duplicate the current block and
   // also miss potential opportunities in case we dont/cant duplicate.
   if (OnlyDest && OnlyDest != MultipleDestSentinel) {
-    if (PredWithKnownDest == (size_t)pred_size(BB)) {
+    if (BB->hasNPredecessors(PredToDestList.size())) {
       bool SeenFirstBranchToOnlyDest = false;
       std::vector <DominatorTree::UpdateType> Updates;
       Updates.reserve(BB->getTerminator()->getNumSuccessors() - 1);
       for (BasicBlock *SuccBB : successors(BB)) {
         if (SuccBB == OnlyDest && !SeenFirstBranchToOnlyDest) {
           SeenFirstBranchToOnlyDest = true; // Don't modify the first branch.
         } else {
           SuccBB->removePredecessor(BB, true); // This is unreachable successor.

llvm/test/Transforms/JumpThreading/pr40992-indirectbr-folding.ll

+; RUN: opt -S < %s -jump-threading | FileCheck %s
+
+; PR40992: Do not incorrectly fold %bb5 into an unconditional br to %bb7.
+;          Also verify we correctly thread %bb1 -> %bb7 when %c is false.
+
+define i32 @jtbr(i1 %v1, i1 %v2, i1 %v3) {
+; CHECK: bb0:
+bb0:
+  br label %bb1
+
+; CHECK: bb1:
+bb1:
+  %c = and i1 %v1, %v2
+  br i1 %c, label %bb2, label %bb5
+
+; CHECK: bb2:
+; CHECK-NEXT: select
+; CHECK-NEXT: indirectbr i8* %ba, [label %bb3, label %bb5]
+bb2:
+  %ba = select i1 %v3, i8* blockaddress(@jtbr, %bb3), i8* blockaddress(@jtbr, %bb4)
+  indirectbr i8* %ba, [label %bb3, label %bb4]
+
+; CHECK: bb3:
+bb3:
+  br label %bb1
+
+; CHECK-NOT: bb4:
+bb4:
+  br label %bb5
+
+; CHECK: bb5:
+bb5:
+  br i1 %c, label %bb6, label %bb7
+
+; CHECK: bb6:
+bb6:
+  ret i32 0
+
+; CHECK: bb7:
+bb7:
+  ret i32 1
+}