[llvm] [LSR][term-fold] Ensure the simple recurrence is reachable from the current loop (PR #83085)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: Patrick O'Neill (patrick-rivos)

<details>
<summary>Changes</summary>

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](https://github.com/llvm/llvm-project/blob/main/llvm/test/Transforms/LoopStrengthReduce/lsr-term-fold.ll#L368) and [SeveralLoopLatch2](https://github.com/llvm/llvm-project/blob/main/llvm/test/Transforms/LoopStrengthReduce/lsr-term-fold.ll#L406) 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:
```c++
  // 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:
```c
// 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
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
+}

``````````

</details>


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