[PM/LoopUnswitch] Begin teaching SimpleLoopUnswitch to use the new update API for dominators rather than doing manual, hacky updates.

[chandlerc]

Summary

This is just the first step, but in some ways the most important as it
moves the non-trivial unswitching to update the domtree rather than
fully recalculating it each time.

Subsequent patches should remove the custom update logic used by the
trivial unswitch and replace it with uses of the update API.

This also fixes a number of bugs I was seeing when testing non-trivial
unswitch due to it querying the quasi-correct dominator tree. Now the
tree is 100% correct and safe to query. That said, there are still more
bugs I can see with non-trivial unswitch just running over the test
suite, so more bugfix patches are needed as well.

[chandlerc]

Rebase and ping.

This is now the only patch needed for the new PM's unswitch to pass the LLVM
test suite with nontrivial unswitching enabled.

There may of course still be bugs here, but I think this is getting pretty
close.

[sanjoy]

llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp

886

The DT update mechanism gets rid of duplicates anyway (see LegalizeUpdates), so SuccSet should not be necessary (unless you expect a lot of duplicates at which point it can serve to save memory).

[fedor.sergeev]

Looks good to me.
Thanks for the effort.
This version finally passes all the tests on our internal pipeline!

llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp

 /// blocks to them but doesn't create the cloned `DominatorTree` structure and
 /// instead the caller must recompute an accurate DT. It *does* correctly
 /// update the `AssumptionCache` provided in `AC`.
 static BasicBlock *buildClonedLoopBlocks(
     Loop &L, BasicBlock *LoopPH, BasicBlock *SplitBB,
     ArrayRef<BasicBlock *> ExitBlocks, BasicBlock *ParentBB,
     BasicBlock *UnswitchedSuccBB, BasicBlock *ContinueSuccBB,
     const SmallPtrSetImpl<BasicBlock *> &SkippedLoopAndExitBlocks,
-    ValueToValueMapTy &VMap, AssumptionCache &AC, DominatorTree &DT,
-    LoopInfo &LI) {
+    ValueToValueMapTy &VMap,
+    SmallVectorImpl<DominatorTree::UpdateType> &DTUpdates, AssumptionCache &AC,
+    DominatorTree &DT, LoopInfo &LI) {
   SmallVector<BasicBlock *, 4> NewBlocks;
   NewBlocks.reserve(L.getNumBlocks() + ExitBlocks.size());
 
   // We will need to clone a bunch of blocks, wrap up the clone operation in
   // a helper.
   auto CloneBlock = [&](BasicBlock *OldBB) {
     // Clone the basic block and insert it before the new preheader.
     BasicBlock *NewBB = CloneBasicBlock(OldBB, VMap, ".us", OldBB->getParent());
     NewBB->moveBefore(LoopPH);
 
     // Record this block and the mapping.
     NewBlocks.push_back(NewBB);
     VMap[OldBB] = NewBB;
 
-    // Add the block to the domtree. We'll move it to the correct position
-    // below.
-    DT.addNewBlock(NewBB, SplitBB);
-
     return NewBB;
   };
 
   // First, clone the preheader.
   auto *ClonedPH = CloneBlock(LoopPH);
 
   // Then clone all the loop blocks, skipping the ones that aren't necessary.
   for (auto *LoopBB : L.blocks())
for (auto *ExitBB : ExitBlocks) {
 
     // Now clone the original exit block.
     auto *ClonedExitBB = CloneBlock(ExitBB);
     assert(ClonedExitBB->getTerminator()->getNumSuccessors() == 1 &&
            "Exit block should have been split to have one successor!");
     assert(ClonedExitBB->getTerminator()->getSuccessor(0) == MergeBB &&
            "Cloned exit block has the wrong successor!");
 
-    // Move the merge block's idom to be the split point as one exit is
-    // dominated by one header, and the other by another, so we know the split
-    // point dominates both. While the dominator tree isn't fully accurate, we
-    // want sub-trees within the original loop to be correctly reflect
-    // dominance within that original loop (at least) and that requires moving
-    // the merge block out of that subtree.
-    // FIXME: This is very brittle as we essentially have a partial contract on
-    // the dominator tree. We really need to instead update it and keep it
-    // valid or stop relying on it.
-    DT.changeImmediateDominator(MergeBB, SplitBB);
-
     // Remap any cloned instructions and create a merge phi node for them.
     for (auto ZippedInsts : llvm::zip_first(
              llvm::make_range(ExitBB->begin(), std::prev(ExitBB->end())),
              llvm::make_range(ClonedExitBB->begin(),
                               std::prev(ClonedExitBB->end())))) {
       Instruction &I = std::get<0>(ZippedInsts);
       Instruction &ClonedI = std::get<1>(ZippedInsts);
 
static BasicBlock *buildClonedLoopBlocks(
   // have spurious incoming values.
   for (auto *LoopBB : L.blocks())
     if (SkippedLoopAndExitBlocks.count(LoopBB))
       for (auto *SuccBB : successors(LoopBB))
         if (auto *ClonedSuccBB = cast_or_null<BasicBlock>(VMap.lookup(SuccBB)))
           for (PHINode &PN : ClonedSuccBB->phis())
             PN.removeIncomingValue(LoopBB, /*DeletePHIIfEmpty*/ false);
 
+  // Record the domtree updates for the new blocks.
+  SmallPtrSet<BasicBlock *, 4> SuccSet;
+  for (auto *ClonedBB : NewBlocks) {
+    for (auto *SuccBB : successors(ClonedBB))
+      if (SuccSet.insert(SuccBB).second)
+        DTUpdates.push_back({DominatorTree::Insert, ClonedBB, SuccBB});
+    SuccSet.clear();
+  }
+
   return ClonedPH;
 }
 
 /// Recursively clone the specified loop and all of its children.
 ///
 /// The target parent loop for the clone should be provided, or can be null if
 /// the clone is a top-level loop. While cloning, all the blocks are mapped
 /// with the provided value map. The entire original loop must be present in
#endif
     NonChildClonedLoops.push_back(cloneLoopNest(
         *ChildL, ExitLoopMap.lookup(ClonedChildHeader), VMap, LI));
   }
 
   // Return the main cloned loop if any.
   return ClonedL;
 }
 
-static void deleteDeadBlocksFromLoop(Loop &L, BasicBlock *DeadSubtreeRoot,
-                                     SmallVectorImpl<BasicBlock *> &ExitBlocks,
-                                     DominatorTree &DT, LoopInfo &LI) {
-  // Walk the dominator tree to build up the set of blocks we will delete here.
-  // The order is designed to allow us to always delete bottom-up and avoid any
-  // dangling uses.
-  SmallSetVector<BasicBlock *, 16> DeadBlocks;
-  DeadBlocks.insert(DeadSubtreeRoot);
-  for (int i = 0; i < (int)DeadBlocks.size(); ++i)
-    for (DomTreeNode *ChildN : *DT[DeadBlocks[i]]) {
-      // FIXME: This assert should pass and that means we don't change nearly
-      // as much below! Consider rewriting all of this to avoid deleting
-      // blocks. They are always cloned before being deleted, and so instead
-      // could just be moved.
-      // FIXME: This in turn means that we might actually be more able to
-      // update the domtree.
-      assert((L.contains(ChildN->getBlock()) ||
-              llvm::find(ExitBlocks, ChildN->getBlock()) != ExitBlocks.end()) &&
-             "Should never reach beyond the loop and exits when deleting!");
-      DeadBlocks.insert(ChildN->getBlock());
-    }
+static void
+deleteDeadBlocksFromLoop(Loop &L,
+                         const SmallVectorImpl<BasicBlock *> &DeadBlocks,
+  SmallPtrSet<BasicBlock *, 16> DeadBlockSet(DeadBlocks.begin(),
+                                             DeadBlocks.end());
 
   // Filter out the dead blocks from the exit blocks list so that it can be
   // used in the caller.
   llvm::erase_if(ExitBlocks,
-                 [&](BasicBlock *BB) { return DeadBlocks.count(BB); });
+                 [&](BasicBlock *BB) { return DeadBlockSet.count(BB); });
 
   // Remove these blocks from their successors.
   for (auto *BB : DeadBlocks)
     for (BasicBlock *SuccBB : successors(BB))
       SuccBB->removePredecessor(BB, /*DontDeleteUselessPHIs*/ true);
 
   // Walk from this loop up through its parents removing all of the dead blocks.
   for (Loop *ParentL = &L; ParentL; ParentL = ParentL->getParentLoop()) {
     for (auto *BB : DeadBlocks)
       ParentL->getBlocksSet().erase(BB);
     llvm::erase_if(ParentL->getBlocksVector(),
-                   [&](BasicBlock *BB) { return DeadBlocks.count(BB); });
+                   [&](BasicBlock *BB) { return DeadBlockSet.count(BB); });
   }
 
   // Now delete the dead child loops. This raw delete will clear them
   // recursively.
   llvm::erase_if(L.getSubLoopsVector(), [&](Loop *ChildL) {
-    if (!DeadBlocks.count(ChildL->getHeader()))
+    if (!DeadBlockSet.count(ChildL->getHeader()))
       return false;
 
     assert(llvm::all_of(ChildL->blocks(),
                         [&](BasicBlock *ChildBB) {
-                          return DeadBlocks.count(ChildBB);
+                          return DeadBlockSet.count(ChildBB);
                         }) &&
            "If the child loop header is dead all blocks in the child loop must "
            "be dead as well!");
     LI.destroy(ChildL);
     return true;
   });
 
-  // Remove the mappings for the dead blocks.
-  for (auto *BB : DeadBlocks)
+  // Remove the loop mappings for the dead blocks and drop all the references
+  // from these blocks to others to handle cyclic references as we start
+  // deleting the blocks themselves.
+  for (auto *BB : DeadBlocks) {
+    // Check that the dominator tree has already been updated.
+    assert(!DT.getNode(BB) && "Should already have cleared domtree!");
     LI.changeLoopFor(BB, nullptr);
-
-  // Drop all the references from these blocks to others to handle cyclic
-  // references as we start deleting the blocks themselves.
-  for (auto *BB : DeadBlocks)
     BB->dropAllReferences();
+  }
 
-  for (auto *BB : llvm::reverse(DeadBlocks)) {
-    DT.eraseNode(BB);
+  // Actually delete the blocks now that they've been fully unhooked from the
+  // IR.
+  for (auto *BB : DeadBlocks)
     BB->eraseFromParent();
-  }
-}
 
 /// Recompute the set of blocks in a loop after unswitching.
 ///
 /// This walks from the original headers predecessors to rebuild the loop. We
 /// take advantage of the fact that new blocks can't have been added, and so we
 /// filter by the original loop's blocks. This also handles potentially
 /// unreachable code that we don't want to explore but might be found examining
 /// the predecessors of the header.
