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

[dmakogon]

Summary

When peeling a loop, we assume that the latch has a br terminator and that all loop exits are either terminated with an unreachable or have a terminating deoptimize call. So when we peel off the 1st iteration, we change the IDom of all loop exits to the peeled copy of NCD(IDom(Exit), Latch). This works now, but if we add logic to support loops with exits that are followed by a block with an unreachable or a terminating deoptimize call, changing the exit's idom wouldn't be enough and DT would be broken.

For example, let Exit1 and Exit2 are loop exits, and each of them unconditionally branches to the same unreachable terminated block. So neither of the exits dominates this unreachable block. If we change the IDoms of the exits to some peeled loop block, we don't update the dominators of the unreachable block. Currently we just don't get to the peeling logic, saying that we can't peel such loops.

Previously we stored exits' IDoms in a map before peeling a loop and then, after peeling off one iteration, we changed their IDoms.
Now we use the same logic not only for exits but for all non-loop blocks dominated by the loop.
So when we add logic to support peeling loops with exits which branch, for example, to an unreachable-terminated block, we would update the IDoms not only for exits, but for their successors.

[mkazantsev]

Looks fine, but let's have more eyes on this.

[fhahn]

llvm/lib/Transforms/Utils/LoopPeel.cpp

836

Passing DT means that cloneLoopBlocks previously could rely on DT being up to date here. That's not true any more, as it is only updated later. Is that a problem? We also need to check all other uses of plain DT in until the updates are applied.

[mkazantsev]

llvm/lib/Transforms/Utils/LoopPeel.cpp

836

We can completely remove DT from the signatures here and pass it updates for filling as well. Also may be faster.

[dmakogon]

Now we pass a DomTreeUpdater to cloneLoopBlocks() for it to store updates there.

[mkazantsev]

llvm/lib/Transforms/Utils/LoopPeel.cpp

16

Pls remove this, DenseMap is already included on line 13.

563

Pls make it const

573

You already have DTU. You can just call insertEdge/deleteEdge on it without any checks. They will not be applied immediately because you are using lazy update strategy. So no need for DTUpdates at all.

835

for (auto *BB : L->blocks())

855

You can just call insertEdge here. DTU with null tree handles it correctly. The application will be lazy. No need to check hasDomTree.

896

Lazy DTU will only apply updates when it goes out of scope, or when it flushes. simplifyLoop below uses DT and needs it to be up to date. I'd suggest calling DTU.flush() here to apply all pending updates and leave assert as is.

[mkazantsev]

llvm/lib/Transforms/Utils/LoopPeel.cpp

573

**applyUpdates, since these two are deprecated. But a separate vector is not needed anyways, you have this lazy vector in DTU.

[dmakogon]

llvm/lib/Transforms/Utils/LoopPeel.cpp

896

Calling DTU.getDomTree flushes all updates to DT

[dmakogon]

Removed DTU.hasDomTree() checks and DTUpdates in cloneLoopBlocks. All updates are applied on DTU

[dmakogon]

llvm/lib/Transforms/Utils/LoopPeel.cpp

563

How are we going to apply updates to it then? DTU::applyUpdates is a non-const function

836

Now we pass this function a DomTreeUpdater

[mkazantsev]

llvm/lib/Transforms/Utils/LoopPeel.cpp

563

LoopBlocksIDoms should be const, not DTU.

[dmakogon]

Made LoopBlocksIDoms const

[mkazantsev]

llvm/lib/Transforms/Utils/LoopPeel.cpp

835

*BB

896

And you only call it in assert, i.e. in debug mode. It will simply not work in release. Please do it properly. I insist it to be flush + old assert.

[dmakogon]

