diff --git a/llvm/lib/Transforms/Scalar/JumpThreading.cpp b/llvm/lib/Transforms/Scalar/JumpThreading.cpp --- a/llvm/lib/Transforms/Scalar/JumpThreading.cpp +++ b/llvm/lib/Transforms/Scalar/JumpThreading.cpp @@ -2239,12 +2239,8 @@ cloneInstructions(PredBB->begin(), PredBB->end(), NewBB, PredPredBB); // Copy the edge probabilities from PredBB to NewBB. - if (HasProfileData) { - SmallVector Probs; - for (BasicBlock *Succ : successors(PredBB)) - Probs.push_back(BPI->getEdgeProbability(PredBB, Succ)); - BPI->setEdgeProbability(NewBB, Probs); - } + if (HasProfileData) + BPI->copyEdgeProbabilities(PredBB, NewBB); // Update the terminator of PredPredBB to jump to NewBB instead of PredBB. // This eliminates predecessors from PredPredBB, which requires us to simplify @@ -2674,6 +2670,8 @@ // Remove the unconditional branch at the end of the PredBB block. OldPredBranch->eraseFromParent(); + if (HasProfileData) + BPI->copyEdgeProbabilities(BB, PredBB); DTU->applyUpdatesPermissive(Updates); ++NumDupes; diff --git a/llvm/test/Transforms/JumpThreading/thread-prob-3.ll b/llvm/test/Transforms/JumpThreading/thread-prob-3.ll new file mode 100644 --- /dev/null +++ b/llvm/test/Transforms/JumpThreading/thread-prob-3.ll @@ -0,0 +1,29 @@ +; RUN: opt -debug-only=branch-prob -jump-threading -S %s 2>&1 | FileCheck %s +; REQUIRES: asserts + +; Make sure that we set edge probabilities for bb2 as we +; call DuplicateCondBranchOnPHIIntoPred(bb3, {bb2}). +; +; CHECK-LABEL: ---- Branch Probability Info : foo +; CHECK: set edge bb2 -> 0 successor probability to 0x7fffffff / 0x80000000 = 100.00% +; CHECK-NEXT: set edge bb2 -> 1 successor probability to 0x00000001 / 0x80000000 = 0.00% +define void @foo(i1 %f0, i1 %f1, i1 %f2) !prof !{!"function_entry_count", i64 0} { +; CHECK-LABEL: @foo( +bb1: + br i1 %f0, label %bb3, label %bb2 + +bb2: +; CHECK: bb2: +; CHECK-NEXT: br i1 %f2, label %exit1, label %unreach + br label %bb3 + +bb3: + %ph = phi i1 [ %f1, %bb1 ], [ %f2, %bb2 ] + br i1 %ph, label %exit1, label %unreach + +exit1: + ret void + +unreach: + unreachable +}