[NFC][SCEV] Generalize monotonicity check for full and limited iteration space

[mkazantsev]

Summary

A piece of logic of isLoopInvariantExitCondDuringFirstIterations is actually
a generalized predicate monotonicity check. This patch moves it into the
corresponding method and generalizes it a bit.

[apilipenko]

Can we split restructuring and generalization (possibly adding a test for the generalization part)? It's hard to assess correctness when you combine the two in one change.

[mkazantsev]

Factored out functional part as https://reviews.llvm.org/D91126

[apilipenko]

llvm/lib/Analysis/ScalarEvolution.cpp

9291–9325

Maybe move ProvedNoWrap into a lambda? This way the checks can be turned into early exits which should make it easier to follow the logic.

[mkazantsev]

Now it's completely an NFC.

llvm/include/llvm/Analysis/ScalarEvolution.h

enum MonotonicPredicateType {
     MonotonicallyIncreasing,
     MonotonicallyDecreasing
   };
 
   /// If, for all loop invariant X, the predicate "LHS `Pred` X" is
   /// monotonically increasing or decreasing, returns
   /// Some(MonotonicallyIncreasing) and Some(MonotonicallyDecreasing)
   /// respectively. If we could not prove either of these facts, returns None.
+  ///
+  /// If NumIter was provided, then we are proving monotonicity during at least
+  /// NumIter first iterations. If it was not provided, then we are proving
+  /// monotonicity on all iteration space.
   Optional<MonotonicPredicateType>
-  getMonotonicPredicateType(const SCEVAddRecExpr *LHS,
-                            ICmpInst::Predicate Pred);
+  getMonotonicPredicateType(const SCEVAddRecExpr *LHS, ICmpInst::Predicate Pred,
+                            Optional<const SCEV *> NumIter = None,
+                            const Instruction *Context = nullptr);
 
   /// Return true if the result of the predicate LHS `Pred` RHS is loop
   /// invariant with respect to L.  Set InvariantPred, InvariantLHS and
   /// InvariantLHS so that InvariantLHS `InvariantPred` InvariantRHS is the
   /// loop invariant form of LHS `Pred` RHS.
   bool isLoopInvariantPredicate(ICmpInst::Predicate Pred, const SCEV *LHS,
                                 const SCEV *RHS, const Loop *L,
                                 ICmpInst::Predicate &InvariantPred,
private:
   /// (resp. `SCEVZeroExtendExpr`).
   template <typename ExtendOpTy>
   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);
 
-  Optional<MonotonicPredicateType>
-  getMonotonicPredicateTypeImpl(const SCEVAddRecExpr *LHS,
-                                ICmpInst::Predicate Pred);
+  Optional<MonotonicPredicateType> getMonotonicPredicateTypeImpl(
+      const SCEVAddRecExpr *LHS, ICmpInst::Predicate Pred,
+      Optional<const SCEV *> NumIter, const Instruction *Context);
 
   /// Return SCEV no-wrap flags that can be proven based on reasoning about
   /// how poison produced from no-wrap flags on this value (e.g. a nuw add)
   /// would trigger undefined behavior on overflow.
   SCEV::NoWrapFlags getNoWrapFlagsFromUB(const Value *V);
 
   /// Return true if the SCEV corresponding to \p I is never poison.  Proving
   /// this is more complex than proving that just \p I is never poison, since

llvm/lib/Analysis/ScalarEvolution.cpp

