[VPlan] Implement VPInterleaveRecipe::computeCost.

[fhahn]

Implement computing costs for VPInterleaveRecipe.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Florian Hahn (fhahn)

Changes

Implement computing costs for VPInterleaveRecipe.

---

Full diff: https://github.com/llvm/llvm-project/pull/106067.diff

2 Files Affected:

• (modified) llvm/lib/Transforms/Vectorize/VPlan.h (+4)
• (modified) llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp (+31-2)

diff --git a/llvm/lib/Transforms/Vectorize/VPlan.h b/llvm/lib/Transforms/Vectorize/VPlan.h
index 36a1aa08654d5b..2ab0a1826bd216 100644
--- a/llvm/lib/Transforms/Vectorize/VPlan.h
+++ b/llvm/lib/Transforms/Vectorize/VPlan.h
@@ -2165,6 +2165,10 @@ class VPInterleaveRecipe : public VPRecipeBase {
   /// Generate the wide load or store, and shuffles.
   void execute(VPTransformState &State) override;
 
+  /// Return the cost of this VPInterleaveRecipe.
+  InstructionCost computeCost(ElementCount VF,
+                              VPCostContext &Ctx) const override;
+
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
   /// Print the recipe.
   void print(raw_ostream &O, const Twine &Indent,
diff --git a/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp b/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp
index fe1325f4163004..b1fcc13dc34717 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp
@@ -269,8 +269,6 @@ void VPRecipeBase::moveBefore(VPBasicBlock &BB,
 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();
   if (auto *WidenMem = dyn_cast<VPWidenMemoryRecipe>(R))
     return &WidenMem->getIngredient();
   return nullptr;
@@ -2627,6 +2625,37 @@ void VPInterleaveRecipe::execute(VPTransformState &State) {
   }
 }
 
+InstructionCost VPInterleaveRecipe::computeCost(ElementCount VF,
+                                                VPCostContext &Ctx) const {
+  Instruction *I = getInsertPos();
+  Type *ValTy = Ctx.Types.inferScalarType(
+      getNumDefinedValues() > 0 ? getVPValue(0) : getStoredValues()[0]);
+  auto *VectorTy = cast<VectorType>(ToVectorTy(ValTy, VF));
+  unsigned AS = getLoadStoreAddressSpace(I);
+  enum TTI::TargetCostKind CostKind = TTI::TCK_RecipThroughput;
+
+  unsigned InterleaveFactor = IG->getFactor();
+  auto *WideVecTy = VectorType::get(ValTy, VF * InterleaveFactor);
+
+  // Holds the indices of existing members in the interleaved group.
+  SmallVector<unsigned, 4> Indices;
+  for (unsigned IF = 0; IF < InterleaveFactor; IF++)
+    if (IG->getMember(IF))
+      Indices.push_back(IF);
+
+  // Calculate the cost of the whole interleaved group.
+  InstructionCost Cost = Ctx.TTI.getInterleavedMemoryOpCost(
+      I->getOpcode(), WideVecTy, IG->getFactor(), Indices, IG->getAlign(), AS,
+      CostKind, getMask(), NeedsMaskForGaps);
+
+  if (!IG->isReverse())
+    return Cost;
+
+  return Cost + IG->getNumMembers() *
+                    Ctx.TTI.getShuffleCost(TargetTransformInfo::SK_Reverse,
+                                           VectorTy, std::nullopt, CostKind, 0);
+}
+
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
 void VPInterleaveRecipe::print(raw_ostream &O, const Twine &Indent,
                                VPSlotTracker &SlotTracker) const {

[rengolin]

Where will this be used, do we already have tests that hit this code?

[rengolin]

This looks a bit hard-coded. It would be better to get the attribute from IG and pass it to getShuffleCost, but this can be done in a future PR.

[fhahn]

IIUC this extra computation could/should also be handled by getInterleavedMemoryOpCost above?

[fhahn]

Where will this be used, do we already have tests that hit this code?

The VPlan-based cost model is used to pick the best vectorization factor already. For recipes that do not implement ::computeCost, we currently fall back to the legacy cost model (via the instruction returned by getInstructionForCost.

With this patch, we use VPInterleaveRecipe::computeCost when computing the cost for all plans when selecting the best VF/Plan (and remove the need to retrieve the underlying instruction via getInstructionForCost)

[fhahn]

IIUC this extra computation could/should also be handled by getInterleavedMemoryOpCost above?

[fhahn]

ping :)

[ElvisWang123]

Remove VPinterleaveRecipe here will prevent -force-target-instr-cost overwriting recipe cost that may be inconsistent to the legacy model.

[fhahn]

The code has been reworked in fa3258e and the changes here are gone.

[fhahn]

Ping. Updated after fa3258e (#109708) landed

[ElvisWang123]

LGTM, thanks