diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp --- a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp +++ b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp @@ -8976,6 +8976,35 @@ continue; } } + + // If the source is a predicate replicate recipe, we need to move the whole + // replicate region, which should only contain a single recipe in the main + // block. + auto *RepR = dyn_cast(Sink); + if (RepR && RepR->isPredicated()) { + auto *Region = + dyn_cast_or_null(RepR->getParent()->getParent()); + assert(Region->isReplicator() && RepR->getParent()->size() == 1 && + "parent must be a replicator with a single recipe"); + auto *SplitBlock = + Target->getParent()->splitAt(std::next(Target->getIterator())); + + auto *Pred = Region->getSinglePredecessor(); + auto *Succ = Region->getSingleSuccessor(); + VPBlockUtils::disconnectBlocks(Pred, Region); + VPBlockUtils::disconnectBlocks(Region, Succ); + VPBlockUtils::connectBlocks(Pred, Succ); + + auto *SplitPred = SplitBlock->getSinglePredecessor(); + + VPBlockUtils::disconnectBlocks(SplitPred, SplitBlock); + VPBlockUtils::connectBlocks(SplitPred, Region); + VPBlockUtils::connectBlocks(Region, SplitBlock); + if (VPBB == SplitPred) + VPBB = SplitBlock; + continue; + } + Sink->moveAfter(Target); } diff --git a/llvm/lib/Transforms/Vectorize/VPlan.h b/llvm/lib/Transforms/Vectorize/VPlan.h --- a/llvm/lib/Transforms/Vectorize/VPlan.h +++ b/llvm/lib/Transforms/Vectorize/VPlan.h @@ -1509,6 +1509,11 @@ void dropAllReferences(VPValue *NewValue) override; + /// Split current block at \p SplitAt by inserting a new block between the + /// current block and its successors and mobing all recipes starting at + /// SplitAt to the new block. Returns the new block. + VPBasicBlock *splitAt(iterator SplitAt); + #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP) /// Print this VPBsicBlock to \p O, prefixing all lines with \p Indent. \p /// SlotTracker is used to print unnamed VPValue's using consequtive numbers. diff --git a/llvm/lib/Transforms/Vectorize/VPlan.cpp b/llvm/lib/Transforms/Vectorize/VPlan.cpp --- a/llvm/lib/Transforms/Vectorize/VPlan.cpp +++ b/llvm/lib/Transforms/Vectorize/VPlan.cpp @@ -403,6 +403,31 @@ } } +VPBasicBlock *VPBasicBlock::splitAt(iterator SplitAt) { + assert(SplitAt->getParent() == this); + + SmallVector Succs(getSuccessors().begin(), + getSuccessors().end()); + // First, disconnect the current block from its successors. + for (VPBlockBase *Succ : Succs) + VPBlockUtils::disconnectBlocks(this, Succ); + + // Create new empty block after the block to split. + auto *SplitBlock = new VPBasicBlock(getName() + ".split"); + VPBlockUtils::insertBlockAfter(SplitBlock, this); + + // Add successors for block to split to new block. + for (VPBlockBase *Succ : Succs) + VPBlockUtils::connectBlocks(SplitBlock, Succ); + + // Finally, move the recipes starting at SplitAt to new block. + for (VPRecipeBase &ToMove : + make_early_inc_range(make_range(SplitAt, this->end()))) + ToMove.moveBefore(*SplitBlock, SplitBlock->end()); + + return SplitBlock; +} + #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP) void VPBasicBlock::print(raw_ostream &O, const Twine &Indent, VPSlotTracker &SlotTracker) const {