[llvm] [VPlan] Add cost for `VPWidenMemoryRecipe` (PR #105614)

[Elvis Wang via llvm-commits]

https://github.com/ElvisWang123 created https://github.com/llvm/llvm-project/pull/105614

In this patch, we add the `computeCost()` function for `VPWidenMemoryRecipe`.

>From c0adb649fd9a1a41b09548c013259a5da382b38f Mon Sep 17 00:00:00 2001
From: Elvis Wang <elvis.wang at sifive.com>
Date: Wed, 21 Aug 2024 20:39:31 -0700
Subject: [PATCH] [VPlan] Add cost for `VPWidenMemoryRecipe`

In this patch, we add the `computeCost()` function for `VPWidenMemoryRecipe`.
---
 llvm/lib/Transforms/Vectorize/VPlan.h         |  7 +++++
 .../lib/Transforms/Vectorize/VPlanRecipes.cpp | 31 +++++++++++++++++++
 2 files changed, 38 insertions(+)

diff --git a/llvm/lib/Transforms/Vectorize/VPlan.h b/llvm/lib/Transforms/Vectorize/VPlan.h
index a99f3882092c2c..1f6de54822ebcf 100644
--- a/llvm/lib/Transforms/Vectorize/VPlan.h
+++ b/llvm/lib/Transforms/Vectorize/VPlan.h
@@ -2529,6 +2529,13 @@ class VPWidenMemoryRecipe : public VPRecipeBase {
     llvm_unreachable("VPWidenMemoryRecipe should not be instantiated.");
   }
 
+  /// Get element Type
+  Type *getElementType() const { return getLoadStoreType(&Ingredient); }
+
+  /// Return the cost of this VPWidenMemoryRecipe.
+  InstructionCost computeCost(ElementCount VF,
+                              VPCostContext &Ctx) const override;
+
   Instruction &getIngredient() const { return Ingredient; }
 };
 
diff --git a/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp b/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp
index c9d603612aecea..d3b36b145d470e 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp
@@ -2084,6 +2084,37 @@ void VPPredInstPHIRecipe::print(raw_ostream &O, const Twine &Indent,
 }
 #endif
 
+InstructionCost VPWidenMemoryRecipe::computeCost(ElementCount VF,
+                                                 VPCostContext &Ctx) const {
+  Instruction *I = getInstructionForCost(this);
+  Type *Ty = ToVectorTy(getElementType(), VF);
+  const Align Alignment = getLoadStoreAlignment(const_cast<Instruction *>(I));
+  const Value *Ptr = getLoadStorePointerOperand(I);
+  unsigned AS = getLoadStoreAddressSpace(const_cast<Instruction *>(I));
+  TTI::TargetCostKind CostKind = TTI::TCK_RecipThroughput;
+
+  if (Consecutive) {
+    InstructionCost Cost = 0;
+    if (IsMasked) {
+      Cost += Ctx.TTI.getMaskedMemoryOpCost(I->getOpcode(), Ty, Alignment, AS,
+                                            CostKind);
+    } else {
+      TTI::OperandValueInfo OpInfo = Ctx.TTI.getOperandInfo(I->getOperand(0));
+      Cost += Ctx.TTI.getMemoryOpCost(I->getOpcode(), Ty, Alignment, AS,
+                                      CostKind, OpInfo, I);
+    }
+    if (Reverse)
+      Cost += Ctx.TTI.getShuffleCost(TargetTransformInfo::SK_Reverse,
+                                     cast<VectorType>(Ty), std::nullopt,
+                                     CostKind, 0);
+
+    return Cost;
+  }
+  return Ctx.TTI.getAddressComputationCost(Ty) +
+         Ctx.TTI.getGatherScatterOpCost(I->getOpcode(), Ty, Ptr, IsMasked,
+                                        Alignment, CostKind, I);
+}
+
 void VPWidenLoadRecipe::execute(VPTransformState &State) {
   auto *LI = cast<LoadInst>(&Ingredient);