Index: llvm/include/llvm/Transforms/Utils/UnrollLoop.h =================================================================== --- llvm/include/llvm/Transforms/Utils/UnrollLoop.h +++ llvm/include/llvm/Transforms/Utils/UnrollLoop.h @@ -79,7 +79,8 @@ bool ForgetAllSCEV; }; -LoopUnrollResult UnrollLoop(Loop *L, UnrollLoopOptions ULO, LoopInfo *LI, +LoopUnrollResult UnrollLoop(Loop *L, BasicBlock *ExitingBlock, + UnrollLoopOptions ULO, LoopInfo *LI, ScalarEvolution *SE, DominatorTree *DT, AssumptionCache *AC, const llvm::TargetTransformInfo *TTI, Index: llvm/lib/Transforms/Scalar/LoopUnrollPass.cpp =================================================================== --- llvm/lib/Transforms/Scalar/LoopUnrollPass.cpp +++ llvm/lib/Transforms/Scalar/LoopUnrollPass.cpp @@ -1162,7 +1162,7 @@ // Unroll the loop. Loop *RemainderLoop = nullptr; LoopUnrollResult UnrollResult = UnrollLoop( - L, + L, ExitingBlock, {UP.Count, TripCount, UP.Force, UP.Runtime, UP.AllowExpensiveTripCount, UseUpperBound, MaxOrZero, TripMultiple, PP.PeelCount, UP.UnrollRemainder, ForgetAllSCEV}, Index: llvm/lib/Transforms/Utils/LoopUnroll.cpp =================================================================== --- llvm/lib/Transforms/Utils/LoopUnroll.cpp +++ llvm/lib/Transforms/Utils/LoopUnroll.cpp @@ -246,21 +246,17 @@ /// 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. +/// ExitingBlock. However, control may exit the loop prior to TripCount +/// iterations via a *different* exit. If control can exit from ExitingBlock +/// prior to TripCount iterations, the PreserveCondBr flag needs to be set. /// -/// PreserveCondBr indicates whether the conditional branch of the LatchBlock +/// PreserveCondBr indicates whether the conditional branch of the ExitingBlock /// 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. +/// Similarly, TripMultiple divides the number of times that the ExitingBlock +/// may execute without exiting the loop. /// /// If AllowRuntime is true then UnrollLoop will consider unrolling loops that /// have a runtime (i.e. not compile time constant) trip count. Unrolling these @@ -279,7 +275,8 @@ /// /// If RemainderLoop is non-null, it will receive the remainder loop (if /// required and not fully unrolled). -LoopUnrollResult llvm::UnrollLoop(Loop *L, UnrollLoopOptions ULO, LoopInfo *LI, +LoopUnrollResult llvm::UnrollLoop(Loop *L, BasicBlock *ExitingBlock, + UnrollLoopOptions ULO, LoopInfo *LI, ScalarEvolution *SE, DominatorTree *DT, AssumptionCache *AC, const TargetTransformInfo *TTI, @@ -352,7 +349,7 @@ // According to our guards and profitability checks the only // meaningful exit should be latch block. Other exits go to deopt, // so we do not worry about them. - BasicBlock *ExitingBlock = L->getLoopLatch(); + ExitingBlock = L->getLoopLatch(); assert(ExitingBlock && "Loop without exiting block?"); assert(L->isLoopExiting(ExitingBlock) && "Latch is not exiting?"); ULO.TripCount = SE->getSmallConstantTripCount(L, ExitingBlock); @@ -387,14 +384,12 @@ // FIXME: The implementation can be extended to work with more complicated // cases, e.g. loops with multiple latches. BranchInst *LatchBI = dyn_cast(LatchBlock->getTerminator()); + bool LatchIsExiting = L->isLoopExiting(LatchBlock); // A conditional branch which exits the loop, which can be optimized to an // unconditional branch in the unrolled loop in some cases. BranchInst *ExitingBI = nullptr; - bool LatchIsExiting = L->isLoopExiting(LatchBlock); - if (LatchIsExiting) - ExitingBI = LatchBI; - else if (BasicBlock *ExitingBlock = L->getExitingBlock()) + if (ExitingBlock) ExitingBI = dyn_cast(ExitingBlock->getTerminator()); if (!LatchBI || (LatchBI->isConditional() && !LatchIsExiting)) { // If the peeling guard is changed this assert may be relaxed or even @@ -526,7 +521,7 @@ SE->forgetTopmostLoop(L); } - if (!LatchIsExiting) + if (LatchBlock != ExitingBlock) ++NumUnrolledNotLatch; Optional ContinueOnTrue = None; BasicBlock *LoopExit = nullptr; @@ -746,13 +741,14 @@ // Connect latches of the unrolled iterations to the headers of the next // iteration. If the latch is also the exiting block, the conditional branch // may have to be preserved. + Optional ContinueOnTrueLatch = L->contains(LatchBI->getSuccessor(0)); for (unsigned i = 0, e = Latches.size(); i != e; ++i) { // The branch destination. unsigned j = (i + 1) % e; BasicBlock *Dest = Headers[j]; bool NeedConditional = LatchIsExiting; - if (LatchIsExiting) { + if (LatchBlock == ExitingBlock) { if (RuntimeTripCount && j != 0) NeedConditional = false; @@ -777,13 +773,13 @@ } } - setDest(Latches[i], Dest, Headers[i], NeedConditional, ContinueOnTrue, + setDest(Latches[i], Dest, Headers[i], NeedConditional, ContinueOnTrueLatch, Dest == LoopExit); } - if (!LatchIsExiting) { - // If the latch is not exiting, we may be able to simplify the conditional - // branches in the unrolled exiting blocks. + if (LatchBlock != ExitingBlock) { + // If we're unrolling against a non-latch exit, we may be able to simplify + // simplify the conditional branches in the unrolled exiting blocks. for (unsigned i = 0, e = ExitingBlocks.size(); i != e; ++i) { // The branch destination. unsigned j = (i + 1) % e; @@ -813,6 +809,9 @@ // When completely unrolling, the last latch becomes unreachable. if (CompletelyUnroll) { + for (BasicBlock *Succ : successors(Latches.back())) + Succ->removePredecessor(Latches.back()); + BranchInst *Term = cast(Latches.back()->getTerminator()); new UnreachableInst(Term->getContext(), Term); Term->eraseFromParent(); @@ -872,7 +871,8 @@ for (BasicBlock *Latch : Latches) { BranchInst *Term = dyn_cast(Latch->getTerminator()); assert((Term || - (CompletelyUnroll && !LatchIsExiting && Latch == Latches.back())) && + (CompletelyUnroll && LatchBlock != ExitingBlock && + Latch == Latches.back())) && "Need a branch as terminator, except when fully unrolling with " "unconditional latch"); if (Term && Term->isUnconditional()) { Index: llvm/lib/Transforms/Utils/LoopUnrollRuntime.cpp =================================================================== --- llvm/lib/Transforms/Utils/LoopUnrollRuntime.cpp +++ llvm/lib/Transforms/Utils/LoopUnrollRuntime.cpp @@ -983,7 +983,7 @@ if (remainderLoop && UnrollRemainder) { LLVM_DEBUG(dbgs() << "Unrolling remainder loop\n"); UnrollResult = - UnrollLoop(remainderLoop, + UnrollLoop(remainderLoop, remainderLoop->getLoopLatch(), {/*Count*/ Count - 1, /*TripCount*/ Count - 1, /*Force*/ false, /*AllowRuntime*/ false, /*AllowExpensiveTripCount*/ false, /*PreserveCondBr*/ true,