[llvm] [VPlan] Return Max from computeCost without underlying instrs. (PR #109708)

[Florian Hahn via llvm-commits]

https://github.com/fhahn created https://github.com/llvm/llvm-project/pull/109708

Update computeCost to return InstructionCost::getMax() for recipes without underlying instructions. Max is used as a sentinel value to handle cases where there is no underlying instruction. At the moment we need to catch those cases when -force-target-instruction-cost is passed to avoid applying the forced cost to auxiliary recipes (like scalar steps). This is needed to match the legacy behavior.

Unfortunately we cannot use InstructionCost::getInvalid, as this is used to indicate that an instruction cannot be legalized (e.g. for scalable vectors).

Alternatively computeCost could return an optional cost.

>From 8117ba1e94dfa4c995d496031b6c9fe989bf9b04 Mon Sep 17 00:00:00 2001
From: Florian Hahn <flo at fhahn.com>
Date: Mon, 23 Sep 2024 20:19:35 +0100
Subject: [PATCH] [VPlan] Return Max from computeCost without underlying
 instrs.

Update computeCost to return InstructionCost::getMax() for recipes
without underlying instructions. Max is used as a sentinel value to
handle cases where there is no underlying instruction. At the moment we
need to catch those cases when -force-target-instruction-cost is passed
to avoid applying the forced cost to auxiliary recipes (like scalar
steps). This is needed to match the legacy behavior.

Unfortunately we cannot use InstructionCost::getInvalid, as this is used
to indicate that an instruction cannot be legalized (e.g. for scalable
vectors).

Alternatively computeCost could return an optional cost.
---
 llvm/lib/Transforms/Vectorize/VPlan.h         |  4 +++-
 .../lib/Transforms/Vectorize/VPlanRecipes.cpp | 20 +++++++++----------
 2 files changed, 12 insertions(+), 12 deletions(-)

diff --git a/llvm/lib/Transforms/Vectorize/VPlan.h b/llvm/lib/Transforms/Vectorize/VPlan.h
index 0632495bc511cd..6f8a44e1cd1802 100644
--- a/llvm/lib/Transforms/Vectorize/VPlan.h
+++ b/llvm/lib/Transforms/Vectorize/VPlan.h
@@ -844,7 +844,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;
 };
diff --git a/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp b/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp
index 3f7ab416e877bc..ca4622b74275e5 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp
@@ -276,14 +276,6 @@ static Instruction *getInstructionForCost(const VPRecipeBase *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;
 }
 
@@ -293,9 +285,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.
+  if (RecipeCost == InstructionCost::getMax())
+    RecipeCost = 0;
 
   LLVM_DEBUG({
     dbgs() << "Cost of " << RecipeCost << " for VF " << VF << ": ";
@@ -315,7 +311,9 @@ InstructionCost VPRecipeBase::computeCost(ElementCount VF,
     // transform, avoid computing their cost multiple times for now.
     Ctx.SkipCostComputation.insert(UI);
   }
-  return UI ? Ctx.getLegacyCost(UI, VF) : 0;
+  // Max cost is used as a sentinel value to detect recipes without underlying
+  // instructions for which no forced target instruction cost should be applied.
+  return UI ? Ctx.getLegacyCost(UI, VF) : InstructionCost::getMax();
 }
 
 FastMathFlags VPRecipeWithIRFlags::getFastMathFlags() const {