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

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 <[email protected]>

---------

Signed-off-by: Patrick O'Neill <[email protected]>

Author: patrick-rivos

llvm/lib/Transforms/Scalar/LoopStrengthReduce.cpp

@@ -6808,6 +6808,10 @@ canFoldTermCondOfLoop(Loop *L, ScalarEvolution &SE, DominatorTree &DT,
   if (!matchSimpleRecurrence(LHS, ToFold, ToFoldStart, ToFoldStep))
     return std::nullopt;
 
+  // Ensure the simple recurrence is a part of the current loop.
+  if (ToFold->getParent() != L->getHeader())
+    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))

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())

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
+}