[llvm] [JumpThreading] Use [BB->SuccIndx] to get probability when updating BB info. (PR #134585)

[via llvm-commits]

https://github.com/tianleliu updated https://github.com/llvm/llvm-project/pull/134585

>From 0c525f6304d9c2c16d8c4c91d8e9d670fc5405a0 Mon Sep 17 00:00:00 2001
From: tianleli <tianle.l.liu at intel.com>
Date: Sun, 6 Apr 2025 15:16:08 +0800
Subject: [PATCH 1/2] [JumpThreading] Use [BB->SuccIndx] to get probability
 when updating BB info.

In case the same src BB targets to the same dest BB in different
conditions/edges, such as switch-cases, we should use prob[SrcBB->SuccIndx]
instead of prob[SrcBB->DstBB].
---
 llvm/lib/Transforms/Scalar/JumpThreading.cpp  |  9 ++--
 .../Transforms/JumpThreading/thread-prob-8.ll | 41 +++++++++++++++++++
 2 files changed, 45 insertions(+), 5 deletions(-)
 create mode 100644 llvm/test/Transforms/JumpThreading/thread-prob-8.ll

diff --git a/llvm/lib/Transforms/Scalar/JumpThreading.cpp b/llvm/lib/Transforms/Scalar/JumpThreading.cpp
index 18d5f201413c8..5c112a429c6bc 100644
--- a/llvm/lib/Transforms/Scalar/JumpThreading.cpp
+++ b/llvm/lib/Transforms/Scalar/JumpThreading.cpp
@@ -2530,17 +2530,16 @@ void JumpThreadingPass::updateBlockFreqAndEdgeWeight(BasicBlock *PredBB,
   // frequency of BB.
   auto BBOrigFreq = BFI->getBlockFreq(BB);
   auto NewBBFreq = BFI->getBlockFreq(NewBB);
-  auto BB2SuccBBFreq = BBOrigFreq * BPI->getEdgeProbability(BB, SuccBB);
   auto BBNewFreq = BBOrigFreq - NewBBFreq;
   BFI->setBlockFreq(BB, BBNewFreq);
 
   // Collect updated outgoing edges' frequencies from BB and use them to update
   // edge probabilities.
   SmallVector<uint64_t, 4> BBSuccFreq;
-  for (BasicBlock *Succ : successors(BB)) {
-    auto SuccFreq = (Succ == SuccBB)
-                        ? BB2SuccBBFreq - NewBBFreq
-                        : BBOrigFreq * BPI->getEdgeProbability(BB, Succ);
+  for (succ_iterator I = succ_begin(BB), E = succ_end(BB); I != E; ++I) {
+    auto BB2SuccBBFreq =
+        BBOrigFreq * BPI->getEdgeProbability(BB, I.getSuccessorIndex());
+    auto SuccFreq = (*I == SuccBB) ? BB2SuccBBFreq - NewBBFreq : BB2SuccBBFreq;
     BBSuccFreq.push_back(SuccFreq.getFrequency());
   }
 
diff --git a/llvm/test/Transforms/JumpThreading/thread-prob-8.ll b/llvm/test/Transforms/JumpThreading/thread-prob-8.ll
new file mode 100644
index 0000000000000..0952b87a35429
--- /dev/null
+++ b/llvm/test/Transforms/JumpThreading/thread-prob-8.ll
@@ -0,0 +1,41 @@
+; RUN: opt -debug-only=branch-prob -passes=jump-threading -S %s 2>&1 | FileCheck %s
+; REQUIRES: asserts
+
+; Make sure that edges' probabilities would not accumulate if they are
+; the same target BB.
+; Edge L0 -> 2 and L0 -> 3 's targets are both L2, but their respective
+; probability should not be L0 -> L2, because prob[L0->L2] equls to
+; prob[L0->2] + prob[L0->3]
+
+; CHECK: Computing probabilities for entry
+; CHECK: eraseBlock L0
+; CHECK-NOT: set edge L0 -> 0 successor probability to 0x12492492 / 0x80000000 = 14.29%
+; CHECK-NOT: set edge L0 -> 1 successor probability to 0x24924925 / 0x80000000 = 28.57%
+; CHECK-NOT: set edge L0 -> 2 successor probability to 0x24924925 / 0x80000000 = 28.57%
+; CHECK-NOT: set edge L0 -> 3 successor probability to 0x24924925 / 0x80000000 = 28.57%
+; CHECK: set edge L0 -> 0 successor probability to 0x1999999a / 0x80000000 = 20.00%
+; CHECK: set edge L0 -> 1 successor probability to 0x33333333 / 0x80000000 = 40.00%
+; CHECK: set edge L0 -> 2 successor probability to 0x1999999a / 0x80000000 = 20.00%
+; CHECK: set edge L0 -> 3 successor probability to 0x1999999a / 0x80000000 = 20.00%
+
+define void @test_switch(i1 %cond, i8 %value) nounwind {
+entry:
+  br i1 %cond, label %L0, label %L4
+L0:
+  %expr = select i1 %cond, i8 1, i8 %value
+  switch i8 %expr, label %L3 [
+    i8 1, label %L1
+    i8 2, label %L2
+    i8 3, label %L2
+  ], !prof !0
+
+L1:
+  ret void
+L2:
+  ret void
+L3:
+  ret void
+L4:
+  br label %L0
+}
+!0 = !{!"branch_weights", i32 1, i32 7, i32 1, i32 1}

>From 15b89f106d6dca542817dec55ecc1b65339f7975 Mon Sep 17 00:00:00 2001
From: tianleli <tianle.l.liu at intel.com>
Date: Tue, 8 Apr 2025 14:51:09 +0800
Subject: [PATCH 2/2] Add branch_weights checking.

---
 llvm/test/Transforms/JumpThreading/thread-prob-8.ll | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/llvm/test/Transforms/JumpThreading/thread-prob-8.ll b/llvm/test/Transforms/JumpThreading/thread-prob-8.ll
index 0952b87a35429..b63c789515966 100644
--- a/llvm/test/Transforms/JumpThreading/thread-prob-8.ll
+++ b/llvm/test/Transforms/JumpThreading/thread-prob-8.ll
@@ -17,7 +17,8 @@
 ; CHECK: set edge L0 -> 1 successor probability to 0x33333333 / 0x80000000 = 40.00%
 ; CHECK: set edge L0 -> 2 successor probability to 0x1999999a / 0x80000000 = 20.00%
 ; CHECK: set edge L0 -> 3 successor probability to 0x1999999a / 0x80000000 = 20.00%
-
+; CHECK-NOT: !0 = !{!"branch_weights", i32 306783378, i32 613566757, i32 613566757, i32 613566757}
+; CHECK: !0 = !{!"branch_weights", i32 429496730, i32 858993459, i32 429496730, i32 429496730}
 define void @test_switch(i1 %cond, i8 %value) nounwind {
 entry:
   br i1 %cond, label %L0, label %L4