[LoopUnroll] Explicitly specify exit to unroll against (NFCI)

[nikic]

Summary

Loop unrolling is usually performed against an exit for which additional information is known, such as a trip count, trip multiple or runtime-enforced trip multiple. This exit receives special treatment, because it can be folded on some or all iterations.

Currently, this exit is not explicitly provided to the unrolling implementation. Instead it assumes that this must be either the unique exit, or the latch exit. For D102635 (where the latch is non-exiting) this results in missed folding opportunities, because unrolling doesn't know which exit the trip count refers to. For D102982 (where the latch can be exiting) it results in miscompiles, because it assumes the wrong exit.

This patch addresses the problem by explicitly providing the ExitingBlock to unrolling. Currently, this shouldn't have any effect, but would have an effect in conjunction with D102635/D102982.

[reames]

Not sure if this would work or not, but if the implementation allows, I might suggest computing the symbolic exit count for each loop upfront (before unrolling starts). Once you have that, there's a couple of cases:

• If Exit EC > UC, branch folds to untaken.
• If Exit EC == UC, then branch folds to untaken on all but one iteration, and taken on that iteration.
• If Exit EC < UC, we can clamp the UC.

Doing this as symbolic comparisons in the unroll implementation (as opposed to passing in a flag) seems likely to be more robust.

[mkazantsev]

llvm/lib/Transforms/Utils/LoopUnroll.cpp

352

I don't really understand the implications of this change. We'll process a different exit in case of non-zero peeled count than passed in param. Is this intentional? Could you please explain?

According to the comment, it looks like we can just assert that ExitingBlock == L->getLoopLatch(), but I never understood subtleties of unrolling code.

[nikic]

I'll experiment with a variant that performs folding of all exits using SCEV information inside unrolling, to possibly remove the need for passing in a trip count entirely.

[nikic]

No longer needed, UnrollLoop() is no longer passed TripCount/TripMultiple for a specific exit.

llvm/include/llvm/Transforms/Utils/UnrollLoop.h

struct UnrollLoopOptions {
   bool PreserveCondBr;
   bool PreserveOnlyFirst;
   unsigned TripMultiple;
   unsigned PeelCount;
   bool UnrollRemainder;
   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,
                             OptimizationRemarkEmitter *ORE, bool PreserveLCSSA,
                             Loop **RemainderLoop = nullptr);
 
 bool UnrollRuntimeLoopRemainder(
     Loop *L, unsigned Count, bool AllowExpensiveTripCount,

llvm/lib/Transforms/Scalar/LoopUnrollPass.cpp

if (TripCount && UP.Count > TripCount)
     UP.Count = TripCount;
 
   // Save loop properties before it is transformed.
   MDNode *OrigLoopID = L->getLoopID();
 
   // 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},
       LI, &SE, &DT, &AC, &TTI, &ORE, PreserveLCSSA, &RemainderLoop);
   if (UnrollResult == LoopUnrollResult::Unmodified)
     return LoopUnrollResult::Unmodified;
 
   if (RemainderLoop) {

llvm/lib/Transforms/Utils/LoopUnroll.cpp

 }
 
 /// 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.
