[LoopUnroll] Fix division by zero #166258
Conversation
jdenny-ornl
requested review from
MatzeB,
arsenm,
fhahn,
jdoerfert and
mtrofin
|
@llvm/pr-subscribers-llvm-transforms Author: Joel E. Denny (jdenny-ornl) ChangesPR #159163's probability computation for epilogue loops does not handle the possibility of an original loop probability of one. Runtime loop unrolling does not make sense for such an infinite loop, and a division by zero results. This patch works around that case. Issue #165998. Full diff: https://github.com/llvm/llvm-project/pull/166258.diff 2 Files Affected:
diff --git a/llvm/lib/Transforms/Utils/LoopUnrollRuntime.cpp b/llvm/lib/Transforms/Utils/LoopUnrollRuntime.cpp
index 1e8f6cc76900c..1599c743c51cf 100644
--- a/llvm/lib/Transforms/Utils/LoopUnrollRuntime.cpp
+++ b/llvm/lib/Transforms/Utils/LoopUnrollRuntime.cpp
@@ -200,8 +200,17 @@ static void ConnectProlog(Loop *L, Value *BECount, unsigned Count,
/// from 0 to \p N more iterations can possibly execute. Among such cases in
/// the original loop (with loop probability \p OriginalLoopProb), what is the
/// probability of executing at least one more iteration?
+///
+/// If \p OriginalLoopProb is 1, then the original loop was somehow determined
+/// to be always infinite. Runtime loop unrolling should be impossible for that
+/// case. What is infinity % UnrollCount? But we might have bad profile data.
+/// In that case, we arbitrarily choose to keep the probability at 1 throughout
+/// the remainder loop.
static BranchProbability
probOfNextInRemainder(BranchProbability OriginalLoopProb, unsigned N) {
+ if (OriginalLoopProb == BranchProbability::getOne())
+ return BranchProbability::getOne();
+
// Each of these variables holds the original loop's probability that the
// number of iterations it will execute is some m in the specified range.
BranchProbability ProbOne = OriginalLoopProb; // 1 <= m
diff --git a/llvm/test/Transforms/LoopUnroll/loop-probability-one.ll b/llvm/test/Transforms/LoopUnroll/loop-probability-one.ll
new file mode 100644
index 0000000000000..bdd2aaf5b7f4e
--- /dev/null
+++ b/llvm/test/Transforms/LoopUnroll/loop-probability-one.ll
@@ -0,0 +1,114 @@
+; Check that a loop probability of one (indicating an always infinite loop) does
+; not crash or otherwise break LoopUnroll behavior when it tries to compute new
+; probabilities from it.
+;
+; Runtime loop unrolling should be impossible for that case. What is infinity %
+; UnrollCount? But we can have bad profile data. In that case, the
+; implementation arbitrarily chooses to keep the probability at 1 throughout the
+; remainder loop.
+
+; DEFINE: %{unroll} = opt < %s -unroll-count=3 -passes=loop-unroll -S
+; DEFINE: %{rt} = %{unroll} -unroll-runtime
+
+; RUN: %{unroll} | FileCheck %s -check-prefix UNROLL
+; RUN: %{rt} -unroll-runtime-epilog=true | FileCheck %s -check-prefix EPILOG
+; RUN: %{rt} -unroll-runtime-epilog=false | FileCheck %s -check-prefix PROLOG
+
+define void @test(i32 %n) {
+entry:
+ br label %loop
+
+loop:
+ %i = phi i32 [ 0, %entry ], [ %inc, %loop ]
+ %inc = add i32 %i, 1
+ %c = icmp slt i32 %inc, %n
+ br i1 %c, label %loop, label %end, !prof !0
+
+end:
+ ret void
+}
+
+
+!0 = !{!"branch_weights", i32 1, i32 0}
+
+; UNROLL: define void @test(i32 %n) {
+; UNROLL: entry:
+; UNROLL: br label %loop
+; UNROLL: loop:
+; UNROLL: br i1 %c, label %loop.1, label %end, !prof !0
+; UNROLL: loop.1:
+; UNROLL: br i1 %c.1, label %loop.2, label %end, !prof !0
+; UNROLL: loop.2:
+; UNROLL: br i1 %c.2, label %loop, label %end, !prof !0, !llvm.loop !1
+; UNROLL-NOT: loop.3
+; UNROLL: end:
+; UNROLL: ret void
+; UNROLL: }
+;
+; Infinite unrolled loop.
+; UNROLL: !0 = !{!"branch_weights", i32 1, i32 0}
+
+; EPILOG: define void @test(i32 %n) {
+; EPILOG: entry:
+; EPILOG: br i1 %{{.*}}, label %loop.epil.preheader, label %entry.new, !prof !0
+; EPILOG: entry.new:
+; EPILOG: br label %loop
+; EPILOG: loop:
+; EPILOG: br i1 %{{.*}}, label %loop, label %end.unr-lcssa, !prof !1
+; EPILOG: end.unr-lcssa:
+; EPILOG: br i1 %{{.*}}, label %loop.epil.preheader, label %end, !prof !1
+; EPILOG: loop.epil.preheader:
+; EPILOG: br label %loop.epil
+; EPILOG: loop.epil:
+; EPILOG: br i1 %{{.*}}, label %loop.epil, label %end.epilog-lcssa, !prof !4
+; EPILOG: end.epilog-lcssa:
+; EPILOG: br label %end
+; EPILOG: end:
+; EPILOG: ret void
+; EPILOG: }
+;
+; Unrolled loop guard: Unrolled loop is always entered.
+; EPILOG: !0 = !{!"branch_weights", i32 0, i32 -2147483648}
+;
+; Unrolled loop latch: Unrolled loop is infinite.
+; Epilogue loop guard: Epilogue loop is always entered if unrolled loop exits.
+; EPILOG: !1 = !{!"branch_weights", i32 -2147483648, i32 0}
+;
+; Epilogue loop latch: Epilogue loop executes both of its 2 iterations.
+; EPILOG: !4 = !{!"branch_weights", i32 1073741824, i32 1073741824}
+
+; PROLOG: define void @test(i32 %n) {
+; PROLOG: entry:
+; PROLOG: br i1 %{{.*}}, label %loop.prol.preheader, label %loop.prol.loopexit, !prof !0
+; PROLOG: loop.prol.preheader:
+; PROLOG: br label %loop.prol
+; PROLOG: loop.prol:
+; PROLOG: br i1 %{{.*}}, label %loop.prol, label %loop.prol.loopexit.unr-lcssa, !prof !1
+; PROLOG: loop.prol.loopexit.unr-lcssa:
+; PROLOG: br label %loop.prol.loopexit
+; PROLOG: loop.prol.loopexit:
+; PROLOG: br i1 %{{.*}}, label %end, label %entry.new, !prof !0
+; PROLOG: entry.new:
+; PROLOG: br label %loop
+; PROLOG: loop:
+; PROLOG: br i1 %{{.*}}, label %loop, label %end.unr-lcssa, !prof !4
+; PROLOG: end.unr-lcssa:
+; PROLOG: br label %end
+; PROLOG: end:
+; PROLOG: ret void
+; PROLOG: }
+;
+; FIXME: Branch weights still need to be fixed in the case of prologues (issue
+; #135812), so !0 and !1 do not yet match their comments below. When we do
+; fix it, this test will hopefully catch any bug like issue #165998, which
+; impacted the case of epilogues.
+;
+; Prologue loop guard: Prologue loop is always entered.
+; Unrolled loop guard: Unrolled loop is always entered.
+; PROLOG: !0 = !{!"branch_weights", i32 1, i32 127}
+;
+; Prologue loop latch: Prologue loop executes both of its 2 iterations.
+; PROLOG: !1 = !{!"branch_weights", i32 0, i32 1}
+;
+; Unrolled loop latch: Unrolled loop is infinite.
+; PROLOG: !4 = !{!"branch_weights", i32 1, i32 0}
|
There was a problem hiding this comment.
Code change LGTM.
Not sure it's worth speculating about the situation in the comment to a random function... I think I remember the sample profile loader avoiding 1:0 probabilities even if we do have not picked up any samples on one side (but it's statistics, and it can happen when there is few/not enough samples for the loop at hand); but given the IR verifier isn't rejecting this case and who knows what other sources or passes produce probability annotationes, seems fair/necessary to handle this case.
| /// | ||
| /// If \p OriginalLoopProb is 1, then the original loop was somehow determined | ||
| /// to be always infinite. Runtime loop unrolling should be impossible for that | ||
| /// case. What is infinity % UnrollCount? But we might have bad profile data. | ||
| /// In that case, we arbitrarily choose to keep the probability at 1 throughout | ||
| /// the remainder loop. | ||
| static BranchProbability | ||
| probOfNextInRemainder(BranchProbability OriginalLoopProb, unsigned N) { | ||
| if (OriginalLoopProb == BranchProbability::getOne()) | ||
| return BranchProbability::getOne(); | ||
|
|
There was a problem hiding this comment.
| /// | |
| /// If \p OriginalLoopProb is 1, then the original loop was somehow determined | |
| /// to be always infinite. Runtime loop unrolling should be impossible for that | |
| /// case. What is infinity % UnrollCount? But we might have bad profile data. | |
| /// In that case, we arbitrarily choose to keep the probability at 1 throughout | |
| /// the remainder loop. | |
| static BranchProbability | |
| probOfNextInRemainder(BranchProbability OriginalLoopProb, unsigned N) { | |
| if (OriginalLoopProb == BranchProbability::getOne()) | |
| return BranchProbability::getOne(); | |
| static BranchProbability | |
| probOfNextInRemainder(BranchProbability OriginalLoopProb, unsigned N) { | |
| // Avoid division by zero if loop exit probability is zero. | |
| if (OriginalLoopProb == BranchProbability::getOne()) | |
| return BranchProbability::getOne(); | |
There was a problem hiding this comment.
Not a strong opinion here, but for my taste I'd rather have the talk about special cases within the function (rather than in the "API" documentation in front of the function) and even then rather discuss in reviews, discourse etc. what the values could mean rather than particularily here.
(But I realize this is a bit of a subjective thing; so do whatever you think is right)
There was a problem hiding this comment.
As I replied to your other comment, I feel like the code comment should stay somewhere in the code. I thought of it as an API-level issue as it explains that/how the function is robust to a case that a caller might be concerned about. I can live with it being an internal comment though if you decide you do feel strongly about it.
But I now realize we should have something saying division by zero is being avoided. I'll add that.
Thanks.
I'm not sure I understand what you mean. I wrote the probOfNextInRemainder function to compute correct probabilities for epilogue remainder loops. I think it is worthwhile to point out that there is no reasonable probability for that case when the original loop is infinite, so the probability we return is an arbitrary choice. |
|
Thanks! I can confirm the problem has been fixed. |
|
¯_(ツ)_/¯ just felt like a lot of text (which to me felt like it may as easily confuse readers as it may help). Is the profile really "bad" in that case? I think having a 1:0 probability is still possible either because it's actually infinite, or because the difference is too small to represent. But as long as the IR verifier accepts it, it's not "bad" in my book... But as I said arguing about source code comments is highly subjective, so do with my comment as you wish or ignore... |
You convinced me. I moved it into the function body. I reworded it to focus more on our attempt to derive a meaningful result from the probability rather than on whether it was truly wrong.
Thanks for raising it. |
And trigger pre-commit CI again because it was failing on a test that surely this PR did not affect.
PR #159163's probability computation for epilogue loops does not handle the possibility of an original loop probability of one. Runtime loop unrolling does not make sense for such an infinite loop, and a division by zero results. This patch works around that case.
Issue #165998.