[LoopUnroll] Reorder code to max dom tree update more obvious [nfc]

[reames]

Summary

This cleans up the unroll action into two phases. Phase 1 does the mechanical act of unrolling, and leaves all conditional branches in place. Phase 2 optimizes away some of the conditional branches and then simplifies the loop. The primary benefit of the reordering is that we can delete some special cases dom tree update logic.

[nikic]

I tried this change as well, so have some compile-time numbers for a similar implementation: https://llvm-compile-time-tracker.com/compare.php?from=90310dfff8fd17b0cabdee1fd72d675e5eb2aa78&to=efa3a5246dbf80fca8c79fb7fd9464ba56f51f86&stat=instructions Which shows that impact is minor, though not entirely free. I dropped the change at the time because I didn't find it particularly helpful for followup work (the code is easy to adjust for multiple exit folding).

That said, this LGTM and at least simplifies the code.

llvm/lib/Transforms/Utils/LoopUnroll.cpp

768

clang-format warning is legit, some extra indentation snuck in here.

llvm/lib/Transforms/Utils/LoopUnroll.cpp

LoopUnrollResult llvm::UnrollLoop(Loop *L, UnrollLoopOptions ULO, LoopInfo *LI,
 
   // 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]);
   }
 
+  // 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 (ULO.Count > 1) {
+    for (auto *BB : OriginalLoopBlocks) {
+      auto *BBDomNode = DT->getNode(BB);
+      SmallVector<BasicBlock *, 16> ChildrenToUpdate;
+      for (auto *ChildDomNode : BBDomNode->children()) {
+        auto *ChildBB = ChildDomNode->getBlock();
+        if (!L->contains(ChildBB))
+          ChildrenToUpdate.push_back(ChildBB);
+      }
+      // The new idom of the block will be the nearest common dominator
+      // of all copies of the previous idom. This is equivalent to the
+      // nearest common dominator of the previous idom and the first latch,
+      // which dominates all copies of the previous idom.
+      BasicBlock *NewIDom = DT->findNearestCommonDominator(BB, LatchBlock);
+      for (auto *ChildBB : ChildrenToUpdate)
+        DT->changeImmediateDominator(ChildBB, NewIDom);
+    }
+  }
+
+  assert(!UnrollVerifyDomtree ||
+         DT->verify(DominatorTree::VerificationLevel::Fast));
+
+  DomTreeUpdater DTU(DT, DomTreeUpdater::UpdateStrategy::Lazy);
+
   if (ExitingBI) {
-    auto SetDest = [](BasicBlock *Src, bool WillExit, bool ExitOnTrue) {
+    auto SetDest = [&](BasicBlock *Src, bool WillExit, bool ExitOnTrue) {
       auto *Term = cast<BranchInst>(Src->getTerminator());
-      BasicBlock *Dest = Term->getSuccessor(ExitOnTrue ^ WillExit);
+      const unsigned Idx = ExitOnTrue ^ WillExit;
+      BasicBlock *Dest = Term->getSuccessor(Idx);
+      BasicBlock *DeadSucc = Term->getSuccessor(1-Idx);
 
       // Remove predecessors from all non-Dest successors.
-      for (BasicBlock *Succ : successors(Src)) {
-        if (Succ == Dest)
-          continue;
-        Succ->removePredecessor(Src, /* KeepOneInputPHIs */ true);
-      }
+      DeadSucc->removePredecessor(Src, /* KeepOneInputPHIs */ true);
 
       // Replace the conditional branch with an unconditional one.
       BranchInst::Create(Dest, Term);
       Term->eraseFromParent();
+
+      DTU.applyUpdates({{DominatorTree::Delete, Src, DeadSucc}});
     };
 
     auto WillExit = [&](unsigned i, unsigned j) -> Optional<bool> {
       if (CompletelyUnroll) {
         if (ULO.PreserveCondBr && j && !(ULO.PreserveOnlyFirst && i != 0))
           return None;
         return j == 0;
       }
 
       if (RuntimeTripCount && j != 0)
         return false;
 
       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.
         return false;
       }
       return None;
     };
 
     // 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;
       Optional<bool> KnownWillExit = WillExit(i, j);
       if (!KnownWillExit)
         continue;
 
       // TODO: Also fold known-exiting branches for non-latch exits.
       if (*KnownWillExit && !LatchIsExiting)
         continue;
 
       SetDest(ExitingBlocks[i], *KnownWillExit, ExitOnTrue);
     }
   }
 
-  // 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 (ULO.Count > 1) {
-    for (auto *BB : OriginalLoopBlocks) {
-      auto *BBDomNode = DT->getNode(BB);
-      SmallVector<BasicBlock *, 16> ChildrenToUpdate;
-      for (auto *ChildDomNode : BBDomNode->children()) {
-        auto *ChildBB = ChildDomNode->getBlock();
-        if (!L->contains(ChildBB))
-          ChildrenToUpdate.push_back(ChildBB);
-      }
-      BasicBlock *NewIDom;
-      if (ExitingBI && BB == ExitingBlocks[0]) {
-        // The latch is special because we emit unconditional branches in
-        // some cases where the original loop contained a conditional branch.
-        // Since the latch is always at the bottom of the loop, if the latch
-        // dominated an exit before unrolling, the new dominator of that exit
-        // must also be a latch.  Specifically, the dominator is the first
-        // latch which ends in a conditional branch, or the last latch if
-        // there is no such latch.
-        // For loops exiting from non latch exiting block, we limit the
-        // branch simplification to single exiting block loops.
-        NewIDom = ExitingBlocks.back();
-        for (unsigned i = 0, e = ExitingBlocks.size(); i != e; ++i) {
-          Instruction *Term = ExitingBlocks[i]->getTerminator();
-          if (isa<BranchInst>(Term) && cast<BranchInst>(Term)->isConditional()) {
-            NewIDom =
-                DT->findNearestCommonDominator(ExitingBlocks[i], Latches[i]);
-            break;
-          }
-        }
-      } else {
-        // The new idom of the block will be the nearest common dominator
-        // of all copies of the previous idom. This is equivalent to the
-        // nearest common dominator of the previous idom and the first latch,
-        // which dominates all copies of the previous idom.
-        NewIDom = DT->findNearestCommonDominator(BB, LatchBlock);
-      }
-      for (auto *ChildBB : ChildrenToUpdate)
-        DT->changeImmediateDominator(ChildBB, NewIDom);
-    }
-  }
-
-  assert(!UnrollVerifyDomtree ||
-         DT->verify(DominatorTree::VerificationLevel::Fast));
-
-  DomTreeUpdater DTU(DT, DomTreeUpdater::UpdateStrategy::Lazy);
 
   // When completely unrolling, the last latch becomes unreachable.
   if (!LatchIsExiting && CompletelyUnroll)
     changeToUnreachable(Latches.back()->getTerminator(), /* UseTrap */ false,
                         PreserveLCSSA, &DTU);
 
   // Merge adjacent basic blocks, if possible.
   for (BasicBlock *Latch : Latches) {