[IndVars] Use exit count up to current block when handing ANDs/ORs

[mkazantsev]

Summary

When analyzing exit conditions joined through AND/OR, use exit count
computed not for current block, but for all blocks from loop start and up
to the current block together. It is not actually required to know that the
current block will exit to skip the last iteration optimistically, it is enough
to know that we will exit on this block or earlier.

[mkazantsev]

This needs rebase afrer underlying code changed a lot.

[mkazantsev]

Rebased

[mkazantsev]

Ping

[mkazantsev]

Ping

llvm/lib/Transforms/Scalar/IndVarSimplify.cpp

createReplacement(ICmpInst *ICmp, const Loop *L, BasicBlock *ExitingBB,
   // Can we prove it to be trivially true?
   if (SE->isKnownPredicateAt(LIP->Pred, LIP->LHS, LIP->RHS, BI))
     return createFoldedExitCond(L, ExitingBB, /*IsTaken*/ false);
   else
     return createInvariantCond(L, ExitingBB, *LIP, Rewriter);
 }
 
 static bool optimizeLoopExitWithUnknownExitCount(
-    const Loop *L, BranchInst *BI, BasicBlock *ExitingBB, const SCEV *MaxIter,
-    bool SkipLastIter, ScalarEvolution *SE, SCEVExpander &Rewriter,
+    const Loop *L, BranchInst *BI, BasicBlock *ExitingBB,
+    const SCEV *CurrMaxIter, const SCEV *MaxIter, bool SkipLastIter,
+    ScalarEvolution *SE, SCEVExpander &Rewriter,
     SmallVectorImpl<WeakTrackingVH> &DeadInsts) {
   assert(
       (L->contains(BI->getSuccessor(0)) != L->contains(BI->getSuccessor(1))) &&
       "Not a loop exit!");
 
   // For branch that stays in loop by TRUE condition, go through AND. For branch
   // that stays in loop by FALSE condition, go through OR. Both gives the
   // similar logic: "stay in loop iff all conditions are true(false)".
if (Curr->hasOneUse())
           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)
+  if (!SkipLastIter && LeafConditions.size() > 1) {
+    auto *CurrBBExitCount = SE->getExitCount(
+        L, ExitingBB, ScalarEvolution::ExitCountKind::SymbolicMaximum);
+    if (!isa<SCEVCouldNotCompute>(CurrBBExitCount)) {
+      auto *ExitCountUpToCurrBB =
+          isa<SCEVCouldNotCompute>(CurrMaxIter)
+              ? CurrBBExitCount
+              : SE->getUMinFromMismatchedTypes(CurrBBExitCount, CurrMaxIter);
+      if (ExitCountUpToCurrBB == MaxIter)
         for (auto *ICmp : LeafConditions) {
           auto EL = SE->computeExitLimitFromCond(L, ICmp, Inverted,
                                                  /*ControlsExit*/ false);
           auto *ExitMax = EL.SymbolicMaxNotTaken;
           if (isa<SCEVCouldNotCompute>(ExitMax))
             continue;
+          auto *MaxUpToCurrentICmp =
+              isa<SCEVCouldNotCompute>(CurrMaxIter)
+                  ? ExitMax
+                  : SE->getUMinFromMismatchedTypes(CurrMaxIter, ExitMax);
           // 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 *WiderType = SE->getWiderType(ExitMax->getType(),
+                                             MaxUpToCurrentICmp->getType());
+          auto *WideMaxUpToCurrentICmp =
+              SE->getNoopOrZeroExtend(MaxUpToCurrentICmp, WiderType);
           auto *WideMaxIter = SE->getNoopOrZeroExtend(MaxIter, WiderType);
-      if (WideExitMax == WideMaxIter)
+          if (WideMaxUpToCurrentICmp == WideMaxIter)
             ICmpsFailingOnLastIter.insert(ICmp);
         }
+    }
+  }
 
   bool Changed = false;
   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;
for (BasicBlock *ExitingBB : ExitingBlocks) {
     const SCEV *ExactExitCount = SE->getExitCount(L, ExitingBB);
     const SCEV *MaxExitCount = SE->getExitCount(
         L, ExitingBB, ScalarEvolution::ExitCountKind::SymbolicMaximum);
     if (isa<SCEVCouldNotCompute>(ExactExitCount)) {
       // 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 SkipLastIter) {
-        return optimizeLoopExitWithUnknownExitCount(L, BI, ExitingBB,
-                                                    MaxBECount, SkipLastIter,
-                                                    SE, Rewriter, DeadInsts);
+        return optimizeLoopExitWithUnknownExitCount(
+            L, BI, ExitingBB, CurrMaxExit, MaxBECount, SkipLastIter, SE,
+            Rewriter, DeadInsts);
       };
 
       // 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 (MaxBECount - 1). However, there is a nasty
       // corner case:
       //
       //   for (i = len; i != 0; i--) { ... check (i ult X) ... }

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

 exit:                                             ; preds = %loop
   ret void
 }
 
 define i32 @test_02(i32 %start, i32 %limit) {
 ; WIDENING_ON-LABEL: @test_02(
 ; 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:    [[RANGE_CHECK_WIDE_FIRST_ITER:%.*]] = icmp ult i64 [[TMP1]], [[TMP2]]
 ; 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:    [[CANONICAL_IV:%.*]] = phi i32 [ [[CANONICAL_IV_NEXT:%.*]], [[BACKEDGE]] ], [ 0, [[BB]] ]
 ; WIDENING_ON-NEXT:    [[EXITCOND:%.*]] = icmp ne i32 [[CANONICAL_IV]], 65635
 ; WIDENING_ON-NEXT:    br i1 [[EXITCOND]], label [[CHECKED:%.*]], label [[FAILED:%.*]]
 ; WIDENING_ON:       checked:
-; 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:    [[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:    [[CANONICAL_IV_NEXT]] = add nuw nsw i32 [[CANONICAL_IV]], 1
 ; WIDENING_ON-NEXT:    br label [[LOOP]]
 ; WIDENING_ON:       exit:
 ; WIDENING_ON-NEXT:    ret i32 0
 ; WIDENING_ON:       failed:
 ; WIDENING_ON-NEXT:    ret i32 1
 ;
 ; WIDENING_OFF-LABEL: @test_02(
 ; WIDENING_OFF-NEXT:  bb:
+; WIDENING_OFF-NEXT:    [[TMP0:%.*]] = add i32 [[START:%.*]], -1
+; WIDENING_OFF-NEXT:    [[RANGE_CHECK_FIRST_ITER:%.*]] = icmp ult i32 [[TMP0]], [[LIMIT:%.*]]
 ; WIDENING_OFF-NEXT:    br label [[LOOP:%.*]]
 ; WIDENING_OFF:       loop:
 ; WIDENING_OFF-NEXT:    [[CANONICAL_IV:%.*]] = phi i32 [ [[CANONICAL_IV_NEXT:%.*]], [[BACKEDGE:%.*]] ], [ 0, [[BB:%.*]] ]
-; 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:    [[EXITCOND:%.*]] = icmp ne i32 [[CANONICAL_IV]], 65635
 ; WIDENING_OFF-NEXT:    br i1 [[EXITCOND]], label [[CHECKED:%.*]], label [[FAILED:%.*]]
 ; WIDENING_OFF:       checked:
 ; 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:    [[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:    [[CANONICAL_IV_NEXT]] = add nuw nsw i32 [[CANONICAL_IV]], 1
 ; WIDENING_OFF-NEXT:    br label [[LOOP]]
 ; WIDENING_OFF:       exit:
 ; WIDENING_OFF-NEXT:    ret i32 0