[IndVars] Apply more optimistic SkipLastIter for AND/OR conditions

[mkazantsev]

Summary

When exit by condition C1 dominates exit by condition C2, and
max symbolic exit count for C1 matches those for loop, we will
apply more optimistic logic to C2 by setting SkipLastIter for it,
meaning that it will do 1 iteration less because the dominating branch
must exit on the last loop iteration.

But when we have a single exit by condition C1 & C2, we cannot
apply the same logic, because there is no dominating condition.

However, if we can prove that the symbolic max exit count of C1 & C2
matches those of C1, it means that for C2 we can assume that it
doesn't matter on the last iteration (because the whole thing is false
because C1 must be false). Therefore, in this situation, we can handle
C2 as if it had SkipLastIter.

This is potentiallty CT-consuming operation, so I've limited it with reasonably
small amount of conditions to deal with.

[mkazantsev]

Rebase needed

[mkazantsev]

Simple rebase.

[mkazantsev]

Rebase2

llvm/include/llvm/Analysis/ScalarEvolution.h

1740 ↗

(On Diff #487371)

Note that this is supposed to be moved to public block, it's just to show that nothing else changes here.

[mkazantsev]

Rebased:

• Split off SCEV access change into NFC
• Introduced option to control max amount of icmps to analyze.

[mkazantsev]

Rebased on top of test var renaming.

[mkazantsev]

ping

[nikic]

llvm/lib/Transforms/Scalar/IndVarSimplify.cpp

1499

I don't really get why we need this quadratic-complexity code. It doesn't look like anything inside here depends on the outer i?

[mkazantsev]

llvm/lib/Transforms/Scalar/IndVarSimplify.cpp

1499

OldCond depends on outer i. The logic here is following: we have n conditions joined through AND. Currenly we are handing i'th condition. If any other condition gives exact same exit count as the whole aggregate, it means that the current condition can be handled w/o last iteration. Scanning these other conditions takes linear time (I guess we could cache something if it's a problem, but they can be rewritten, so it's safer not to).

[nikic]

llvm/lib/Transforms/Scalar/IndVarSimplify.cpp

1499

It should be sufficient to go through all the conditions once, and just remember which ones match the maximum. Then it should be possible to use skip last iter if either a) there is more than one such condition or b) the one such condition is not the one currently processed. Probably, for all practical purposes it's fine to just ignore a) and only store a single index for the condition that matches the max exit count.

1499

And if we do that, then I don't think we need the MaxJoinedICmpsToAnalyze limit either, because the complexity is linear.

[mkazantsev]

llvm/lib/Transforms/Scalar/IndVarSimplify.cpp

1499

I was thinking about something like this actually. The fact that stopped me is that, after rewriting, the property of "exits on last iteration" will disappear. Imagine A & B & C, all three exiting on last iteration. We could replace A basing on B, B basing on C and C basing on A. The simplistic approach I am currently using prevents that, however if we sustain a set of all conditions exiting on last iter, that should work. Let me try it out.

[mkazantsev]

Removed quadratic compexity, now it's linear.

[mkazantsev]

ping

[nikic]

LGTM

llvm/lib/Transforms/Scalar/IndVarSimplify.cpp

1489

Would just WideExitMax == WideMaxIter be sufficient? Don't think we usually use isKnownPredicate for equality comparisons and rely on canonicalization instead.

[nikic]

(patch description needs an update, last paragraph can be dropped)

[mkazantsev]

llvm/lib/Transforms/Scalar/IndVarSimplify.cpp

1489

I think it should be, in any practical sense.

llvm/lib/Transforms/Scalar/IndVarSimplify.cpp

do {
     // Go through AND/OR conditions. Collect leaf ICMPs. We only care about
     // those with one use, to avoid instruction duplication.
     if (Curr->hasOneUse())
       if (!GoThrough(Curr))
         if (auto *ICmp = dyn_cast<ICmpInst>(Curr))
           LeafConditions.push_back(ICmp);
   } while (!Worklist.empty());
 
