[PATCH] D111611: [NFC] [LoopPeel] Update IDoms of non-loop blocks dominated by the loop

[Dmitry Makogon via Phabricator via llvm-commits]

dmakogon updated this revision to Diff 381218.
dmakogon retitled this revision from "[NFC] [LoopPeel] Change the way DT is updated for loop exits" to "[NFC] [LoopPeel] Update IDoms of non-loop blocks dominated by the loop".
dmakogon edited the summary of this revision.
dmakogon added a comment.

As pointed out by @nikic the proposed solution which used `DomTreeUpdater` was incorrect. 
Now we remember IDoms of non-loop blocks which are dominated by some loop block. 
Their IDom would be the nearest common dominator of their loop dominator and the latch. 
The same logic was already used in `peelLoop` function, but only for the loop exits. 
With this patch we also cover other blocks that are dominated by the loop.

The new logic is copied from the LoopUnroll pass <https://github.com/llvm/llvm-project/blob/3f0b178de21ee82791a6ebe198314f14c0287a44/llvm/lib/Transforms/Utils/LoopUnroll.cpp#L643-L658>. It's probably better to create an utility function for this, but need to check if it would work in the loop unroll code.


Repository:
  rG LLVM Github Monorepo

CHANGES SINCE LAST ACTION
  https://reviews.llvm.org/D111611/new/

https://reviews.llvm.org/D111611

Files:
  llvm/lib/Transforms/Utils/LoopPeel.cpp


Index: llvm/lib/Transforms/Utils/LoopPeel.cpp
===================================================================
--- llvm/lib/Transforms/Utils/LoopPeel.cpp
+++ llvm/lib/Transforms/Utils/LoopPeel.cpp
@@ -732,34 +732,27 @@
   SmallVector<std::pair<BasicBlock *, BasicBlock *>, 4> ExitEdges;
   L->getExitEdges(ExitEdges);
 
-  DenseMap<BasicBlock *, BasicBlock *> ExitIDom;
+  // Remember dominators of blocks we might reach through exits to change them
+  // later. Immediate dominator of such block might change, because we add more
+  // routes which can lead to the exit: we can reach it from the peeled
+  // iterations too.
+  DenseMap<BasicBlock *, BasicBlock *> NonLoopBlocksIDom;
   if (DT) {
-    // We'd like to determine the idom of exit block after peeling one
-    // iteration.
-    // Let Exit is exit block.
-    // Let ExitingSet - is a set of predecessors of Exit block. They are exiting
-    // blocks.
-    // Let Latch' and ExitingSet' are copies after a peeling.
-    // We'd like to find an idom'(Exit) - idom of Exit after peeling.
-    // It is an evident that idom'(Exit) will be the nearest common dominator
-    // of ExitingSet and ExitingSet'.
-    // idom(Exit) is a nearest common dominator of ExitingSet.
-    // idom(Exit)' is a nearest common dominator of ExitingSet'.
-    // Taking into account that we have a single Latch, Latch' will dominate
-    // Header and idom(Exit).
-    // So the idom'(Exit) is nearest common dominator of idom(Exit)' and Latch'.
-    // All these basic blocks are in the same loop, so what we find is
-    // (nearest common dominator of idom(Exit) and Latch)'.
-    // In the loop below we remember nearest common dominator of idom(Exit) and
-    // Latch to update idom of Exit later.
-    assert(L->hasDedicatedExits() && "No dedicated exits?");
-    for (auto Edge : ExitEdges) {
-      if (ExitIDom.count(Edge.second))
-        continue;
-      BasicBlock *BB = DT->findNearestCommonDominator(
-          DT->getNode(Edge.second)->getIDom()->getBlock(), Latch);
-      assert(L->contains(BB) && "IDom is not in a loop");
-      ExitIDom[Edge.second] = BB;
+    for (auto *BB : L->blocks()) {
+      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, Latch);
+      for (auto *ChildBB : ChildrenToUpdate)
+        NonLoopBlocksIDom[ChildBB] = NewIDom;
     }
   }
 
@@ -848,13 +841,11 @@
     remapInstructionsInBlocks(NewBlocks, VMap);
 
     if (DT) {
-      // Latches of the cloned loops dominate over the loop exit, so idom of the
-      // latter is the first cloned loop body, as original PreHeader dominates
-      // the original loop body.
+      // Update IDoms of the blocks reachable through exits.
       if (Iter == 0)
-        for (auto Exit : ExitIDom)
-          DT->changeImmediateDominator(Exit.first,
-                                       cast<BasicBlock>(LVMap[Exit.second]));
+        for (auto BBIDom : NonLoopBlocksIDom)
+          DT->changeImmediateDominator(BBIDom.first,
+                                       cast<BasicBlock>(LVMap[BBIDom.second]));
 #ifdef EXPENSIVE_CHECKS
       assert(DT->verify(DominatorTree::VerificationLevel::Fast));
 #endif


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D111611.381218.patch
Type: text/x-patch
Size: 3739 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20211021/3e46f8ff/attachment.bin>