4 files changed, 36 insertions, 53 deletions
diff --git a/llvm/lib/Analysis/ScalarEvolution.cpp b/llvm/lib/Analysis/ScalarEvolution.cpp
index 22e361c22de..9e4eb111c2e 100644
--- a/llvm/lib/Analysis/ScalarEvolution.cpp
+++ b/llvm/lib/Analysis/ScalarEvolution.cpp
@@ -1274,25 +1274,13 @@ static const SCEV *getPreStartForExtend(const SCEVAddRecExpr *AR, Type *Ty,
   const SCEVAddRecExpr *PreAR = dyn_cast<SCEVAddRecExpr>(
       SE->getAddRecExpr(PreStart, Step, L, SCEV::FlagAnyWrap));
 
-  // WARNING: FIXME: the optimization below assumes that a sign/zero-overflowing
-  // nsw/nuw operation is undefined behavior.  This is strictly more aggressive
-  // than the interpretation of nsw in other parts of LLVM (for instance, they
-  // may unconditionally hoist nsw/nuw arithmetic through control flow).  This
-  // logic needs to be revisited once we have a consistent semantics for poison
-  // values.
-  //
-  // "{S,+,X} is <nsw>/<nuw>" and "{S,+,X} is evaluated at least once" implies
-  // "S+X does not sign/unsign-overflow" (we'd have undefined behavior if it
-  // did).  If `L->getExitingBlock() == L->getLoopLatch()` then `PreAR` (=
-  // {S,+,X}<nsw>/<nuw>) is evaluated every-time `AR` (= {S+X,+,X}) is
-  // evaluated, and hence within `AR` we are safe to assume that "S+X" will not
-  // sign/unsign-overflow.
+  // "{S,+,X} is <nsw>/<nuw>" and "the backedge is taken at least once" implies
+  // "S+X does not sign/unsign-overflow".
   //
 
-  BasicBlock *ExitingBlock = L->getExitingBlock();
-  BasicBlock *LatchBlock = L->getLoopLatch();
-  if (PreAR && PreAR->getNoWrapFlags(WrapType) && ExitingBlock != nullptr &&
-      ExitingBlock == LatchBlock)
+  const SCEV *BECount = SE->getBackedgeTakenCount(L);
+  if (PreAR && PreAR->getNoWrapFlags(WrapType) &&
+      !isa<SCEVCouldNotCompute>(BECount) && SE->isKnownPositive(BECount))
     return PreStart;
 
   // 2. Direct overflow check on the step operation's expression.
diff --git a/llvm/test/Analysis/ScalarEvolution/infer-prestart-no-wrap.ll b/llvm/test/Analysis/ScalarEvolution/infer-prestart-no-wrap.ll
index 45b3e4c10b3..c9689f7fe14 100644
--- a/llvm/test/Analysis/ScalarEvolution/infer-prestart-no-wrap.ll
+++ b/llvm/test/Analysis/ScalarEvolution/infer-prestart-no-wrap.ll
@@ -6,12 +6,14 @@ define void @infer.sext.0(i1* %c, i32 %start) {
   br label %loop
 
  loop:
+  %counter = phi i32 [ 0, %entry ], [ %counter.inc, %loop ]
   %idx = phi i32 [ %start, %entry ], [ %idx.inc, %loop ]
   %idx.inc = add nsw i32 %idx, 1
   %idx.inc.sext = sext i32 %idx.inc to i64
 ; CHECK: %idx.inc.sext = sext i32 %idx.inc to i64
 ; CHECK-NEXT: -->  {(1 + (sext i32 %start to i64)),+,1}<nsw><%loop>
-  %condition = load volatile i1* %c
+  %condition = icmp eq i32 %counter, 1
+  %counter.inc = add i32 %counter, 1
   br i1 %condition, label %exit, label %loop
 
  exit:
@@ -24,12 +26,14 @@ define void @infer.zext.0(i1* %c, i32 %start) {
   br label %loop
 
  loop:
+  %counter = phi i32 [ 0, %entry ], [ %counter.inc, %loop ]
   %idx = phi i32 [ %start, %entry ], [ %idx.inc, %loop ]
   %idx.inc = add nuw i32 %idx, 1
   %idx.inc.sext = zext i32 %idx.inc to i64
 ; CHECK: %idx.inc.sext = zext i32 %idx.inc to i64
 ; CHECK-NEXT: -->  {(1 + (zext i32 %start to i64)),+,1}<nuw><%loop>
-  %condition = load volatile i1* %c
+  %condition = icmp eq i32 %counter, 1
+  %counter.inc = add i32 %counter, 1
   br i1 %condition, label %exit, label %loop
 
  exit:
diff --git a/llvm/test/Analysis/ScalarEvolution/pr22641.ll b/llvm/test/Analysis/ScalarEvolution/pr22641.ll
new file mode 100644
index 00000000000..3b55afe0845
--- /dev/null
+++ b/llvm/test/Analysis/ScalarEvolution/pr22641.ll
@@ -0,0 +1,25 @@
+; RUN: opt -analyze -scalar-evolution < %s | FileCheck %s
+
+define i1 @main(i16 %a) {
+; CHECK-LABEL: Classifying expressions for: @main
+entry:
+  br label %body
+
+body:
+  %dec2 = phi i16 [ %a, %entry ], [ %dec, %cond ]
+  %dec = add i16 %dec2, -1
+  %conv2 = zext i16 %dec2 to i32
+  %conv = zext i16 %dec to i32
+; CHECK:   %conv = zext i16 %dec to i32
+; CHECK-NEXT: -->  {(zext i16 (-1 + %a) to i32),+,65535}<nuw><%body>
+; CHECK-NOT:  -->  {(65535 + (zext i16 %a to i32)),+,65535}<nuw><%body>
+
+  br label %cond
+
+cond:
+  br i1 false, label %body, label %exit
+
+exit:
+  %ret = icmp ne i32 %conv, 0
+  ret i1 %ret
+}
diff --git a/llvm/test/Analysis/ScalarEvolution/scev-prestart-nowrap.ll b/llvm/test/Analysis/ScalarEvolution/scev-prestart-nowrap.ll
index 5dbb6e6d03e..3ca32bdd9a1 100644
--- a/llvm/test/Analysis/ScalarEvolution/scev-prestart-nowrap.ll
+++ b/llvm/test/Analysis/ScalarEvolution/scev-prestart-nowrap.ll
@@ -80,37 +80,3 @@ define i64 @bad.1(i32 %start, i32 %low.limit, i32 %high.limit, i1* %unknown) {
  early.exit:
   ret i64 %postinc.sext
 }
-
-
-; WARNING: FIXME: it is safe to make the inference demonstrated here
-; only if we assume `add nsw` has undefined behavior if the result
-; sign-overflows; and this interpretation is stronger than what most
-; of LLVM assumes.  This test here only serves as a documentation of
-; current behavior and will need to be revisited once we've decided
-; upon a consistent semantics for nsw (and nuw) arithetic operations.
-;
-define i64 @good(i32 %start, i32 %low.limit, i32 %high.limit) {
-; CHECK-LABEL: Classifying expressions for: @good
- entry:
-  %postinc.start = add i32 %start, 1
-  br label %loop
-
- loop:
-  %idx = phi i32 [ %start, %entry ], [ %idx.inc, %loop ]
-  %postinc = phi i32 [ %postinc.start, %entry ], [ %postinc.inc, %loop ]
-  %postinc.inc = add nsw i32 %postinc, 1
-  %postinc.sext = sext i32 %postinc to i64
-; CHECK: %postinc.sext = sext i32 %postinc to i64
-; CHECK-NEXT: -->  {(1 + (sext i32 %start to i64)),+,1}<nsw><%loop>
-
-  %break.early = icmp slt i32 %postinc, %low.limit
-  %idx.inc = add nsw i32 %idx, 1
-  %cmp = icmp slt i32 %idx.inc, %high.limit
-  br i1 %cmp, label %loop, label %exit
-
- exit:
-  ret i64 0
-
- early.exit:
-  ret i64 %postinc.sext
-}