[SCEV] Use knowledge of stride to prove loops finite for LT exit count computation

[reames]

Summary

This came up in the post commit review discussion of D103834.

The basic idea is that the general case was using a loop being finite by assumption to prove that step was non-zero. Non-zero combined with NoWrap flags in turns proves the loop must exit before overflow occurs. This change adds an explicit check for when the loop isn't known mustprogress, but we can directly prove that the step isn't zero. This will mostly help non-C/C++ family languages, though it may benefit some C language test cases as well.

[efriedma]

I don't see any obvious issue here. I'd like to work through the known miscompiles before we start expanding this, though.

[mkazantsev]

Looks fine, but agreed with Eli. Please don't land it for some more days to find and fix miscompiles from underlying patches (if any).

llvm/lib/Analysis/ScalarEvolution.cpp

11479

Please refactor as

if (PredicatedIV || !NoWrap || isKnownNonPositive(Stride))
     return getCouldNotCompute();
bool StrideNonZero = loopIsFiniteByAssumption || isLoopEntryGuardedByCond ...
 if (!StrideNonZero )
     return getCouldNotCompute();

The dom tree walking is expensive, let's only do it if easy checks failed.

[reames]

I agree with the request to hold until latent miscompile in underlying code is addressed. I will hold this patch until that happens, and then land without further review. Max, good suggestion on condition ordering, will address in landed patch.

[bjope]

In D104066#2813641, @reames wrote:

I agree with the request to hold until latent miscompile in underlying code is addressed. I will hold this patch until that happens, and then land without further review. Max, good suggestion on condition ordering, will address in landed patch.

I guess it will be "safe" to land this downstream meanwhile (as everyone thinks this looks good).

Out-of-curiosity. What are these known/latent miscompiles that is being referred to? Is there a bug report somewhere?

[nikic]

@bjope See https://reviews.llvm.org/D101722#2809588 and discussion on D104140.

[reames]

Code as posted is incorrect. A loop being finite only proves stride != 0 if the backedge is taken. If the backedge is not taken, stride may still be zero.

[efriedma]

If the stride is zero, and RHS is invariant, the comparison feeding the backedge is invariant. If we don't exit via any other exit, that implies the backedge-taken count is either zero or infinity.

So if we know the loop is finite, the RHS is invariant, we don't exit via any other exit, and the IV doesn't overflow, we can divide strides into three cases:

1. Stride is zero. The backedge-taken count is zero, as noted above.
2. Stride is negative. I think we can also prove the backedge-taken count is zero here?
3. Stride is positive. We can use the standard formula.

[reames]

JFYI, this idea is on hold until some of the previous patches land. The interplay between this and some of the others is delicate and I'm worried about introducing bugs because I tried to keep too many things in my head at once. We should probably also do a bit of cleanup first - e.g. lift the RHS invariant check.

llvm/lib/Analysis/ScalarEvolution.cpp

if (!PositiveStride) {
     // The backedge taken count for such loops is evaluated as -
     // (max(end, start + stride) - start - 1) /u stride
     //
     // The additional preconditions that we need to check to prove correctness
     // of the above formula is as follows -
     //
     // a) IV is either nuw or nsw depending upon signedness (indicated by the
     //    NoWrap flag).
-    // b) loop is single exit with no side effects.
-    //
+    // b) s is non-zero (and thus the loop is non-infinite)
     //
     // Precondition a) implies that if the stride is negative, this is a single
     // trip loop. The backedge taken count formula reduces to zero in this case.
     //
-    // Precondition b) implies that the unknown stride cannot be zero otherwise
-    // we have UB.
+    // Precondition b) implies that the loop can't be infinite (as a non-zero
+    // stride requires the comparison to eventually produce poison and we know
+    // this condition value reaches the branch exit)
     //
     // The positive stride case is the same as isKnownPositive(Stride) returning
     // true (original behavior of the function).
     //
     // We want to make sure that the stride is truly unknown as there are edge
     // cases where ScalarEvolution propagates no wrap flags to the
     // post-increment/decrement IV even though the increment/decrement operation
     // itself is wrapping. The computed backedge taken count may be wrong in
     // such cases. This is prevented by checking that the stride is not known to
     // be either positive or non-positive. For example, no wrap flags are
     // propagated to the post-increment IV of this loop with a trip count of 2 -
     //
     // unsigned char i;
     // for(i=127; i<128; i+=129)
     //   A[i] = i;
     //
-    if (PredicatedIV || !NoWrap || isKnownNonPositive(Stride) ||
-        !loopIsFiniteByAssumption(L))
+
+    bool StrideNonZero = loopIsFiniteByAssumption(L) ||
+      isLoopEntryGuardedByCond(L, ICmpInst::ICMP_NE, Stride,
+                               getZero(Stride->getType()));
+
+    if (PredicatedIV || !NoWrap || isKnownNonPositive(Stride) || !StrideNonZero)
       return getCouldNotCompute();
   } else if (!Stride->isOne() && !NoWrap) {
     auto isUBOnWrap = [&]() {
       // Can we prove this loop *must* be UB if overflow of IV occurs?
       // Reasoning goes as follows:
       // * Suppose the IV did self wrap.
       // * If Stride evenly divides the iteration space, then once wrap
       //   occurs, the loop must revisit the same values.

llvm/test/Analysis/ScalarEvolution/trip-count-unknown-stride.ll

for.body:                                         ; preds = %entry, %for.body
   %add = add nsw i32 %i.05, %s
   %cmp = icmp slt i32 %add, %n
   br i1 %cmp, label %for.body, label %for.end, !llvm.loop !8
 
 for.end:                                          ; preds = %for.body, %entry
   ret void
 }
 
+; Without explicit mustprogress, but the combination of nsw flags, the
+; condition being branched on, and the fact the step is known non-zero
+; is enough to prove this loop finite.
+
+; CHECK: Determining loop execution counts for: @non_zero_stride
+; CHECK: Loop %for.body4: backedge-taken count is ((-1 + (-1 * %arg) + ((%arg + %div) smax %n)) /u %div)
+; CHECK: Loop %for.body4: max backedge-taken count is -1
+
+define void @non_zero_stride(i64 %arg, i64 %div, i64 %n) {
+entry:
+  %assume = icmp sgt i64 %div, 0
+  call void @llvm.assume(i1 %assume)
+  br label %for.body4
+
+for.body4:                                        ; preds = %for.cond1.preheader, %for.body4
+  %i.020 = phi i64 [ %add5, %for.body4 ], [ %arg, %entry ]
+  %add5 = add nsw i64 %i.020, %div
+  %cmp2 = icmp slt i64 %add5, %n
+  br i1 %cmp2, label %for.body4, label %exit
+
+exit:
+  ret void
+}
+
+declare void @llvm.assume(i1)
+
 !8 = distinct !{!8, !9}
 !9 = !{!"llvm.loop.mustprogress"}