[SCEV] Make exact taken count calculation more optimistic

[mkazantsev]

Summary

Currently, getExact fails if it sees two exit counts in different blocks. There is
no solid reason to do so, given that we only calculate exact non-taken count
for exiting blocks that dominate latch. Using this fact, we can simply take min
out of all exits of all blocks to get the exact taken count.

This patch makes the calculation more optimistic with enforcing our assumption
with asserts. It allows us to calculate exact backedge taken count in trivial loops
like

for (int i = 0; i < 100; i++) {
  if (i > 50) break;
  . . .
}

[fhahn]

lib/Analysis/ScalarEvolution.cpp

6672

With exits you mean "exiting blocks", rather than "exit blocks", right? (http://lists.llvm.org/pipermail/llvm-dev/2017-July/115962.html)

Maybe it would be clearer to explicitly use "exiting blocks" instead of exits here (and throughout the patch) :)

[mkazantsev]

lib/Analysis/ScalarEvolution.cpp

6672

Sure, I will fix it. :)

6672

Actually I can do it as a separate NFC. This is not the only place where the term is used.

[fhahn]

lib/Analysis/ScalarEvolution.cpp

6672

Sounds good to me.

[fhahn]

LGTM, with clearing up the exiting/exit block thing.

Please wait a bit with committing, in case the other reviewers want to raise additional concerns

[fhahn]

lib/Analysis/ScalarEvolution.cpp

6657

Please also update the comment here, stating when multiple exits are OK.

[fhahn]

lib/Analysis/ScalarEvolution.cpp

6673

Another thing to double check. Are you sure we only compute exit counts when we have a single loop latch that dominates all exiting nodes? I had a quick look but could not find the place where we ensure that (computeBackedgeTakenCount states that the loop latch may be null). In that case we might miss a case when all exit counts are equal.

[mkazantsev]

lib/Analysis/ScalarEvolution.cpp

6657

Will do before commit.

6673

The latch should be dominated by all exiting blocks, yes. It follows from current implementation of computeBackedgeTakenCount . It currently has *very* strict limitations that in practice say that the exit should actually be the latch (its successors being the loop's header and its chain of single predecessors also leads to header), so this requirement exists but implicit. If you look at https://reviews.llvm.org/D44677, here I tried to explain that the current requirement set is actually noting but a small subset of all situations when all exits dominate the latch. And in this patch, this requirement replaces all odd logic and becomes explicit.

[mkazantsev]

lib/Analysis/ScalarEvolution.cpp

6673

Besides, single latch is a canonical form, and any multiple-latch loop can be easily converted to it (and LLVM does it).

[fhahn]

lib/Analysis/ScalarEvolution.cpp

6673

Right , thanks for making this clearer and more explicit in this patch and D44677 :)

[mkazantsev]

llvm/trunk/lib/Analysis/ScalarEvolution.cpp

6675 ↗