Moved DTU.getDomTree out of assert, now all further uses of DT (if it's present) will see the updated tree not only in debug mode

[dmakogon]

fixed small mistake

[mkazantsev]

Fine by me now, thanks! Let's see if someone else has concerns.

[nikic]

llvm/lib/Transforms/Utils/LoopPeel.cpp

601

This code doesn't make sense to me. DT->addNewBlock() directly adds a new edge to the DT. DTU.applyUpdates() applies CFG updates to the DT. You seem to be passing a DT update to a CFG API here. Maybe that happens to work, but this is not how this is supposed to be used.

I'm not sure using DTU in this code is really useful.

[mkazantsev]

llvm/lib/Transforms/Utils/LoopPeel.cpp

601

Let me try to explain. We are going to expand the logic of exit detection on D110922. The current implementation relies on fact that we only need to update idoms of loop blocks and their immediate exits to get the right tree. However, this assumption will break once we handle exit blocks that may have successors, meaning idom for those successors may change. Imagine the case:

loop:
  br loop2, exit1

loop2:
  br loop3, exit 2

loop3:
  br loop, exit 3

exit1:
  br post_exit

exit2:
  br post_exit

exit3:
  br post_exit

After all it is done, idom of post_exit will change. Meaning that we should not just change idoms of exits, but also their successors. Now this situation is impossible, but it becomes possible as the scope of the optimization expands. We want to make current implementation independent of the scope (in fact, this patch is just a split-out piece of that one). That's why we need to use insert updates in the caller method.

Now, lazy strategy (at least in theory) should require less tree traversal than eager one. That's why we use lazy DTU in the caller.

Finally, we use DTU in the outer method and we have this cloneLoopBlocks here in that, I agree, we could leave it be with DT. However, the code written as is both reads and writes DT. What is the interaction between this and pending lazy updates, isn't very clear. I guess it just happens to be OK, but since we are to use DTU anyways, why not just use it everywhere for uniformity reasons? Just to avoid the situation that part of updates are applied, part is pending, and we also get idom from this half-updated tree. It should *happen* to work, but it doesn't really look nearly safe.

So the prime motivation for this was "ok, we are going to use lazy DTU anyways, let's be consistent and get all updates through it". It means that cloneLoopBlocks should be deprived of direct DT access, to avoid temptation to read something from it, knowing that not all updates are applied. This is why no DT is here, and there is DTU in its stead. And now, DTU simply doesn't have addNewBlock.

Does it all make more sense now?

[nikic]

llvm/lib/Transforms/Utils/LoopPeel.cpp

601

I understand the motivation for this change -- my problem is that the code misuses the DomTreeUpdater API. There is (generally) no CFG edge between VMap[IDom] and NewBB, but you are pretending there is. Maybe this happens to work because of implementation details, but it's not the correct way to use the API.

You'd have to actually insert all new control flow edges for the cloned blocks, not just the idom edges in the DT. It's one of the cases where directly updating the DT is simpler than going through DTU.

FYI the corresponding code for unrolling is https://github.com/llvm/llvm-project/blob/3f0b178de21ee82791a6ebe198314f14c0287a44/llvm/lib/Transforms/Utils/LoopUnroll.cpp#L643-L658 -- probably the same code would work for peeling as well.

[mkazantsev]

llvm/lib/Transforms/Utils/LoopPeel.cpp

601

I got your point now. Reverting the patch.

[mkazantsev]

llvm/lib/Transforms/Utils/LoopPeel.cpp

601

I've re-read this code, and now I got why I had this confusion. The deprecated methods insertEdge and deleteEdge did check that the corresponding edge does [not] exist in CFG by the moment of this update. I was certain that lazy mass update also preserves this behavior, but it doesn't. This is highly confusing. Not sure how it happens to pass the tests, but it does.

Dima will come up with alternative approach.

[dmakogon]

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. It's probably better to create an utility function for this, but need to check if it would work in the loop unroll code.

[mkazantsev]

Looks fine for me, but let's give other reviewers time to raise their concerns. And please run some extensive fuzzer testing for it.

[nikic]

LGTM

[dmakogon]

In D111611#3084537, @mkazantsev wrote:

Looks fine for me, but let's give other reviewers time to raise their concerns. And please run some extensive fuzzer testing for it.

The testing finished successfully.

llvm/lib/Transforms/Utils/LoopPeel.cpp

 //===- LoopPeel.cpp -------------------------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // Loop Peeling Utilities.
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Transforms/Utils/LoopPeel.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Analysis/Loads.h"
 #include "llvm/Analysis/LoopInfo.h"
 #include "llvm/Analysis/LoopIterator.h"
 #include "llvm/Analysis/ScalarEvolution.h"
 #include "llvm/Analysis/ScalarEvolutionExpressions.h"
 #include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/Dominators.h"

 /// \param LoopBlocks A helper for DFS-traversal of the loop.
 /// \param LVMap A value-map that maps instructions from the original loop to
 /// instructions in the last peeled-off iteration.
 static void cloneLoopBlocks(
     Loop *L, unsigned IterNumber, BasicBlock *InsertTop, BasicBlock *InsertBot,
     SmallVectorImpl<std::pair<BasicBlock *, BasicBlock *>> &ExitEdges,
     SmallVectorImpl<BasicBlock *> &NewBlocks, LoopBlocksDFS &LoopBlocks,
     ValueToValueMapTy &VMap, ValueToValueMapTy &LVMap, DominatorTree *DT,
     LoopInfo *LI, ArrayRef<MDNode *> LoopLocalNoAliasDeclScopes) {
   BasicBlock *Header = L->getHeader();
   BasicBlock *Latch = L->getLoopLatch();
   BasicBlock *PreHeader = L->getLoopPreheader();
 
   Function *F = Header->getParent();
   LoopBlocksDFS::RPOIterator BlockBegin = LoopBlocks.beginRPO();
   LoopBlocksDFS::RPOIterator BlockEnd = LoopBlocks.endRPO();
   Loop *ParentLoop = L->getParentLoop();
 
   // For each block in the original loop, create a new copy,
   // and update the value map with the newly created values.
   for (LoopBlocksDFS::RPOIterator BB = BlockBegin; BB != BlockEnd; ++BB) {
     BasicBlock *NewBB = CloneBasicBlock(*BB, VMap, ".peel", F);
     NewBlocks.push_back(NewBB);
 
     // If an original block is an immediate child of the loop L, its copy
     // is a child of a ParentLoop after peeling. If a block is a child of
     // a nested loop, it is handled in the cloneLoop() call below.
if (DT) {
         // VMap must contain entry for IDom, as the iteration order is RPO.
         DT->addNewBlock(NewBB, cast<BasicBlock>(VMap[IDom->getBlock()]));
       }
     }
   }
 
   {
     // Identify what other metadata depends on the cloned version. After
     // cloning, replace the metadata with the corrected version for both
     // memory instructions and noalias intrinsics.
     std::string Ext = (Twine("Peel") + Twine(IterNumber)).str();
     cloneAndAdaptNoAliasScopes(LoopLocalNoAliasDeclScopes, NewBlocks,
                                Header->getContext(), Ext);
   }
 
   // Recursively create the new Loop objects for nested loops, if any,
   // to preserve LoopInfo.
bool llvm::peelLoop(Loop *L, unsigned PeelCount, LoopInfo *LI,
   LoopBlocks.perform(LI);
 
   BasicBlock *Header = L->getHeader();
   BasicBlock *PreHeader = L->getLoopPreheader();
   BasicBlock *Latch = L->getLoopLatch();
   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;
     }
   }
 
   Function *F = Header->getParent();
 
   // Set up all the necessary basic blocks. It is convenient to split the
   // preheader into 3 parts - two blocks to anchor the peeled copy of the loop
   // body, and a new preheader for the "real" loop.