bool ScalarEvolution::isKnownOnEveryIteration(ICmpInst::Predicate Pred,
                                               const SCEV *RHS) {
   const Loop *L = LHS->getLoop();
   return isLoopEntryGuardedByCond(L, Pred, LHS->getStart(), RHS) &&
          isLoopBackedgeGuardedByCond(L, Pred, LHS->getPostIncExpr(*this), RHS);
 }
 
 Optional<ScalarEvolution::MonotonicPredicateType>
 ScalarEvolution::getMonotonicPredicateType(const SCEVAddRecExpr *LHS,
-                                           ICmpInst::Predicate Pred) {
-  auto Result = getMonotonicPredicateTypeImpl(LHS, Pred);
+                                           ICmpInst::Predicate Pred,
+                                           Optional<const SCEV *> NumIter,
+                                           const Instruction *Context) {
+  assert((!NumIter || !isa<SCEVCouldNotCompute>(*NumIter)) &&
+         "provided number of iterations must be computable!");
+  auto Result = getMonotonicPredicateTypeImpl(LHS, Pred, NumIter, Context);
 
 #ifndef NDEBUG
   // Verify an invariant: inverting the predicate should turn a monotonically
   // increasing change to a monotonically decreasing one, and vice versa.
   if (Result) {
-    auto ResultSwapped =
-        getMonotonicPredicateTypeImpl(LHS, ICmpInst::getSwappedPredicate(Pred));
+    auto ResultSwapped = getMonotonicPredicateTypeImpl(
+        LHS, ICmpInst::getSwappedPredicate(Pred), NumIter, Context);
 
     assert(ResultSwapped.hasValue() && "should be able to analyze both!");
     assert(ResultSwapped.getValue() != Result.getValue() &&
            "monotonicity should flip as we flip the predicate");
   }
 #endif
 
   return Result;
 }
 
 Optional<ScalarEvolution::MonotonicPredicateType>
 ScalarEvolution::getMonotonicPredicateTypeImpl(const SCEVAddRecExpr *LHS,
-                                               ICmpInst::Predicate Pred) {
+                                               ICmpInst::Predicate Pred,
+                                               Optional<const SCEV *> NumIter,
+                                               const Instruction *Context) {
   // A zero step value for LHS means the induction variable is essentially a
   // loop invariant value. We don't really depend on the predicate actually
   // flipping from false to true (for increasing predicates, and the other way
   // around for decreasing predicates), all we care about is that *if* the
   // predicate changes then it only changes from false to true.
   //
   // A zero step value in itself is not very useful, but there may be places
   // where SCEV can prove X >= 0 but not prove X > 0, so it is helpful to be
   // as general as possible.
 
   // Only handle LE/LT/GE/GT predicates.
   if (!ICmpInst::isRelational(Pred))
     return None;
 
   bool IsGreater = ICmpInst::isGE(Pred) || ICmpInst::isGT(Pred);
   assert((IsGreater || ICmpInst::isLE(Pred) || ICmpInst::isLT(Pred)) &&
          "Should be greater or less!");
+  bool IsUnsigned = ICmpInst::isUnsigned(Pred);
+  assert((IsUnsigned || ICmpInst::isSigned(Pred)) &&
+         "Should be either signed or unsigned!");
+
+  // Check if we can prove no-wrap in the relevant range.
+  bool ProvedNoWrap =
+      IsUnsigned ? LHS->hasNoUnsignedWrap() : LHS->hasNoSignedWrap();
+  const SCEV *Step = LHS->getStepRecurrence(*this);
+  bool IsStepNonNegative = isKnownNonNegative(Step);
+  bool IsStepNonPositive = isKnownNonPositive(Step);
 
-  // Check that AR does not wrap.
-  if (ICmpInst::isUnsigned(Pred)) {
-    if (!LHS->hasNoUnsignedWrap())
-      return None;
+  // If we only care about first several iterations, check that start <= end for
+  // non-negative step and start >= end for non-positive step.
+  if (!ProvedNoWrap && NumIter && (IsStepNonNegative || IsStepNonPositive)) {
+    bool HasNoSelfWrap = LHS->hasNoSelfWrap();
+    if (!HasNoSelfWrap)
+      // If num iter has same type as the AddRec, and step is +/- 1, even max
+      // possible number of iterations is not enough to self-wrap.
+      if (NumIter.getValue()->getType() == LHS->getType())
+        if (Step == getOne(LHS->getType()) ||
+            Step == getMinusOne(LHS->getType()))
+          HasNoSelfWrap = true;
+    if (HasNoSelfWrap) {
+      const SCEV *Start = LHS->getStart();
+      const SCEV *End = LHS->evaluateAtIteration(*NumIter, *this);
+      ICmpInst::Predicate NoOverflowPred;
+      if (IsStepNonNegative)
+        NoOverflowPred = ICmpInst::ICMP_SLE;
+      else
+        NoOverflowPred = ICmpInst::ICMP_SGE;
+      if (IsUnsigned)
+        NoOverflowPred = ICmpInst::getUnsignedPredicate(NoOverflowPred);
+      if (isKnownPredicateAt(NoOverflowPred, Start, End, Context))
+        ProvedNoWrap = true;
+    }
+  }
+
+  // If nothing worked, bail.
+  if (!ProvedNoWrap)
+
+  if (IsUnsigned) {
     return IsGreater ? MonotonicallyIncreasing : MonotonicallyDecreasing;
   } else {
-    assert(ICmpInst::isSigned(Pred) &&
-           "Relational predicate is either signed or unsigned!");
-    if (!LHS->hasNoSignedWrap())
-      return None;
-
-    const SCEV *Step = LHS->getStepRecurrence(*this);
-
-    if (isKnownNonNegative(Step))
+    if (IsStepNonNegative)
       return IsGreater ? MonotonicallyIncreasing : MonotonicallyDecreasing;
 
-    if (isKnownNonPositive(Step))
+    if (IsStepNonPositive)
       return !IsGreater ? MonotonicallyIncreasing : MonotonicallyDecreasing;
 
     return None;
   }
 }
 
 bool ScalarEvolution::isLoopInvariantPredicate(
     ICmpInst::Predicate Pred, const SCEV *LHS, const SCEV *RHS, const Loop *L,
if (!isLoopInvariant(RHS, L)) {
     if (!isLoopInvariant(LHS, L))
       return false;
 
     std::swap(LHS, RHS);
     Pred = ICmpInst::getSwappedPredicate(Pred);
   }
 
   auto *AR = dyn_cast<SCEVAddRecExpr>(LHS);
-  // TODO: Lift affinity limitation in the future.
-  if (!AR || AR->getLoop() != L || !AR->isAffine())
-    return false;
-
-  // The predicate must be relational (i.e. <, <=, >=, >).
-  if (!ICmpInst::isRelational(Pred))
-    return false;
-
-  // TODO: Support steps other than +/- 1.
-  const SCEV *Step = AR->getOperand(1);
-  auto *One = getOne(Step->getType());
-  auto *MinusOne = getNegativeSCEV(One);
-  if (Step != One && Step != MinusOne)
+  if (!AR || AR->getLoop() != L)
     return false;
 
-  // Type mismatch here means that MaxIter is potentially larger than max
-  // unsigned value in start type, which mean we cannot prove no wrap for the
-  // indvar.
-  if (AR->getType() != MaxIter->getType())
+  // The predicate must be monotonic (i.e. <, <=, >=, >).
+  if (!getMonotonicPredicateType(AR, Pred, MaxIter, Context))
     return false;
 
   // Value of IV on suggested last iteration.
   const SCEV *Last = AR->evaluateAtIteration(MaxIter, *this);
   // Does it still meet the requirement?
   if (!isKnownPredicateAt(Pred, Last, RHS, Context))
     return false;
-  // Because step is +/- 1 and MaxIter has same type as Start (i.e. it does
-  // not exceed max unsigned value of this type), this effectively proves
-  // that there is no wrap during the iteration. To prove that there is no
-  // signed/unsigned wrap, we need to check that
-  // Start <= Last for step = 1 or Start >= Last for step = -1.
-  ICmpInst::Predicate NoOverflowPred =
-      CmpInst::isSigned(Pred) ? ICmpInst::ICMP_SLE : ICmpInst::ICMP_ULE;
-  if (Step == MinusOne)
-    NoOverflowPred = CmpInst::getSwappedPredicate(NoOverflowPred);
-  const SCEV *Start = AR->getStart();
-  if (!isKnownPredicateAt(NoOverflowPred, Start, Last, Context))
-    return false;
 
   // Everything is fine.
   InvariantPred = Pred;
-  InvariantLHS = Start;
+  InvariantLHS = AR->getStart();
   InvariantRHS = RHS;
   return true;
 }
 
 bool ScalarEvolution::isKnownPredicateViaConstantRanges(
     ICmpInst::Predicate Pred, const SCEV *LHS, const SCEV *RHS) {
   if (HasSameValue(LHS, RHS))
     return ICmpInst::isTrueWhenEqual(Pred);