[VPlan] Sink retrieving legacy costs to more specific computeCost impls.

llvm/lib/Transforms/Vectorize/VPlan.h

@@ -828,7 +828,9 @@ class VPRecipeBase : public ilist_node_with_parent<VPRecipeBase, VPBasicBlock>,
 protected:
   /// Compute the cost of this recipe either using a recipe's specialized
   /// implementation or using the legacy cost model and the underlying
-  /// instructions.
+  /// instructions. Returns InstructionCost::max() if the cost of this recipe
+  /// should be ignored. Forced target instruction cost is not applied for such
+  /// recipes.
   virtual InstructionCost computeCost(ElementCount VF,
                                       VPCostContext &Ctx) const;
 };

llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp

@@ -275,20 +275,15 @@ void VPRecipeBase::moveBefore(VPBasicBlock &BB,
 }
 
 /// Return the underlying instruction to be used for computing \p R's cost via
-/// the legacy cost model. Return nullptr if there's no suitable instruction.
+/// the legacy cost model. Return nullptr if there's no suitable instruction or
+/// computeCost is already implemented for the recipe and there is no need for
+/// the underlying instruction, i.e. it does not need to be skipped for cost
+/// computations.
 static Instruction *getInstructionForCost(const VPRecipeBase *R) {
   if (auto *S = dyn_cast<VPSingleDefRecipe>(R))
     return dyn_cast_or_null<Instruction>(S->getUnderlyingValue());
   if (auto *IG = dyn_cast<VPInterleaveRecipe>(R))
     return IG->getInsertPos();
-  // Currently the legacy cost model only calculates the instruction cost with
-  // underlying instruction. Removing the WidenMem here will prevent
-  // force-target-instruction-cost overwriting the cost of recipe with
-  // underlying instruction which is inconsistent with the legacy model.
-  // TODO: Remove WidenMem from this function when we don't need to compare to
-  // the legacy model.
-  if (auto *WidenMem = dyn_cast<VPWidenMemoryRecipe>(R))
-    return &WidenMem->getIngredient();
   return nullptr;
 }
 
@@ -298,9 +293,13 @@ InstructionCost VPRecipeBase::cost(ElementCount VF, VPCostContext &Ctx) {
     return 0;
 
   InstructionCost RecipeCost = computeCost(VF, Ctx);
-  if (UI && ForceTargetInstructionCost.getNumOccurrences() > 0 &&
-      RecipeCost.isValid())
+  if (ForceTargetInstructionCost.getNumOccurrences() > 0 &&
+      (RecipeCost.isValid() && RecipeCost != InstructionCost::getMax()))
     RecipeCost = InstructionCost(ForceTargetInstructionCost);
+  // Max cost is used as a sentinel value to detect recipes without underlying
+  // instructions for which no forced target instruction cost should be applied.
+  else if (RecipeCost == InstructionCost::getMax())
+    RecipeCost = 0;
 
   LLVM_DEBUG({
     dbgs() << "Cost of " << RecipeCost << " for VF " << VF << ": ";
@@ -312,15 +311,19 @@ InstructionCost VPRecipeBase::cost(ElementCount VF, VPCostContext &Ctx) {
 InstructionCost VPRecipeBase::computeCost(ElementCount VF,
                                           VPCostContext &Ctx) const {
   // Compute the cost for the recipe falling back to the legacy cost model using
-  // the underlying instruction. If there is no underlying instruction, returns
-  // 0.
+  // the underlying instruction. If there is no underlying instruction or the
+  // cost is computed by the recipe's computeCost, returns
+  // InstructionCost::getMax. It is used as  a sentinel value to detect recipes
+  // without underlying instructions for which no forced target instruction cost
+  // should be applied.
+
   Instruction *UI = getInstructionForCost(this);
   if (UI && isa<VPReplicateRecipe>(this)) {
     // VPReplicateRecipe may be cloned as part of an existing VPlan-to-VPlan
     // transform, avoid computing their cost multiple times for now.
     Ctx.SkipCostComputation.insert(UI);
   }
-  return UI ? Ctx.getLegacyCost(UI, VF) : 0;
+  return UI ? Ctx.getLegacyCost(UI, VF) : InstructionCost::getMax();
 }
 
 FastMathFlags VPRecipeWithIRFlags::getFastMathFlags() const {