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,11 +351,12 @@ /// 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, + // 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 adjustRecipesForInLoopReductions(VPBasicBlock *LatchVPBB, VPlanPtr &Plan, VPRecipeBuilder &RecipeBuilder, ElementCount MinVF); }; 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 @@ -9401,21 +9401,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}); - } - } + adjustRecipesForInLoopReductions(VPBB, Plan, RecipeBuilder, Range.Start); VPlanTransforms::sinkScalarOperands(*Plan); VPlanTransforms::mergeReplicateRegions(*Plan); @@ -9471,12 +9457,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. +// 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::adjustRecipesForInLoopReductions( - VPlanPtr &Plan, VPRecipeBuilder &RecipeBuilder, ElementCount MinVF) { + VPBasicBlock *LatchVPBB, VPlanPtr &Plan, VPRecipeBuilder &RecipeBuilder, + ElementCount MinVF) { for (auto &Reduction : CM.getInLoopReductionChains()) { PHINode *Phi = Reduction.first; RecurrenceDescriptor &RdxDesc = Legal->getReductionVars()[Phi]; @@ -9533,6 +9521,24 @@ 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. + // Note that this is currently only used to reflect the def-use chains in + // VPlan. fixReduction will introduce the appropriate selects and update the + // users outside the loop without relying on the select recipe. + 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)