[JumpThreading] Skip unconditional PredBB when threading jumps through two basic blocks

[MaskRay]

Summary

Fixes https://bugs.llvm.org/show_bug.cgi?id=44922 (caused by 4698bf145d583e26ed438026ef7fde031ef322b1)

ThreadThroughTwoBasicBlocks assumes PredBBBranch is conditional.

AddPHINodeEntriesForMappedBlock(PredBBBranch->getSuccessor(1), PredBB, NewBB,
                                ValueMapping);

can segfault.

[kazu]

Thank you for catching this! I meant to handle PredBB ending with a conditional branch only, but I wasn't aware that PredBB could end with an unconditional branch if its address is taken.

llvm/test/Transforms/JumpThreading/thread-two-bbs6.ll

8

Nit: May I suggest to remove dso_local if it's not critical to the test?

[MaskRay]

Delete dso_local

[MaskRay]

Improve bb names

[MaskRay]

Skip unconditional PredBB

[kazu]

LGTM.

llvm/lib/Transforms/Scalar/JumpThreading.cpp

2116

Could you say "conditional Branch" instead of "Branch"? Thanks!

[MaskRay]

Improve comments

[kees]

Thank you! I can confirm this fixes the problems I saw building the Linux kernel with CONFIG_UBSAN=y.

llvm/lib/Transforms/Scalar/JumpThreading.cpp

bool JumpThreadingPass::MaybeThreadThroughTwoBasicBlocks(BasicBlock *BB,
   if (!CondBr)
     return false;
 
   // BB must have exactly one predecessor.
   BasicBlock *PredBB = BB->getSinglePredecessor();
   if (!PredBB)
     return false;
 
   // Require that PredBB end with a Branch.  If PredBB ends with an
   // unconditional branch, we should be merging PredBB and BB instead.  For
   // simplicity, we don't deal with a switch.
   BranchInst *PredBBBranch = dyn_cast<BranchInst>(PredBB->getTerminator());
   if (!PredBBBranch)
     return false;
 
   // If PredBB has exactly one incoming edge, we don't gain anything by copying
   // PredBB.
void JumpThreadingPass::ThreadThroughTwoBasicBlocks(BasicBlock *PredPredBB,
   for (unsigned i = 0, e = PredPredTerm->getNumSuccessors(); i != e; ++i)
     if (PredPredTerm->getSuccessor(i) == PredBB) {
       PredBB->removePredecessor(PredPredBB, true);
       PredPredTerm->setSuccessor(i, NewBB);
     }
 
   AddPHINodeEntriesForMappedBlock(PredBBBranch->getSuccessor(0), PredBB, NewBB,
                                   ValueMapping);
-  AddPHINodeEntriesForMappedBlock(PredBBBranch->getSuccessor(1), PredBB, NewBB,
-                                  ValueMapping);
+  if (PredBBBranch->isConditional())
+    AddPHINodeEntriesForMappedBlock(PredBBBranch->getSuccessor(1), PredBB,
+                                    NewBB, ValueMapping);
 
   DTU->applyUpdatesPermissive(
       {{DominatorTree::Insert, NewBB, CondBr->getSuccessor(0)},
        {DominatorTree::Insert, NewBB, CondBr->getSuccessor(1)},
        {DominatorTree::Insert, PredPredBB, NewBB},
        {DominatorTree::Delete, PredPredBB, PredBB}});
 
   UpdateSSA(PredBB, NewBB, ValueMapping);

llvm/test/Transforms/JumpThreading/thread-two-bbs6.ll

+;; Test that we can handle unconditional PredBB when threading jumps through two
+;; successive basic blocks.
+; RUN: opt -S -passes='function(jump-threading)' < %s | FileCheck %s
+
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+define dso_local i32 @f(i32* %0) {
+; CHECK-LABEL: @f
+; CHECK: br i1 %good, label %bb, label %cont1.thread
+entry:
+  %size = call i64 @get_size(i32* %0)
+  %good = icmp ugt i64 %size, 3
+  br i1 %good, label %bb, label %pred.bb
+
+; CHECK:      bb:
+; CHECK-NEXT:  %v = load i32, i32* %0
+; CHECK-NEXT:  br label %cont1
+; CHECK:      cont1.thread:
+; CHECK:       call void @effect()
+; CHECK-NEXT:  %v1 = load i32, i32* %0
+; CHECK-NEXT:  call void @effect1(i8* blockaddress(@f, %cont1))
+; CHECK:      cont1:
+; CHECK:       call void @effect1(i8* blockaddress(@f, %cont1))
+pred.bb:                                             ; preds = %entry
+  call void @effect()
+  br label %bb
+bb:                                                  ; preds = %pred.bb, %entry
+  %v = load i32, i32* %0
+  br label %cont1
+
+cont1:                                               ; preds = %bb
+  call void @effect1(i8* blockaddress(@f, %cont1))
+  br i1 %good, label %cont3, label %cont2
+
+cont2:                                               ; preds = %cont1
+  br i1 %good, label %exit, label %cont3
+cont3:                                               ; preds = %cont2
+  br label %exit
+exit:                                                ; preds = %cont3, %cont2
+  ret i32 %v
+}
+
+declare i64 @get_size(i32*)
+declare void @effect()
+declare void @effect1(i8*)