[LoopUnroll] Use tripcount from exiting header, if latch not exiting.

[fhahn]

Summary

Currently we use the trip count from the exiting latch or from the
unique exit block. If the loop is not rotated, the latch won't exit but
the header will.

This patch falls back to using the trip count from the header, if the
header is exiting and the latch is not.

This fixes cases where we failed to unroll with -Oz, like
https://godbolt.org/z/fP6sna8qK

bool foo(int *ptr, int limit) {
    #pragma clang loop unroll(full)
    for (unsigned int i = 0; i < 4; i++) {
        if (ptr[i] > limit)
        return false;
        ptr[i]++;
    }
    return true;
}

[Meinersbur]

LGTM

However, I wonder why it is not using the overloads without BB arguments. At the end, unrolling would be only interested in some max trip count imposed by any exit, not of a specific one.

[nikic]

In D102635#2763879, @Meinersbur wrote:

However, I wonder why it is not using the overloads without BB arguments. At the end, unrolling would be only interested in some max trip count imposed by any exit, not of a specific one.

Full unrolling can operate in a number of different modes. Max tripcount is one of them, but is actually disabled on most targets (maybe it shouldn't be...) The one this targets is an exact tripcount on one exit, which is generally enabled. The exact tripcount case will always eliminate branches for one exit (though there may be others), while max tripcount will effectively retain all branches.

Unfortunately the loop unrolling implementation currently has some implicit contracts about what kinds of loops it can unroll -- e.g. it would be great if this worked even if the latch were exiting (but had no exact exit count itself), but that does wouldn't work out of the box.

[reames]

This is semantically incorrect, and can not land. There is an important distinction between a maximum trip count (which this code would be correct for), and an exact trip count (which is what this code appears to need). Consider the case where the header exits after 20 iterations, and another non-latch exit exits after 10. Fully unrolling by 20 would be incorrect. (Since control flow is elided and we'd execute too many iterations.)

If you want to achieve the same effect correctly, you should update ScalarEvolution::getSmallConstantTripCount(L) to do what the block variant does, but with the getExitCount(L) version. The change is straight forward, but you will need to audit callers to ensure they're not relying on the single exit block post-condition.

[Meinersbur]

In D102635#2764456, @reames wrote:

This is semantically incorrect, and can not land. There is an important distinction between a maximum trip count (which this code would be correct for), and an exact trip count (which is what this code appears to need). Consider the case where the header exits after 20 iterations, and another non-latch exit exits after 10. Fully unrolling by 20 would be incorrect. (Since control flow is elided and we'd execute too many iterations.)

Does this mean only the branch in the latch is elided by (full) unrolling? The situation could also be that the latch exits after 20 iterations. I was not expecting that there is something special about the latch when exiting.

[reames]

In D102635#2764588, @Meinersbur wrote:

In D102635#2764456, @reames wrote:

This is semantically incorrect, and can not land. There is an important distinction between a maximum trip count (which this code would be correct for), and an exact trip count (which is what this code appears to need). Consider the case where the header exits after 20 iterations, and another non-latch exit exits after 10. Fully unrolling by 20 would be incorrect. (Since control flow is elided and we'd execute too many iterations.)

Does this mean only the branch in the latch is elided by (full) unrolling? The situation could also be that the latch exits after 20 iterations. I was not expecting that there is something special about the latch when exiting.

When we full unroll, we don't duplicate exit blocks. (I haven't checked how we handle single vs multiple exits, but this patch presupposes we handle multiple.) If we unroll by 20 *without* duplicating the earlier exit control flow, we have a miscompile.

[fhahn]

In D102635#2764588, @Meinersbur wrote:

In D102635#2764456, @reames wrote:

This is semantically incorrect, and can not land. There is an important distinction between a maximum trip count (which this code would be correct for), and an exact trip count (which is what this code appears to need). Consider the case where the header exits after 20 iterations, and another non-latch exit exits after 10. Fully unrolling by 20 would be incorrect. (Since control flow is elided and we'd execute too many iterations.)

Does this mean only the branch in the latch is elided by (full) unrolling? The situation could also be that the latch exits after 20 iterations. I was not expecting that there is something special about the latch when exiting.

Thanks for taking a look! LoopUnroll expects TripCount to be an upper bound on the number of times the terminator we simplify during unrolling executes. At the moment, that's either the latch block terminator, if it exits or the unique exiting block's termaintor. Otherwise the exiting branches should be retained. I've also added another test that has an early exit depending on the IV as well. It gets simplified, but only due to simplifications after unrolling.

If there's a single exit block or the latch is exiting, we use TripCount to decide whether the loop is fully unrolled. If that's the case, we simplify the branches in the latch/unique exiting block. If there the latch is not exiting and there's no unique exiting block, we won't simplify the exiting branches to unconditional branches (the code to do so is around https://github.com/llvm/llvm-project/blob/main/llvm/lib/Transforms/Utils/LoopUnroll.cpp#L787, but ExitingBlocks will be empty, if there's no single exiting branch instruction).

I also had another look at some of the comments and it seems like the one for llvm::UnrollLoop was not updated to reflect the fact that loops with multiple exiting blocks are also supported. I tried to update the comment to reflect the current requirement. I might have missed something (there' been a few other changes since I added support for unrolling header-exiting loops and quite a few cases are handled in the code), but I think it should more accurately reflect what's going on. I think it might make sense to split off the comment change, but I included it here to start with as it is directly related to the discussion here.

[reames]

@fhahn I'm willing to defer to you here. If you're really sure the TripCount is always interpreted as an upper bound, you can land and move forward. From your description, it really sounds like this code is duplicating a lot of logic already existing in SCEV. As a cleanup, it would be nice if we named the variables appropriately and used scev's existing distinction between exact and max trip counts.

[mkazantsev]

Just curious: how come that the loop wasn't rotated properly before unrolling? Looks like it's a pass ordering problem.

[fhahn]

In D102635#2764658, @reames wrote:

@fhahn I'm willing to defer to you here. If you're really sure the TripCount is always interpreted as an upper bound, you can land and move forward. From your description, it really sounds like this code is duplicating a lot of logic already existing in SCEV. As a cleanup, it would be nice if we named the variables appropriately and used scev's existing distinction between exact and max trip counts.

After taking a closer look, I think we might be able to indeed re-use parts of SCEV in some parts. The variable names & co could also be clarified & improved. I'll look into some follow-up patches.

In D102635#2765051, @mkazantsev wrote:

Just curious: how come that the loop wasn't rotated properly before unrolling? Looks like it's a pass ordering problem.

The case I am looking at uses -Oz, where loop-rotate is very conservative when it comes to duplicating code and does not rotate the loop (e.g. https://godbolt.org/z/fP6sna8qK). With all other optimization levels, the loop gets rotated.

[mkazantsev]

Looks ok by me, but please make sure you address Philip's concerns.

[fhahn]

More general refactoring/improvements mean that this patch is no longer needed. Thanks!

llvm/lib/Transforms/Scalar/LoopUnrollPass.cpp

static LoopUnrollResult tryToUnrollLoop(
   unsigned TripCount = 0;
   unsigned TripMultiple = 1;
   // If there are multiple exiting blocks but one of them is the latch, use the
   // latch for the trip count estimation. Otherwise insist on a single exiting
   // block for the trip count estimation.
   BasicBlock *ExitingBlock = L->getLoopLatch();
   if (!ExitingBlock || !L->isLoopExiting(ExitingBlock))
     ExitingBlock = L->getExitingBlock();
+  // If the latch is not exiting and there's no single exiting block, check if
+  // the header is exiting and use it to estimate the trip count. This can
+  // happen when unrolling non-rotated loops.
+  if (!ExitingBlock && L->isLoopExiting(L->getHeader()))
+    ExitingBlock = L->getHeader();
   if (ExitingBlock) {
     TripCount = SE.getSmallConstantTripCount(L, ExitingBlock);
     TripMultiple = SE.getSmallConstantTripMultiple(L, ExitingBlock);
   }
 
   // If the loop contains a convergent operation, the prelude we'd add
   // to do the first few instructions before we hit the unrolled loop
   // is unsafe -- it adds a control-flow dependency to the convergent

llvm/lib/Transforms/Utils/LoopUnroll.cpp

void llvm::simplifyLoopAfterUnroll(Loop *L, bool SimplifyIVs, LoopInfo *LI,
   // appropriate.
 }
 
 /// Unroll the given loop by Count. The loop must be in LCSSA form.  Unrolling
 /// can only fail when the loop's latch block is not terminated by a conditional
 /// branch instruction. However, if the trip count (and multiple) are not known,
 /// loop unrolling will mostly produce more code that is no faster.
 ///
-/// TripCount is the upper bound of the iteration on which control exits
-/// LatchBlock. Control may exit the loop prior to TripCount iterations either
-/// via an early branch in other loop block or via LatchBlock terminator. This
-/// is relaxed from the general definition of trip count which is the number of
-/// times the loop header executes. Note that UnrollLoop assumes that the loop
-/// counter test is in LatchBlock in order to remove unnecesssary instances of
-/// the test.  If control can exit the loop from the LatchBlock's terminator
-/// prior to TripCount iterations, flag PreserveCondBr needs to be set.
+/// TripCount is the upper bound of the iteration on which control exits via
+/// latch block (if it exits), the header (if it exits) or the unique exiting
+/// block of the loop, if there is one. Control may exit the loop prior to
+/// TripCount iterations via any of its exits. This is relaxed from the general
+/// definition of trip count which is the number of times the loop header
+/// executes. Note that UnrollLoop assumes that the loop counter test is in
+/// LatchBlock or in the unique exiting block in order to remove unnecesssary
+/// instances of the test. Otherwise the conditional branches are preserved,
+/// pending simplification after unrolling. If control can exit the loop from
+/// the LatchBlock's terminator prior to TripCount iterations, flag
+/// PreserveCondBr needs to be set.
 ///
 /// PreserveCondBr indicates whether the conditional branch of the LatchBlock
 /// needs to be preserved.  It is needed when we use trip count upper bound to
 /// fully unroll the loop. If PreserveOnlyFirst is also set then only the first
 /// conditional branch needs to be preserved.
 ///
 /// Similarly, TripMultiple divides the number of times that the LatchBlock may
 /// execute without exiting the loop.

llvm/test/Transforms/LoopUnroll/unroll-header-exiting-with-phis-multiple-exiting-blocks.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt -loop-unroll -S %s | FileCheck %s
 
 ; Loop with multiple exiting blocks, where the header exits but not the latch,
 ; e.g. because it has not been rotated.
 define i16 @full_unroll_multiple_exiting_blocks(i16* %A, i16 %x, i16 %y) {
 ; CHECK-LABEL: @full_unroll_multiple_exiting_blocks(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    br label [[HEADER:%.*]]
 ; CHECK:       header:
-; CHECK-NEXT:    [[RES:%.*]] = phi i16 [ 123, [[ENTRY:%.*]] ], [ [[RES_NEXT:%.*]], [[LATCH:%.*]] ]
-; CHECK-NEXT:    [[I_0:%.*]] = phi i64 [ 0, [[ENTRY]] ], [ [[INC9:%.*]], [[LATCH]] ]
-; CHECK-NEXT:    [[PTR:%.*]] = getelementptr inbounds i16, i16* [[A:%.*]], i64 [[I_0]]
-; CHECK-NEXT:    [[LV:%.*]] = load i16, i16* [[PTR]], align 2
-; CHECK-NEXT:    [[RES_NEXT]] = add i16 [[RES]], [[LV]]
-; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i64 [[I_0]], 3
-; CHECK-NEXT:    br i1 [[CMP]], label [[EXITING_1:%.*]], label [[EXIT:%.*]]
+; CHECK-NEXT:    [[LV:%.*]] = load i16, i16* [[A:%.*]], align 2
+; CHECK-NEXT:    [[RES_NEXT:%.*]] = add i16 123, [[LV]]
+; CHECK-NEXT:    br i1 true, label [[EXITING_1:%.*]], label [[EXIT:%.*]]
 ; CHECK:       exiting.1:
 ; CHECK-NEXT:    [[EC_1:%.*]] = icmp eq i16 [[LV]], [[X:%.*]]
 ; CHECK-NEXT:    br i1 [[EC_1]], label [[EXIT]], label [[EXITING_2:%.*]]
 ; CHECK:       exiting.2:
 ; CHECK-NEXT:    [[EC_2:%.*]] = icmp eq i16 [[LV]], [[Y:%.*]]
-; CHECK-NEXT:    br i1 [[EC_2]], label [[EXIT]], label [[LATCH]]
+; CHECK-NEXT:    br i1 [[EC_2]], label [[EXIT]], label [[LATCH:%.*]]
 ; CHECK:       latch:
-; CHECK-NEXT:    [[INC9]] = add i64 [[I_0]], 1
-; CHECK-NEXT:    br label [[HEADER]]
+; CHECK-NEXT:    [[PTR_1:%.*]] = getelementptr inbounds i16, i16* [[A]], i64 1
+; CHECK-NEXT:    [[LV_1:%.*]] = load i16, i16* [[PTR_1]], align 2
+; CHECK-NEXT:    [[RES_NEXT_1:%.*]] = add i16 [[RES_NEXT]], [[LV_1]]
+; CHECK-NEXT:    br i1 true, label [[EXITING_1_1:%.*]], label [[EXIT]]
 ; CHECK:       exit:
-; CHECK-NEXT:    [[RES_LCSSA:%.*]] = phi i16 [ [[RES_NEXT]], [[HEADER]] ], [ 0, [[EXITING_1]] ], [ 1, [[EXITING_2]] ]
+; CHECK-NEXT:    [[RES_LCSSA:%.*]] = phi i16 [ [[RES_NEXT]], [[HEADER]] ], [ 0, [[EXITING_1]] ], [ 1, [[EXITING_2]] ], [ [[RES_NEXT_1]], [[LATCH]] ], [ 0, [[EXITING_1_1]] ], [ 1, [[EXITING_2_1:%.*]] ], [ [[RES_NEXT_2:%.*]], [[LATCH_1:%.*]] ], [ 0, [[EXITING_1_2:%.*]] ], [ 1, [[EXITING_2_2:%.*]] ], [ [[RES_NEXT_3:%.*]], [[LATCH_2:%.*]] ], [ 0, [[EXITING_1_3:%.*]] ], [ 1, [[EXITING_2_3:%.*]] ]
 ; CHECK-NEXT:    ret i16 [[RES_LCSSA]]
+; CHECK:       exiting.1.1:
+; CHECK-NEXT:    [[EC_1_1:%.*]] = icmp eq i16 [[LV_1]], [[X]]
+; CHECK-NEXT:    br i1 [[EC_1_1]], label [[EXIT]], label [[EXITING_2_1]]
+; CHECK:       exiting.2.1:
+; CHECK-NEXT:    [[EC_2_1:%.*]] = icmp eq i16 [[LV_1]], [[Y]]
+; CHECK-NEXT:    br i1 [[EC_2_1]], label [[EXIT]], label [[LATCH_1]]
+; CHECK:       latch.1:
+; CHECK-NEXT:    [[PTR_2:%.*]] = getelementptr inbounds i16, i16* [[A]], i64 2
+; CHECK-NEXT:    [[LV_2:%.*]] = load i16, i16* [[PTR_2]], align 2
+; CHECK-NEXT:    [[RES_NEXT_2]] = add i16 [[RES_NEXT_1]], [[LV_2]]
+; CHECK-NEXT:    br i1 true, label [[EXITING_1_2]], label [[EXIT]]
+; CHECK:       exiting.1.2:
+; CHECK-NEXT:    [[EC_1_2:%.*]] = icmp eq i16 [[LV_2]], [[X]]
+; CHECK-NEXT:    br i1 [[EC_1_2]], label [[EXIT]], label [[EXITING_2_2]]
+; CHECK:       exiting.2.2:
+; CHECK-NEXT:    [[EC_2_2:%.*]] = icmp eq i16 [[LV_2]], [[Y]]
+; CHECK-NEXT:    br i1 [[EC_2_2]], label [[EXIT]], label [[LATCH_2]]
+; CHECK:       latch.2:
+; CHECK-NEXT:    [[PTR_3:%.*]] = getelementptr inbounds i16, i16* [[A]], i64 3
+; CHECK-NEXT:    [[LV_3:%.*]] = load i16, i16* [[PTR_3]], align 2
+; CHECK-NEXT:    [[RES_NEXT_3]] = add i16 [[RES_NEXT_2]], [[LV_3]]
+; CHECK-NEXT:    br i1 false, label [[EXITING_1_3]], label [[EXIT]]
+; CHECK:       exiting.1.3:
+; CHECK-NEXT:    [[EC_1_3:%.*]] = icmp eq i16 [[LV_3]], [[X]]
+; CHECK-NEXT:    br i1 [[EC_1_3]], label [[EXIT]], label [[EXITING_2_3]]
+; CHECK:       exiting.2.3:
+; CHECK-NEXT:    [[EC_2_3:%.*]] = icmp eq i16 [[LV_3]], [[Y]]
+; CHECK-NEXT:    br i1 [[EC_2_3]], label [[EXIT]], label [[LATCH_3:%.*]]
+; CHECK:       latch.3:
+; CHECK-NEXT:    unreachable
 ;
 entry:
   br label %header
 
 header:
   %res = phi i16 [ 123, %entry ], [ %res.next, %latch ]
   %i.0 = phi i64 [ 0, %entry ], [ %inc9, %latch ]
   %ptr = getelementptr inbounds i16, i16* %A, i64 %i.0

 }
 
 
 define i16 @full_unroll_multiple_exiting_blocks_2(i16* %A, i16 %x, i16 %y) {
 ; CHECK-LABEL: @full_unroll_multiple_exiting_blocks_2(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    br label [[HEADER:%.*]]
 ; CHECK:       header:
-; CHECK-NEXT:    [[RES:%.*]] = phi i16 [ 123, [[ENTRY:%.*]] ], [ [[RES_NEXT:%.*]], [[LATCH:%.*]] ]
-; CHECK-NEXT:    [[I_0:%.*]] = phi i64 [ 0, [[ENTRY]] ], [ [[INC9:%.*]], [[LATCH]] ]
-; CHECK-NEXT:    [[PTR:%.*]] = getelementptr inbounds i16, i16* [[A:%.*]], i64 [[I_0]]
-; CHECK-NEXT:    [[LV:%.*]] = load i16, i16* [[PTR]], align 2
-; CHECK-NEXT:    [[RES_NEXT]] = add i16 [[RES]], [[LV]]
-; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i64 [[I_0]], 3
-; CHECK-NEXT:    br i1 [[CMP]], label [[EXITING_1:%.*]], label [[EXIT:%.*]]
+; CHECK-NEXT:    [[LV:%.*]] = load i16, i16* [[A:%.*]], align 2
+; CHECK-NEXT:    [[RES_NEXT:%.*]] = add i16 123, [[LV]]
+; CHECK-NEXT:    br i1 true, label [[EXITING_1:%.*]], label [[EXIT:%.*]]
 ; CHECK:       exiting.1:
-; CHECK-NEXT:    [[EC_1:%.*]] = icmp eq i64 [[I_0]], 1
-; CHECK-NEXT:    br i1 [[EC_1]], label [[EXIT]], label [[EXITING_2:%.*]]
+; CHECK-NEXT:    br i1 false, label [[EXIT]], label [[EXITING_2:%.*]]
 ; CHECK:       exiting.2:
 ; CHECK-NEXT:    [[EC_2:%.*]] = icmp eq i16 [[LV]], [[Y:%.*]]
-; CHECK-NEXT:    br i1 [[EC_2]], label [[EXIT]], label [[LATCH]]
+; CHECK-NEXT:    br i1 [[EC_2]], label [[EXIT]], label [[LATCH:%.*]]
 ; CHECK:       latch:
-; CHECK-NEXT:    [[INC9]] = add i64 [[I_0]], 1
-; CHECK-NEXT:    br label [[HEADER]]
+; CHECK-NEXT:    [[PTR_1:%.*]] = getelementptr inbounds i16, i16* [[A]], i64 1
+; CHECK-NEXT:    [[LV_1:%.*]] = load i16, i16* [[PTR_1]], align 2
+; CHECK-NEXT:    [[RES_NEXT_1:%.*]] = add i16 [[RES_NEXT]], [[LV_1]]
+; CHECK-NEXT:    br i1 true, label [[EXITING_1_1:%.*]], label [[EXIT]]
 ; CHECK:       exit:
-; CHECK-NEXT:    [[RES_LCSSA:%.*]] = phi i16 [ [[RES_NEXT]], [[HEADER]] ], [ 0, [[EXITING_1]] ], [ 1, [[EXITING_2]] ]
+; CHECK-NEXT:    [[RES_LCSSA:%.*]] = phi i16 [ [[RES_NEXT]], [[HEADER]] ], [ 0, [[EXITING_1]] ], [ 1, [[EXITING_2]] ], [ [[RES_NEXT_1]], [[LATCH]] ], [ 0, [[EXITING_1_1]] ], [ 1, [[EXITING_2_1:%.*]] ], [ [[RES_NEXT_2:%.*]], [[LATCH_1:%.*]] ], [ 0, [[EXITING_1_2:%.*]] ], [ 1, [[EXITING_2_2:%.*]] ], [ [[RES_NEXT_3:%.*]], [[LATCH_2:%.*]] ], [ 0, [[EXITING_1_3:%.*]] ], [ 1, [[EXITING_2_3:%.*]] ]
 ; CHECK-NEXT:    ret i16 [[RES_LCSSA]]
+; CHECK:       exiting.1.1:
+; CHECK-NEXT:    br i1 true, label [[EXIT]], label [[EXITING_2_1]]
+; CHECK:       exiting.2.1:
+; CHECK-NEXT:    [[EC_2_1:%.*]] = icmp eq i16 [[LV_1]], [[Y]]
+; CHECK-NEXT:    br i1 [[EC_2_1]], label [[EXIT]], label [[LATCH_1]]
+; CHECK:       latch.1:
+; CHECK-NEXT:    [[PTR_2:%.*]] = getelementptr inbounds i16, i16* [[A]], i64 2
+; CHECK-NEXT:    [[LV_2:%.*]] = load i16, i16* [[PTR_2]], align 2
+; CHECK-NEXT:    [[RES_NEXT_2]] = add i16 [[RES_NEXT_1]], [[LV_2]]
+; CHECK-NEXT:    br i1 true, label [[EXITING_1_2]], label [[EXIT]]
+; CHECK:       exiting.1.2:
+; CHECK-NEXT:    br i1 false, label [[EXIT]], label [[EXITING_2_2]]
+; CHECK:       exiting.2.2:
+; CHECK-NEXT:    [[EC_2_2:%.*]] = icmp eq i16 [[LV_2]], [[Y]]
+; CHECK-NEXT:    br i1 [[EC_2_2]], label [[EXIT]], label [[LATCH_2]]
+; CHECK:       latch.2:
+; CHECK-NEXT:    [[PTR_3:%.*]] = getelementptr inbounds i16, i16* [[A]], i64 3
+; CHECK-NEXT:    [[LV_3:%.*]] = load i16, i16* [[PTR_3]], align 2
+; CHECK-NEXT:    [[RES_NEXT_3]] = add i16 [[RES_NEXT_2]], [[LV_3]]
+; CHECK-NEXT:    br i1 false, label [[EXITING_1_3]], label [[EXIT]]
+; CHECK:       exiting.1.3:
+; CHECK-NEXT:    br i1 false, label [[EXIT]], label [[EXITING_2_3]]
+; CHECK:       exiting.2.3:
+; CHECK-NEXT:    [[EC_2_3:%.*]] = icmp eq i16 [[LV_3]], [[Y]]
+; CHECK-NEXT:    br i1 [[EC_2_3]], label [[EXIT]], label [[LATCH_3:%.*]]
+; CHECK:       latch.3:
+; CHECK-NEXT:    unreachable
 ;
 entry:
   br label %header
 
 header:
   %res = phi i16 [ 123, %entry ], [ %res.next, %latch ]
   %i.0 = phi i64 [ 0, %entry ], [ %inc9, %latch ]
   %ptr = getelementptr inbounds i16, i16* %A, i64 %i.0