[llvm] [JumpThreading] Don't divide by all-zero branch weights (PR #221846)

[Timur Baidusenov via llvm-commits]

https://github.com/bai-tim created https://github.com/llvm/llvm-project/pull/221846

Fixes #192165.

>From 5a72d91b06b35a460c429f5a58bca3c4c4dac082 Mon Sep 17 00:00:00 2001
From: Timur Baidusenov <timurbaidusenov at gmail.com>
Date: Mon, 7 Sep 2026 21:34:48 +0300
Subject: [PATCH] [JumpThreading] Don't divide by all-zero branch weights

Fixes #192165.
---
 llvm/lib/Transforms/Scalar/JumpThreading.cpp  | 35 ++++++++++---------
 .../unfold-select-zero-branch-weights.ll      | 27 ++++++++++++++
 2 files changed, 46 insertions(+), 16 deletions(-)
 create mode 100644 llvm/test/Transforms/JumpThreading/unfold-select-zero-branch-weights.ll

diff --git a/llvm/lib/Transforms/Scalar/JumpThreading.cpp b/llvm/lib/Transforms/Scalar/JumpThreading.cpp
index 7a0542aac83fc..5e65db3eb495d 100644
--- a/llvm/lib/Transforms/Scalar/JumpThreading.cpp
+++ b/llvm/lib/Transforms/Scalar/JumpThreading.cpp
@@ -3056,22 +3056,25 @@ bool JumpThreadingPass::tryToUnfoldSelectInCurrBB(BasicBlock *BB) {
       assert(Extracted);
       uint64_t Denominator =
           sum_of(llvm::map_range(BW, StaticCastTo<uint64_t>));
-      assert(Denominator > 0 &&
-             "At least one of the branch probabilities should be non-zero");
-      BranchProbability TrueProb =
-          BranchProbability::getBranchProbability(BW[0], Denominator);
-      BranchProbability FalseProb =
-          BranchProbability::getBranchProbability(BW[1], Denominator);
-      SmallVector<BranchProbability, 2> BP = {TrueProb, FalseProb};
-
-      if (BPI)
-        BPI->setEdgeProbability(BB, BP);
-
-      if (BFI) {
-        auto BBOrigFreq = BFI->getBlockFreq(BB);
-        auto NewBBFreq = BBOrigFreq * TrueProb;
-        BFI->setBlockFreq(NewBB, NewBBFreq);
-        BFI->setBlockFreq(SplitBB, BBOrigFreq);
+      // All-zero branch weights are valid IR, and they say that both sides are
+      // equally likely rather than giving a probability to derive. Leave BPI
+      // and BFI as they are instead of dividing by zero.
+      if (Denominator > 0) {
+        BranchProbability TrueProb =
+            BranchProbability::getBranchProbability(BW[0], Denominator);
+        BranchProbability FalseProb =
+            BranchProbability::getBranchProbability(BW[1], Denominator);
+        SmallVector<BranchProbability, 2> BP = {TrueProb, FalseProb};
+
+        if (BPI)
+          BPI->setEdgeProbability(BB, BP);
+
+        if (BFI) {
+          auto BBOrigFreq = BFI->getBlockFreq(BB);
+          auto NewBBFreq = BBOrigFreq * TrueProb;
+          BFI->setBlockFreq(NewBB, NewBBFreq);
+          BFI->setBlockFreq(SplitBB, BBOrigFreq);
+        }
       }
     }
     SI->eraseFromParent();
diff --git a/llvm/test/Transforms/JumpThreading/unfold-select-zero-branch-weights.ll b/llvm/test/Transforms/JumpThreading/unfold-select-zero-branch-weights.ll
new file mode 100644
index 0000000000000..90d610d22946c
--- /dev/null
+++ b/llvm/test/Transforms/JumpThreading/unfold-select-zero-branch-weights.ll
@@ -0,0 +1,27 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
+; RUN: opt -passes="require<branch-prob>,require<block-freq>,jump-threading" -S < %s | FileCheck %s
+
+define i16 @zero_branch_weights(i1 %c) {
+; CHECK-LABEL: define i16 @zero_branch_weights(
+; CHECK-SAME: i1 [[C:%.*]]) {
+; CHECK-NEXT:  [[ENTRY:.*]]:
+; CHECK-NEXT:    br i1 [[C]], label %[[EXIT:.*]], label %[[EXIT]]
+; CHECK:       [[EXIT]]:
+; CHECK-NEXT:    [[P:%.*]] = phi i16 [ 0, %[[ENTRY]] ], [ 0, %[[ENTRY]] ]
+; CHECK-NEXT:    ret i16 [[P]]
+
+entry:
+  br i1 %c, label %sel, label %exit
+
+sel:
+  %iv = phi i32 [ 0, %entry ]
+  %cond = icmp eq i32 %iv, 0
+  %s = select i1 %cond, i16 0, i16 0, !prof !0
+  br label %exit
+
+exit:
+  %p = phi i16 [ 0, %sel ], [ 0, %entry ]
+  ret i16 %p
+}
+
+!0 = !{!"branch_weights", i32 0, i32 0}