static bool unswitchInvariantBranch(
   // between the unswitched versions, and we will have a new preheader for the
   // original loop.
   BasicBlock *SplitBB = L.getLoopPreheader();
   BasicBlock *LoopPH = SplitEdge(SplitBB, L.getHeader(), &DT, &LI);
 
   // Keep a mapping for the cloned values.
   ValueToValueMapTy VMap;
 
+  // Keep track of the dominator tree updates needed.
+  SmallVector<DominatorTree::UpdateType, 4> DTUpdates;
+
   // Build the cloned blocks from the loop.
   auto *ClonedPH = buildClonedLoopBlocks(
       L, LoopPH, SplitBB, ExitBlocks, ParentBB, UnswitchedSuccBB,
-      ContinueSuccBB, SkippedLoopAndExitBlocks, VMap, AC, DT, LI);
-
-  // Build the cloned loop structure itself. This may be substantially
-  // different from the original structure due to the simplified CFG. This also
-  // handles inserting all the cloned blocks into the correct loops.
-  SmallVector<Loop *, 4> NonChildClonedLoops;
-  Loop *ClonedL =
-      buildClonedLoops(L, ExitBlocks, VMap, LI, NonChildClonedLoops);
+      ContinueSuccBB, SkippedLoopAndExitBlocks, VMap, DTUpdates, AC, DT, LI);
 
   // Remove the parent as a predecessor of the unswitched successor.
   UnswitchedSuccBB->removePredecessor(ParentBB, /*DontDeleteUselessPHIs*/ true);
 
   // Now splice the branch from the original loop and use it to select between
   // the two loops.
   SplitBB->getTerminator()->eraseFromParent();
   SplitBB->getInstList().splice(SplitBB->end(), ParentBB->getInstList(), BI);
   BI.setSuccessor(ClonedSucc, ClonedPH);
   BI.setSuccessor(1 - ClonedSucc, LoopPH);
 
   // Create a new unconditional branch to the continuing block (as opposed to
   // the one cloned).
   BranchInst::Create(ContinueSuccBB, ParentBB);
 
