Address comments and add isInLoopReduction().

Author: ElvisWang123

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

@@ -7137,6 +7137,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<Instruction *>(UI), VF, VectorTy,
+                               TTI::TCK_RecipThroughput)
+      .has_value();
+}
+
 InstructionCost
 LoopVectorizationPlanner::precomputeCosts(VPlan &Plan, ElementCount VF,
                                           VPCostContext &CostCtx) const {

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

llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp

@@ -2074,18 +2074,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<VectorType>(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);