[LoopUnroll] Fix assert fail on zeroed branch weights

[jdenny-ornl]

BranchProbability fails an assert when its denominator is zero.

Reported at #159163 (review).

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Joel E. Denny (jdenny-ornl)

Changes

BranchProbability fails an assert when its denominator is zero.

Reported at <#159163 (review)>.

---

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

2 Files Affected:

• (modified) llvm/lib/Transforms/Utils/LoopUtils.cpp (+4-1)
• (added) llvm/test/Transforms/LoopUnroll/zeroed-branch-weights.ll (+30)

diff --git a/llvm/lib/Transforms/Utils/LoopUtils.cpp b/llvm/lib/Transforms/Utils/LoopUtils.cpp
index 8be471bee5579..6e60b94be78e3 100644
--- a/llvm/lib/Transforms/Utils/LoopUtils.cpp
+++ b/llvm/lib/Transforms/Utils/LoopUtils.cpp
@@ -992,9 +992,12 @@ BranchProbability llvm::getBranchProbability(BranchInst *B,
   uint64_t Weight0, Weight1;
   if (!extractBranchWeights(*B, Weight0, Weight1))
     return BranchProbability::getUnknown();
+  uint64_t Denominator = Weight0 + Weight1;
+  if (Denominator == 0)
+    return BranchProbability::getUnknown();
   if (!ForFirstTarget)
     std::swap(Weight0, Weight1);
-  return BranchProbability::getBranchProbability(Weight0, Weight0 + Weight1);
+  return BranchProbability::getBranchProbability(Weight0, Denominator);
 }
 
 bool llvm::setBranchProbability(BranchInst *B, BranchProbability P,
diff --git a/llvm/test/Transforms/LoopUnroll/zeroed-branch-weights.ll b/llvm/test/Transforms/LoopUnroll/zeroed-branch-weights.ll
new file mode 100644
index 0000000000000..4d378b0d22f7d
--- /dev/null
+++ b/llvm/test/Transforms/LoopUnroll/zeroed-branch-weights.ll
@@ -0,0 +1,30 @@
+; Check that zeroed branch weights do not crash or otherwise break basic
+; LoopUnroll behavior when it tries to compute a probability from them.
+
+; RUN: opt < %s -S -unroll-count=2 -passes='loop-unroll' 2>&1 | FileCheck %s
+
+define void @test() {
+entry:
+  br label %loop
+
+loop:
+  br i1 false, label %end, label %loop, !prof !0
+
+end:
+  ret void
+}
+
+!0 = !{!"branch_weights", i32 0, i32 0}
+
+; CHECK: define void @test() {
+; CHECK: entry:
+; CHECK:   br label %loop
+; CHECK: loop:
+; CHECK:   br i1 false, label %end, label %loop.1, !prof !0
+; CHECK: loop.1:
+; CHECK:   br i1 false, label %end, label %loop, !prof !0, !llvm.loop !1
+; CHECK-NOT: loop.2
+; CHECK: end:
+; CHECK:   ret void
+; CHECK: }
+; CHECK: !0 = !{!"branch_weights", i32 0, i32 0}

[fhahn]

LGTM, thanks.

[fhahn]

@jdenny-ornl It looks like there are other, similar failures reported, including #165998 so a more comprehensive fix may be needed.

As this has been crashing for a few days now, it may be good to revert back to a known good state

[jdenny-ornl]

@jdenny-ornl It looks like there are other, similar failures reported, including #165998 so a more comprehensive fix may be needed.

As this has been crashing for a few days now, it may be good to revert back to a known good state

Thanks for suggesting that. I'm looking now. If I cannot track this down within a few hours, I will start rolling back the patch series. Is that reasonable?

[wlei-llvm]

I guess we alternatively can avoid to create all zero branch_weights metadata by setting the ElideAllZero parameter in setBranchWeights , in this case we can save some metadata memory.