[SCEV] Make howFarToZero max backedge-taken count check for precondition

[efriedma]

Summary

Refines max backedge-taken count if a loop like "for (int i = 0; i != n; ++i) { /* body */ }" is rotated.

[sanjoy]

lgtm

lib/Analysis/ScalarEvolution.cpp

7226 ↗

(On Diff #83874)

Use APInt::umin here?

llvm/trunk/lib/Analysis/ScalarEvolution.cpp

ScalarEvolution::howFarToZero(const SCEV *V, const Loop *L, bool ControlsExit,
   bool CountDown = StepC->getAPInt().isNegative();
   const SCEV *Distance = CountDown ? Start : getNegativeSCEV(Start);
 
   // Handle unitary steps, which cannot wraparound.
   // 1*N = -Start; -1*N = Start (mod 2^BW), so:
   //   N = Distance (as unsigned)
   if (StepC->getValue()->equalsInt(1) || StepC->getValue()->isAllOnesValue()) {
     APInt MaxBECount = getUnsignedRange(Distance).getUnsignedMax();
+
+    // When a loop like "for (int i = 0; i != n; ++i) { /* body */ }" is rotated,
+    // we end up with a loop whose backedge-taken count is n - 1.  Detect this
+    // case, and see if we can improve the bound.
+    //
+    // Explicitly handling this here is necessary because getUnsignedRange
+    // isn't context-sensitive; it doesn't know that we only care about the
+    // range inside the loop.
+    const SCEV *Zero = getZero(Distance->getType());
+    const SCEV *One = getOne(Distance->getType());
+    const SCEV *DistancePlusOne = getAddExpr(Distance, One);
+    if (isLoopEntryGuardedByCond(L, ICmpInst::ICMP_NE, DistancePlusOne, Zero)) {
+      // If Distance + 1 doesn't overflow, we can compute the maximum distance
+      // as "unsigned_max(Distance + 1) - 1".
+      ConstantRange CR = getUnsignedRange(DistancePlusOne);
+      MaxBECount = APIntOps::umin(MaxBECount, CR.getUnsignedMax() - 1);
+    }
     return ExitLimit(Distance, getConstant(MaxBECount), false, Predicates);
   }
 
   // As a special case, handle the instance where Step is a positive power of
   // two. In this case, determining whether Step divides Distance evenly can be
   // done by counting and comparing the number of trailing zeros of Step and
   // Distance.
   if (!CountDown) {

llvm/trunk/test/Analysis/ScalarEvolution/max-trip-count.ll

for.body:
   %add = add nsw i32 %i, 1
   %cmp = icmp ne i32 %add, %masked
   br i1 %cmp, label %for.body, label %bar.exit
 
 bar.exit:
   ret i32 0
 }
 
-; TODO: Improve count based on guard.
 ; CHECK-LABEL: @ne_max_trip_count_3
-; CHECK: Loop %for.body: max backedge-taken count is -1
+; CHECK: Loop %for.body: max backedge-taken count is 6
 define i32 @ne_max_trip_count_3(i32 %n) {
 entry:
   %masked = and i32 %n, 7
   %guard = icmp eq i32 %masked, 0
   br i1 %guard, label %exit, label %for.preheader
 
 for.preheader:
   br label %for.body
 
 for.body:
   %i = phi i32 [ 0, %for.preheader ], [ %add, %for.body ]
   %add = add nsw i32 %i, 1
   %cmp = icmp ne i32 %add, %masked
   br i1 %cmp, label %for.body, label %loop.exit
 
 loop.exit:
   br label %exit
 
 exit:
   ret i32 0
 }
 
-; TODO: Improve count based on guard.
 ; CHECK-LABEL: @ne_max_trip_count_4
-; CHECK: Loop %for.body: max backedge-taken count is -1
+; CHECK: Loop %for.body: max backedge-taken count is -2
 define i32 @ne_max_trip_count_4(i32 %n) {
 entry:
   %guard = icmp eq i32 %n, 0
   br i1 %guard, label %exit, label %for.preheader
 
 for.preheader:
   br label %for.body