[LoopPredication] Fix division by zero in case of zero branch weights

[danilaml]

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:

llvm-project/llvm/lib/Analysis/BranchProbabilityInfo.cpp

Lines 430 to 434 in 7472490

	if (WeightSum == 0 || ReachableIdxs.size() == 0) {
	for (unsigned I = 0, E = TI->getNumSuccessors(); I != E; ++I)
	Weights[I] = 1;
	WeightSum = TI->getNumSuccessors();
	}

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Changes

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:

llvm-project/llvm/lib/Analysis/BranchProbabilityInfo.cpp

Lines 430 to 434 in 7472490

	if (WeightSum == 0 || ReachableIdxs.size() == 0) {
	for (unsigned I = 0, E = TI->getNumSuccessors(); I != E; ++I)
	Weights[I] = 1;
	WeightSum = TI->getNumSuccessors();
	}

Full diff: https://github.com/llvm/llvm-project/pull/66506.diff

2 Files Affected:

• (modified) llvm/lib/Transforms/Scalar/LoopPredication.cpp (+3)
• (modified) llvm/test/Transforms/LoopPredication/pr66382.ll (+20-1)

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
 

[dexonsmith]

It looks to me like the test case could probably be reduced. E.g., is the call to @llvm.experimental.guard necessary to trigger the bug? Do you need all of the attributes? Can you reduce the number of basic blocks?

I can't add comments to most of the file though because you've already committed an XFAILed version. (Since this is an assertion failure, does it even fail consistently when assertions are off? I could imagine it XPASS-ing on a non-assertion bot.)

[dexonsmith]

I don't think we usually land failing tests in tree (unless policies have changed? I'm not doing a ton of reviewing these days...). This makes it a bit harder to comment on them in the review.

[danilaml]

I've seen it done before. Regarding assertions - without them the crash would be just division by zero. Otherwise, the buildbots would complain about XFAIL passing.

[dexonsmith]

It's not clear to me if this is checking anything relevant about the expected output of the pass, in the context of the interpretation of the branch weights as an even distribution. If it is, can you explain?

If not, is there a way to do that? Maybe you can observe somehow that the branch weights are correctly interpreted as "even" by looking at DEBUG output (the DEBUG_TYPE for this pass is loop-predication), or maybe STATISTIC?

For example, I see that there's a command-line option -loop-predication-latch-probability-scale, which is a scaling factor applied to the latch probability. This affects how the branch weights are used, ultimately changing the return of LoopPredication::isLoopProfitableToPredicate. Can you construct a test case where, if the branch probability is even (the correct interpretation of branch_weights of 0), then two RUN lines with different -loop-predication-latch-probability-scale will give you different output for STATISTIC and/or DEBUG? If so, then we could have see CHECK lines on the STATISTIC/DEBUG output.

[danilaml]

Right now, it just checks that there are simply no crashes. experimental_guard is required otherwise the LoopPredication pass would not run

// There is nothing to do if the module doesn't use guards
  auto *GuardDecl =
      M->getFunction(Intrinsic::getName(Intrinsic::experimental_guard));
  bool HasIntrinsicGuards = GuardDecl && !GuardDecl->use_empty();
  auto *WCDecl = M->getFunction(
      Intrinsic::getName(Intrinsic::experimental_widenable_condition));
  bool HasWidenableConditions =
      PredicateWidenableBranchGuards && WCDecl && !WCDecl->use_empty();
  if (!HasIntrinsicGuards && !HasWidenableConditions)
    return false;

The test was reduced with bugpoint and llvm-reduce, so I don't know if it can be meaningfully reduced further (also why it doesn't really do much if there is no crash).

Trying to make it output something more meaningful would likely make the test bigger (and would require assertions).

[danilaml]

@dexonsmith I've added a test that would test that zero weights would be treated as if there was no profile, and it also tests scale factor as a byproduct. I left the original test case since I think it's worth having a small reduced regression test for the original issue.

[dexonsmith]

LGTM!