[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

9499–9533

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!");
 
-  // Check that AR does not wrap.
-  if (ICmpInst::isUnsigned(Pred)) {
-    if (!LHS->hasNoUnsignedWrap())
-      return None;
-    return IsGreater ? MonotonicallyIncreasing : MonotonicallyDecreasing;
-  } else {
-    assert(ICmpInst::isSigned(Pred) &&
-           "Relational predicate is either signed or unsigned!");
-    if (!LHS->hasNoSignedWrap())
-      return None;
+  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.
 
-    const SCEV *Step = LHS->getStepRecurrence(*this);
+  bool IsStepNonNegative = isKnownNonNegative(Step);
+  bool IsStepNonPositive = isKnownNonPositive(Step);
+  // We need to know which direction the iteration is going.
+  if (!IsStepNonNegative && !IsStepNonPositive)
+    return None;
+
+  auto ProvedNoWrap = [&]() {
+    // If the AddRec already has the flag, we are done.
+    if (IsUnsigned ? LHS->hasNoUnsignedWrap() : LHS->hasNoSignedWrap())
+      return true;
+
+    if (!NumIter)
+      return false;
+    // We could not prove no-wrap on all iteration space. Can we prove it for
+    // first iterations? In order to achieve it, check that:
+    // 1. The addrec does not self-wrap;
+    // 2. start <= end for non-negative step and start >= end for non-positive
+    // step.
+    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)
+      return false;
+    const SCEV *Start = LHS->getStart();
+    const SCEV *End = LHS->evaluateAtIteration(*NumIter, *this);
+    ICmpInst::Predicate NoOverflowPred =
+        IsStepNonNegative ? ICmpInst::ICMP_SLE : ICmpInst::ICMP_SGE;
+    if (IsUnsigned)
+      NoOverflowPred = ICmpInst::getUnsignedPredicate(NoOverflowPred);
+    return isKnownPredicateAt(NoOverflowPred, Start, End, Context);
+  };
 
-    if (isKnownNonNegative(Step))
+  // If nothing worked, bail.
+  if (!ProvedNoWrap())
+    return None;
+
+  if (IsUnsigned)
+    return IsGreater ? MonotonicallyIncreasing : MonotonicallyDecreasing;
+  else {
+    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,

 }
 
 bool ScalarEvolution::isLoopInvariantExitCondDuringFirstIterations(
     ICmpInst::Predicate Pred, const SCEV *LHS, const SCEV *RHS, const Loop *L,
     const Instruction *Context, const SCEV *MaxIter,
     ICmpInst::Predicate &InvariantPred, const SCEV *&InvariantLHS,
     const SCEV *&InvariantRHS) {
   // Try to prove the following set of facts:
-  // - The predicate is monotonic.
+  // - The predicate is monotonic in the iteration space.
   // - If the check does not fail on the 1st iteration:
-  //   - No overflow will happen during first MaxIter iterations;
   //   - It will not fail on the MaxIter'th iteration.
   // If the check does fail on the 1st iteration, we leave the loop and no
   // other checks matter.
 
   // If there is a loop-invariant, force it into the RHS, otherwise bail out.
   if (!isLoopInvariant(RHS, L)) {
     if (!isLoopInvariant(LHS, L))
       return false;
 
     std::swap(LHS, RHS);
     Pred = ICmpInst::getSwappedPredicate(Pred);
   }
 
   auto *AR = dyn_cast<SCEVAddRecExpr>(LHS);
   if (!AR || AR->getLoop() != L)
     return false;
 
-  // The predicate must be relational (i.e. <, <=, >=, >).
-  if (!ICmpInst::isRelational(Pred))
-    return false;
-
-  const SCEV *Step = AR->getStepRecurrence(*this);
-  bool IsStepNonPositive = isKnownNonPositive(Step);
-  if (!IsStepNonPositive && !isKnownNonNegative(Step))
-    return false;
-  bool HasNoSelfWrap = AR->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 (MaxIter->getType() == AR->getType())
-      if (Step == getOne(AR->getType()) || Step == getMinusOne(AR->getType()))
-        HasNoSelfWrap = true;
-  // Only proceed with non-self-wrapping ARs.
-  if (!HasNoSelfWrap)
+  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;
-  // We know that the addrec does not have a self-wrap. To prove that there is
-  // no signed/unsigned wrap, we need to check that
-  // Start <= Last for positive step or Start >= Last for negative step. Either
-  // works for zero step.
-  ICmpInst::Predicate NoOverflowPred =
-      CmpInst::isSigned(Pred) ? ICmpInst::ICMP_SLE : ICmpInst::ICMP_ULE;
-  if (IsStepNonPositive)
-    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);