[PATCH] D138132: [JumpThreading] Preserve profile metadata during select unfolding

[Roman Paukner via Phabricator via llvm-commits]

rp updated this revision to Diff 476840.
rp added a comment.

rebase


Repository:
  rG LLVM Github Monorepo

CHANGES SINCE LAST ACTION
  https://reviews.llvm.org/D138132/new/

https://reviews.llvm.org/D138132

Files:
  llvm/lib/Transforms/Scalar/JumpThreading.cpp
  llvm/test/Transforms/JumpThreading/select-unfold-prof.ll


Index: llvm/test/Transforms/JumpThreading/select-unfold-prof.ll
===================================================================
--- /dev/null
+++ 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}
Index: llvm/lib/Transforms/Scalar/JumpThreading.cpp
===================================================================
--- llvm/lib/Transforms/Scalar/JumpThreading.cpp
+++ 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<BranchProbability, 2> 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();


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D138132.476840.patch
Type: text/x-patch
Size: 3870 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20221121/cc497b27/attachment.bin>