+/// 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
 /// loops require a unroll "prologue" that runs "RuntimeTripCount % Count"
 /// iterations before branching into the unrolled loop.  UnrollLoop will not
 /// runtime-unroll the loop if computing RuntimeTripCount will be expensive and
 /// AllowExpensiveTripCount is false.
 ///
 /// If we want to perform PGO-based loop peeling, PeelCount is set to the
 /// number of iterations we want to peel off.
 ///
 /// The LoopInfo Analysis that is passed will be kept consistent.
 ///
 /// This utility preserves LoopInfo. It will also preserve ScalarEvolution and
 /// DominatorTree if they are non-null.
 ///
 /// 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,
                                   OptimizationRemarkEmitter *ORE,
                                   bool PreserveLCSSA, Loop **RemainderLoop) {
 
   if (!L->getLoopPreheader()) {
     LLVM_DEBUG(dbgs() << "  Can't unroll; loop preheader-insertion failed.\n");
if (ULO.PeelCount) {
     Peeled = peelLoop(L, ULO.PeelCount, LI, SE, DT, AC, PreserveLCSSA);
 
     // Successful peeling may result in a change in the loop preheader/trip
     // counts. If we later unroll the loop, we want these to be updated.
     if (Peeled) {
       // 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);
       ULO.TripMultiple = SE->getSmallConstantTripMultiple(L, ExitingBlock);
     }
   }
 
   // All these values should be taken only after peeling because they might have
LoopUnrollResult llvm::UnrollLoop(Loop *L, BasicBlock *ExitingBlock,
 
   // The current loop unroll pass can unroll loops that have
   // (1) single latch; and
   // (2a) latch is unconditional; or
   // (2b) latch is conditional and is an exiting block
   // FIXME: The implementation can be extended to work with more complicated
   // cases, e.g. loops with multiple latches.
   BranchInst *LatchBI = dyn_cast<BranchInst>(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<BranchInst>(ExitingBlock->getTerminator());
   if (!LatchBI || (LatchBI->isConditional() && !LatchIsExiting)) {
     // If the peeling guard is changed this assert may be relaxed or even
     // deleted.
     assert(!Peeled && "Peeling guard changed!");
     LLVM_DEBUG(
         dbgs() << "Can't unroll; a conditional latch must exit the loop");
     return LoopUnrollResult::Unmodified;
LoopUnrollResult llvm::UnrollLoop(Loop *L, BasicBlock *ExitingBlock,
   // and this is what we need here.
   if (SE) {
     if (ULO.ForgetAllSCEV)
       SE->forgetAllLoops();
     else
       SE->forgetTopmostLoop(L);
   }
 
-  if (!LatchIsExiting)
+  if (LatchBlock != ExitingBlock)
     ++NumUnrolledNotLatch;
   Optional<bool> ContinueOnTrue = None;
   BasicBlock *LoopExit = nullptr;
   if (ExitingBI) {
     ContinueOnTrue = L->contains(ExitingBI->getSuccessor(0));
     LoopExit = ExitingBI->getSuccessor(*ContinueOnTrue);
   }
 
if (NeedConditional) {
       BranchInst::Create(Dest, Term);
       Term->eraseFromParent();
     }
   };
 
   // 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<bool> 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;
 
       // For a complete unroll, make the last iteration end with a branch
       // to the exit block.
       if (CompletelyUnroll) {
         if (j == 0)
           Dest = LoopExit;
         // If using trip count upper bound to completely unroll, we need to
         // keep the conditional branch except the last one because the loop
         // may exit after any iteration.
         assert(NeedConditional &&
                "NeedCondition cannot be modified by both complete "
                "unrolling and runtime unrolling");
         NeedConditional =
             (ULO.PreserveCondBr && j && !(ULO.PreserveOnlyFirst && i != 0));
       } else if (j != BreakoutTrip &&
                  (ULO.TripMultiple == 0 || j % ULO.TripMultiple != 0)) {
         // If we know the trip count or a multiple of it, we can safely use an
         // unconditional branch for some iterations.
         NeedConditional = false;
       }
     }
 
-    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;
       bool NeedConditional = true;
 
       if (RuntimeTripCount && j != 0)
         NeedConditional = false;
 
for (unsigned i = 0, e = ExitingBlocks.size(); i != e; ++i) {
       if (NeedConditional)
         continue;
       setDest(ExitingBlocks[i], ExitingSucc[i], ExitingSucc[i], NeedConditional,
               None, false);
     }
 
     // When completely unrolling, the last latch becomes unreachable.
     if (CompletelyUnroll) {
+      for (BasicBlock *Succ : successors(Latches.back()))
+        Succ->removePredecessor(Latches.back());
+
       BranchInst *Term = cast<BranchInst>(Latches.back()->getTerminator());
       new UnreachableInst(Term->getContext(), Term);
       Term->eraseFromParent();
     }
   }
 
   // Update dominators of blocks we might reach through exits.
   // Immediate dominator of such block might change, because we add more
LoopUnrollResult llvm::UnrollLoop(Loop *L, BasicBlock *ExitingBlock,
 
   assert(!DT || !UnrollVerifyDomtree ||
          DT->verify(DominatorTree::VerificationLevel::Fast));
 
   DomTreeUpdater DTU(DT, DomTreeUpdater::UpdateStrategy::Lazy);
   // Merge adjacent basic blocks, if possible.
   for (BasicBlock *Latch : Latches) {
     BranchInst *Term = dyn_cast<BranchInst>(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()) {
       BasicBlock *Dest = Term->getSuccessor(0);
       BasicBlock *Fold = Dest->getUniquePredecessor();
       if (MergeBlockIntoPredecessor(Dest, &DTU, LI)) {
         // Dest has been folded into Fold. Update our worklists accordingly.
         std::replace(Latches.begin(), Latches.end(), Dest, Fold);

llvm/lib/Transforms/Utils/LoopUnrollRuntime.cpp

if (OtherExits.size() > 0) {
     if (remainderLoop)
       formDedicatedExitBlocks(remainderLoop, DT, LI, nullptr, PreserveLCSSA);
   }
 
   auto UnrollResult = LoopUnrollResult::Unmodified;
   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,
                     /*PreserveOnlyFirst*/ false, /*TripMultiple*/ 1,
                     /*PeelCount*/ 0, /*UnrollRemainder*/ false, ForgetAllSCEV},
                    LI, SE, DT, AC, TTI, /*ORE*/ nullptr, PreserveLCSSA);
   }
 
   if (ResultLoop && UnrollResult != LoopUnrollResult::FullyUnrolled)
     *ResultLoop = remainderLoop;
   NumRuntimeUnrolled++;
   return true;
 }