diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorizationPlanner.h b/llvm/lib/Transforms/Vectorize/LoopVectorizationPlanner.h --- a/llvm/lib/Transforms/Vectorize/LoopVectorizationPlanner.h +++ b/llvm/lib/Transforms/Vectorize/LoopVectorizationPlanner.h @@ -351,13 +351,14 @@ /// legal to vectorize the loop. This method creates VPlans using VPRecipes. void buildVPlansWithVPRecipes(ElementCount MinVF, ElementCount MaxVF); - /// Adjust the recipes for any inloop reductions. The chain of instructions - /// leading from the loop exit instr to the phi need to be converted to - /// reductions, with one operand being vector and the other being the scalar - /// reduction chain. - void adjustRecipesForInLoopReductions(VPlanPtr &Plan, - VPRecipeBuilder &RecipeBuilder, - ElementCount MinVF); + // Adjust the recipes for reductions. For in-loop reductions the chain of + // instructions leading from the loop exit instr to the phi need to be + // converted to reductions, with one operand being vector and the other being + // the scalar reduction chain. For other reductions, a select is introduced + // between the phi and live-out recipes when folding the tail. + void adjustRecipesForReductions(VPBasicBlock *LatchVPBB, VPlanPtr &Plan, + VPRecipeBuilder &RecipeBuilder, + ElementCount MinVF); }; } // namespace llvm 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 @@ -4295,7 +4295,7 @@ Instruction *LoopExitInst = RdxDesc.getLoopExitInstr(); setDebugLocFromInst(ReductionStartValue); - VPValue *LoopExitInstDef = State.Plan->getVPValue(LoopExitInst); + VPValue *LoopExitInstDef = PhiR->getBackedgeValue(); // This is the vector-clone of the value that leaves the loop. Type *VecTy = State.get(LoopExitInstDef, 0)->getType(); @@ -9438,21 +9438,7 @@ } // Adjust the recipes for any inloop reductions. - adjustRecipesForInLoopReductions(Plan, RecipeBuilder, Range.Start); - - // Finally, if tail is folded by masking, introduce selects between the phi - // and the live-out instruction of each reduction, at the end of the latch. - if (CM.foldTailByMasking() && !Legal->getReductionVars().empty()) { - Builder.setInsertPoint(VPBB); - auto *Cond = RecipeBuilder.createBlockInMask(OrigLoop->getHeader(), Plan); - for (auto &Reduction : Legal->getReductionVars()) { - if (CM.isInLoopReduction(Reduction.first)) - continue; - VPValue *Phi = Plan->getOrAddVPValue(Reduction.first); - VPValue *Red = Plan->getOrAddVPValue(Reduction.second.getLoopExitInstr()); - Builder.createNaryOp(Instruction::Select, {Cond, Red, Phi}); - } - } + adjustRecipesForReductions(VPBB, Plan, RecipeBuilder, Range.Start); VPlanTransforms::sinkScalarOperands(*Plan); VPlanTransforms::mergeReplicateRegions(*Plan); @@ -9508,12 +9494,14 @@ return Plan; } -// Adjust the recipes for any inloop reductions. The chain of instructions -// leading from the loop exit instr to the phi need to be converted to -// reductions, with one operand being vector and the other being the scalar -// reduction chain. -void LoopVectorizationPlanner::adjustRecipesForInLoopReductions( - VPlanPtr &Plan, VPRecipeBuilder &RecipeBuilder, ElementCount MinVF) { +// Adjust the recipes for reductions. For in-loop reductions the chain of +// instructions leading from the loop exit instr to the phi need to be converted +// to reductions, with one operand being vector and the other being the scalar +// reduction chain. For other reductions, a select is introduced between the phi +// and live-out recipes when folding the tail. +void LoopVectorizationPlanner::adjustRecipesForReductions( + VPBasicBlock *LatchVPBB, VPlanPtr &Plan, VPRecipeBuilder &RecipeBuilder, + ElementCount MinVF) { for (auto &Reduction : CM.getInLoopReductionChains()) { PHINode *Phi = Reduction.first; RecurrenceDescriptor &RdxDesc = Legal->getReductionVars()[Phi]; @@ -9570,6 +9558,21 @@ Chain = R; } } + + // If tail is folded by masking, introduce selects between the phi + // and the live-out instruction of each reduction, at the end of the latch. + if (CM.foldTailByMasking()) { + for (VPRecipeBase &R : Plan->getEntry()->getEntryBasicBlock()->phis()) { + VPReductionPHIRecipe *PhiR = dyn_cast(&R); + if (!PhiR || PhiR->isInLoop()) + continue; + Builder.setInsertPoint(LatchVPBB); + VPValue *Cond = + RecipeBuilder.createBlockInMask(OrigLoop->getHeader(), Plan); + VPValue *Red = PhiR->getBackedgeValue(); + Builder.createNaryOp(Instruction::Select, {Cond, Red, PhiR}); + } + } } #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)