bool llvm::peelLoop(Loop *L, unsigned PeelCount, LoopInfo *LI,
   SmallVector<MDNode *, 6> LoopLocalNoAliasDeclScopes;
   identifyNoAliasScopesToClone(L->getBlocks(), LoopLocalNoAliasDeclScopes);
 
   // For each peeled-off iteration, make a copy of the loop.
   for (unsigned Iter = 0; Iter < PeelCount; ++Iter) {
     SmallVector<BasicBlock *, 8> NewBlocks;
     ValueToValueMapTy VMap;
 
     cloneLoopBlocks(L, Iter, InsertTop, InsertBot, ExitEdges, NewBlocks,
                     LoopBlocks, VMap, LVMap, DT, LI,
                     LoopLocalNoAliasDeclScopes);
 
     // Remap to use values from the current iteration instead of the
     // previous one.
     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
     }
 
     auto *LatchBRCopy = cast<BranchInst>(VMap[LatchBR]);
     updateBranchWeights(InsertBot, LatchBRCopy, ExitWeight, FallThroughWeight);
     // Remove Loop metadata from the latch branch instruction
     // because it is not the Loop's latch branch anymore.
     LatchBRCopy->setMetadata(LLVMContext::MD_loop, nullptr);
 
     InsertTop = InsertBot;
     InsertBot = SplitBlock(InsertBot, InsertBot->getTerminator(), DT, LI);
     InsertBot->setName(Header->getName() + ".peel.next");
 
#endif
 
   if (Loop *ParentLoop = L->getParentLoop())
     L = ParentLoop;
 
   // We modified the loop, update SE.
   SE->forgetTopmostLoop(L);
 
   // Finally DomtTree must be correct.
   assert(DT->verify(DominatorTree::VerificationLevel::Fast));
 
   // FIXME: Incrementally update loop-simplify
   simplifyLoop(L, DT, LI, SE, AC, nullptr, PreserveLCSSA);
 
   NumPeeled++;
 
   return true;
 }