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 @@ -2759,8 +2759,24 @@ // Create a conditional branch and update PHI nodes. auto *BI = BranchInst::Create(NewBB, BB, SI->getCondition(), Pred); BI->applyMergedLocation(PredTerm->getDebugLoc(), SI->getDebugLoc()); + BI->copyMetadata(*SI, {LLVMContext::MD_prof}); SIUse->setIncomingValue(Idx, SI->getFalseValue()); SIUse->addIncoming(SI->getTrueValue(), NewBB); + // Set the block frequency of NewBB. + if (HasProfileData) { + uint64_t TrueWeight, FalseWeight; + if (extractBranchWeights(*SI, TrueWeight, FalseWeight) && + (TrueWeight + FalseWeight) != 0) { + SmallVector BP; + BP.emplace_back(BranchProbability(TrueWeight, TrueWeight + FalseWeight)); + BP.emplace_back(BranchProbability(FalseWeight, TrueWeight + FalseWeight)); + BPI->setEdgeProbability(Pred, BP); + } + + auto NewBBFreq = + BFI->getBlockFreq(Pred) * BPI->getEdgeProbability(Pred, NewBB); + BFI->setBlockFreq(NewBB, NewBBFreq.getFrequency()); + } // The select is now dead. SI->eraseFromParent(); diff --git a/llvm/test/Transforms/JumpThreading/select-unfold-prof.ll b/llvm/test/Transforms/JumpThreading/select-unfold-prof.ll new file mode 100644 --- /dev/null +++ b/llvm/test/Transforms/JumpThreading/select-unfold-prof.ll @@ -0,0 +1,58 @@ +; NOTE: Assertions have been autogenerated by utils/update_test_checks.py +; RUN: opt -S -jump-threading -debug-only=branch-prob < %s 2>&1 | FileCheck %s + +; Jump threading of a branch with select as a condition. +; Prof atrribute attached to select should be preserved on conditional branch. + +declare void @foo() + +define void @unfold2(i32 %x, i32 %y) nounwind !prof !1 { +; CHECK-LABEL: ---- Branch Probability Info : unfold2 ---- +; CHECK: set edge cond.false -> 0 successor probability to 0x0ccccccd / 0x80000000 = 10.00% +; CHECK: set edge cond.false -> 1 successor probability to 0x73333333 / 0x80000000 = 90.00% +; CHECK-LABEL: @unfold2( +; CHECK-NEXT: entry: +; CHECK-NEXT: [[SUB:%.*]] = sub nsw i32 [[X:%.*]], [[Y:%.*]] +; CHECK-NEXT: [[CMP:%.*]] = icmp sgt i32 [[SUB]], 10 +; CHECK-NEXT: br i1 [[CMP]], label [[IF_END:%.*]], label [[COND_FALSE:%.*]] +; CHECK: cond.false: +; CHECK-NEXT: [[ADD:%.*]] = add nsw i32 [[X]], [[Y]] +; CHECK-NEXT: [[CMP1:%.*]] = icmp sgt i32 [[ADD]], 10 +; CHECK-NEXT: br i1 [[CMP1]], label [[IF_THEN:%.*]], label [[COND_END4:%.*]], !prof [[PROF1:![0-9]+]] +; CHECK: cond.end4: +; CHECK-NEXT: [[COND5:%.*]] = phi i32 [ [[ADD]], [[COND_FALSE]] ] +; CHECK-NEXT: [[CMP6:%.*]] = icmp eq i32 [[COND5]], 0 +; CHECK-NEXT: br i1 [[CMP6]], label [[IF_THEN]], label [[IF_END]] +; CHECK: if.then: +; CHECK-NEXT: call void @foo() +; CHECK-NEXT: br label [[IF_END]] +; CHECK: if.end: +; CHECK-NEXT: ret void +; +entry: + %sub = sub nsw i32 %x, %y + %cmp = icmp sgt i32 %sub, 10 + br i1 %cmp, label %cond.end4, label %cond.false + +cond.false: ; preds = %entry + %add = add nsw i32 %x, %y + %cmp1 = icmp sgt i32 %add, 10 + %add. = select i1 %cmp1, i32 0, i32 %add, !prof !0 + br label %cond.end4 + +cond.end4: ; preds = %entry, %cond.false + %cond5 = phi i32 [ %add., %cond.false ], [ %sub, %entry ] + %cmp6 = icmp eq i32 %cond5, 0 + br i1 %cmp6, label %if.then, label %if.end + +if.then: ; preds = %cond.end4 + call void @foo() + br label %if.end + +if.end: ; preds = %if.then, %cond.end4 + ret void + +} + +!0 = !{!"branch_weights", i32 1, i32 9} +!1 = !{!"function_entry_count", i64 1984}