diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp index fbc82fb9696e..3aff0a02e059 100644 --- a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp +++ b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp @@ -8038,6 +8038,17 @@ VPValue *VPRecipeBuilder::createEdgeMask(BasicBlock *Src, BasicBlock *Dst, return EdgeMaskCache[Edge] = EdgeMask; } +VPValue *VPRecipeBuilder::getEdgeMask(BasicBlock *Src, BasicBlock *Dst) const { + assert(is_contained(predecessors(Dst), Src) && "Invalid edge"); + + // Look for cached value. + std::pair Edge(Src, Dst); + EdgeMaskCacheTy::const_iterator ECEntryIt = EdgeMaskCache.find(Edge); + assert(ECEntryIt != EdgeMaskCache.end() && + "looking up mask for edge which has not been created"); + return ECEntryIt->second; +} + void VPRecipeBuilder::createHeaderMask(VPlan &Plan) { BasicBlock *Header = OrigLoop->getHeader(); @@ -8728,10 +8739,13 @@ LoopVectorizationPlanner::tryToBuildVPlanWithVPRecipes(VFRange &Range) { DFS.perform(LI); VPBasicBlock *VPBB = HeaderVPBB; - bool NeedsMasks = CM.foldTailByMasking() || - any_of(OrigLoop->blocks(), [this](BasicBlock *BB) { - return Legal->blockNeedsPredication(BB); - }); + BasicBlock *HeaderBB = OrigLoop->getHeader(); + bool NeedsMasks = + CM.foldTailByMasking() || + any_of(OrigLoop->blocks(), [this, HeaderBB](BasicBlock *BB) { + bool NeedsBlends = BB != HeaderBB && !BB->phis().empty(); + return Legal->blockNeedsPredication(BB) || NeedsBlends; + }); for (BasicBlock *BB : make_range(DFS.beginRPO(), DFS.endRPO())) { // Relevant instructions from basic block BB will be grouped into VPRecipe // ingredients and fill a new VPBasicBlock. @@ -8750,7 +8764,7 @@ LoopVectorizationPlanner::tryToBuildVPlanWithVPRecipes(VFRange &Range) { Instruction *Instr = &I; SmallVector Operands; auto *Phi = dyn_cast(Instr); - if (Phi && Phi->getParent() == OrigLoop->getHeader()) { + if (Phi && Phi->getParent() == HeaderBB) { Operands.push_back(Plan->getVPValueOrAddLiveIn( Phi->getIncomingValueForBlock(OrigLoop->getLoopPreheader()))); } else { diff --git a/llvm/lib/Transforms/Vectorize/VPRecipeBuilder.h b/llvm/lib/Transforms/Vectorize/VPRecipeBuilder.h index 4b3143aead46..5645cfa1473a 100644 --- a/llvm/lib/Transforms/Vectorize/VPRecipeBuilder.h +++ b/llvm/lib/Transforms/Vectorize/VPRecipeBuilder.h @@ -148,6 +148,10 @@ public: /// and DST. VPValue *createEdgeMask(BasicBlock *Src, BasicBlock *Dst, VPlan &Plan); + /// A helper that returns the previously computed predicate of the edge + /// between SRC and DST. + VPValue *getEdgeMask(BasicBlock *Src, BasicBlock *Dst) const; + /// Mark given ingredient for recording its recipe once one is created for /// it. void recordRecipeOf(Instruction *I) {