[IndVars] Use knowledge about execution on last iteration when removing checks

[mkazantsev]

Summary

If we know that some check will not be executed on the last iteration, we can use this
fact to eliminate its check.

[mkazantsev]

Rebased.

[reames]

Marking as Needs Changes purely to indicate I'm holding off on reviewing this in depth until changes this is built upon have all landed. Please refresh at that point.

Max, have you looked at SCEV->isLoopInvariantPredicate? It tries to handle some of the same reasoning. It might be useable directly here, or if not, can the logic maybe be consolidated in one place rather than two? I have not looked at this in depth yet, but that'll be the question I'm asking myself when reviewing this in detail.

[mkazantsev]

Opening after underlying patches have been merged.

[ebrevnov]

llvm/lib/Transforms/Scalar/IndVarSimplify.cpp

2309

I think we better avoid adding SkipLastIter parameter since we should be able to cover all cases with existing MaxIter and it just unnecessarily complicates the API.

2340

Let's move this to the caller.

2470

I suggest to adjust MaxIter here. It should greatly simplify the logic in this loop.

[mkazantsev]

Move last iter adjustment into lambda.

[ebrevnov]

llvm/lib/Transforms/Scalar/IndVarSimplify.cpp

2462

Please rework the comment so it is crystal clear that we actually call OptimizeCond for both MaxIter and MaxIter -1. And the reason is implementation limitation in SCEV which is unable to prove that MaxIter-1 < MaxIter.

[mkazantsev]

Tried to make the comment more clear.

[ebrevnov]

LGTM. But let chance for others to comment. Thanks.

[ebrevnov]

llvm/lib/Transforms/Scalar/IndVarSimplify.cpp

2429–2436

NIT1: please list all captured values explicitly. Maybe capture SkipLastIter by value?
NIT2: please think of moving handling of SkipLastIter case to this lambda. In this case we will "hide" complexity from the main loop.

2433

This lambda may be called twice with SkipLastIter equal to true for each exit. As a result MaxIter may be decremented by 2 for particular exit.

[mkazantsev]

llvm/lib/Transforms/Scalar/IndVarSimplify.cpp

2429–2436

As for nit 1: will do;
As for nit 2: this logic is going to grow more complex in further patch (see dependency), so, to avoid copy-paste, I'd rather keep it as is.

2433

This cannot happen because MaxIter is local.

[ebrevnov]

LGTM

llvm/lib/Transforms/Scalar/IndVarSimplify.cpp

