[LSR][term-fold] Ensure the simple recurrence is reachable from the current loop

[patrick-rivos]

If the phi node found by matchSimpleRecurrence is not from the current
loop, then isAlmostDeadIV panics. With this patch we bail out early.

Signed-off-by: Patrick O'Neill patrick@rivosinc.com

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Patrick O'Neill (patrick-rivos)

Changes

If the phi node is unreachable from the current loop, then isAlmostDeadIV panics. With this patch we bail out early.

In some cases like SeveralLoopLatch and SeveralLoopLatch2 we want to optimize using a simple recurrence from another (reachable) BB. This patch does not impact those cases.

An alternative solution would be to delete:

  // If that IV isn't dead after we rewrite the exit condition in terms of
  // another IV, there's no point in doing the transform.
  if (!isAlmostDeadIV(ToFold, LoopLatch, TermCond))
    return std::nullopt;

since removing it has no impact on the testsuite and resolves this issue.

This issue was found via fuzzer with this testcase:

// clang red.c -O2 -S -march=rv64gc
struct a {
  unsigned b;
};
int *c;
static struct a d;
int e;
void f() {
  struct a *g[18] = {&d};
  d.b = 8;
  for (; d.b >= 4; ++d.b)
    *c = 0;
  do {
    e++;
  } while (d.b);
}

https://godbolt.org/z/9GGq9Kcfx

This is my first LLVM bugfix so beyond the regular review if there are any nits about patch/commit message/testsuite etiquette please let me know.

---

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

3 Files Affected:

• (modified) llvm/lib/Transforms/Scalar/LoopStrengthReduce.cpp (+6)
• (modified) llvm/lib/Transforms/Utils/LoopUtils.cpp (+1)
• (added) llvm/test/Transforms/LoopStrengthReduce/lsr-unreachable-bb-phi-node.ll (+40)

diff --git a/llvm/lib/Transforms/Scalar/LoopStrengthReduce.cpp b/llvm/lib/Transforms/Scalar/LoopStrengthReduce.cpp
index 08021f3ba853e8..62c271494cf2d7 100644
--- a/llvm/lib/Transforms/Scalar/LoopStrengthReduce.cpp
+++ b/llvm/lib/Transforms/Scalar/LoopStrengthReduce.cpp
@@ -6808,6 +6808,12 @@ canFoldTermCondOfLoop(Loop *L, ScalarEvolution &SE, DominatorTree &DT,
   if (!matchSimpleRecurrence(LHS, ToFold, ToFoldStart, ToFoldStep))
     return std::nullopt;
 
+  // If ToFold does not have an incoming value from LoopLatch then the simple
+  // recurrence is from a prior loop unreachable from the loop we're currently
+  // considering.
+  if (ToFold->getBasicBlockIndex(LoopLatch) == -1)
+    return std::nullopt;
+
   // If that IV isn't dead after we rewrite the exit condition in terms of
   // another IV, there's no point in doing the transform.
   if (!isAlmostDeadIV(ToFold, LoopLatch, TermCond))
diff --git a/llvm/lib/Transforms/Utils/LoopUtils.cpp b/llvm/lib/Transforms/Utils/LoopUtils.cpp
index a4fdc1f8c12e50..7491a99b03f66e 100644
--- a/llvm/lib/Transforms/Utils/LoopUtils.cpp
+++ b/llvm/lib/Transforms/Utils/LoopUtils.cpp
@@ -468,6 +468,7 @@ llvm::collectChildrenInLoop(DomTreeNode *N, const Loop *CurLoop) {
 
 bool llvm::isAlmostDeadIV(PHINode *PN, BasicBlock *LatchBlock, Value *Cond) {
   int LatchIdx = PN->getBasicBlockIndex(LatchBlock);
+  assert(LatchIdx != -1 && "LatchBlock is not a case in this PHINode");
   Value *IncV = PN->getIncomingValue(LatchIdx);
 
   for (User *U : PN->users())
diff --git a/llvm/test/Transforms/LoopStrengthReduce/lsr-unreachable-bb-phi-node.ll b/llvm/test/Transforms/LoopStrengthReduce/lsr-unreachable-bb-phi-node.ll
new file mode 100644
index 00000000000000..1454535b52bccb
--- /dev/null
+++ b/llvm/test/Transforms/LoopStrengthReduce/lsr-unreachable-bb-phi-node.ll
@@ -0,0 +1,40 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt < %s -loop-reduce -S -lsr-term-fold | FileCheck %s
+
+; This test used to crash due to matchSimpleRecurrence matching the simple
+; recurrence in pn-loop when evaluating unrelated-loop. Since unrelated-loop
+; cannot jump to pn-node isAlmostDeadIV panics.
+define void @phi_node_different_bb() {
+; CHECK-LABEL: @phi_node_different_bb(
+; CHECK-NEXT:    br label [[PN_LOOP:%.*]]
+; CHECK:       pn-loop:
+; CHECK-NEXT:    [[TMP1:%.*]] = phi i32 [ 1, [[TMP0:%.*]] ], [ [[TMP2:%.*]], [[PN_LOOP]] ]
+; CHECK-NEXT:    [[TMP2]] = add i32 [[TMP1]], 1
+; CHECK-NEXT:    [[TMP3:%.*]] = icmp ugt i32 [[TMP2]], 1
+; CHECK-NEXT:    br i1 [[TMP3]], label [[PN_LOOP]], label [[UNRELATED_LOOP_PREHEADER:%.*]]
+; CHECK:       unrelated-loop.preheader:
+; CHECK-NEXT:    br label [[UNRELATED_LOOP:%.*]]
+; CHECK:       unrelated-loop:
+; CHECK-NEXT:    [[TMP4:%.*]] = icmp eq i32 [[TMP2]], 0
+; CHECK-NEXT:    br i1 [[TMP4]], label [[END:%.*]], label [[UNRELATED_LOOP]]
+; CHECK:       end:
+; CHECK-NEXT:    ret void
+;
+  br label %pn-loop
+
+pn-loop:                                          ; preds = %pn-loop, %0
+  %1 = phi i32 [ 1, %0 ], [ %2, %pn-loop ]
+  %2 = add i32 %1, 1
+  %3 = icmp ugt i32 %2, 1
+  br i1 %3, label %pn-loop, label %unrelated-loop.preheader
+
+unrelated-loop.preheader:                         ; preds = %pn-loop
+  br label %unrelated-loop
+
+unrelated-loop:                                   ; preds = %unrelated-loop, %unrelated-loop.preheader
+  %4 = icmp eq i32 %2, 0
+  br i1 %4, label %end, label %unrelated-loop
+
+end:                                              ; preds = %unrelated-loop
+  ret void
+}

[preames]

A more idiomatic check here would be: L->getHeader() != ToFold->getParent(). This directly checks the property you're comment indicates.

[patrick-rivos]

Updated.

Previously I had tried comparing with the LoopLatch which interfered with the SeveralLoopLatch testcase. I'm going to read up on how getHeader differs from LoopLatch. Thanks!

[preames]

Your comment about unreachability is a bit misleading - the reachability of the other loop isn't the key detail. Instead, maybe something along the lines of: Make sure the recurrence matched is part of the current loop of interest.

[patrick-rivos]

Updated comment and order of comparison to match the new comment.

Also updated PR message to be the squashed commit message.

[preames]

LGTM - Do you want me to land it for you?

[patrick-rivos]

Yes please - I don't have permissions