+  // Before we update the dominator tree, collect the dead blocks if we're going
+  // to end up deleting the unswitched successor.
+  SmallVector<BasicBlock *, 16> DeadBlocks;
+  if (DeleteUnswitchedSucc) {
+    DeadBlocks.push_back(UnswitchedSuccBB);
+    for (int i = 0; i < (int)DeadBlocks.size(); ++i) {
+      // If we reach an exit block, stop recursing as the unswitched loop will
+      // end up reaching the merge block which we make the successor of the
+      // exit.
+      if (ExitBlockSet.count(DeadBlocks[i]))
+        continue;
+
+      // Insert the children that are within the loop or exit block set. Other
+      // children may reach out of the loop. While we don't expect these to be
+      // dead (as the unswitched clone should reach them) we don't try to prove
+      // that here.
+      for (DomTreeNode *ChildN : *DT[DeadBlocks[i]])
+        if (L.contains(ChildN->getBlock()) ||
+            ExitBlockSet.count(ChildN->getBlock()))
+          DeadBlocks.push_back(ChildN->getBlock());
+    }
+  }
+
+  // Add the remaining edges to our updates and apply them to get an up-to-date
+  // dominator tree. Note that this will cause the dead blocks above to be
+  // unreachable and no longer in the dominator tree.
+  DTUpdates.push_back({DominatorTree::Delete, ParentBB, UnswitchedSuccBB});
+  DTUpdates.push_back({DominatorTree::Insert, SplitBB, ClonedPH});
+  DT.applyUpdates(DTUpdates);
+
+  // Build the cloned loop structure itself. This may be substantially
+  // different from the original structure due to the simplified CFG. This also
+  // handles inserting all the cloned blocks into the correct loops.
+  SmallVector<Loop *, 4> NonChildClonedLoops;
+  Loop *ClonedL =
+      buildClonedLoops(L, ExitBlocks, VMap, LI, NonChildClonedLoops);
+
   // Delete anything that was made dead in the original loop due to
   // unswitching.
