[LVI] Disable at-use reasoning across phi nodes (PR60629)

For phi nodes, while we can make use of the edge condition for the
incoming value, we shouldn't look past the phi node to look for
further conditions, because we might be reasoning about values
from two different cycle iterations (which will have the same
SSA value).

To handle this more specifically we would have to detect cycles,
and there doesn't seem to be any motivating case for that at this
point.

2 files changed, 53 insertions, 6 deletions

diff --git a/llvm/lib/Analysis/LazyValueInfo.cpp b/llvm/lib/Analysis/LazyValueInfo.cpp
index f1587cecf9fb..0e9fab667e6e 100644
--- a/llvm/lib/Analysis/LazyValueInfo.cpp
+++ b/llvm/lib/Analysis/LazyValueInfo.cpp
@@ -1673,11 +1673,6 @@ ConstantRange LazyValueInfo::getConstantRangeAtUse(const Use &U,
       // TODO: Use non-local query?
       CondVal =
           getEdgeValueLocal(V, PHI->getIncomingBlock(*CurrU), PHI->getParent());
-    } else if (!isSafeToSpeculativelyExecute(CurrI)) {
-      // Stop walking if we hit a non-speculatable instruction. Even if the
-      // result is only used under a specific condition, executing the
-      // instruction itself may cause side effects or UB already.
-      break;
     }
     if (CondVal && CondVal->isConstantRange())
       CR = CR.intersectWith(CondVal->getConstantRange());
@@ -1685,7 +1680,13 @@ ConstantRange LazyValueInfo::getConstantRangeAtUse(const Use &U,
     // Only follow one-use chain, to allow direct intersection of conditions.
     // If there are multiple uses, we would have to intersect with the union of
     // all conditions at different uses.
-    if (!CurrI->hasOneUse())
+    // Stop walking if we hit a non-speculatable instruction. Even if the
+    // result is only used under a specific condition, executing the
+    // instruction itself may cause side effects or UB already.
+    // This also disallows looking through phi nodes: If the phi node is part
+    // of a cycle, we might end up reasoning about values from different cycle
+    // iterations (PR60629).
+    if (!CurrI->hasOneUse() || !isSafeToSpeculativelyExecute(CurrI))
       break;
     CurrU = &*CurrI->use_begin();
   }
diff --git a/llvm/test/Analysis/LazyValueAnalysis/pr60629.ll b/llvm/test/Analysis/LazyValueAnalysis/pr60629.ll
new file mode 100644
index 000000000000..7cc314e693a4
--- /dev/null
+++ b/llvm/test/Analysis/LazyValueAnalysis/pr60629.ll
@@ -0,0 +1,46 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -passes=correlated-propagation -S %s | FileCheck %s
+
+; Check that shift is not removed by CVP because of incorrect range returned by LVI::getConstantRangeAtUse.
+; https://github.com/llvm/llvm-project/issues/60629
+
+define i32 @test(i32 %arg) {
+; CHECK-LABEL: @test(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br label [[LOOP:%.*]]
+; CHECK:       loop:
+; CHECK-NEXT:    [[IV:%.*]] = phi i32 [ 127, [[ENTRY:%.*]] ], [ [[IV_NEXT:%.*]], [[DO_SHIFT:%.*]] ]
+; CHECK-NEXT:    [[SHIFTED_IV:%.*]] = phi i32 [ 0, [[ENTRY]] ], [ [[SHIFT:%.*]], [[DO_SHIFT]] ]
+; CHECK-NEXT:    [[IV_NEXT]] = add i32 [[IV]], 1
+; CHECK-NEXT:    [[LOOP_COND:%.*]] = icmp sgt i32 [[IV]], 127
+; CHECK-NEXT:    br i1 [[LOOP_COND]], label [[EXIT:%.*]], label [[DO_SHIFT]]
+; CHECK:       do.shift:
+; CHECK-NEXT:    [[SHIFT]] = ashr i32 [[IV]], [[ARG:%.*]]
+; CHECK-NEXT:    br label [[LOOP]]
+; CHECK:       bb1:
+; CHECK-NEXT:    br label [[EXIT]]
+; CHECK:       exit:
+; CHECK-NEXT:    [[RETVAL:%.*]] = phi i32 [ 0, [[BB1:%.*]] ], [ [[SHIFTED_IV]], [[LOOP]] ]
+; CHECK-NEXT:    ret i32 [[RETVAL]]
+;
+entry:
+  br label %loop
+
+loop:                                              ; preds = %do.shift, %entry
+  %iv = phi i32 [ 127, %entry ], [ %iv.next, %do.shift ]
+  %shifted.iv = phi i32 [ 0, %entry ], [ %shift, %do.shift ]
+  %iv.next = add i32 %iv, 1
+  %loop_cond = icmp sgt i32 %iv, 127
+  br i1 %loop_cond, label %exit, label %do.shift
+
+do.shift:                                          ; preds = %loop
+  %shift = ashr i32 %iv, %arg
+  br label %loop
+
+bb1:                                               ; No predecessors!
+  br label %exit
+
+exit:                                              ; preds = %loop, %bb1
+  %retval = phi i32 [ 0, %bb1 ], [ %shifted.iv, %loop ]
+  ret i32 %retval
+}