+  // If the current basic block has the same exit count as the whole loop, and
+  // it consists of multiple icmp's, try to collect all icmp's that give exact
+  // same exit count. For all other icmp's, we could use one less iteration,
+  // because their value on the last iteration doesn't really matter.
+  SmallPtrSet<ICmpInst *, 4> ICmpsFailingOnLastIter;
+  if (!SkipLastIter && LeafConditions.size() > 1 &&
+      SE->getExitCount(L, ExitingBB,
+                       ScalarEvolution::ExitCountKind::SymbolicMaximum) ==
+          MaxIter)
+    for (auto *ICmp : LeafConditions) {
+      auto EL = SE->computeExitLimitFromCond(L, ICmp, Inverted,
+                                             /*ControlsExit*/ false);
+      auto *ExitMax = EL.SymbolicMaxNotTaken;
+      if (isa<SCEVCouldNotCompute>(ExitMax))
+        continue;
+      // They could be of different types (specifically this happens after
+      // IV widening).
+      auto *WiderType =
+          SE->getWiderType(ExitMax->getType(), MaxIter->getType());
+      auto *WideExitMax = SE->getNoopOrZeroExtend(ExitMax, WiderType);
+      auto *WideMaxIter = SE->getNoopOrZeroExtend(MaxIter, WiderType);
+      if (WideExitMax == WideMaxIter)
+        ICmpsFailingOnLastIter.insert(ICmp);
+    }
+
   bool Changed = false;
