[LoopUnroll] Clean up exit folding (NFC)

[nikic]

Summary

This does some non-functional cleanup of exit folding during unrolling. The two main changes are:

• First rewrite latch->header edges, which is unrelated to exit folding.
• Combine folding for latch and non-latch exits. After the previous change, these use the same logic, with the caveat that for non-latch exits we currently only fold "known non-exit" cases, but not "known exit" cases.

I think this helps a lot to clarify this code and prepare it for future changes.

[reames]

LGTM, this looks correct to me, and definitely helps to simplify the code!

llvm/lib/Transforms/Utils/LoopUnroll.cpp

if (SE) {
     if (ULO.ForgetAllSCEV)
       SE->forgetAllLoops();
     else
       SE->forgetTopmostLoop(L);
   }
 
   if (!LatchIsExiting)
     ++NumUnrolledNotLatch;
-  Optional<bool> ContinueOnTrue = None;
-  BasicBlock *LoopExit = nullptr;
-  if (ExitingBI) {
-    ContinueOnTrue = L->contains(ExitingBI->getSuccessor(0));
-    LoopExit = ExitingBI->getSuccessor(*ContinueOnTrue);
-  }
 
   // For the first iteration of the loop, we should use the precloned values for
   // PHI nodes.  Insert associations now.
   ValueToValueMapTy LastValueMap;
   std::vector<PHINode*> OrigPHINode;
   for (BasicBlock::iterator I = Header->begin(); isa<PHINode>(I); ++I) {
     OrigPHINode.push_back(cast<PHINode>(I));
   }
 
   std::vector<BasicBlock *> Headers;
   std::vector<BasicBlock *> ExitingBlocks;
-  std::vector<BasicBlock *> ExitingSucc;
   std::vector<BasicBlock *> Latches;
   Headers.push_back(Header);
   Latches.push_back(LatchBlock);
-  if (ExitingBI) {
+  if (ExitingBI)
     ExitingBlocks.push_back(ExitingBI->getParent());
-    ExitingSucc.push_back(ExitingBI->getSuccessor(!(*ContinueOnTrue)));
-  }
 
   // The current on-the-fly SSA update requires blocks to be processed in
   // reverse postorder so that LastValueMap contains the correct value at each
   // exit.
   LoopBlocksDFS DFS(L);
   DFS.perform(LI);
 
   // Stash the DFS iterators before adding blocks to the loop.
for (LoopBlocksDFS::RPOIterator BB = BlockBegin; BB != BlockEnd; ++BB) {
       // we can insert the proper branches later.
       if (*BB == Header)
         Headers.push_back(New);
       if (*BB == LatchBlock)
         Latches.push_back(New);
 
       // Keep track of the exiting block and its successor block contained in
       // the loop for the current iteration.
-      if (ExitingBI) {
+      if (ExitingBI)
         if (*BB == ExitingBlocks[0])
           ExitingBlocks.push_back(New);
-        if (*BB == ExitingSucc[0])
-          ExitingSucc.push_back(New);
-      }
 
       NewBlocks.push_back(New);
       UnrolledLoopBlocks.push_back(New);
 
       // Update DomTree: since we just copy the loop body, and each copy has a
       // dedicated entry block (copy of the header block), this header's copy
       // dominates all copied blocks. That means, dominance relations in the
       // copied body are the same as in the original body.
if (CompletelyUnroll) {
         if (L->contains(InValI))
           InVal = LastValueMap[InVal];
       }
       assert(Latches.back() == LastValueMap[LatchBlock] && "bad last latch");
       PN->addIncoming(InVal, Latches.back());
     }
   }
 
-  auto setDest = [](BasicBlock *Src, BasicBlock *Dest, BasicBlock *BlockInLoop,
-                    bool NeedConditional, Optional<bool> ContinueOnTrue,
-                    bool IsDestLoopExit) {
+  // Connect latches of the unrolled iterations to the headers of the next
+  // iteration. Currently they point to the header of the same iteration.
+  for (unsigned i = 0, e = Latches.size(); i != e; ++i) {
+    unsigned j = (i + 1) % e;
+    Latches[i]->getTerminator()->replaceSuccessorWith(Headers[i], Headers[j]);
+  }
+
+  if (ExitingBI) {
+    auto SetDest = [](BasicBlock *Src, bool WillExit, bool ExitOnTrue) {
       auto *Term = cast<BranchInst>(Src->getTerminator());
-    if (NeedConditional) {
-      // Update the conditional branch's successor for the following
-      // iteration.
-      assert(ContinueOnTrue.hasValue() &&
-             "Expecting valid ContinueOnTrue when NeedConditional is true");
-      Term->setSuccessor(!(*ContinueOnTrue), Dest);
-    } else {
-      // Remove phi operands at this loop exit
-      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);
-        }
+      BasicBlock *Dest = Term->getSuccessor(ExitOnTrue ^ WillExit);
+
+      // Remove predecessors from all non-Dest successors.
+      for (BasicBlock *Succ : successors(Src)) {
+        if (Succ == Dest)
+        Succ->removePredecessor(Src, /* KeepOneInputPHIs */ true);
       }
+
       // Replace the conditional branch with an unconditional one.
       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.
-  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;
+    auto WillExit = [&](unsigned i, unsigned j) -> Optional<bool> {
+      if (CompletelyUnroll) {
+        if (ULO.PreserveCondBr && j && !(ULO.PreserveOnlyFirst && i != 0))
+          return None;
+        return j == 0;
+      }
 
-    if (LatchIsExiting) {
       if (RuntimeTripCount && j != 0)
-        NeedConditional = false;
+        return 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 (j != BreakoutTrip &&
         // 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,
-            Dest == LoopExit);
+        return false;
       }
+      return None;
+    };
 
-  if (!LatchIsExiting) {
-    // If the latch is not exiting, we may be able to simplify the conditional
-    // branches in the unrolled exiting blocks.
+    // Fold branches for iterations where we know that they will exit or not
+    // exit.
+    bool ExitOnTrue = !L->contains(ExitingBI->getSuccessor(0));
     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;
-
-      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;
+      Optional<bool> KnownWillExit = WillExit(i, j);
+      if (!KnownWillExit)
+        continue;
 
-      // 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)
+      // TODO: Also fold known-exiting branches for non-latch exits.
+      if (*KnownWillExit && !LatchIsExiting)
         continue;
-      setDest(ExitingBlocks[i], ExitingSucc[i], ExitingSucc[i], NeedConditional,
-              None, false);
-    }
+
+      SetDest(ExitingBlocks[i], *KnownWillExit, ExitOnTrue);
+    }
 
-    // When completely unrolling, the last latch becomes unreachable.
-    if (CompletelyUnroll) {
-      BranchInst *Term = cast<BranchInst>(Latches.back()->getTerminator());
-      new UnreachableInst(Term->getContext(), Term);
-      Term->eraseFromParent();
-    }
-  }
+  if (!LatchIsExiting && CompletelyUnroll) {
 
   // 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.
   if (DT && ULO.Count > 1) {
     for (auto *BB : OriginalLoopBlocks) {
       auto *BBDomNode = DT->getNode(BB);