diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp index 001c8987667df8..866568cb9b2a8a 100644 --- a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp +++ b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp @@ -7103,6 +7103,14 @@ bool VPCostContext::skipCostComputation(Instruction *UI, bool IsVector) const { SkipCostComputation.contains(UI); } +bool VPCostContext::isInLoopReduction(const Instruction *UI, ElementCount VF, + Type *VectorTy) const { + return CM + .getReductionPatternCost(const_cast(UI), VF, VectorTy, + TTI::TCK_RecipThroughput) + .has_value(); +} + InstructionCost LoopVectorizationPlanner::precomputeCosts(VPlan &Plan, ElementCount VF, VPCostContext &CostCtx) const { diff --git a/llvm/lib/Transforms/Vectorize/VPlan.h b/llvm/lib/Transforms/Vectorize/VPlan.h index 05d0da691354c0..41ae3a184091a4 100644 --- a/llvm/lib/Transforms/Vectorize/VPlan.h +++ b/llvm/lib/Transforms/Vectorize/VPlan.h @@ -718,6 +718,10 @@ struct VPCostContext { /// Return true if the cost for \p UI shouldn't be computed, e.g. because it /// has already been pre-computed. bool skipCostComputation(Instruction *UI, bool IsVector) const; + + /// Return true if the \p UI is part of the in-loop reduction. + bool isInLoopReduction(const Instruction *UI, ElementCount VF, + Type *VectorTy) const; }; /// VPRecipeBase is a base class modeling a sequence of one or more output IR diff --git a/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp b/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp index cc0231ce766128..8217b8cda5118b 100644 --- a/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp +++ b/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp @@ -2015,18 +2015,24 @@ void VPReductionEVLRecipe::execute(VPTransformState &State) { InstructionCost VPReductionRecipe::computeCost(ElementCount VF, VPCostContext &Ctx) const { RecurKind RdxKind = RdxDesc.getRecurrenceKind(); - // TODO: Support any-of reduction and in-loop reduction - assert(!RecurrenceDescriptor::isAnyOfRecurrenceKind(RdxKind) && - "Not support any-of reduction in VPlan-based cost model currently."); - - Type *ElementTy = Ctx.Types.inferScalarType(this->getVPSingleValue()); - assert(ElementTy->getTypeID() == RdxDesc.getRecurrenceType()->getTypeID() && - "Infered type and recurrence type mismatch."); - + Type *ElementTy = Ctx.Types.inferScalarType(this); auto *VectorTy = cast(ToVectorTy(ElementTy, VF)); TTI::TargetCostKind CostKind = TTI::TCK_RecipThroughput; unsigned Opcode = RdxDesc.getOpcode(); + // TODO: Support any-of reduction and in-loop reductions. + assert( + (!RecurrenceDescriptor::isAnyOfRecurrenceKind(RdxKind) || + ForceTargetInstructionCost.getNumOccurrences() > 0) && + "Any-of reduction not implemented in VPlan-based cost model currently."); + assert( + (!Ctx.isInLoopReduction(getUnderlyingInstr(), VF, VectorTy) || + ForceTargetInstructionCost.getNumOccurrences() > 0) && + "In-loop reduction not implemented in VPlan-based cost model currently."); + + assert(ElementTy->getTypeID() == RdxDesc.getRecurrenceType()->getTypeID() && + "Infered type and recurrence type mismatch."); + // Cost = Reduction cost + BinOp cost InstructionCost Cost = Ctx.TTI.getArithmeticInstrCost(Opcode, ElementTy, CostKind);