-  for (auto *OldCond : LeafConditions)
+  for (auto *OldCond : LeafConditions) {
+    // Skip last iteration for this icmp under one of two conditions:
+    // - We do it for all conditions;
+    // - There is another ICmp that would fail on last iter, so this one doesn't
+    // really matter.
+    bool OptimisticSkipLastIter = SkipLastIter;
+    if (!OptimisticSkipLastIter) {
+      if (ICmpsFailingOnLastIter.size() > 1)
+        OptimisticSkipLastIter = true;
+      else if (ICmpsFailingOnLastIter.size() == 1)
+        OptimisticSkipLastIter = !ICmpsFailingOnLastIter.count(OldCond);
+    }
     if (auto Replaced =
             createReplacement(OldCond, L, ExitingBB, MaxIter, Inverted,
-                              SkipLastIter, SE, Rewriter)) {
+                              OptimisticSkipLastIter, SE, Rewriter)) {
       Changed = true;
       auto *NewCond = *Replaced;
       if (auto *NCI = dyn_cast<Instruction>(NewCond)) {
         NCI->setName(OldCond->getName() + ".first_iter");
         NCI->moveBefore(cast<Instruction>(OldCond));
       }
       LLVM_DEBUG(dbgs() << "Unknown exit count: Replacing " << *OldCond
                         << " with " << *NewCond << "\n");
       assert(OldCond->hasOneUse() && "Must be!");
       OldCond->replaceAllUsesWith(NewCond);
       DeadInsts.push_back(OldCond);
-    }
+      // Make sure we no longer consider this condition as failing on last
+      // iteration.
+      ICmpsFailingOnLastIter.erase(OldCond);
+    }
   return Changed;
 }
 
 bool IndVarSimplify::canonicalizeExitCondition(Loop *L) {
   // Note: This is duplicating a particular part on SimplifyIndVars reasoning.
   // We need to duplicate it because given icmp zext(small-iv), C, IVUsers
   // never reaches the icmp since the zext doesn't fold to an AddRec unless
   // it already has flags.  The alternative to this would be to extending the

llvm/test/Transforms/IndVarSimplify/X86/widening-vs-and-elimination.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt -S -passes=indvars < %s | FileCheck %s --check-prefixes=WIDENING_ON
 ; RUN: opt -S -passes=indvars -indvars-widen-indvars=false < %s | FileCheck %s --check-prefixes=WIDENING_OFF
 
 target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128-ni:1-p2:32:8:8:32-ni:2"
 target triple = "x86_64-unknown-linux-gnu"
 
 define void @test_01(i32 %start, i32 %limit) {
 ; WIDENING_ON-LABEL: @test_01(
 ; WIDENING_ON-NEXT:  bb:
 ; WIDENING_ON-NEXT:    [[TMP0:%.*]] = zext i32 [[START:%.*]] to i64
+; WIDENING_ON-NEXT:    [[TMP1:%.*]] = add nsw i64 [[TMP0]], -1
+; WIDENING_ON-NEXT:    [[TMP2:%.*]] = zext i32 [[LIMIT:%.*]] to i64
 ; WIDENING_ON-NEXT:    br label [[LOOP:%.*]]
 ; WIDENING_ON:       loop:
 ; WIDENING_ON-NEXT:    [[INDVARS_IV:%.*]] = phi i64 [ [[INDVARS_IV_NEXT:%.*]], [[BACKEDGE:%.*]] ], [ [[TMP0]], [[BB:%.*]] ]
-; WIDENING_ON-NEXT:    [[TMP1:%.*]] = add nsw i64 [[INDVARS_IV]], -1
 ; WIDENING_ON-NEXT:    [[INDVARS_IV_NEXT]] = add nsw i64 [[INDVARS_IV]], -1
 ; WIDENING_ON-NEXT:    [[NOT_ZERO:%.*]] = icmp ne i64 [[INDVARS_IV]], 0
-; WIDENING_ON-NEXT:    [[TMP2:%.*]] = zext i32 [[LIMIT:%.*]] to i64
-; WIDENING_ON-NEXT:    [[RANGE_CHECK_WIDE:%.*]] = icmp ult i64 [[TMP1]], [[TMP2]]
-; WIDENING_ON-NEXT:    [[AND:%.*]] = and i1 [[NOT_ZERO]], [[RANGE_CHECK_WIDE]]
+; WIDENING_ON-NEXT:    [[RANGE_CHECK_WIDE_FIRST_ITER:%.*]] = icmp ult i64 [[TMP1]], [[TMP2]]
+; WIDENING_ON-NEXT:    [[AND:%.*]] = and i1 [[NOT_ZERO]], [[RANGE_CHECK_WIDE_FIRST_ITER]]
 ; WIDENING_ON-NEXT:    br i1 [[AND]], label [[BACKEDGE]], label [[EXIT:%.*]]
 ; WIDENING_ON:       backedge:
 ; WIDENING_ON-NEXT:    br label [[LOOP]]
 ; WIDENING_ON:       exit:
 ; WIDENING_ON-NEXT:    ret void
 ;
 ; WIDENING_OFF-LABEL: @test_01(
 ; WIDENING_OFF-NEXT:  bb:
+; WIDENING_OFF-NEXT:    [[TMP0:%.*]] = add i32 [[START:%.*]], -1
 ; WIDENING_OFF-NEXT:    br label [[LOOP:%.*]]
 ; WIDENING_OFF:       loop:
-; WIDENING_OFF-NEXT:    [[IV:%.*]] = phi i32 [ [[IV_NEXT:%.*]], [[BACKEDGE:%.*]] ], [ [[START:%.*]], [[BB:%.*]] ]
+; WIDENING_OFF-NEXT:    [[IV:%.*]] = phi i32 [ [[IV_NEXT:%.*]], [[BACKEDGE:%.*]] ], [ [[START]], [[BB:%.*]] ]
 ; WIDENING_OFF-NEXT:    [[IV_NEXT]] = add i32 [[IV]], -1
 ; WIDENING_OFF-NEXT:    [[NOT_ZERO:%.*]] = icmp ne i32 [[IV]], 0
-; WIDENING_OFF-NEXT:    [[RANGE_CHECK:%.*]] = icmp ult i32 [[IV_NEXT]], [[LIMIT:%.*]]
-; WIDENING_OFF-NEXT:    [[AND:%.*]] = and i1 [[NOT_ZERO]], [[RANGE_CHECK]]
+; WIDENING_OFF-NEXT:    [[RANGE_CHECK_FIRST_ITER:%.*]] = icmp ult i32 [[TMP0]], [[LIMIT:%.*]]
+; WIDENING_OFF-NEXT:    [[AND:%.*]] = and i1 [[NOT_ZERO]], [[RANGE_CHECK_FIRST_ITER]]
 ; WIDENING_OFF-NEXT:    br i1 [[AND]], label [[BACKEDGE]], label [[EXIT:%.*]]
 ; WIDENING_OFF:       backedge:
 ; WIDENING_OFF-NEXT:    [[ZEXT:%.*]] = zext i32 [[IV_NEXT]] to i64
 ; WIDENING_OFF-NEXT:    [[GEP:%.*]] = getelementptr inbounds i32, ptr addrspace(1) poison, i64 [[ZEXT]]
 ; WIDENING_OFF-NEXT:    br label [[LOOP]]
 ; WIDENING_OFF:       exit:
 ; WIDENING_OFF-NEXT:    ret void
 ;

llvm/test/Transforms/IndVarSimplify/turn-to-invariant.ll

 
 failed_1:
   ret i32 -1
 
 failed_2:
   ret i32 -2
 }
 
-; TODO: This test is equivalent to @test_simple_case, with only difference
-;       that both checks are merged together into one 'and' check. This
-;       should not prevent turning the range check into invariant.
-;       https://alive2.llvm.org/ce/z/G-2ERB
+; This test is equivalent to @test_simple_case, with only difference
+; that both checks are merged together into one 'and' check. This
+; should not prevent turning the range check into invariant.
+; https://alive2.llvm.org/ce/z/G-2ERB
 define i32 @test_and_conditions(i32 %start, i32 %len) {
 ; CHECK-LABEL: @test_and_conditions(
 ; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = add i32 [[START:%.*]], -1
 ; CHECK-NEXT:    br label [[LOOP:%.*]]
 ; CHECK:       loop:
-; CHECK-NEXT:    [[IV:%.*]] = phi i32 [ [[START:%.*]], [[ENTRY:%.*]] ], [ [[IV_NEXT:%.*]], [[BACKEDGE:%.*]] ]
+; CHECK-NEXT:    [[IV:%.*]] = phi i32 [ [[START]], [[ENTRY:%.*]] ], [ [[IV_NEXT:%.*]], [[BACKEDGE:%.*]] ]
 ; CHECK-NEXT:    [[ZERO_CHECK:%.*]] = icmp ne i32 [[IV]], 0
-; CHECK-NEXT:    [[IV_MINUS_1:%.*]] = add i32 [[IV]], -1
-; CHECK-NEXT:    [[RANGE_CHECK:%.*]] = icmp ult i32 [[IV_MINUS_1]], [[LEN:%.*]]
-; CHECK-NEXT:    [[BOTH_CHECKS:%.*]] = and i1 [[ZERO_CHECK]], [[RANGE_CHECK]]
+; CHECK-NEXT:    [[RANGE_CHECK_FIRST_ITER:%.*]] = icmp ult i32 [[TMP0]], [[LEN:%.*]]
+; CHECK-NEXT:    [[BOTH_CHECKS:%.*]] = and i1 [[ZERO_CHECK]], [[RANGE_CHECK_FIRST_ITER]]
 ; CHECK-NEXT:    br i1 [[BOTH_CHECKS]], label [[BACKEDGE]], label [[FAILED:%.*]]
 ; CHECK:       backedge:
 ; CHECK-NEXT:    [[IV_NEXT]] = add i32 [[IV]], -1
 ; CHECK-NEXT:    [[LOOP_COND:%.*]] = call i1 @cond()
 ; CHECK-NEXT:    br i1 [[LOOP_COND]], label [[DONE:%.*]], label [[LOOP]]
 ; CHECK:       done:
 ; CHECK-NEXT:    [[IV_LCSSA1:%.*]] = phi i32 [ [[IV]], [[BACKEDGE]] ]
 ; CHECK-NEXT:    ret i32 [[IV_LCSSA1]]

 
 done:
   ret i32 %iv
 
 failed:
   ret i32 -3
 }
 
-; TODO: Same as test_and_conditions, but swapped exit block branches,
-;       and condition is expressed as OR. Still optimizable.
+; Same as test_and_conditions, but swapped exit block branches,
+; and condition is expressed as OR. Still optimizable.
 define i32 @test_and_conditions_inverse(i32 %start, i32 %len) {
 ; CHECK-LABEL: @test_and_conditions_inverse(
 ; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = add i32 [[START:%.*]], -1
 ; CHECK-NEXT:    br label [[LOOP:%.*]]
 ; CHECK:       loop:
-; CHECK-NEXT:    [[IV:%.*]] = phi i32 [ [[START:%.*]], [[ENTRY:%.*]] ], [ [[IV_NEXT:%.*]], [[BACKEDGE:%.*]] ]
+; CHECK-NEXT:    [[IV:%.*]] = phi i32 [ [[START]], [[ENTRY:%.*]] ], [ [[IV_NEXT:%.*]], [[BACKEDGE:%.*]] ]
 ; CHECK-NEXT:    [[ZERO_CHECK_FAILED:%.*]] = icmp eq i32 [[IV]], 0
-; CHECK-NEXT:    [[IV_MINUS_1:%.*]] = add i32 [[IV]], -1
-; CHECK-NEXT:    [[RANGE_CHECK_FAILED:%.*]] = icmp uge i32 [[IV_MINUS_1]], [[LEN:%.*]]
-; CHECK-NEXT:    [[EITHER_CHECK:%.*]] = or i1 [[ZERO_CHECK_FAILED]], [[RANGE_CHECK_FAILED]]
+; CHECK-NEXT:    [[RANGE_CHECK_FAILED_FIRST_ITER:%.*]] = icmp uge i32 [[TMP0]], [[LEN:%.*]]
+; CHECK-NEXT:    [[EITHER_CHECK:%.*]] = or i1 [[ZERO_CHECK_FAILED]], [[RANGE_CHECK_FAILED_FIRST_ITER]]
 ; CHECK-NEXT:    br i1 [[EITHER_CHECK]], label [[FAILED:%.*]], label [[BACKEDGE]]
 ; CHECK:       backedge:
 ; CHECK-NEXT:    [[IV_NEXT]] = add i32 [[IV]], -1
 ; CHECK-NEXT:    [[LOOP_COND:%.*]] = call i1 @cond()
 ; CHECK-NEXT:    br i1 [[LOOP_COND]], label [[DONE:%.*]], label [[LOOP]]
 ; CHECK:       done:
 ; CHECK-NEXT:    [[IV_LCSSA1:%.*]] = phi i32 [ [[IV]], [[BACKEDGE]] ]
 ; CHECK-NEXT:    ret i32 [[IV_LCSSA1]]

 
 failed_1:
   ret i32 -1
 
 failed_2:
   ret i32 -2
 }
 
-; TODO: Same as test_and_conditions, but with logical AND.
+; Same as test_and_conditions, but with logical AND.
 define i32 @test_and_conditions_logical_and(i32 %start, i32 %len) {
 ; CHECK-LABEL: @test_and_conditions_logical_and(
 ; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = add i32 [[START:%.*]], -1
 ; CHECK-NEXT:    br label [[LOOP:%.*]]
 ; CHECK:       loop:
-; CHECK-NEXT:    [[IV:%.*]] = phi i32 [ [[START:%.*]], [[ENTRY:%.*]] ], [ [[IV_NEXT:%.*]], [[BACKEDGE:%.*]] ]
+; CHECK-NEXT:    [[IV:%.*]] = phi i32 [ [[START]], [[ENTRY:%.*]] ], [ [[IV_NEXT:%.*]], [[BACKEDGE:%.*]] ]
 ; CHECK-NEXT:    [[ZERO_CHECK:%.*]] = icmp ne i32 [[IV]], 0
-; CHECK-NEXT:    [[IV_MINUS_1:%.*]] = add i32 [[IV]], -1
-; CHECK-NEXT:    [[RANGE_CHECK:%.*]] = icmp ult i32 [[IV_MINUS_1]], [[LEN:%.*]]
-; CHECK-NEXT:    [[BOTH_CHECKS:%.*]] = select i1 [[ZERO_CHECK]], i1 [[RANGE_CHECK]], i1 false
+; CHECK-NEXT:    [[RANGE_CHECK_FIRST_ITER:%.*]] = icmp ult i32 [[TMP0]], [[LEN:%.*]]
+; CHECK-NEXT:    [[BOTH_CHECKS:%.*]] = select i1 [[ZERO_CHECK]], i1 [[RANGE_CHECK_FIRST_ITER]], i1 false
 ; CHECK-NEXT:    br i1 [[BOTH_CHECKS]], label [[BACKEDGE]], label [[FAILED:%.*]]
 ; CHECK:       backedge:
 ; CHECK-NEXT:    [[IV_NEXT]] = add i32 [[IV]], -1
 ; CHECK-NEXT:    [[LOOP_COND:%.*]] = call i1 @cond()
 ; CHECK-NEXT:    br i1 [[LOOP_COND]], label [[DONE:%.*]], label [[LOOP]]
 ; CHECK:       done:
 ; CHECK-NEXT:    [[IV_LCSSA1:%.*]] = phi i32 [ [[IV]], [[BACKEDGE]] ]
 ; CHECK-NEXT:    ret i32 [[IV_LCSSA1]]

 
 done:
   ret i32 %iv
 
 failed:
   ret i32 -3
 }
 
-; TODO: Same as test_and_conditions_inverse, but with logical OR.
+; Same as test_and_conditions_inverse, but with logical OR.
 define i32 @test_and_conditions_inverse_logical_or(i32 %start, i32 %len) {
 ; CHECK-LABEL: @test_and_conditions_inverse_logical_or(
 ; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = add i32 [[START:%.*]], -1
 ; CHECK-NEXT:    br label [[LOOP:%.*]]
 ; CHECK:       loop:
-; CHECK-NEXT:    [[IV:%.*]] = phi i32 [ [[START:%.*]], [[ENTRY:%.*]] ], [ [[IV_NEXT:%.*]], [[BACKEDGE:%.*]] ]
+; CHECK-NEXT:    [[IV:%.*]] = phi i32 [ [[START]], [[ENTRY:%.*]] ], [ [[IV_NEXT:%.*]], [[BACKEDGE:%.*]] ]
 ; CHECK-NEXT:    [[ZERO_CHECK_FAILED:%.*]] = icmp eq i32 [[IV]], 0
-; CHECK-NEXT:    [[IV_MINUS_1:%.*]] = add i32 [[IV]], -1
-; CHECK-NEXT:    [[RANGE_CHECK_FAILED:%.*]] = icmp uge i32 [[IV_MINUS_1]], [[LEN:%.*]]
-; CHECK-NEXT:    [[EITHER_CHECK:%.*]] = select i1 [[ZERO_CHECK_FAILED]], i1 true, i1 [[RANGE_CHECK_FAILED]]
+; CHECK-NEXT:    [[RANGE_CHECK_FAILED_FIRST_ITER:%.*]] = icmp uge i32 [[TMP0]], [[LEN:%.*]]
+; CHECK-NEXT:    [[EITHER_CHECK:%.*]] = select i1 [[ZERO_CHECK_FAILED]], i1 true, i1 [[RANGE_CHECK_FAILED_FIRST_ITER]]
 ; CHECK-NEXT:    br i1 [[EITHER_CHECK]], label [[FAILED:%.*]], label [[BACKEDGE]]
 ; CHECK:       backedge:
 ; CHECK-NEXT:    [[IV_NEXT]] = add i32 [[IV]], -1
 ; CHECK-NEXT:    [[LOOP_COND:%.*]] = call i1 @cond()
 ; CHECK-NEXT:    br i1 [[LOOP_COND]], label [[DONE:%.*]], label [[LOOP]]
 ; CHECK:       done:
 ; CHECK-NEXT:    [[IV_LCSSA1:%.*]] = phi i32 [ [[IV]], [[BACKEDGE]] ]
 ; CHECK-NEXT:    ret i32 [[IV_LCSSA1]]