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

[Florian Hahn via llvm-commits]

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

>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 1/2] [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 {

>From 0bf27606a01112300f237f4ea11287b1ebff6f87 Mon Sep 17 00:00:00 2001
From: Florian Hahn <flo at fhahn.com>
Date: Sun, 6 Oct 2024 13:49:03 +0100
Subject: [PATCH 2/2] !fixup address latest comments, thanks!

---
 llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp | 17 +++++++++++------
 1 file changed, 11 insertions(+), 6 deletions(-)

diff --git a/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp b/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp
index c7e94336b59a35..56b9f76e1def41 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp
@@ -275,7 +275,10 @@ 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());
@@ -295,7 +298,7 @@ InstructionCost VPRecipeBase::cost(ElementCount VF, VPCostContext &Ctx) {
     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())
+  else if (RecipeCost == InstructionCost::getMax())
     RecipeCost = 0;
 
   LLVM_DEBUG({
@@ -308,16 +311,18 @@ 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);
   }
-  // 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();
 }