bool IndVarSimplify::sinkUnusedInvariants(Loop *L) {
 
   return MadeAnyChanges;
 }
 
 // Returns true if the condition of \p BI being checked is invariant and can be
 // proved to be trivially true during at least first \p MaxIter iterations.
 static bool isTrivialCond(const Loop *L, BranchInst *BI, ScalarEvolution *SE,
                           bool ProvingLoopExit, const SCEV *MaxIter) {
   ICmpInst::Predicate Pred;
   Value *LHS, *RHS;
   using namespace PatternMatch;
   BasicBlock *TrueSucc, *FalseSucc;
   if (!match(BI, m_Br(m_ICmp(Pred, m_Value(LHS), m_Value(RHS)),
                       m_BasicBlock(TrueSucc), m_BasicBlock(FalseSucc))))
     return false;
 
   assert((L->contains(TrueSucc) != L->contains(FalseSucc)) &&
if (SE->isKnownPredicateAt(Pred, LHSS, RHSS, BI))
     return true;
 
   if (ProvingLoopExit)
     return false;
 
   ICmpInst::Predicate InvariantPred;
   const SCEV *InvariantLHS, *InvariantRHS;
 
   // Check if there is a loop-invariant predicate equivalent to our check.
   if (!SE->isLoopInvariantExitCondDuringFirstIterations(
            Pred, LHSS, RHSS, L, BI, MaxIter, InvariantPred, InvariantLHS,
            InvariantRHS))
     return false;
 
   // Can we prove it to be trivially true?
   return SE->isKnownPredicateAt(InvariantPred, InvariantLHS, InvariantRHS, BI);
 }
auto FoldExit = [&](BasicBlock *ExitingBB, bool IsTaken) {
     auto *NewCond = ConstantInt::get(OldCond->getType(),
                                      IsTaken ? ExitIfTrue : !ExitIfTrue);
     BI->setCondition(NewCond);
     if (OldCond->use_empty())
       DeadInsts.emplace_back(OldCond);
   };
 
   bool Changed = false;
+  bool SkipLastIter = false;
   SmallSet<const SCEV*, 8> DominatingExitCounts;
   for (BasicBlock *ExitingBB : ExitingBlocks) {
     const SCEV *ExitCount = SE->getExitCount(L, ExitingBB);
     if (isa<SCEVCouldNotCompute>(ExitCount)) {
       // Okay, we do not know the exit count here. Can we at least prove that it
       // will remain the same within iteration space?
       auto *BI = cast<BranchInst>(ExitingBB->getTerminator());
-      auto OptimizeCond = [&](bool Inverted) {
-        if (isTrivialCond(L, BI, SE, Inverted, MaxExitCount)) {
+      auto OptimizeCond = [this, L, BI, ExitingBB, MaxExitCount, &FoldExit](
+          bool Inverted, bool SkipLastIter) {
+        const SCEV *MaxIter = MaxExitCount;
+        if (SkipLastIter) {
+          const SCEV *One = SE->getOne(MaxIter->getType());
+          MaxIter = SE->getMinusSCEV(MaxIter, One);
+        }
+        if (isTrivialCond(L, BI, SE, Inverted, MaxIter)) {
           FoldExit(ExitingBB, Inverted);
           return true;
         }
         return false;
       };
-      if (OptimizeCond(false) || OptimizeCond(true))
+
+      // TODO: We might have proved that we can skip the last iteration for
+      // this check. In this case, we only want to check the condition on the
+      // pre-last iteration (MaxExitCount - 1). However, there is a nasty
+      // corner case:
+      //
+      //   for (i = len; i != 0; i--) { ... check (i ult X) ... }
+      //
+      // If we could not prove that len != 0, then we also could not prove that
+      // (len - 1) is not a UINT_MAX. If we simply query (len - 1), then
+      // OptimizeCond will likely not prove anything for it, even if it could
+      // prove the same fact for len.
+      //
+      // As a temporary solution, we query both last and pre-last iterations in
+      // hope that we will be able to prove triviality for at least one of
+      // them. We can stop querying MaxExitCount for this case once SCEV
+      // understands that (MaxExitCount - 1) will not overflow here.
+      if (OptimizeCond(false, false) || OptimizeCond(true, false))
+        Changed = true;
+      else if (SkipLastIter)
+        if (OptimizeCond(false, true) || OptimizeCond(true, true))
           Changed = true;
       continue;
     }
 
+    if (MaxExitCount == ExitCount)
+      // If the loop has more than 1 iteration, all further checks will be
+      // executed 1 iteration less.
+      SkipLastIter = true;
+
     // If we know we'd exit on the first iteration, rewrite the exit to
     // reflect this.  This does not imply the loop must exit through this
     // exit; there may be an earlier one taken on the first iteration.
     // TODO: Given we know the backedge can't be taken, we should go ahead
     // and break it.  Or at least, kill all the header phis and simplify.
     if (ExitCount->isZero()) {
       FoldExit(ExitingBB, true);
       Changed = true;

llvm/test/Transforms/IndVarSimplify/predicated_ranges.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt -indvars -S < %s | FileCheck %s
 ; RUN: opt -passes=indvars -S < %s | FileCheck %s
 
-; TODO: should be able to remove the range check basing on the following facts:
+; Check that we are able to remove the range check basing on the following facts:
 ; 0 <= len <= MAX_INT [1];
 ; iv starts from len and goes down stopping at zero and [1], therefore
 ;   0 <= iv <= len [2];
 ; 3. In range_check_block, iv != 0 and [2], therefore
 ;   1 <= iv <= len [3];
 ; 4. iv.next = iv - 1 and [3], therefore
 ;   0 <= iv.next < len.
 define void @test_predicated_simple_unsigned(i32* %p, i32* %arr) {
 ; CHECK-LABEL: @test_predicated_simple_unsigned(
 ; CHECK-NEXT:  preheader:
 ; CHECK-NEXT:    [[LEN:%.*]] = load i32, i32* [[P:%.*]], align 4, [[RNG0:!range !.*]]
 ; CHECK-NEXT:    br label [[LOOP:%.*]]
 ; CHECK:       loop:
 ; CHECK-NEXT:    [[IV:%.*]] = phi i32 [ [[LEN]], [[PREHEADER:%.*]] ], [ [[IV_NEXT:%.*]], [[BACKEDGE:%.*]] ]
 ; CHECK-NEXT:    [[ZERO_COND:%.*]] = icmp eq i32 [[IV]], 0
 ; CHECK-NEXT:    br i1 [[ZERO_COND]], label [[EXIT:%.*]], label [[RANGE_CHECK_BLOCK:%.*]]
 ; CHECK:       range_check_block:
 ; CHECK-NEXT:    [[IV_NEXT]] = sub i32 [[IV]], 1
-; CHECK-NEXT:    [[RANGE_CHECK:%.*]] = icmp ult i32 [[IV_NEXT]], [[LEN]]
-; CHECK-NEXT:    br i1 [[RANGE_CHECK]], label [[BACKEDGE]], label [[FAIL:%.*]]
+; CHECK-NEXT:    br i1 true, label [[BACKEDGE]], label [[FAIL:%.*]]
 ; CHECK:       backedge:
 ; CHECK-NEXT:    [[EL_PTR:%.*]] = getelementptr i32, i32* [[P]], i32 [[IV]]
 ; CHECK-NEXT:    [[EL:%.*]] = load i32, i32* [[EL_PTR]], align 4
 ; CHECK-NEXT:    [[LOOP_COND:%.*]] = icmp eq i32 [[EL]], 0
 ; CHECK-NEXT:    br i1 [[LOOP_COND]], label [[LOOP]], label [[EXIT]]
 ; CHECK:       exit:
 ; CHECK-NEXT:    ret void
 ; CHECK:       fail: