[LoopUnroll] Directly update DT instead of DTU.

[fhahn]

Summary

The scope of DT updates are very limited when unrolling loops: the DT
should only need updating for

• new blocks added
• exiting blocks we simplified branches

This can be done manually without too much extra work.
MergeBlockIntoPredecessor and changeToUnreachable also need to be
updated to support direct DT updates.

This fixes excessive time spent in DTU for same cases. In an internal
example, time spent in LoopUnroll with this patch goes from ~200s to 2s.

It also is slightly positive for CTMark:

• NewPM-O3: -0.13%
• NewPM-ReleaseThinLTO: -0.11%
• NewPM-ReleaseLTO-g: -0.13%

Notable improvements are mafft (~ -0.50%) and lencod (~ -0.30%), with no
workload regressed.

https://llvm-compile-time-tracker.com/compare.php?from=78a9ee7834331fb4360457cc565fa36f5452f7e0&to=687e08d011b0dc6d3edd223612761e44225c7537&stat=instructions:u

[nikic]

llvm/include/llvm/Transforms/Utils/BasicBlockUtils.h

100

Here and elsewhere, please add a comment -- normally, if there's a DT rather than DTU argument, I'd expect the DT to be used rather than updated.

llvm/lib/Transforms/Utils/BasicBlockUtils.cpp

239

Do we need to check whether PredNode / BBNode are null? A bit degenerate, but we might be merging two unreachable blocks, no?

[fhahn]

Document DT argument, check if DT nodes are null (=unreachable)

[fhahn]

llvm/include/llvm/Transforms/Utils/BasicBlockUtils.h

100

Good point, added a comment here and for changeToUnreachable

llvm/lib/Transforms/Utils/BasicBlockUtils.cpp

239

Added a check. I think it should be sufficient to check if PRedNode is reachable (!= nullptr), as BB must have PredBB as unique predecessor.

Also added a check to updateDominatorTreeUsingPredecessors.

[nikic]

llvm/lib/Transforms/Utils/Local.cpp

3530 ↗

(On Diff #488227)

What about the case where BB has a (now also unreachable) successor? Don't we have to erase all those nodes as well?

[fhahn]

Update to remove all nodes connected to node that became unreachable.

[fhahn]

llvm/lib/Transforms/Utils/Local.cpp

3530 ↗

(On Diff #488227)

I am not sure what the exact invariant in the DT is, but I updated the code to remove all connected nodes to be on the safe side. Same in MergeBlockIntoPredecessor

[asbirlea]

I'm curious where all the time is spent in the 200s to 2s case. Is Eager strategy just as slow?

[fhahn]

In D141487#4045554, @asbirlea wrote:

I'm curious where all the time is spent in the 200s to 2s case. Is Eager strategy just as slow?

The DTU here was using the lazy strategy.

A profile is below (note that this is on a different (slower) machine, so it takes longer than 200s :)) . AFACIT most of the time is spent discovering reachability after deleting edges. I think the main reason manual updates are much quicker here is that we can use the loop structure to avoid determining reachability; after updating the DT for nodes connected to the exits, no other nodes in the DT should need updating.

8.06 min   99.9%	0 s	 	             llvm::UnrollLoop(llvm::Loop*, llvm::UnrollLoopOptions, llvm::LoopInfo*, llvm::ScalarEvolution*, llvm::DominatorTree*, llvm::AssumptionCache*, llvm::TargetTransformInfo const*, llvm::OptimizationRemarkEmitter*, bool, llvm::Loop**)
8.04 min   99.8%	0 s	 	              llvm::DomTreeUpdater::getDomTree()
8.04 min   99.7%	0 s	 	               llvm::DomTreeUpdater::applyDomTreeUpdates()
8.04 min   99.7%	0 s	 	                llvm::DominatorTreeBase<llvm::BasicBlock, false>::applyUpdates(llvm::ArrayRef<llvm::cfg::Update<llvm::BasicBlock*> >)
8.04 min   99.7%	0 s	 	                 void llvm::DomTreeBuilder::ApplyUpdates<llvm::DominatorTreeBase<llvm::BasicBlock, false> >(llvm::DominatorTreeBase<llvm::BasicBlock, false>&, llvm::GraphDiff<llvm::DominatorTreeBase<llvm::BasicBlock, false>::NodePtr, llvm::DominatorTreeBase<llvm::BasicBlock, false>::IsPostDominator>&, llvm::GraphDiff<llvm::DominatorTreeBase<llvm::BasicBlock, false>::NodePtr, llvm::DominatorTreeBase<llvm::BasicBlock, false>::IsPostDominator>*)
8.04 min   99.7%	2.00 ms	 	                  llvm::DomTreeBuilder::SemiNCAInfo<llvm::DominatorTreeBase<llvm::BasicBlock, false> >::ApplyUpdates(llvm::DominatorTreeBase<llvm::BasicBlock, false>&, llvm::GraphDiff<llvm::BasicBlock*, false>&, llvm::GraphDiff<llvm::BasicBlock*, false>*)
8.04 min   99.7%	3.00 ms	 	                   llvm::DomTreeBuilder::SemiNCAInfo<llvm::DominatorTreeBase<llvm::BasicBlock, false> >::ApplyNextUpdate(llvm::DominatorTreeBase<llvm::BasicBlock, false>&, llvm::DomTreeBuilder::SemiNCAInfo<llvm::DominatorTreeBase<llvm::BasicBlock, false> >::BatchUpdateInfo&)
7.10 min   88.0%	6.00 ms	 	                    llvm::DomTreeBuilder::SemiNCAInfo<llvm::DominatorTreeBase<llvm::BasicBlock, false> >::DeleteEdge(llvm::DominatorTreeBase<llvm::BasicBlock, false>&, llvm::DomTreeBuilder::SemiNCAInfo<llvm::DominatorTreeBase<llvm::BasicBlock, false> >::BatchUpdateInfo*, llvm::BasicBlock*, llvm::BasicBlock*)
7.09 min   88.0%	5.00 ms	 	                     llvm::DomTreeBuilder::SemiNCAInfo<llvm::DominatorTreeBase<llvm::BasicBlock, false> >::DeleteReachable(llvm::DominatorTreeBase<llvm::BasicBlock, false>&, llvm::DomTreeBuilder::SemiNCAInfo<llvm::DominatorTreeBase<llvm::BasicBlock, false> >::BatchUpdateInfo*, llvm::DomTreeNodeBase<llvm::BasicBlock>*, llvm::DomTreeNodeBase<llvm::BasicBlock>*)
4.61 min   57.2%	2.88 s	 	                      unsigned int llvm::DomTreeBuilder::SemiNCAInfo<llvm::DominatorTreeBase<llvm::BasicBlock, false> >::runDFS<false, llvm::DomTreeBuilder::SemiNCAInfo<llvm::DominatorTreeBase<llvm::BasicBlock, false> >::DeleteReachable(llvm::DominatorTreeBase<llvm::BasicBlock, false>&, llvm::DomTreeBuilder::SemiNCAInfo<llvm::DominatorTreeBase<llvm::BasicBlock, false> >::BatchUpdateInfo*, llvm::DomTreeNodeBase<llvm::BasicBlock>*, llvm::DomTreeNodeBase<llvm::BasicBlock>*)::'lambda'(llvm::BasicBlock*, llvm::BasicBlock*)>(llvm::BasicBlock*, unsigned int, llvm::DomTreeBuilder::SemiNCAInfo<llvm::DominatorTreeBase<llvm::BasicBlock, false> >::DeleteReachable(llvm::DominatorTreeBase<llvm::BasicBlock, false>&, llvm::DomTreeBuilder::SemiNCAInfo<llvm::DominatorTreeBase<llvm::BasicBlock, false> >::BatchUpdateInfo*, llvm::DomTreeNodeBase<llvm::BasicBlock>*, llvm::DomTreeNodeBase<llvm::BasicBlock>*)::'lambda'(llvm::BasicBlock*, llvm::BasicBlock*), unsigned int, llvm::DenseMap<llvm::BasicBlock*, unsigned int, llvm::DenseMapInfo<llvm::BasicBlock*, void>, llvm::detail::DenseMapPair<llvm::BasicBlock*, unsigned int> > const*)
1.52 min   18.9%	5.86 s	 	                      llvm::DomTreeBuilder::SemiNCAInfo<llvm::DominatorTreeBase<llvm::BasicBlock, false> >::runSemiNCA(llvm::DominatorTreeBase<llvm::BasicBlock, false>&, unsigned int)
41.53 s    8.5%	1.07 s	 	                      llvm::DomTreeBuilder::SemiNCAInfo<llvm::DominatorTreeBase<llvm::BasicBlock, false> >::reattachExistingSubtree(llvm::DominatorTreeBase<llvm::BasicBlock, false>&, llvm::DomTreeNodeBase<llvm::BasicBlock>*)
15.22 s    3.1%	4.00 ms	 	                      llvm::DomTreeBuilder::SemiNCAInfo<llvm::DominatorTreeBase<llvm::BasicBlock, false> >::~SemiNCAInfo()
452.00 ms    0.0%	0 s	 	                      0xfffffffffffffffe
64.00 ms    0.0%	64.00 ms	 	                      DYLD-STUB$$llvm::DominatorTreeBase<llvm::BasicBlock, false>::getNode(llvm::BasicBlock const*) const
31.00 ms    0.0%	0 s	 	                      llvm::DomTreeBuilder::SemiNCAInfo<llvm::DominatorTreeBase<llvm::BasicBlock, false> >::SemiNCAInfo(llvm::DomTreeBuilder::SemiNCAInfo<llvm::DominatorTreeBase<llvm::BasicBlock, false> >::BatchUpdateInfo*)
29.00 ms    0.0%	29.00 ms	 	                      DYLD-STUB$$llvm::SmallVectorBase<unsigned int>::empty() const
3.00 ms    0.0%	0 s	 	                      llvm::DominatorTreeBase<llvm::BasicBlock, false>::findNearestCommonDominator(llvm::BasicBlock*, llvm::BasicBlock*) const
1.00 ms    0.0%	1.00 ms	 	                      DYLD-STUB$$llvm::DomTreeNodeBase<llvm::BasicBlock>::getLevel() const
1.00 ms    0.0%	1.00 ms	 	                      DYLD-STUB$$llvm::DomTreeNodeBase<llvm::BasicBlock>::setIDom(llvm::DomTreeNodeBase<llvm::BasicBlock>*)
1.00 ms    0.0%	0 s	 	                      llvm::DominatorTreeBase<llvm::BasicBlock, false>::getNode(llvm::BasicBlock const*) const
90.00 ms    0.0%	1.00 ms	 	                     llvm::DomTreeBuilder::SemiNCAInfo<llvm::DominatorTreeBase<llvm::BasicBlock, false> >::DeleteEdge(llvm::DominatorTreeBase<llvm::BasicBlock, false>&, llvm::DomTreeBuilder::SemiNCAInfo<llvm::DominatorTreeBase<llvm::BasicBlock, false> >::BatchUpdateInfo*, llvm::BasicBlock*, llvm::BasicBlock*)::'lambda'(llvm::BasicBlock*, llvm::BasicBlock*)::operator()(llvm::BasicBlock*, llvm::BasicBlock*) const
35.00 ms    0.0%	1.00 ms	 	                     llvm::DomTreeBuilder::SemiNCAInfo<llvm::DominatorTreeBase<llvm::BasicBlock, false> >::DeleteUnreachable(llvm::DominatorTreeBase<llvm::BasicBlock, false>&, llvm::DomTreeBuilder::SemiNCAInfo<llvm::DominatorTreeBase<llvm::BasicBlock, false> >::BatchUpdateInfo*, llvm::DomTreeNodeBase<llvm::BasicBlock>*)
24.00 ms    0.0%	1.00 ms	 	                     llvm::DomTreeBuilder::SemiNCAInfo<llvm::DominatorTreeBase<llvm::BasicBlock, false> >::HasProperSupport(llvm::DominatorTreeBase<llvm::BasicBlock, false>&, llvm::DomTreeBuilder::SemiNCAInfo<llvm::DominatorTreeBase<llvm::BasicBlock, false> >::BatchUpdateInfo*, llvm::DomTreeNodeBase<llvm::BasicBlock>*)
23.00 ms    0.0%	0 s	 	                     llvm::DominatorTreeBase<llvm::BasicBlock, false>::findNearestCommonDominator(llvm::BasicBlock*, llvm::BasicBlock*) const
11.00 ms    0.0%	2.00 ms	 	                     llvm::DominatorTreeBase<llvm::BasicBlock, false>::getNode(llvm::BasicBlock const*) const
1.00 ms    0.0%	1.00 ms	 	                     llvm::DomTreeNodeBase<llvm::BasicBlock>::getIDom() const
56.59 s   11.7%	1.00 ms	 	                    llvm::DomTreeBuilder::SemiNCAInfo<llvm::DominatorTreeBase<llvm::BasicBlock, false> >::InsertEdge(llvm::DominatorTreeBase<llvm::BasicBlock, false>&, llvm::DomTreeBuilder::SemiNCAInfo<llvm::DominatorTreeBase<llvm::BasicBlock, false> >::BatchUpdateInfo*, llvm::BasicBlock*, llvm::BasicBlock*)
54.00 ms    0.0%	7.00 ms	 	                    llvm::GraphDiff<llvm::BasicBlock*, false>::popUpdateForIncrementalUpdates()
1.00 ms    0.0%	1.00 ms	 	                    llvm::cfg::Update<llvm::BasicBlock*>::getFrom() const

[kuhar]

Some thoughts:

1. If possible, it'd be great if we could have a repro to use as a benchmark for the updater. And experiment with. In the past, I found that manually instrumenting the construction & updater algorithms with event counters (e.g., number of nodes visited by DFS, number of updates in-flight, initial/final tree size) to give more insight than perf traces.
2. I find it very hard to prove that manual updates are correct. We found a lot of very rare corner cases that required combination of things like infinite loops, multi-edges, irreducible loops, etc., that broke seemingly correct manual update functions. This makes me strongly prefer not to add more manual update code, unless we have a very good justification. Here, 200s -> 2s seems convincing for me if we cannot find other workarounds.
3. I haven't looked into this specific case closely, but I'll list some generic alternative ideas:
   • See if we can determine to recalculate in this case and if this would be faster than running the updater.
   • Detect the unreachable node case in the updater and handle it there.
4. Could you add assertions to verify the tree just before and after applying the updates manually? This can be guarded with EXPENSIVE_CHECKS.

llvm/include/llvm/Transforms/Utils/BasicBlockUtils.h

93

nit: s/use it to update the dominator tree directly/update it directly

llvm/include/llvm/Transforms/Utils/Local.h

345 ↗

(On Diff #488282)

nit: also here

501 ↗

(On Diff #488282)

nit: s/it's/its

501–503 ↗

(On Diff #488282)

It's not clear to me what the exact assumptions are here. I'm thinking about things like predecessors having edges to other predecessors and other corner cases.

In addition, could you add some ascii diagrams to illustrate these two cases (new/deleted)?

llvm/lib/Transforms/Utils/BasicBlockUtils.cpp

242

nit: either BBNode->children() in the loop directly or use llvm::to_vector?

llvm/lib/Transforms/Utils/LoopUnroll.cpp

699

nit: also here

805–809

nit: could you add comments with parameter names?

[kuhar]

llvm/include/llvm/Transforms/Utils/Local.h

504 ↗

(On Diff #488282)

Also, could we have some unit tests? I'm thinking about both having some concrete examples that we can look at and possibly death test to make sure that the assertions catch what we wanted to assume.

[fhahn]

Thanks for taking a look!

In D141487#4048701, @kuhar wrote:

Some thoughts:

1. If possible, it'd be great if we could have a repro to use as a benchmark for the updater. And experiment with. In the past, I found that manually instrumenting the construction & updater algorithms with event counters (e.g., number of nodes visited by DFS, number of updates in-flight, initial/final tree size) to give more insight than perf traces.

I'll work with the person who reported the issue to see if it is possible to come up with a case that can be easily shared publicly.

The general structure is a function with ~70 smallish consecutive loops, where each of them can be fully unrolled (with trip counts of 32 or somethingn.

2. I find it very hard to prove that manual updates are correct. We found a lot of very rare corner cases that required combination of things like infinite loops, multi-edges, irreducible loops, etc., that broke seemingly correct manual update functions. This makes me strongly prefer not to add more manual update code, unless we have a very good justification. Here, 200s -> 2s seems convincing for me if we cannot find other workarounds.

Agreed that it is very difficult to make sure manual updates are correct in general. In the loop-unroll case, the types of changes we make are limited, which hopefully makes it more tractable.

3. I haven't looked into this specific case closely, but I'll list some generic alternative ideas:
   • See if we can determine to recalculate in this case and if this would be faster than running the updater.

Recalculating the DTs from scratch is much quicker than DTU (~18s), but still noticeable slower than manual updates (~2s).

1. • Detect the unreachable node case in the updater and handle it there.
2. Could you add assertions to verify the tree just before and after applying the updates manually? This can be guarded with EXPENSIVE_CHECKS.

There are already DT verification assertions just before the begin of the manual updates and at the end. I can add some in-between as well if you think that's beneficial, but I am a bit worried about doing so too liberally even with expensive checks.

llvm/include/llvm/Transforms/Utils/BasicBlockUtils.h

93

I simplified it a bit further: If \p DT is not nullptr, update it directly;

llvm/include/llvm/Transforms/Utils/Local.h

345 ↗

(On Diff #488282)

updated, thanks!

501 ↗

(On Diff #488282)

fixed, thanks!

501–503 ↗

(On Diff #488282)

I updated the comment to hopefully be more clear. The key thing is that it assumes the dominator tree to be up-to-date for all predecessors:

The function assumes the dominator tree is valid for all predecessors and sets \p BB's immediate dominator to the nearest common dominator of all predecessors. \p BB cannot be a predecessor of itself.

I also added an assertion that BB cannot have itself as predecessor; in that instance, the assumption that the DT is up-to-date for all predecessors doesn't hold, as we just asked to update the DT for BB.

In addition, could you add some ascii diagrams to illustrate these two cases (new/deleted)?

I am not sure if that would help too much as the main requirement is that the DT is up-to-date for the predecessors, hopefully the comment update adds enough info.

504 ↗

(On Diff #488282)

I added some small example tests in llvm/unittests/Transforms/Utils/LocalTest.cpp. Please let me know if you have other cases in mind. I also built various external binaries and clang bootstrap with DT verification and saw no issues.

llvm/lib/Transforms/Utils/BasicBlockUtils.cpp

242

Changing the immediate dominator of a child of BB might invalidate the children iterator I think, so we need a vector. Updated to use to_vector, thanks for point out that utility, I wasn't aware of it.

llvm/lib/Transforms/Utils/LoopUnroll.cpp

699

updated, thanks!

805–809

Done, thanks!

[kuhar]

Recalculating the DTs from scratch is much quicker than DTU (~18s), but still noticeable slower than manual updates (~2s).

It would be great if we could detect those cases and teach DTU to trigger this recalculation. Would 18s be an acceptable fix for your use case?

There are already DT verification assertions just before the begin of the manual updates and at the end. I can add some in-between as well if you think that's beneficial, but I am a bit worried about doing so too liberally even with expensive checks.

Ack.

llvm/unittests/Transforms/Utils/LocalTest.cpp

1092 ↗

(On Diff #489015)

FYI, I added a utility for writing similar tests in the past (CFGBuilder), not sure if we could link with it here and if it would simplify the test code

[fhahn]

I took a step back and removed the need to add and use a new updateDominatorTreeUsingPredecessors helper. I discovered cases in general where we miss updates. While this was not an issue for the LoopUnroll test case, it seems dangerous to expose it.

The current patch switches to a hybrid approach: if there's a single exiting node, it should be sufficient to directly update the immediate dominators of the original exiting block and move them to the first actual exiting block of the unrolled loop. for all other cases, use the DTU.

After D141810, only MergeBlockIntoPredecessor needs udpating, and here the update should also be a straight-forward transfer of immediate dominators. The latest version achieves the same CTMark gains, while it should be much more limited in scope.

In D141487#4051794, @kuhar wrote:

Recalculating the DTs from scratch is much quicker than DTU (~18s), but still noticeable slower than manual updates (~2s).

It would be great if we could detect those cases and teach DTU to trigger this recalculation. Would 18s be an acceptable fix for your use case?

It would be an improvement, but LLVM performs loop-unrolling very aggressively and I think it is very desirable to make sure the code to update the DT here is as fast as possible. The latest version of the patch is even slightly faster.

I also experimented with cut-offs for the DTU based on the number of nodes it needs to traverse in its DFS, which yielded good results for the slow case, but was negative on CTMark :(

[fhahn]

llvm/lib/Transforms/Utils/Local.cpp

3530 ↗

(On Diff #488227)

I undid the change in MergeBlockIntoPredecessor. It is definitely not needed there, as the successors of the original BB are transferred to the predecessor.

[nikic]

llvm/lib/Transforms/Utils/BasicBlockUtils.cpp

243

Doesn't this need to check whether the old idom was BB? Couldn't it be higher up?

[fhahn]

llvm/lib/Transforms/Utils/BasicBlockUtils.cpp

243

I might be missing something here, but BBNode->children() iterates over the nodes which have BB as idom.

[kuhar]

This looks fine to me but I haven't looked carefully at the surrounding code. Ideally, I think this should albo be reviewed by someone familiar with the loop unrolling implementation.

llvm/lib/Transforms/Utils/LoopUnroll.cpp

774

nit: we can use structured bindings here: auto& [Something, ExitInfo] = *ExitInfos.begin();

[fhahn]

Thanks for taking a look! Updated to use structured bindings. I am planning to land this tomorrow unless there are additional comments.

[aeubanks]

this causes opt -passes=loop-unroll-full to crash on

define void @foo() {
bb:
  br label %bb1
    
bb1:                                              ; preds = %bb1, %bb1, %bb
  switch i1 true, label %bb1 [
    i1 true, label %bb2
    i1 false, label %bb1
  ]
    
 bb2:                                              ; preds = %bb1
  ret void
}

so I've reverted this for now

[raj.khem]

In D141487#4067417, @aeubanks wrote:

this causes opt -passes=loop-unroll-full to crash on

define void @foo() {
bb:
  br label %bb1
    
bb1:                                              ; preds = %bb1, %bb1, %bb
  switch i1 true, label %bb1 [
    i1 true, label %bb2
    i1 false, label %bb1
  ]
    
 bb2:                                              ; preds = %bb1
  ret void
}

so I've reverted this for now

I also reduced a clang segfault to this commit while compiling nettle library for RISCV64, The test script and case are here https://uclibc.org/~kraj/sexp-conv-5be224.zip

[fhahn]

Thanks for the revert + test case! The issue was that some Exiting blocks may not be added to ExitInfos if they have a non-branch terminator. I added an extra check for ExitInfo's size and recommitted the patch. A test case has been added in f92b35392ed8e4631

[fhahn]

Recommitted as https://reviews.llvm.org/rG12dd3a7b54cd

llvm/include/llvm/Transforms/Utils/BasicBlockUtils.h

 
 /// Attempts to merge a block into its predecessor, if possible. The return
 /// value indicates success or failure.
 /// By default do not merge blocks if BB's predecessor has multiple successors.
 /// If PredecessorWithTwoSuccessors = true, the blocks can only be merged
 /// if BB's Pred has a branch to BB and to AnotherBB, and BB has a single
 /// successor Sing. In this case the branch will be updated with Sing instead of
 /// BB, and BB will still be merged into its predecessor and removed.
+/// If \p DT is not nullptr, update it directly; in that case, DTU must be
+/// nullptr.
 bool MergeBlockIntoPredecessor(BasicBlock *BB, DomTreeUpdater *DTU = nullptr,
                                LoopInfo *LI = nullptr,
                                MemorySSAUpdater *MSSAU = nullptr,
                                MemoryDependenceResults *MemDep = nullptr,
-                               bool PredecessorWithTwoSuccessors = false);
+                               bool PredecessorWithTwoSuccessors = false,
+                               DominatorTree *DT = nullptr);
 
 /// Merge block(s) sucessors, if possible. Return true if at least two
 /// of the blocks were merged together.
 /// In order to merge, each block must be terminated by an unconditional
 /// branch. If L is provided, then the blocks merged into their predecessors
 /// must be in L. In addition, This utility calls on another utility:
 /// MergeBlockIntoPredecessor. Blocks are successfully merged when the call to
 /// MergeBlockIntoPredecessor returns true.

llvm/lib/Transforms/Utils/BasicBlockUtils.cpp

if (PHINode *PN = dyn_cast_or_null<PHINode>(PHIs[i].operator Value*()))
       Changed |= RecursivelyDeleteDeadPHINode(PN, TLI, MSSAU);
 
   return Changed;
 }
 
 bool llvm::MergeBlockIntoPredecessor(BasicBlock *BB, DomTreeUpdater *DTU,
                                      LoopInfo *LI, MemorySSAUpdater *MSSAU,
                                      MemoryDependenceResults *MemDep,
-                                     bool PredecessorWithTwoSuccessors) {
+                                     bool PredecessorWithTwoSuccessors,
+                                     DominatorTree *DT) {
   if (BB->hasAddressTaken())
     return false;
 
   // Can't merge if there are multiple predecessors, or no predecessors.
   BasicBlock *PredBB = BB->getUniquePredecessor();
   if (!PredBB) return false;
 
   // Don't break self-loops.
bool llvm::MergeBlockIntoPredecessor(BasicBlock *BB, DomTreeUpdater *DTU,
   if (isa<PHINode>(BB->front())) {
     for (PHINode &PN : BB->phis())
       if (!isa<PHINode>(PN.getIncomingValue(0)) ||
           cast<PHINode>(PN.getIncomingValue(0))->getParent() != BB)
         IncomingValues.push_back(PN.getIncomingValue(0));
     FoldSingleEntryPHINodes(BB, MemDep);
   }
 
+  if (DT) {
+    assert(!DTU && "cannot use both DT and DTU for updates");
+    DomTreeNode *PredNode = DT->getNode(PredBB);
+    DomTreeNode *BBNode = DT->getNode(BB);
+    if (PredNode) {
+      assert(BBNode && "PredNode unreachable but BBNode reachable?");
+      for (DomTreeNode *C : to_vector(BBNode->children()))
+        C->setIDom(PredNode);
+    }
+  }
   // DTU update: Collect all the edges that exit BB.
   // These dominator edges will be redirected from Pred.
   std::vector<DominatorTree::UpdateType> Updates;
   if (DTU) {
+    assert(!DT && "cannot use both DT and DTU for updates");
     // To avoid processing the same predecessor more than once.
     SmallPtrSet<BasicBlock *, 8> SeenSuccs;
     SmallPtrSet<BasicBlock *, 2> SuccsOfPredBB(succ_begin(PredBB),
                                                succ_end(PredBB));
     Updates.reserve(Updates.size() + 2 * succ_size(BB) + 1);
     // Add insert edges first. Experimentally, for the particular case of two
     // blocks that can be merged, with a single successor and single predecessor
     // respectively, it is beneficial to have all insert updates first. Deleting
if (LI)
     LI->removeBlock(BB);
 
   if (MemDep)
     MemDep->invalidateCachedPredecessors();
 
   if (DTU)
     DTU->applyUpdates(Updates);
 
+  if (DT) {
+    assert(succ_empty(BB) &&
+           "successors should have been transferred to PredBB");
+    DT->eraseNode(BB);
+  }
+
   // Finally, erase the old block and update dominator info.
   DeleteDeadBlock(BB, DTU);
 
   return true;
 }
 
 bool llvm::MergeBlockSuccessorsIntoGivenBlocks(
     SmallPtrSetImpl<BasicBlock *> &MergeBlocks, Loop *L, DomTreeUpdater *DTU,

llvm/lib/Transforms/Utils/LoopUnroll.cpp

LoopUnrollResult llvm::UnrollLoop(Loop *L, UnrollLoopOptions ULO, LoopInfo *LI,
   if (MaxTripCount && ULO.Count > MaxTripCount)
     ULO.Count = MaxTripCount;
 
   struct ExitInfo {
     unsigned TripCount;
     unsigned TripMultiple;
     unsigned BreakoutTrip;
     bool ExitOnTrue;
+    BasicBlock *FirstExitingBlock = nullptr;
     SmallVector<BasicBlock *> ExitingBlocks;
   };
   DenseMap<BasicBlock *, ExitInfo> ExitInfos;
   SmallVector<BasicBlock *, 4> ExitingBlocks;
   L->getExitingBlocks(ExitingBlocks);
   for (auto *ExitingBlock : ExitingBlocks) {
     // The folding code is not prepared to deal with non-branch instructions
     // right now.
for (auto *BB : OriginalLoopBlocks) {
       for (auto *ChildBB : ChildrenToUpdate)
         DT->changeImmediateDominator(ChildBB, NewIDom);
     }
   }
 
   assert(!UnrollVerifyDomtree ||
          DT->verify(DominatorTree::VerificationLevel::Fast));
 
-  DomTreeUpdater DTU(DT, DomTreeUpdater::UpdateStrategy::Lazy);
-
+  SmallVector<DominatorTree::UpdateType> DTUpdates;
   auto SetDest = [&](BasicBlock *Src, bool WillExit, bool ExitOnTrue) {
     auto *Term = cast<BranchInst>(Src->getTerminator());
     const unsigned Idx = ExitOnTrue ^ WillExit;
     BasicBlock *Dest = Term->getSuccessor(Idx);
     BasicBlock *DeadSucc = Term->getSuccessor(1-Idx);
 
     // Remove predecessors from all non-Dest successors.
     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}});
+    DTUpdates.emplace_back(DominatorTree::Delete, Src, DeadSucc);
   };
 
   auto WillExit = [&](const ExitInfo &Info, unsigned i, unsigned j,
                       bool IsLatch) -> std::optional<bool> {
     if (CompletelyUnroll) {
       if (PreserveOnlyFirst) {
         if (i == 0)
           return std::nullopt;
         return j == 0;
if (j != Info.BreakoutTrip &&
       // unconditional branch for some iterations.
       return false;
     }
     return std::nullopt;
   };
 
   // Fold branches for iterations where we know that they will exit or not
   // exit.
-  for (const auto &Pair : ExitInfos) {
-    const ExitInfo &Info = Pair.second;
+  for (auto &Pair : ExitInfos) {
+    ExitInfo &Info = Pair.second;
     for (unsigned i = 0, e = Info.ExitingBlocks.size(); i != e; ++i) {
       // The branch destination.
       unsigned j = (i + 1) % e;
       bool IsLatch = Pair.first == LatchBlock;
       std::optional<bool> KnownWillExit = WillExit(Info, i, j, IsLatch);
-      if (!KnownWillExit)
+      if (!KnownWillExit) {
+        if (!Info.FirstExitingBlock)
+          Info.FirstExitingBlock = Info.ExitingBlocks[i];
         continue;
+      }
 
       // We don't fold known-exiting branches for non-latch exits here,
       // because this ensures that both all loop blocks and all exit blocks
       // remain reachable in the CFG.
       // TODO: We could fold these branches, but it would require much more
       // sophisticated updates to LoopInfo.
-      if (*KnownWillExit && !IsLatch)
+      if (*KnownWillExit && !IsLatch) {
+        if (!Info.FirstExitingBlock)
+          Info.FirstExitingBlock = Info.ExitingBlocks[i];
         continue;
+      }
 
       SetDest(Info.ExitingBlocks[i], *KnownWillExit, Info.ExitOnTrue);
     }
   }
 
+  DomTreeUpdater DTU(DT, DomTreeUpdater::UpdateStrategy::Lazy);
+  DomTreeUpdater *DTUToUse = &DTU;
+  if (ExitingBlocks.size() == 1) {
+    // Manually update the DT if there's a single exiting node. In that case
+    // there's a single exit node and it is sufficient to update the nodes
+    // immediately dominated by the original exiting block. They will become
+    // dominated by the first exiting block that leaves the loop after
+    // unrolling. Note that the CFG inside the loop does not change, so there's
+    // no need to update the DT inside the unrolled loop.
+    DTUToUse = nullptr;
+    auto &[OriginalExit, Info] = *ExitInfos.begin();
+    if (!Info.FirstExitingBlock)
+      Info.FirstExitingBlock = Info.ExitingBlocks.back();
+    for (auto *C : to_vector(DT->getNode(OriginalExit)->children())) {
+      if (L->contains(C->getBlock()))
+        continue;
+      C->setIDom(DT->getNode(Info.FirstExitingBlock));
+    }
+  } else {
+    DTU.applyUpdates(DTUpdates);
+  }
+
   // When completely unrolling, the last latch becomes unreachable.
   if (!LatchIsExiting && CompletelyUnroll) {
     // There is no need to update the DT here, because there must be a unique
     // latch. Hence if the latch is not exiting it must directly branch back to
     // the original loop header and does not dominate any nodes.
     assert(LatchBlock->getSingleSuccessor() && "Loop with multiple latches?");
     changeToUnreachable(Latches.back()->getTerminator(), PreserveLCSSA);
   }
 
   // Merge adjacent basic blocks, if possible.
   for (BasicBlock *Latch : Latches) {
     BranchInst *Term = dyn_cast<BranchInst>(Latch->getTerminator());
     assert((Term ||
             (CompletelyUnroll && !LatchIsExiting && Latch == Latches.back())) &&
            "Need a branch as terminator, except when fully unrolling with "
            "unconditional latch");
     if (Term && Term->isUnconditional()) {
       BasicBlock *Dest = Term->getSuccessor(0);
       BasicBlock *Fold = Dest->getUniquePredecessor();
-      if (MergeBlockIntoPredecessor(Dest, &DTU, LI)) {
+      if (MergeBlockIntoPredecessor(Dest, /*DTU=*/DTUToUse, LI,
+                                    /*MSSAU=*/nullptr, /*MemDep=*/nullptr,
+                                    /*PredecessorWithTwoSuccessors=*/false,
+                                    DTUToUse ? nullptr : DT)) {
         // Dest has been folded into Fold. Update our worklists accordingly.
         std::replace(Latches.begin(), Latches.end(), Dest, Fold);
         llvm::erase_value(UnrolledLoopBlocks, Dest);
       }
     }
   }
+
+  if (DTUToUse) {
     // Apply updates to the DomTree.
     DT = &DTU.getDomTree();
-
+  }
   assert(!UnrollVerifyDomtree ||
          DT->verify(DominatorTree::VerificationLevel::Fast));
 
   // At this point, the code is well formed.  We now simplify the unrolled loop,
   // doing constant propagation and dead code elimination as we go.
   simplifyLoopAfterUnroll(L, !CompletelyUnroll && ULO.Count > 1, LI, SE, DT, AC,
                           TTI);