[SCEV] Generalize no-self-wrap check in isLoopInvariantExitCondDuringFirstIterations

[mkazantsev]

Summary

Lift limitation on step being +/- 1. In fact, the only thing it is needed for
is proving no-self-wrap. We can instead check this flag directly.

Theoretically it can increase the scope of the transform, but I could not
construct such test easily.

[apilipenko]

llvm/lib/Analysis/ScalarEvolution.cpp

9589–9591

Nit. IsStepNonNegative local variable can be removed.

[mkazantsev]

This was completely illegal. Need to find a way to revert it now.

llvm/lib/Analysis/ScalarEvolution.cpp

bool ScalarEvolution::isLoopInvariantExitCondDuringFirstIterations(
   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;
 
-  // TODO: Support steps other than +/- 1.
   const SCEV *Step = AR->getStepRecurrence(*this);
-  auto *One = getOne(Step->getType());
-  auto *MinusOne = getNegativeSCEV(One);
-  if (Step != One && Step != MinusOne)
-    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())
+  bool IsStepNonNegative = isKnownNonNegative(Step);
+  bool IsStepNonPositive = isKnownNonPositive(Step);
+  if (!IsStepNonNegative && !IsStepNonPositive)
+    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)
     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.
+  // 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 (Step == MinusOne)
+  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;