[llvm] [LoopPredication] Fix division by zero in case of zero branch weights (PR #66506)

[Danila Malyutin via llvm-commits]

https://github.com/danilaml created https://github.com/llvm/llvm-project/pull/66506

Treat the case where all branch weights are zero as if there was no profile.
Fixes #66382

Should be in line with the approach taken by BPI: https://github.com/llvm/llvm-project/blob/74724902ba2fc0dc2734cdb5fae15cb338f1b635/llvm/lib/Analysis/BranchProbabilityInfo.cpp#L430-L434

>From 03614b095177946544a18086a31030bdf8787c5b Mon Sep 17 00:00:00 2001
From: Danila Malyutin <dmalyutin at azul.com>
Date: Fri, 15 Sep 2023 15:13:54 +0300
Subject: [PATCH] [LoopPredication] Fix division by zero in case of zero branch
 weights

Treat the case where all branch weights are zero as if there was no profile.
Fixes #66382
---
 .../lib/Transforms/Scalar/LoopPredication.cpp |  3 +++
 .../Transforms/LoopPredication/pr66382.ll     | 21 ++++++++++++++++++-
 2 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/Transforms/Scalar/LoopPredication.cpp b/llvm/lib/Transforms/Scalar/LoopPredication.cpp
index a58ab093a1f75d3..55079b4a42d2fae 100644
--- a/llvm/lib/Transforms/Scalar/LoopPredication.cpp
+++ b/llvm/lib/Transforms/Scalar/LoopPredication.cpp
@@ -967,6 +967,9 @@ bool LoopPredication::isLoopProfitableToPredicate() {
           Numerator += Weight;
         Denominator += Weight;
       }
+      // If all weights are zero act as if there was no profile data
+      if (Denominator == 0)
+        return BranchProbability::getBranchProbability(1, NumSucc);
       return BranchProbability::getBranchProbability(Numerator, Denominator);
     } else {
       assert(LatchBlock != ExitingBlock &&
diff --git a/llvm/test/Transforms/LoopPredication/pr66382.ll b/llvm/test/Transforms/LoopPredication/pr66382.ll
index 3ac4cac0615f464..f9a14d470453cf0 100644
--- a/llvm/test/Transforms/LoopPredication/pr66382.ll
+++ b/llvm/test/Transforms/LoopPredication/pr66382.ll
@@ -1,4 +1,4 @@
-; XFAIL: *
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 3
 ; RUN: opt -S -loop-predication-skip-profitability-checks=false -passes='require<scalar-evolution>,loop-mssa(loop-predication)' %s | FileCheck %s
 
 target triple = "x86_64-unknown-linux-gnu"
@@ -6,7 +6,26 @@ target triple = "x86_64-unknown-linux-gnu"
 ; Function Attrs: nocallback nofree nosync willreturn
 declare void @llvm.experimental.guard(i1, ...) #0
 
+; Check that LoopPredication doesn't crash on all-zero branch weights
 define void @foo() {
+; CHECK-LABEL: define void @foo() {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br label [[HEADER:%.*]]
+; CHECK:       Header:
+; CHECK-NEXT:    [[J2:%.*]] = phi i64 [ 0, [[ENTRY:%.*]] ], [ [[J_NEXT:%.*]], [[LATCH:%.*]] ]
+; CHECK-NEXT:    call void (i1, ...) @llvm.experimental.guard(i1 false, i32 0) [ "deopt"() ]
+; CHECK-NEXT:    [[J_NEXT]] = add i64 [[J2]], 1
+; CHECK-NEXT:    br i1 false, label [[LATCH]], label [[EXIT:%.*]]
+; CHECK:       Latch:
+; CHECK-NEXT:    [[SPECULATE_TRIP_COUNT:%.*]] = icmp ult i64 [[J2]], 0
+; CHECK-NEXT:    br i1 [[SPECULATE_TRIP_COUNT]], label [[HEADER]], label [[COMMON_RET_LOOPEXIT:%.*]], !prof [[PROF0:![0-9]+]]
+; CHECK:       common.ret.loopexit:
+; CHECK-NEXT:    br label [[COMMON_RET:%.*]]
+; CHECK:       common.ret:
+; CHECK-NEXT:    ret void
+; CHECK:       exit:
+; CHECK-NEXT:    br label [[COMMON_RET]]
+;
 entry:
   br label %Header