(On Diff #139893)

Actually I start thinking that this may be an over-conservative assert. Imagine the situation:

do { // Loop 1
  do { // Loop 2
     if (invariant_cond) exit
   } while (...)
} while(...)

If we unswitch Loop 2 by invariant condition, we end up with something like

do { // Loop 1
  if (invariant_cond) {
    do { // Loop 2
     } while (...)
   } else {
    do { // Loop 2 copy
     } while (...)
   }
} while(...)

Exits in copies of loop2 no longer dominate the latch of loop1. I didn't comprehend it from the very beginning, but this may be an issue.

include/llvm/Analysis/ScalarEvolution.h

public:
     /// Predicates. A run-time check needs to be emitted for the SCEV
     /// predicate in order for the answer to be valid.
     ///
     /// Note that we should always know if we need to pass a predicate
     /// argument or not from the way the ExitCounts vector was computed.
     /// If we allowed SCEV predicates to be generated when populating this
     /// vector, this information can contain them and therefore a
     /// SCEVPredicate argument should be added to getExact.
-    const SCEV *getExact(ScalarEvolution *SE,
+    const SCEV *getExact(const Loop *L, ScalarEvolution *SE,
                          SCEVUnionPredicate *Predicates = nullptr) const;
 
     /// Return the number of times this loop exit may fall through to the back
     /// edge, or SCEVCouldNotCompute. The loop is guaranteed not to exit via
     /// this block before this number of iterations, but may exit via another
     /// block.
     const SCEV *getExact(BasicBlock *ExitingBlock, ScalarEvolution *SE) const;
 

lib/Analysis/ScalarEvolution.cpp

 const SCEV *ScalarEvolution::getExitCount(const Loop *L,
                                           BasicBlock *ExitingBlock) {
   return getBackedgeTakenInfo(L).getExact(ExitingBlock, this);
 }
 
 const SCEV *
 ScalarEvolution::getPredicatedBackedgeTakenCount(const Loop *L,
                                                  SCEVUnionPredicate &Preds) {
-  return getPredicatedBackedgeTakenInfo(L).getExact(this, &Preds);
+  return getPredicatedBackedgeTakenInfo(L).getExact(L, this, &Preds);
 }
 
 const SCEV *ScalarEvolution::getBackedgeTakenCount(const Loop *L) {
-  return getBackedgeTakenInfo(L).getExact(this);
+  return getBackedgeTakenInfo(L).getExact(L, this);
 }
 
 /// Similar to getBackedgeTakenCount, except return the least SCEV value that is
 /// known never to be less than the actual backedge taken count.
 const SCEV *ScalarEvolution::getMaxBackedgeTakenCount(const Loop *L) {
   return getBackedgeTakenInfo(L).getMax(this);
 }
 
ScalarEvolution::getBackedgeTakenInfo(const Loop *L) {
   if (!Pair.second)
     return Pair.first->second;
 
   // computeBackedgeTakenCount may allocate memory for its result. Inserting it
   // into the BackedgeTakenCounts map transfers ownership. Otherwise, the result
   // must be cleared in this scope.
   BackedgeTakenInfo Result = computeBackedgeTakenCount(L);
 
-  if (Result.getExact(this) != getCouldNotCompute()) {
-    assert(isLoopInvariant(Result.getExact(this), L) &&
+  if (Result.getExact(L, this) != getCouldNotCompute()) {
+    assert(isLoopInvariant(Result.getExact(L, this), L) &&
            isLoopInvariant(Result.getMax(this), L) &&
            "Computed backedge-taken count isn't loop invariant for loop!");
     ++NumTripCountsComputed;
   }
   else if (Result.getMax(this) == getCouldNotCompute() &&
            isa<PHINode>(L->getHeader()->begin())) {
     // Only count loops that have phi nodes as not being computable.
     ++NumTripCountsNotComputed;
while (!Worklist.empty()) {
     }
 
     PushDefUseChildren(I, Worklist);
   }
 }
 
 /// Get the exact loop backedge taken count considering all loop exits. A
 /// computable result can only be returned for loops with a single exit.
 /// Returning the minimum taken count among all exits is incorrect because one
 /// of the loop's exit limit's may have been skipped. howFarToZero assumes that
 /// the limit of each loop test is never skipped. This is a valid assumption as
 /// long as the loop exits via that test. For precise results, it is the
 /// caller's responsibility to specify the relevant loop exit using
 /// getExact(ExitingBlock, SE).
 const SCEV *
-ScalarEvolution::BackedgeTakenInfo::getExact(ScalarEvolution *SE,
+ScalarEvolution::BackedgeTakenInfo::getExact(const Loop *L, ScalarEvolution *SE,
                                              SCEVUnionPredicate *Preds) const {
   // If any exits were not computable, the loop is not computable.
   if (!isComplete() || ExitNotTaken.empty())
     return SE->getCouldNotCompute();
 
   const SCEV *BECount = nullptr;
+  const BasicBlock *Latch = L->getLoopLatch();
+  // All exits we have collected must dominate the only latch.
+  if (!Latch)
+    return SE->getCouldNotCompute();
+
+  // All exits we have gathered dominate loop's latch, so exact trip count is
+  // simply a minimum out of all these calculated exit counts.
   for (auto &ENT : ExitNotTaken) {
     assert(ENT.ExactNotTaken != SE->getCouldNotCompute() && "bad exit SCEV");
+    assert(SE->DT.dominates(ENT.ExitingBlock, Latch) &&
+           "We should only have known counts for exits that dominate latch!");
 
     if (!BECount)
       BECount = ENT.ExactNotTaken;
     else if (BECount != ENT.ExactNotTaken)
-      return SE->getCouldNotCompute();
+      BECount = SE->getUMinFromMismatchedTypes(BECount, ENT.ExactNotTaken);
+
     if (Preds && !ENT.hasAlwaysTruePredicate())
       Preds->add(ENT.Predicate.get());
 
     assert((Preds || ENT.hasAlwaysTruePredicate()) &&
            "Predicate should be always true!");
   }
 
   assert(BECount && "Invalid not taken count for loop exit");

test/Analysis/ScalarEvolution/exact_iter_count.ll

+; RUN: opt < %s -scalar-evolution -analyze | FileCheck %s
+
+; One side exit dominating the latch, exact backedge taken count is known.
+define void @test_01() {
+
+; CHECK-LABEL: Determining loop execution counts for: @test_01
+; CHECK-NEXT:  Loop %loop: <multiple exits> backedge-taken count is 50
+
+entry:
+  br label %loop
+
+loop:
+  %iv = phi i32 [ 0, %entry ], [ %iv.next, %backedge ]
+  %side.cond = icmp slt i32 %iv, 50
+  br i1 %side.cond, label %backedge, label %side.exit
+
+backedge:
+  %iv.next = add i32 %iv, 1
+  %loop.cond = icmp slt i32 %iv, 100
+  br i1 %loop.cond, label %loop, label %exit
+
+exit:
+  ret void
+
+side.exit:
+  ret void
+}

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

 bar.exit:                                         ; preds = %for.cond.i, %for.body.i
   ret i32 0
 }
 
 ; This loop has two must-exits, both of which dominate the latch. The
 ; MaxBECount should be the minimum of them.
 ;
 ; CHECK-LABEL: @two_mustexit
-; CHECK: Loop %for.body.i: <multiple exits> Unpredictable backedge-taken count.
+; CHECK: Loop %for.body.i: <multiple exits> backedge-taken count is 1
 ; CHECK: Loop %for.body.i: max backedge-taken count is 1
 define i32 @two_mustexit() {
 entry:
   store i32 -1, i32* @a, align 4
   br label %for.body.i
 
 for.body.i:                                       ; preds = %for.cond.i, %entry
   %storemerge1.i = phi i32 [ -1, %entry ], [ %add.i.i, %for.cond.i ]

test/Analysis/ScalarEvolution/trip-count14.ll

 if.end:
   %arrayidx = getelementptr i32, i32* %p, i32 %i.0
   store i32 %i.0, i32* %arrayidx, align 4
   %inc = add i32 %i.0, 1
   %cmp1 = icmp slt i32 %i.0, %add
   br i1 %cmp1, label %do.body, label %do.end ; taken either 0 or 2 times
 
 ; CHECK-LABEL: Determining loop execution counts for: @s32_max2_unpredictable_exit
-; CHECK-NEXT: Loop %do.body: <multiple exits> Unpredictable backedge-taken count.
+; CHECK-NEXT: Loop %do.body: <multiple exits> backedge-taken count is (-1 + (-1 * ((-1 + (-1 * ((2 + %n) smax %n)) + %n) umax (-1 + (-1 * %x) + %n))))
 ; CHECK-NEXT: Loop %do.body: max backedge-taken count is 2{{$}}
 
 do.end:
   ret void
 }
 
 define void @u32_max1(i32 %n, i32* %p) {
 entry:

 if.end:
   %arrayidx = getelementptr i32, i32* %p, i32 %i.0
   store i32 %i.0, i32* %arrayidx, align 4
   %inc = add i32 %i.0, 1
   %cmp1 = icmp ult i32 %i.0, %add
   br i1 %cmp1, label %do.body, label %do.end ; taken either 0 or 2 times
 
 ; CHECK-LABEL: Determining loop execution counts for: @u32_max2_unpredictable_exit
-; CHECK-NEXT: Loop %do.body: <multiple exits> Unpredictable backedge-taken count.
+; CHECK-NEXT: Loop %do.body: <multiple exits> backedge-taken count is (-1 + (-1 * ((-1 + (-1 * ((2 + %n) umax %n)) + %n) umax (-1 + (-1 * %x) + %n))))
 ; CHECK-NEXT: Loop %do.body: max backedge-taken count is 2{{$}}
 
 do.end:
   ret void
 }

test/Transforms/IndVarSimplify/loop_evaluate10.ll

 ; RUN: opt < %s -indvars -S | FileCheck %s
 ;
 ; This loop has multiple exits, and the value of %b1 depends on which
 ; exit is taken. Indvars should correctly compute the exit values.
 ;
-; XFAIL: *
-; Indvars does not currently replace loop invariant values unless all
-; loop exits have the same exit value. We could handle some cases,
-; such as this, by making getSCEVAtScope() sensitive to a particular
-; loop exit.  See PR11388.
 
 target datalayout = "e-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v64:64:64-v128:128:128-a0:0:64-s0:64:64-f80:128:128"
 target triple = "x86_64-pc-linux-gnu"
   %struct..0anon = type <{ i8, [3 x i8] }>
 
 define i32 @main() nounwind {
 ; CHECK-LABEL: @main(
 ; CHECK: %b.1 = phi i32 [ 2, %bb ], [ 1, %bb2 ]

test/Transforms/LoopSimplify/preserve-scev.ll

 ; CHECK: Loop %while.cond191: Unpredictable predicated backedge-taken count.
 ; CHECK: Loop %while.cond191.outer: <multiple exits> Unpredictable backedge-taken count.
 ; CHECK: Loop %while.cond191.outer: Unpredictable max backedge-taken count.
 ; CHECK: Loop %while.cond191.outer: Unpredictable predicated backedge-taken count.
 ;
 ; After simplifying, the max backedge count is refined.
 ; Second SCEV print:
 ; CHECK-LABEL: Determining loop execution counts for: @mergeExit
-; CHECK: Loop %while.cond191: <multiple exits> Unpredictable backedge-taken count.
+; CHECK: Loop %while.cond191: <multiple exits> backedge-taken count is 0
 ; CHECK: Loop %while.cond191: max backedge-taken count is 0
-; CHECK: Loop %while.cond191: Unpredictable predicated backedge-taken count.
+; CHECK: Loop %while.cond191: Predicated backedge-taken count is 0
 ; CHECK: Loop %while.cond191.outer: <multiple exits> Unpredictable backedge-taken count.
 ; CHECK: Loop %while.cond191.outer: Unpredictable max backedge-taken count.
 ; CHECK: Loop %while.cond191.outer: Unpredictable predicated backedge-taken count.
 define void @mergeExit(i32 %MapAttrCount) nounwind uwtable ssp {
 entry:
   br i1 undef, label %if.then124, label %if.end126
 
 if.then124:                                       ; preds = %entry