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

[Philip Reames via llvm-commits]

Author: Philip Reames
Date: 2021-06-03T10:19:56-07:00
New Revision: bb5e1c6dcb78435f8cdcaa9eb1c5af39244df491

URL: https://github.com/llvm/llvm-project/commit/bb5e1c6dcb78435f8cdcaa9eb1c5af39244df491
DIFF: https://github.com/llvm/llvm-project/commit/bb5e1c6dcb78435f8cdcaa9eb1c5af39244df491.diff

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

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.

Differential Revision: https://reviews.llvm.org/D103561

Added: 
    

Modified: 
    llvm/lib/Transforms/Utils/LoopUnroll.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/Utils/LoopUnroll.cpp b/llvm/lib/Transforms/Utils/LoopUnroll.cpp
index 8ae9095302d16..3fdcd0d08e91b 100644
--- a/llvm/lib/Transforms/Utils/LoopUnroll.cpp
+++ b/llvm/lib/Transforms/Utils/LoopUnroll.cpp
@@ -708,21 +708,49 @@ LoopUnrollResult llvm::UnrollLoop(Loop *L, UnrollLoopOptions ULO, LoopInfo *LI,
     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> {
@@ -762,55 +790,6 @@ LoopUnrollResult llvm::UnrollLoop(Loop *L, UnrollLoopOptions ULO, LoopInfo *LI,
     }
   }
 
-  // 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)