[IndVars] Replace checks with invariants if we cannot remove them

[mkazantsev]

Summary

If we cannot prove that the check is trivially true, but can prove that it either
fails on the 1st iteration or never fails, we can replace it with first iteration check.

[mkazantsev]

Rebased, ready for review.

[anna]

Just a stylistic comment. You can have the InvariantCond as a struct with the fields: InvariantPred, InvariantLHS, InvariantRHS. I think this will make code more readable (for example, use the struct in ReplaceInvariantCond(ExitingBB, InvariantCond).

llvm/lib/Transforms/Scalar/IndVarSimplify.cpp

1347

I maybe missing something. Where are we checking the "never fails in any iteration other than the first" ?

1425

Nit: ReplaceWithInvariantCond.

[skatkov]

llvm/lib/Transforms/Scalar/IndVarSimplify.cpp

1347

SE->isLoopInvariantExitCondDuringFirstIterations ?

[mkazantsev]

llvm/lib/Transforms/Scalar/IndVarSimplify.cpp

1347

Yes, inside of isLoopInvariantExitCondDuringFirstIterations we effectively check that comparison against IV may be replaced with comparison against start.

[skatkov]

lgtm with Anna's nit + please do a rebase.

[mkazantsev]

(InvariantPred, InvariantLHS, InvariantRHS) API is used in SCEV. I will follow-up with refactoring changing it whevever it is used.

llvm/lib/Transforms/Scalar/IndVarSimplify.cpp

while (I != Preheader->begin()) {
     InsertPt = ToMove->getIterator();
   }
 
   return MadeAnyChanges;
 }
 
 enum ExitCondAnalysisResult {
   CanBeRemoved,
+  CanBeReplacedWithInvariant,
   CannotOptimize
 };
 
 /// If the condition of BI is trivially true during at least first MaxIter
 /// iterations, return CanBeRemoved.
+/// If the condition is equivalent to loop-invariant condition expressed as
+/// 'InvariantLHS `InvariantPred` InvariantRHS', fill them into respective
+/// output parameters and return CanBeReplacedWithInvariant.
 /// Otherwise, return CannotOptimize.
-static ExitCondAnalysisResult analyzeCond(const Loop *L, BranchInst *BI,
-                                          ScalarEvolution *SE,
-                                          bool ProvingLoopExit,
-                                          const SCEV *MaxIter) {
+static ExitCondAnalysisResult
+analyzeCond(const Loop *L, BranchInst *BI, ScalarEvolution *SE,
+            bool ProvingLoopExit, const SCEV *MaxIter,
+            ICmpInst::Predicate &InvariantPred, const SCEV *&InvariantLHS,
+            const SCEV *&InvariantRHS) {
   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 CannotOptimize;
 
analyzeCond(const Loop *L, BranchInst *BI, ScalarEvolution *SE,
   const SCEV *RHSS = SE->getSCEVAtScope(RHS, L);
   // Can we prove it to be trivially true?
   if (SE->isKnownPredicateAt(Pred, LHSS, RHSS, BI))
     return CanBeRemoved;
 
   if (ProvingLoopExit)
     return CannotOptimize;
 
-  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 CannotOptimize;
 
   // Can we prove it to be trivially true?
   if (SE->isKnownPredicateAt(InvariantPred, InvariantLHS, InvariantRHS, BI))
     return CanBeRemoved;
-  return CannotOptimize;
+  return CanBeReplacedWithInvariant;
 }
 
 bool IndVarSimplify::optimizeLoopExits(Loop *L, SCEVExpander &Rewriter) {
   SmallVector<BasicBlock*, 16> ExitingBlocks;
   L->getExitingBlocks(ExitingBlocks);
 
   // Remove all exits which aren't both rewriteable and execute on every
   // iteration.
auto FoldExit = [&](BasicBlock *ExitingBB, bool IsTaken) {
     BranchInst *BI = cast<BranchInst>(ExitingBB->getTerminator());
     bool ExitIfTrue = !L->contains(*succ_begin(ExitingBB));
     auto *OldCond = BI->getCondition();
     auto *NewCond = ConstantInt::get(OldCond->getType(),
                                      IsTaken ? ExitIfTrue : !ExitIfTrue);
     ReplaceExitCond(BI, NewCond);
   };
 
+  auto ReplaceWithInvariantCond = [&](
+      BasicBlock *ExitingBB, ICmpInst::Predicate InvariantPred,
+      const SCEV *InvariantLHS, const SCEV *InvariantRHS) {
+    BranchInst *BI = cast<BranchInst>(ExitingBB->getTerminator());
+    Rewriter.setInsertPoint(BI);
+    auto *LHSV = Rewriter.expandCodeFor(InvariantLHS);
+    auto *RHSV = Rewriter.expandCodeFor(InvariantRHS);
+    IRBuilder<> Builder(BI);
+    auto *NewCond = Builder.CreateICmp(InvariantPred, LHSV, RHSV,
+                                       BI->getCondition()->getName());
+    ReplaceExitCond(BI, NewCond);
+  };
+
   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 = [this, L, BI, ExitingBB, MaxExitCount, &FoldExit](
-          bool Inverted, bool SkipLastIter) {
+      auto OptimizeCond = [this, L, BI, ExitingBB, MaxExitCount, &FoldExit,
+                           &ReplaceWithInvariantCond](bool Inverted,
+                                                      bool SkipLastIter) {
         const SCEV *MaxIter = MaxExitCount;
         if (SkipLastIter) {
           const SCEV *One = SE->getOne(MaxIter->getType());
           MaxIter = SE->getMinusSCEV(MaxIter, One);
         }
-        switch (analyzeCond(L, BI, SE, Inverted, MaxIter)) {
+        ICmpInst::Predicate InvariantPred;
+        const SCEV *InvariantLHS, *InvariantRHS;
+        switch (analyzeCond(L, BI, SE, Inverted, MaxIter, InvariantPred,
+                            InvariantLHS, InvariantRHS)) {
         case CanBeRemoved:
           FoldExit(ExitingBB, Inverted);
           return true;
+        case CanBeReplacedWithInvariant: {
+          ReplaceWithInvariantCond(ExitingBB, InvariantPred, InvariantLHS,
+                                   InvariantRHS);
+          return true;
+        }
         case CannotOptimize:
           return false;
         }
         llvm_unreachable("Unknown case!");
       };
 
       // 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

llvm/test/Transforms/IndVarSimplify/predicated_ranges.ll

 exit:
   ret void
 }
 
 define void @test_can_predicate_simple_unsigned(i32* %p, i32* %arr) {
 ; CHECK-LABEL: @test_can_predicate_simple_unsigned(
 ; CHECK-NEXT:  preheader:
 ; CHECK-NEXT:    [[LEN:%.*]] = load i32, i32* [[P:%.*]], align 4
+; CHECK-NEXT:    [[TMP0:%.*]] = add i32 [[LEN]], -1
 ; 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:    [[RANGE_CHECK1:%.*]] = icmp ult i32 [[TMP0]], [[LEN]]
+; CHECK-NEXT:    br i1 [[RANGE_CHECK1]], 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: