[LoopUnroll] Simplify latch/header block handling (NFC).

[fhahn]

Summary

I think the current code dealing with connecting the unrolled iterations
is a bit more complicated than necessary currently. To connect the
unrolled iterations, we have to update the unrolled latch blocks to
branch to the header of the next unrolled iteration.

We need to do this regardless whether the latch is exiting or not.

Additionally, we try to turn the conditional branch in the exiting block
to an unconditional one. This is an optimization only; alternatively we
could leave the conditional branches in place and rely on other passes
to simplify the conditions.

Logically, this is a separate step from connecting the latches to the
headers, but it is convenient to fold them into the same loop, if the
latch is also exiting. For headers (or other non-latch exiting blocks,
this is done separately).

Hopefully the patch with additional comments makes things a bit clearer.

[efriedma]

Makes sense.

llvm/lib/Transforms/Utils/LoopUnroll.cpp

718–719

While you're here, can you change this so the "LoopExit" and "ContinueOnTrue" are passed in as arguments? I think changing the meaning based on LatchIsExiting is confusing.

[fhahn]

Move ContinueOnTrue, LoopExit captures to explicit argument list. Also turned LoopExit into a bool (IsDestLoopExit), which makes it clear that this is always false in the Header case.

[fhahn]

llvm/lib/Transforms/Utils/LoopUnroll.cpp

718–719

Done thanks. I've turned LoopExit into a bool (IsDestLoopExit). That should make the fact it is always false for the header case more explicit.

[efriedma]

LGTM

llvm/lib/Transforms/Utils/LoopUnroll.cpp

if (CompletelyUnroll) {
       if (Instruction *InValI = dyn_cast<Instruction>(InVal)) {
         if (L->contains(InValI))
           InVal = LastValueMap[InVal];
       }
       assert(Latches.back() == LastValueMap[LatchBlock] && "bad last latch");
       PN->addIncoming(InVal, Latches.back());
     }
   }
 
-  auto setDest = [LoopExit, ContinueOnTrue](BasicBlock *Src, BasicBlock *Dest,
-                                            BasicBlock *BlockInLoop,
-                                            bool NeedConditional) {
+  auto setDest = [](BasicBlock *Src, BasicBlock *Dest, BasicBlock *BlockInLoop,
+                    bool NeedConditional, bool ContinueOnTrue,
+                    bool IsDestLoopExit) {
     auto *Term = cast<BranchInst>(Src->getTerminator());
     if (NeedConditional) {
       // Update the conditional branch's successor for the following
       // iteration.
       Term->setSuccessor(!ContinueOnTrue, Dest);
     } else {
       // Remove phi operands at this loop exit
-      if (Dest != LoopExit) {
+      if (!IsDestLoopExit) {
         BasicBlock *BB = Src;
         for (BasicBlock *Succ : successors(BB)) {
           // Preserve the incoming value from BB if we are jumping to the block
           // in the current loop.
           if (Succ == BlockInLoop)
             continue;
           for (PHINode &Phi : Succ->phis())
             Phi.removeIncomingValue(BB, false);
         }
       }
       // Replace the conditional branch with an unconditional one.
       BranchInst::Create(Dest, Term);
       Term->eraseFromParent();
     }
   };
 
-  // Now that all the basic blocks for the unrolled iterations are in place,
-  // set up the branches to connect them.
-  if (LatchIsExiting) {
-    // Set up latches to branch to the new header in the unrolled iterations or
-    // the loop exit for the last latch in a fully unrolled loop.
-    for (unsigned i = 0, e = Latches.size(); i != e; ++i) {
-      // The branch destination.
-      unsigned j = (i + 1) % e;
-      BasicBlock *Dest = Headers[j];
-      bool NeedConditional = true;
+  // 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.
+    bool NeedConditional = LatchIsExiting;
 
-      if (RuntimeTripCount && j != 0) {
+    if (LatchIsExiting) {
+      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.
+        // 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);
-    }
-  } else {
-    // Setup headers to branch to their new successors in the unrolled
-    // iterations.
+    setDest(Latches[i], Dest, Headers[i], NeedConditional, ContinueOnTrue,
+            Dest == LoopExit);
+
+  if (!LatchIsExiting) {
+    // If the latch is not exiting, we may be able to simplify the conditional
+    // branches in the unrolled exiting blocks.
     for (unsigned i = 0, e = Headers.size(); i != e; ++i) {
       // The branch destination.
       unsigned j = (i + 1) % e;
-      BasicBlock *Dest = HeaderSucc[i];
       bool NeedConditional = true;
 
       if (RuntimeTripCount && j != 0)
         NeedConditional = false;
 
       if (CompletelyUnroll)
         // We cannot drop the conditional branch for the last condition, as we
         // may have to execute the loop body depending on the condition.
         NeedConditional = j == 0 || ULO.PreserveCondBr;
       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(Headers[i], Dest, HeaderSucc[i], NeedConditional);
+      // Conditional branches from non-latch exiting block have successors
+      // either in the same loop iteration or outside the loop. The branches are
+      // already correct.
+      if (NeedConditional)
+        continue;
+      setDest(Headers[i], HeaderSucc[i], HeaderSucc[i], NeedConditional,
+              ContinueOnTrue, false);
     }
 
-    // Set up latches to branch to the new header in the unrolled iterations or
-    // the loop exit for the last latch in a fully unrolled loop.
-
-    for (unsigned i = 0, e = Latches.size(); i != e; ++i) {
-      // The original branch was replicated in each unrolled iteration.
-      BranchInst *Term = cast<BranchInst>(Latches[i]->getTerminator());
-
-      // The branch destination.
-      unsigned j = (i + 1) % e;
-      BasicBlock *Dest = Headers[j];
-
-      // When completely unrolling, the last latch becomes unreachable.
-      if (CompletelyUnroll && j == 0)
-        new UnreachableInst(Term->getContext(), Term);
-      else
-        // Replace the conditional branch with an unconditional one.
-        BranchInst::Create(Dest, Term);
-
+    if (CompletelyUnroll) {
+      BranchInst *Term = cast<BranchInst>(Latches.back()->getTerminator());
       Term->eraseFromParent();
     }
   }
 
   // Update dominators of blocks we might reach through exits.
   // Immediate dominator of such block might change, because we add more
   // routes which can lead to the exit: we can now reach it from the copied
   // iterations too.