diff --git a/llvm/lib/Transforms/Vectorize/VPlan.h b/llvm/lib/Transforms/Vectorize/VPlan.h index 18f5f13073aa63..0770ff032aa6a0 100644 --- a/llvm/lib/Transforms/Vectorize/VPlan.h +++ b/llvm/lib/Transforms/Vectorize/VPlan.h @@ -4113,17 +4113,29 @@ class VPBlockUtils { IfFalse->setParent(BlockPtr->getParent()); } - /// Connect VPBlockBases \p From and \p To bi-directionally. Append \p To to - /// the successors of \p From and \p From to the predecessors of \p To. Both - /// VPBlockBases must have the same parent, which can be null. Both - /// VPBlockBases can be already connected to other VPBlockBases. - static void connectBlocks(VPBlockBase *From, VPBlockBase *To) { + /// Connect VPBlockBases \p From and \p To bi-directionally. If \p PredIdx is + /// -1, append \p From to the predecessors of \p To, otherwise set \p To's + /// predecessor at \p PredIdx to \p From. If \p SuccIdx is -1, append \p To to + /// the successors of \p From, otherwise set \p From's successor at \p SuccIdx + /// to \p To. Both VPBlockBases must have the same parent, which can be null. + /// Both VPBlockBases can be already connected to other VPBlockBases. + static void connectBlocks(VPBlockBase *From, VPBlockBase *To, + unsigned PredIdx = -1u, unsigned SuccIdx = -1u) { assert((From->getParent() == To->getParent()) && "Can't connect two block with different parents"); - assert(From->getNumSuccessors() < 2 && + assert((SuccIdx != -1u || From->getNumSuccessors() < 2) && "Blocks can't have more than two successors."); - From->appendSuccessor(To); - To->appendPredecessor(From); + assert((PredIdx != -1u || To->getNumPredecessors() < 2) && + "Blocks can't have more than two predecessors."); + if (SuccIdx == -1u) + From->appendSuccessor(To); + else + From->getSuccessors()[SuccIdx] = To; + + if (PredIdx == -1u) + To->appendPredecessor(From); + else + To->getPredecessors()[PredIdx] = From; } /// Disconnect VPBlockBases \p From and \p To bi-directionally. Remove \p To diff --git a/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp b/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp index ea8845eaa75d4d..70163aacb19d9d 100644 --- a/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp +++ b/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp @@ -360,9 +360,8 @@ static void addReplicateRegions(VPlan &Plan) { // Record predicated instructions for above packing optimizations. VPBlockBase *Region = createReplicateRegion(RepR, Plan); Region->setParent(CurrentBlock->getParent()); - VPBlockUtils::disconnectBlocks(CurrentBlock, SplitBlock); - VPBlockUtils::connectBlocks(CurrentBlock, Region); - VPBlockUtils::connectBlocks(Region, SplitBlock); + VPBlockUtils::connectBlocks(CurrentBlock, Region, -1, 0); + VPBlockUtils::connectBlocks(Region, SplitBlock, 0, -1); } }