-  if (DeleteUnswitchedSucc)
-    deleteDeadBlocksFromLoop(L, UnswitchedSuccBB, ExitBlocks, DT, LI);
+  if (!DeadBlocks.empty())
+    deleteDeadBlocksFromLoop(L, DeadBlocks, ExitBlocks, DT, LI);
 
   SmallVector<Loop *, 4> HoistedLoops;
   bool IsStillLoop = rebuildLoopAfterUnswitch(L, ExitBlocks, LI, HoistedLoops);
 
-  // This will have completely invalidated the dominator tree. We can't easily
-  // bound how much is invalid because in some cases we will refine the
-  // predecessor set of exit blocks of the loop which can move large unrelated
-  // regions of code into a new subtree.
-  //
-  // FIXME: Eventually, we should use an incremental update utility that
-  // leverages the existing information in the dominator tree (and potentially
-  // the nature of the change) to more efficiently update things.
-  DT.recalculate(*SplitBB->getParent());
+  // This transformation has a high risk of corrupting the dominator tree, and
+  // the below steps to rebuild loop structures will result in hard to debug
+  // errors in that case so verify that the dominator tree is sane first.
+  // FIXME: Remove this when the bugs stop showing up and rely on existing
+  // verification steps.
+  assert(DT.verify(DominatorTree::VerificationLevel::Fast));
 
   // We can change which blocks are exit blocks of all the cloned sibling
   // loops, the current loop, and any parent loops which shared exit blocks
   // with the current loop. As a consequence, we need to re-form LCSSA for
   // them. But we shouldn't need to re-form LCSSA for any child loops.
   // FIXME: This could be made more efficient by tracking which exit blocks are
   // new, and focusing on them, but that isn't likely to be necessary.
   //