10 changes: 7 additions & 3 deletions llvm/lib/Transforms/Vectorize/VPlan.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -465,7 +465,8 @@ void VPlan::execute(VPTransformState *State) {

// We do not attempt to preserve DT for outer loop vectorization currently.
if (!EnableVPlanNativePath)
updateDominatorTree(State->DT, VectorPreHeaderBB, VectorLatchBB);
updateDominatorTree(State->DT, VectorPreHeaderBB, VectorLatchBB,
L->getExitBlock());
}

#if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
Expand All @@ -474,10 +475,10 @@ void VPlan::dump() const { dbgs() << *this << '\n'; }
#endif

void VPlan::updateDominatorTree(DominatorTree *DT, BasicBlock *LoopPreHeaderBB,
BasicBlock *LoopLatchBB) {
BasicBlock *LoopLatchBB,
BasicBlock *LoopExitBB) {
BasicBlock *LoopHeaderBB = LoopPreHeaderBB->getSingleSuccessor();
assert(LoopHeaderBB && "Loop preheader does not have a single successor.");
DT->addNewBlock(LoopHeaderBB, LoopPreHeaderBB);
// The vector body may be more than a single basic-block by this point.
// Update the dominator tree information inside the vector body by propagating
// it from header to latch, expecting only triangular control-flow, if any.
Expand Down Expand Up @@ -508,6 +509,9 @@ void VPlan::updateDominatorTree(DominatorTree *DT, BasicBlock *LoopPreHeaderBB,
DT->addNewBlock(InterimSucc, BB);
DT->addNewBlock(PostDomSucc, BB);
}
// Latch block is a new dominator for the loop exit.
DT->changeImmediateDominator(LoopExitBB, LoopLatchBB);
assert(DT->verify(DominatorTree::VerificationLevel::Fast));
}

const Twine VPlanPrinter::getUID(const VPBlockBase *Block) {
Expand Down
4 changes: 2 additions & 2 deletions llvm/lib/Transforms/Vectorize/VPlan.h
Original file line number Diff line number Diff line change
Expand Up @@ -1424,9 +1424,9 @@ class VPlan {
private:
/// Add to the given dominator tree the header block and every new basic block
/// that was created between it and the latch block, inclusive.
static void updateDominatorTree(DominatorTree *DT,
static void updateDominatorTree(DominatorTree *DT, BasicBlock *LoopLatchBB,
BasicBlock *LoopPreHeaderBB,
BasicBlock *LoopLatchBB);
BasicBlock *LoopExitBB);
};

/// VPlanPrinter prints a given VPlan to a given output stream. The printing is
Expand Down