3 files changed, 95 insertions, 0 deletions

diff --git a/llvm/include/llvm/Analysis/ScalarEvolution.h b/llvm/include/llvm/Analysis/ScalarEvolution.h
index 7650c73f454..1aa3092dc05 100644
--- a/llvm/include/llvm/Analysis/ScalarEvolution.h
+++ b/llvm/include/llvm/Analysis/ScalarEvolution.h
@@ -900,6 +900,9 @@ namespace llvm {
     bool proveNoWrapByVaryingStart(const SCEV *Start, const SCEV *Step,
                                    const Loop *L);
 
+    /// Try to prove NSW or NUW on \p AR relying on ConstantRange manipulation.
+    SCEV::NoWrapFlags proveNoWrapViaConstantRanges(const SCEVAddRecExpr *AR);
+
     bool isMonotonicPredicateImpl(const SCEVAddRecExpr *LHS,
                                   ICmpInst::Predicate Pred, bool &Increasing);
 
diff --git a/llvm/lib/Analysis/ScalarEvolution.cpp b/llvm/lib/Analysis/ScalarEvolution.cpp
index 0365862cc16..2f0347771ba 100644
--- a/llvm/lib/Analysis/ScalarEvolution.cpp
+++ b/llvm/lib/Analysis/ScalarEvolution.cpp
@@ -1454,6 +1454,11 @@ const SCEV *ScalarEvolution::getZeroExtendExpr(const SCEV *Op,
       unsigned BitWidth = getTypeSizeInBits(AR->getType());
       const Loop *L = AR->getLoop();
 
+      if (!AR->hasNoUnsignedWrap()) {
+        auto NewFlags = proveNoWrapViaConstantRanges(AR);
+        const_cast<SCEVAddRecExpr *>(AR)->setNoWrapFlags(NewFlags);
+      }
+
       // If we have special knowledge that this addrec won't overflow,
       // we don't need to do any further analysis.
       if (AR->hasNoUnsignedWrap())
@@ -1663,6 +1668,11 @@ const SCEV *ScalarEvolution::getSignExtendExpr(const SCEV *Op,
       unsigned BitWidth = getTypeSizeInBits(AR->getType());
       const Loop *L = AR->getLoop();
 
+      if (!AR->hasNoSignedWrap()) {
+        auto NewFlags = proveNoWrapViaConstantRanges(AR);
+        const_cast<SCEVAddRecExpr *>(AR)->setNoWrapFlags(NewFlags);
+      }
+
       // If we have special knowledge that this addrec won't overflow,
       // we don't need to do any further analysis.
       if (AR->hasNoSignedWrap())
@@ -3757,6 +3767,37 @@ private:
 };
 } // end anonymous namespace
 
+SCEV::NoWrapFlags
+ScalarEvolution::proveNoWrapViaConstantRanges(const SCEVAddRecExpr *AR) {
+  if (!AR->isAffine())
+    return SCEV::FlagAnyWrap;
+
+  typedef OverflowingBinaryOperator OBO;
+  SCEV::NoWrapFlags Result = SCEV::FlagAnyWrap;
+
+  if (!AR->hasNoSignedWrap()) {
+    ConstantRange AddRecRange = getSignedRange(AR);
+    ConstantRange IncRange = getSignedRange(AR->getStepRecurrence(*this));
+
+    auto NSWRegion = ConstantRange::makeGuaranteedNoWrapRegion(
+        Instruction::Add, IncRange, OBO::NoSignedWrap);
+    if (NSWRegion.contains(AddRecRange))
+      Result = ScalarEvolution::setFlags(Result, SCEV::FlagNSW);
+  }
+
+  if (!AR->hasNoUnsignedWrap()) {
+    ConstantRange AddRecRange = getUnsignedRange(AR);
+    ConstantRange IncRange = getUnsignedRange(AR->getStepRecurrence(*this));
+
+    auto NUWRegion = ConstantRange::makeGuaranteedNoWrapRegion(
+        Instruction::Add, IncRange, OBO::NoUnsignedWrap);
+    if (NUWRegion.contains(AddRecRange))
+      Result = ScalarEvolution::setFlags(Result, SCEV::FlagNUW);
+  }
+
+  return Result;
+}
+
 const SCEV *ScalarEvolution::createAddRecFromPHI(PHINode *PN) {
   const Loop *L = LI.getLoopFor(PN->getParent());
   if (!L || L->getHeader() != PN->getParent())
diff --git a/llvm/test/Analysis/ScalarEvolution/increasing-and-decreasing-range.ll b/llvm/test/Analysis/ScalarEvolution/increasing-or-decreasing-iv.ll
index 5a83306cc06..36dcfd3a549 100644
--- a/llvm/test/Analysis/ScalarEvolution/increasing-and-decreasing-range.ll
+++ b/llvm/test/Analysis/ScalarEvolution/increasing-or-decreasing-iv.ll
@@ -86,3 +86,54 @@ loop:
 leave:
   ret void
 }
+
+define void @f2(i1 %c) {
+; CHECK-LABEL: Classifying expressions for: @f2
+entry:
+  %start = select i1 %c, i32 127, i32 0
+  %step  = select i1 %c, i32 -1,  i32 1
+  br label %loop
+
+loop:
+  %loop.iv = phi i32 [ 0, %entry ], [ %loop.iv.inc, %loop ]
+  %iv = phi i32 [ %start, %entry ], [ %iv.next, %loop ]
+  %iv.sext = sext i32 %iv to i64
+  %iv.next = add i32 %iv, %step
+; CHECK:  %iv.sext = sext i32 %iv to i64
+; CHECK-NEXT:  -->  {(sext i32 %start to i64),+,(sext i32 %step to i64)}<nsw><%loop>
+  %loop.iv.inc = add i32 %loop.iv, 1
+  %be.cond = icmp ne i32 %loop.iv.inc, 128
+  br i1 %be.cond, label %loop, label %leave
+
+leave:
+  ret void
+}
+
+define void @f3(i1 %c) {
+; CHECK-LABEL: Classifying expressions for: @f3
+entry:
+
+; NB! the i16 type (as opposed to i32), the choice of the constant 509
+; and the trip count are all related and not arbitrary.  We want an
+; add recurrence that will look like it can unsign-overflow *unless*
+; SCEV is able to see the correlation between the two selects feeding
+; into the initial value and the step increment.
+
+  %start = select i1 %c, i16 1000, i16 0
+  %step  = select i1 %c, i16 1,  i16 509
+  br label %loop
+
+loop:
+  %loop.iv = phi i16 [ 0, %entry ], [ %loop.iv.inc, %loop ]
+  %iv = phi i16 [ %start, %entry ], [ %iv.next, %loop ]
+  %iv.zext = zext i16 %iv to i64
+; CHECK:  %iv.zext = zext i16 %iv to i64
+; CHECK-NEXT:  -->  {(zext i16 %start to i64),+,(zext i16 %step to i64)}<nuw><%loop>
+  %iv.next = add i16 %iv, %step
+  %loop.iv.inc = add i16 %loop.iv, 1
+  %be.cond = icmp ne i16 %loop.iv.inc, 128
+  br i1 %be.cond, label %loop, label %leave
+
+leave:
+  ret void
+}