[IndVars] Eliminate loop exits with equivalent exit counts

[reames]

Summary

We can end up with two loop exits whose exit counts are equivalent, but whose textual representation is different and non-obvious. For the sub-case where we have a series of exits which dominate one another (common), eliminate any exits which would iterate *after* a previous exit on the exiting iteration.

As noted in the TODO being removed, I'd always thought this was a good idea, but I've now seen this in a real workload as well. This needs to be rebased on D68956 as both need the same dominance order check.

[ebrevnov]

lib/Transforms/Scalar/IndVarSimplify.cpp

2758 ↗

(On Diff #225121)

That would be a fourth loop filtering out "bad" exits before actual work. And the method becomes pretty long. I think we need to refactor in this or consequent patch.

2793 ↗

(On Diff #225121)

Two notes here. 1) Any reason not to use isKnownViaNonRecursiveReasoning or isKnownPredicate? 2) We can actually eliminate an exit if its exit count not less than exit count of one of its predecessors.

None of these is blocking the current patch though.

[reames]

Rebase over landed loop pred change, and some refactoring to address a few concerns brought up in the review.

[reames]

lib/Transforms/Scalar/IndVarSimplify.cpp

2758 ↗

(On Diff #225121)

What do you think of the committed NFC reorgs? If you have suggestions, I'm happy to further tweak.

2793 ↗

(On Diff #225121)

2. I believe you're referring not to the EQ case, but to the provably greater then case just above right?

1. I could generalize, but I don't have an example which illustrates the difference. SCEV tries to canonicalize at construction (as demonstrated by some of the test changes). If you don't mind, I'd like to wait until we have a motivating example before using a more expensive API.

[nikic]

lib/Transforms/Scalar/IndVarSimplify.cpp

2804 ↗

(On Diff #225465)

This code is already repeated three times now, extract?

[reames]

lib/Transforms/Scalar/IndVarSimplify.cpp

2804 ↗

(On Diff #225465)

Not quite. We're folding the branch either to taken or untaken, but yes, it could be factored out. Will do, and rebase.

[reames]

Rebase on requested refactor.

[nikic]

lib/Transforms/Scalar/IndVarSimplify.cpp

2796 ↗

(On Diff #225548)

Now that we're traversing the exits in dominating order, I think we should consider a slightly different overall approach:

Instead of computing umin({ exit counts }) upfront and then checking for umin({ exit counts }) < this exit count, we can instead build this up incrementally, i.e. take the umin between the current max exit count and the exit count of the currently considered exit at each iteration. This will allow us to check for umin({ already seen exit counts }) <= this exit count (note the <=!) instead.

This should implicitly handle the case of exact equalities, but may also be more powerful in general because we're checking for a weaker predicate.

[reames]

lib/Transforms/Scalar/IndVarSimplify.cpp

2796 ↗

(On Diff #225548)

The suggested approach doesn't handle one of the most common cases. If the latch exit is in fact the minimum exit, we want to be able to discharge earlier exits. (i.e. the difference between the dominating set and the total sets is important)

I do agree that your phrasing *might* be more powerful for eliminating exits which SCEV can prove are <= dominating checks (but not either < or ==). However, I don't have a test case which demonstrates that difference, and until I do, I'd like to avoid further complicating the code. (Assuming I'd need both approaches due to the gap described above.)

[nikic]

lib/Transforms/Scalar/IndVarSimplify.cpp

2796 ↗

(On Diff #225548)

For an example, consider the existing @ult test just using a ule predicate instead: https://godbolt.org/z/j-rZeb This should eliminate the second exit, but doesn't, as we have neither the < predicate nor the strict scev identity (but do have a <= predicate).

I didn't quite follow what you mean regarding the latch exit and why it is special. Is there a test case that shows the issue?

[nikic]

LG

lib/Transforms/Scalar/IndVarSimplify.cpp

2796 ↗

(On Diff #225548)

Okay, I see what you mean now. Not specifically the latch, just any future exit. It does seem like we'd have to combine both approaches to get the full coverage, so let's go with your current variant for now.

llvm/lib/Transforms/Scalar/IndVarSimplify.cpp

bool IndVarSimplify::optimizeLoopExits(Loop *L, SCEVExpander &Rewriter) {
   if (ExitingBlocks.empty())
     return false;
   
   // Get a symbolic upper bound on the loop backedge taken count.  
   const SCEV *MaxExitCount = getMaxBackedgeTakenCount(*SE, *DT, L);
   if (isa<SCEVCouldNotCompute>(MaxExitCount))
     return false;
 
+  // Visit our exit blocks in order of dominance.  We know from the fact that
+  // all exits (left) are analyzeable that the must be a total dominance order
+  // between them as each must dominate the latch.  The visit order only
+  // matters for the provably equal case.  
+  llvm::sort(ExitingBlocks,
+             [&](BasicBlock *A, BasicBlock *B) {
+               // std::sort sorts in ascending order, so we want the inverse of
+               // the normal dominance relation.
+               if (DT->properlyDominates(A, B)) return true;
+               if (DT->properlyDominates(B, A)) return false;
+               llvm_unreachable("expected total dominance order!");
+             });
+#ifdef ASSERT
+  for (unsigned i = 1; i < ExitingBlocks.size(); i++) {
+    assert(DT->dominates(ExitingBlocks[i-1], ExitingBlocks[i]));
+  }
+#endif
+  
   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);
     BI->setCondition(NewCond);
     if (OldCond->use_empty())
       DeadInsts.push_back(OldCond);
   };
 
   bool Changed = false;
+  SmallSet<const SCEV*, 8> DominatingExitCounts;
   for (BasicBlock *ExitingBB : ExitingBlocks) {
     const SCEV *ExitCount = SE->getExitCount(L, ExitingBB);
     assert(!isa<SCEVCouldNotCompute>(ExitCount) && "checked above");
     
     // 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
for (BasicBlock *ExitingBB : ExitingBlocks) {
     // one?
     if (SE->isLoopEntryGuardedByCond(L, CmpInst::ICMP_ULT,
                                      MaxExitCount, ExitCount)) {
       FoldExit(ExitingBB, false);
       Changed = true;
       continue;
     }
 
-    // TODO: If we can prove that the exiting iteration is equal to the exit
-    // count for this exit and that no previous exit oppurtunities exist within
-    // the loop, then we can discharge all other exits.  (May fall out of
-    // previous TODO.)
+    // As we run, keep track of which exit counts we've encountered.  If we
+    // find a duplicate, we've found an exit which would have exited on the
+    // exiting iteration, but (from the visit order) strictly follows another
+    // which does the same and is thus dead. 
+    if (!DominatingExitCounts.insert(ExitCount).second) {
+      FoldExit(ExitingBB, false);
+      Changed = true;
+      continue;
+    }
   }
   return Changed;
 }
 
 bool IndVarSimplify::predicateLoopExits(Loop *L, SCEVExpander &Rewriter) {
   SmallVector<BasicBlock*, 16> ExitingBlocks;
   L->getExitingBlocks(ExitingBlocks);
 

llvm/test/Transforms/IndVarSimplify/eliminate-exit.ll

 latch:
   call void @side_effect()
   %cmp2 = icmp ult i64 %iv, %n
   br i1 %cmp2, label %loop, label %exit
 exit:
   ret void
 }
 
+define void @mixed_width(i32 %len) {
+; CHECK-LABEL: @mixed_width(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[LEN_ZEXT:%.*]] = zext i32 [[LEN:%.*]] to i64
+; CHECK-NEXT:    br label [[LOOP:%.*]]
+; CHECK:       loop:
+; CHECK-NEXT:    [[IV:%.*]] = phi i64 [ 0, [[ENTRY:%.*]] ], [ [[IV_NEXT:%.*]], [[BACKEDGE:%.*]] ]
+; CHECK-NEXT:    [[IV_NEXT]] = add nuw nsw i64 [[IV]], 1
+; CHECK-NEXT:    [[CMP1:%.*]] = icmp ult i64 [[IV]], [[LEN_ZEXT]]
+; CHECK-NEXT:    br i1 [[CMP1]], label [[BACKEDGE]], label [[EXIT:%.*]]
+; CHECK:       backedge:
+; CHECK-NEXT:    call void @side_effect()
+; CHECK-NEXT:    br i1 true, label [[LOOP]], label [[EXIT]]
+; CHECK:       exit:
+; CHECK-NEXT:    ret void
+;
+entry:
+  %len.zext = zext i32 %len to i64
+  br label %loop
+loop:
+  %iv = phi i64 [0, %entry], [%iv.next, %backedge]
+  %iv2 = phi i32 [0, %entry], [%iv2.next, %backedge]
+  %iv.next = add i64 %iv, 1
+  %iv2.next = add i32 %iv2, 1
+  %cmp1 = icmp ult i64 %iv, %len.zext
+  br i1 %cmp1, label %backedge, label %exit
+
+backedge:
+  call void @side_effect()
+  %cmp2 = icmp ult i32 %iv2, %len
+  br i1 %cmp2, label %loop, label %exit
+exit:
+  ret void
+}
 
 declare void @side_effect()

llvm/test/Transforms/IndVarSimplify/loop-predication.ll

 ; CHECK-LABEL: @duplicate_checks(
 ; CHECK-NEXT:  loop.preheader:
 ; CHECK-NEXT:    [[TMP0:%.*]] = icmp ugt i32 [[N:%.*]], 1
 ; CHECK-NEXT:    [[UMAX:%.*]] = select i1 [[TMP0]], i32 [[N]], i32 1
 ; CHECK-NEXT:    [[TMP1:%.*]] = add i32 [[UMAX]], -1
 ; CHECK-NEXT:    [[TMP2:%.*]] = icmp ult i32 [[LENGTH:%.*]], [[TMP1]]
 ; CHECK-NEXT:    [[UMIN:%.*]] = select i1 [[TMP2]], i32 [[LENGTH]], i32 [[TMP1]]
 ; CHECK-NEXT:    [[TMP3:%.*]] = icmp ne i32 [[LENGTH]], [[UMIN]]
-; CHECK-NEXT:    [[TMP4:%.*]] = icmp ne i32 [[LENGTH]], [[UMIN]]
 ; CHECK-NEXT:    br label [[LOOP:%.*]]
 ; CHECK:       loop:
 ; CHECK-NEXT:    [[LOOP_ACC:%.*]] = phi i32 [ [[LOOP_ACC_NEXT:%.*]], [[GUARDED1:%.*]] ], [ 0, [[LOOP_PREHEADER:%.*]] ]
 ; CHECK-NEXT:    [[I:%.*]] = phi i32 [ [[I_NEXT:%.*]], [[GUARDED1]] ], [ 0, [[LOOP_PREHEADER]] ]
 ; CHECK-NEXT:    br i1 [[TMP3]], label [[GUARDED:%.*]], label [[DEOPT:%.*]], !prof !0
 ; CHECK:       deopt:
 ; CHECK-NEXT:    call void @prevent_merging()
 ; CHECK-NEXT:    ret i32 -1
 ; CHECK:       guarded:
 ; CHECK-NEXT:    [[I_I64:%.*]] = zext i32 [[I]] to i64
 ; CHECK-NEXT:    [[ARRAY_1_I_PTR:%.*]] = getelementptr inbounds i32, i32* [[ARRAY_1:%.*]], i64 [[I_I64]]
 ; CHECK-NEXT:    [[ARRAY_1_I:%.*]] = load i32, i32* [[ARRAY_1_I_PTR]], align 4
 ; CHECK-NEXT:    [[LOOP_ACC_1:%.*]] = add i32 [[LOOP_ACC]], [[ARRAY_1_I]]
-; CHECK-NEXT:    br i1 [[TMP4]], label [[GUARDED1]], label [[DEOPT2:%.*]], !prof !0
+; CHECK-NEXT:    br i1 true, label [[GUARDED1]], label [[DEOPT2:%.*]], !prof !0
 ; CHECK:       deopt2:
 ; CHECK-NEXT:    call void @prevent_merging()
 ; CHECK-NEXT:    ret i32 -1
 ; CHECK:       guarded1:
 ; CHECK-NEXT:    [[ARRAY_3_I_PTR:%.*]] = getelementptr inbounds i32, i32* [[ARRAY_3:%.*]], i64 [[I_I64]]
 ; CHECK-NEXT:    [[ARRAY_3_I:%.*]] = load i32, i32* [[ARRAY_3_I_PTR]], align 4
 ; CHECK-NEXT:    [[LOOP_ACC_NEXT]] = add i32 [[LOOP_ACC_1]], [[ARRAY_3_I]]
 ; CHECK-NEXT:    [[I_NEXT]] = add nuw i32 [[I]], 1

 exit:
   %result = phi i32 [ %loop.acc.next, %guarded ], [0, %entry]
   ret i32 %result
 }
 
 ; If we have a dominating exit (exit1) which can't be itself rewritten, we
 ; can't rewrite a later exit (exit2).  Doing so would cause the loop to exit
 ; from the exit2 when it should have exited from exit1.
-define i32 @neg_dominating_exit(i32* %array, i32 %length, i32 %n) {
+define i32 @neg_dominating_exit(i32* %array, i32 %length, i32 %length2, i32 %n) {
 ; CHECK-LABEL: @neg_dominating_exit(
 ; CHECK-NEXT:  loop.preheader:
 ; CHECK-NEXT:    br label [[LOOP:%.*]]
 ; CHECK:       loop:
 ; CHECK-NEXT:    [[LOOP_ACC:%.*]] = phi i32 [ [[LOOP_ACC_NEXT:%.*]], [[GUARDED2:%.*]] ], [ 0, [[LOOP_PREHEADER:%.*]] ]
 ; CHECK-NEXT:    [[I:%.*]] = phi i32 [ [[I_NEXT:%.*]], [[GUARDED2]] ], [ 0, [[LOOP_PREHEADER]] ]
 ; CHECK-NEXT:    [[WITHIN_BOUNDS:%.*]] = icmp ult i32 [[I]], [[LENGTH:%.*]]
 ; CHECK-NEXT:    br i1 [[WITHIN_BOUNDS]], label [[GUARDED:%.*]], label [[DEOPT:%.*]], !prof !0
 ; CHECK:       deopt:
 ; CHECK-NEXT:    [[RESULT:%.*]] = phi i32 [ [[LOOP_ACC]], [[LOOP]] ]
 ; CHECK-NEXT:    call void @prevent_merging()
 ; CHECK-NEXT:    ret i32 [[RESULT]]
 ; CHECK:       guarded:
-; CHECK-NEXT:    [[WITHIN_BOUNDS2:%.*]] = icmp ult i32 [[I]], [[LENGTH]]
+; CHECK-NEXT:    [[WITHIN_BOUNDS2:%.*]] = icmp ult i32 [[I]], [[LENGTH2:%.*]]
 ; CHECK-NEXT:    br i1 [[WITHIN_BOUNDS2]], label [[GUARDED2]], label [[DEOPT2:%.*]], !prof !0
 ; CHECK:       deopt2:
 ; CHECK-NEXT:    call void @prevent_merging()
 ; CHECK-NEXT:    ret i32 -1
 ; CHECK:       guarded2:
 ; CHECK-NEXT:    [[I_I64:%.*]] = zext i32 [[I]] to i64
 ; CHECK-NEXT:    [[ARRAY_I_PTR:%.*]] = getelementptr inbounds i32, i32* [[ARRAY:%.*]], i64 [[I_I64]]
 ; CHECK-NEXT:    [[ARRAY_I:%.*]] = load i32, i32* [[ARRAY_I_PTR]], align 4
loop:                                             ; preds = %guarded, %loop.preheader
   br i1 %within.bounds, label %guarded, label %deopt, !prof !0
 
 deopt:                                            ; preds = %loop
   %result = phi i32 [ %loop.acc, %loop ]
   call void @prevent_merging()
   ret i32 %result
 
 guarded:                                          ; preds = %loop
-  %within.bounds2 = icmp ult i32 %i, %length
+  %within.bounds2 = icmp ult i32 %i, %length2
   br i1 %within.bounds2, label %guarded2, label %deopt2, !prof !0
 
 deopt2:                                            ; preds = %loop
   call void @prevent_merging()
   ret i32 -1
 
 guarded2:                                          ; preds = %loop
   %i.i64 = zext i32 %i to i64

llvm/test/Transforms/IndVarSimplify/pr38674.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt -S -indvars < %s | FileCheck %s
 
 ; Check that we don't reuse %zext instead of %inc11 for LCSSA Phi node. Case
 ; with constants SCEV.
 
 define i32 @test_01() {
 ; CHECK-LABEL: @test_01(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    br label [[FOR_COND1_PREHEADER:%.*]]
 ; CHECK:       for.cond1.preheader:
 ; CHECK-NEXT:    br label [[FOR_COND4_PREHEADER:%.*]]
 ; CHECK:       for.cond4.preheader:
 ; CHECK-NEXT:    [[ZEXT:%.*]] = zext i16 1 to i32
 ; CHECK-NEXT:    br label [[FOR_BODY6:%.*]]
 ; CHECK:       for.cond4:
-; CHECK-NEXT:    [[CMP5:%.*]] = icmp ult i32 [[INC:%.*]], 2
-; CHECK-NEXT:    br i1 [[CMP5]], label [[FOR_BODY6]], label [[FOR_END:%.*]]
+; CHECK-NEXT:    br i1 true, label [[FOR_BODY6]], label [[FOR_END:%.*]]
 ; CHECK:       for.body6:
-; CHECK-NEXT:    [[IV:%.*]] = phi i32 [ 0, [[FOR_COND4_PREHEADER]] ], [ [[INC]], [[FOR_COND4:%.*]] ]
+; CHECK-NEXT:    [[IV:%.*]] = phi i32 [ 0, [[FOR_COND4_PREHEADER]] ], [ [[INC:%.*]], [[FOR_COND4:%.*]] ]
 ; CHECK-NEXT:    [[TMP0:%.*]] = icmp eq i32 [[IV]], [[ZEXT]]
 ; CHECK-NEXT:    [[INC]] = add nuw nsw i32 [[IV]], 1
 ; CHECK-NEXT:    br i1 [[TMP0]], label [[RETURN_LOOPEXIT:%.*]], label [[FOR_COND4]]
 ; CHECK:       for.end:
 ; CHECK-NEXT:    br i1 false, label [[FOR_COND4_PREHEADER]], label [[FOR_END9:%.*]]
 ; CHECK:       for.end9:
 ; CHECK-NEXT:    br i1 false, label [[FOR_COND1_PREHEADER]], label [[RETURN_LOOPEXIT3:%.*]]
 ; CHECK:       return.loopexit: