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 @@ -1584,11 +1584,6 @@ return InLoopReductionChains; } - /// Returns true if the Phi is part of an inloop reduction. - bool isInLoopReduction(PHINode *Phi) const { - return InLoopReductionChains.count(Phi); - } - /// Estimate cost of an intrinsic call instruction CI if it were vectorized /// with factor VF. Return the cost of the instruction, including /// scalarization overhead if it's needed. @@ -4281,7 +4276,7 @@ TrackingVH ReductionStartValue = RdxDesc.getRecurrenceStartValue(); Instruction *LoopExitInst = RdxDesc.getLoopExitInstr(); setDebugLocFromInst(Builder, ReductionStartValue); - bool IsInLoopReductionPhi = Cost->isInLoopReduction(OrigPhi); + bool IsInLoopReductionPhi = isa(PhiR->getBackedgeValue()); VPValue *LoopExitInstDef = State.Plan->getVPValue(LoopExitInst); // This is the vector-clone of the value that leaves the loop. @@ -4692,7 +4687,8 @@ if (RdxDesc || Legal->isFirstOrderRecurrence(P)) { Value *Iden = nullptr; bool ScalarPHI = - (State.VF.isScalar()) || Cost->isInLoopReduction(cast(PN)); + (State.VF.isScalar()) || + (RdxDesc && isa(PhiR->getBackedgeValue())); Type *VecTy = ScalarPHI ? PN->getType() : VectorType::get(PN->getType(), State.VF); @@ -9053,16 +9049,29 @@ // 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)) + if (CM.foldTailByMasking()) { + VPValue *Cond = nullptr; + unsigned NumRegularReductions = 0; + for (VPRecipeBase &R : Plan->getEntry()->getEntryBasicBlock()->phis()) { + VPWidenPHIRecipe *PhiR = dyn_cast(&R); + if (!PhiR || !PhiR->getRecurrenceDescriptor() || + isa(PhiR->getBackedgeValue())) continue; - VPValue *Phi = Plan->getOrAddVPValue(Reduction.first); - VPValue *Red = Plan->getOrAddVPValue(Reduction.second.getLoopExitInstr()); + + if (!Cond) { + Builder.setInsertPoint(VPBB); + Cond = RecipeBuilder.createBlockInMask(OrigLoop->getHeader(), Plan); + } + VPValue *Phi = PhiR; + VPValue *Red = Plan->getOrAddVPValue( + PhiR->getRecurrenceDescriptor()->getLoopExitInstr()); Builder.createNaryOp(Instruction::Select, {Cond, Red, Phi}); + NumRegularReductions++; } + assert(!Range.Start.isVector() || + NumRegularReductions == Legal->getReductionVars().size() - + CM.getInLoopReductionChains().size() && + "failed to identify an in-loop reduction in VPlan"); } std::string PlanName;