diff --git a/llvm/lib/Transforms/Utils/LoopUnroll.cpp b/llvm/lib/Transforms/Utils/LoopUnroll.cpp --- a/llvm/lib/Transforms/Utils/LoopUnroll.cpp +++ b/llvm/lib/Transforms/Utils/LoopUnroll.cpp @@ -708,21 +708,49 @@ 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 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(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 { @@ -762,55 +790,6 @@ } } - // 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 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(Term) && cast(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)