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 @@ -2236,6 +2236,14 @@ DenseMap ValueMapping = 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); + } + // Update the terminator of PredPredBB to jump to NewBB instead of PredBB. // This eliminates predecessors from PredPredBB, which requires us to simplify // any PHI nodes in PredBB. @@ -2424,8 +2432,15 @@ if (HasProfileData) // Update frequencies between Pred -> NewBB. NewBBFreq += FreqMap.lookup(Pred); } - if (HasProfileData) // Apply the summed frequency to NewBB. + if (HasProfileData) { + // Apply the summed frequency to NewBB. BFI->setBlockFreq(NewBB, NewBBFreq.getFrequency()); + + // NewBB has exactly one successor. + SmallVector BBSuccProbs; + BBSuccProbs.push_back(BranchProbability::getOne()); + BPI->setEdgeProbability(NewBB, BBSuccProbs); + } } DTU->applyUpdatesPermissive(Updates); @@ -2498,6 +2513,11 @@ // Update edge probabilities in BPI. BPI->setEdgeProbability(BB, BBSuccProbs); + // NewBB has exactly one successor. + SmallVector NewBBSuccProbs; + NewBBSuccProbs.push_back(BranchProbability::getOne()); + BPI->setEdgeProbability(NewBB, NewBBSuccProbs); + // Update the profile metadata as well. // // Don't do this if the profile of the transformed blocks was statically @@ -2708,6 +2728,13 @@ PHINode *Phi = dyn_cast(BI); ++BI) if (Phi != SIUse) Phi->addIncoming(Phi->getIncomingValueForBlock(Pred), NewBB); + + if (HasProfileData) { + // NewBB has exactly one successor. + SmallVector BBSuccProbs; + BBSuccProbs.push_back(BranchProbability::getOne()); + BPI->setEdgeProbability(NewBB, BBSuccProbs); + } } bool JumpThreadingPass::TryToUnfoldSelect(SwitchInst *SI, BasicBlock *BB) { diff --git a/llvm/test/Transforms/JumpThreading/thread-prob-1.ll b/llvm/test/Transforms/JumpThreading/thread-prob-1.ll new file mode 100644 --- /dev/null +++ b/llvm/test/Transforms/JumpThreading/thread-prob-1.ll @@ -0,0 +1,42 @@ +; RUN: opt -debug-only=branch-prob -jump-threading -S %s 2>&1 | FileCheck %s + +; Make sure that we set the branch probability for the newly created +; basic block. + +define void @foo(i1 %arg1, i1 %arg2, i32 %arg3) !prof !0 !PGOFuncName !1 { +entry: + call void @bar(i32 0) + br i1 %arg1, label %bb3, label %bb1, !prof !2 + +bb1: + call void @bar(i32 1) + br i1 %arg2, label %bb2, label %bb3, !prof !3 + +bb2: + call void @bar(i32 2) + br label %bb3 + +bb3: +; CHECK: set edge bb3.thr_comm -> 0 successor probability to 0x80000000 / 0x80000000 +%ptr = phi i32 [ 0, %bb1 ], [ 0, %entry ], [ %arg3, %bb2 ] + call void @bar(i32 3) + %bool = icmp eq i32 %ptr, 0 + br i1 %bool, label %exit, label %bb4, !prof !4 +; CHECK: set edge bb3.thread -> 0 successor probability to 0x80000000 / 0x80000000 + +bb4: + call void @bar(i32 %ptr) + br label %exit + +exit: + ret void +} + +declare void @bar(i32) + +!0 = !{!"function_entry_count", i64 15985} +!1 = !{!"foo:foo"} +!2 = !{!"branch_weights", i32 15973, i32 36865} +!3 = !{!"branch_weights", i32 2957, i32 5798} +!4 = !{!"branch_weights", i32 1807, i32 35058} +!5 = !{!"branch_weights", i32 38, i32 287958} diff --git a/llvm/test/Transforms/JumpThreading/thread-prob-2.ll b/llvm/test/Transforms/JumpThreading/thread-prob-2.ll new file mode 100644 --- /dev/null +++ b/llvm/test/Transforms/JumpThreading/thread-prob-2.ll @@ -0,0 +1,32 @@ +; RUN: opt -debug-only=branch-prob -jump-threading -S %s 2>&1 | FileCheck %s + +; Make sure that we set the branch probability for the newly created +; basic block. + +define void @foo(i32 %v0, i1 %arg2) !prof !0 !PGOFuncName !1 { +entry: + %bool1 = icmp eq i32 %v0, 0 + br i1 %bool1, label %bb2, label %bb1, !prof !2 + +bb1: + %sel = select i1 %arg2, i32 %v0, i32 0, !prof !3 + br label %bb2 +; CHECK: set edge select.unfold -> 0 successor probability to 0x80000000 / 0x80000000 + +bb2: + %phi = phi i32 [ %sel, %bb1 ], [ 0, %entry ] + %bool2 = icmp eq i32 %phi, 0 + br i1 %bool2, label %exit, label %bb3, !prof !4 + +bb3: + br label %exit + +exit: + ret void +} + +!0 = !{!"function_entry_count", i64 15985} +!1 = !{!"foo.cpp:foo"} +!2 = !{!"branch_weights", i32 0, i32 36865} +!3 = !{!"branch_weights", i32 35058, i32 1807} +!4 = !{!"branch_weights", i32 1807, i32 35058} 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,50 @@ +; RUN: opt -debug-only=branch-prob -jump-threading -S %s 2>&1 | FileCheck %s + +; Make sure that we set the branch probability for the newly created +; basic block. + +define void @foo(i1 %cond1, i1 %cond2) !prof !0 !PGOFuncName !1 { +entry: + br i1 %cond1, label %bb.f1, label %bb.f2, !prof !2 + +bb.f1: + call void @f1() + br label %bb.cond2 + +bb.f2: + call void @f2() + br label %bb.cond2 + +bb.cond2: + br i1 %cond2, label %exit, label %bb.cond1again, !prof !3 +; CHECK: set edge bb.cond2.thread -> 0 successor probability to 0x79b9d244 / 0x80000000 +; CHECK: set edge bb.cond2.thread -> 1 successor probability to 0x06462dbc / 0x80000000 = 4.90 + +bb.cond1again: + br i1 %cond1, label %bb.f3, label %bb.f4, !prof !4 + +bb.f3: + call void @f3() + br label %exit + +bb.f4: + call void @f4() + br label %exit + +exit: + ret void +} + +declare void @f1() + +declare void @f2() + +declare void @f3() + +declare void @f4() + +!0 = !{!"function_entry_count", i64 15985} +!1 = !{!"foo.cpp:foo"} +!2 = !{!"branch_weights", i32 0, i32 36865} +!3 = !{!"branch_weights", i32 35058, i32 1807} +!4 = !{!"branch_weights", i32 1807, i32 35058}