[RFC][LV] VPlan-based cost model

[arcbbb]

Summary

This patch follows D89322 to add an initial skeleton of vplan-based cost model.

This difference is that instead of incorporating a cost() interface to VPRecipes,
all cost implementations are put together in VPlanCostModel.

This allows VPlanCostModel to concentrate on assigning costs to vplan,
thus seprating the cost model code from the vplan IR, similar to LLVM IR cost
modeling.

During the transition, it will still use the legacy model to obtain cost until
all cost calculation for recipes are implemented.

Please let me know if you agree with the main idea of this patch.
If there is a general consensus, I'll proceed to implement the cost for the
other recipes for review.

[wangpc]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

8709

Is it always true even for scalar VF? Why can't we have the same signature as CM.expectedCost?

llvm/lib/Transforms/Vectorize/VPlanCostModel.h

31

RVL->VF?

[arcbbb]

• Address comments

[arcbbb]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

8709

VectorizationCostTy is not exported in the header, so I change the interface trying to decouple the dependency.

llvm/lib/Transforms/Vectorize/VPlanCostModel.h

31

Fixed. Thanks!

[rengolin]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

8707

It's costly to create this object in every inner loop.

This should either be hoisted to the outer loop, or outside of the loop altogether if you allow expectedCost to have the Plan as an argument, since it's the only function in VPlanCostModel that uses it.

[arcbbb]

• Move VPCM object outside of the loop

[arcbbb]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

8707

Good point!

[nikolaypanchenko]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

8707

Cost model may eventually have a state to track context information. Even though it's possible to have a single instance of VPCM, but this is going to add extra boilerplate to manage that state.

[rengolin]

I agree with the general idea and the implementation seems reasonable, but I haven't worked in this part of the code for a long time, so would like folks like @fhahn and @Ayal to approve this one.

I'd like to see tests with the new flag. Perhaps reusing existing tests with a new RUN line and checking that the decisions are the same. After all, we don't want to change the behaviour with this patch, just the infrastructure.

[arcbbb]

Here is going to be read-only, so I moved this patch to github pr67647

I'd like to see tests with the new flag. Perhaps reusing existing tests with a new RUN line and checking that the decisions are the same. After all, we don't want to change the behaviour with this patch, just the infrastructure.

Sure, I'll incorporate tests with the new flag and try existing tests with a new RUN line to ensure that the decisions remain consistent.

llvm/lib/Transforms/Vectorize/CMakeLists.txt

 add_llvm_component_library(LLVMVectorize
   LoadStoreVectorizer.cpp
   LoopVectorizationLegality.cpp
   LoopVectorize.cpp
   SLPVectorizer.cpp
   Vectorize.cpp
   VectorCombine.cpp
   VPlan.cpp
+  VPlanCostModel.cpp
   VPlanHCFGBuilder.cpp
   VPlanRecipes.cpp
   VPlanSLP.cpp
   VPlanTransforms.cpp
   VPlanVerifier.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${LLVM_MAIN_INCLUDE_DIR}/llvm/Transforms

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

 //  Vectorizing Compilers.
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Transforms/Vectorize/LoopVectorize.h"
 #include "LoopVectorizationPlanner.h"
 #include "VPRecipeBuilder.h"
 #include "VPlan.h"
+#include "VPlanCostModel.h"
 #include "VPlanHCFGBuilder.h"
 #include "VPlanTransforms.h"
 #include "llvm/ADT/APInt.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/DenseMapInfo.h"
 #include "llvm/ADT/Hashing.h"
 #include "llvm/ADT/MapVector.h"

 
 namespace llvm {
 cl::opt<bool> EnableVPlanNativePath(
     "enable-vplan-native-path", cl::Hidden,
     cl::desc("Enable VPlan-native vectorization path with "
              "support for outer loop vectorization."));
 }
 
+cl::opt<bool> CostUsingVPlan("vplan-use-vplan-cost-model", cl::init(false),
+                             cl::Hidden,
+                             cl::desc("Enable VPlan based costing path. To "
+                                      "become the default in the future."));
+
 // This flag enables the stress testing of the VPlan H-CFG construction in the
 // VPlan-native vectorization path. It must be used in conjuction with
 // -enable-vplan-native-path. -vplan-verify-hcfg can also be used to enable the
 // verification of the H-CFGs built.
 static cl::opt<bool> VPlanBuildStressTest(
     "vplan-build-stress-test", cl::init(false), cl::Hidden,
     cl::desc(
         "Build VPlan for every supported loop nest in the function and bail "

 /// LoopVectorizationCostModel - estimates the expected speedups due to
 /// vectorization.
 /// In many cases vectorization is not profitable. This can happen because of
 /// a number of reasons. In this class we mainly attempt to predict the
 /// expected speedup/slowdowns due to the supported instruction set. We use the
 /// TargetTransformInfo to query the different backends for the cost of
 /// different operations.
 class LoopVectorizationCostModel {
+  friend class VPlanCostModel;
+
 public:
   LoopVectorizationCostModel(ScalarEpilogueLowering SEL, Loop *L,
                              PredicatedScalarEvolution &PSE, LoopInfo *LI,
                              LoopVectorizationLegality *Legal,
                              const TargetTransformInfo &TTI,
                              const TargetLibraryInfo *TLI, DemandedBits *DB,
                              AssumptionCache *AC,
                              OptimizationRemarkEmitter *ORE, const Function *F,
VPRecipeBuilder::tryToCreateWidenRecipe(Instruction *Instr,
   if (auto *SI = dyn_cast<SelectInst>(Instr)) {
     return toVPRecipeResult(new VPWidenSelectRecipe(
         *SI, make_range(Operands.begin(), Operands.end())));
   }
 
   return toVPRecipeResult(tryToWiden(Instr, Operands, VPBB, Plan));
 }
 
+Type *VPlanCostModel::truncateToMinimalBitwidth(Type *ValTy,
+                                                Instruction *I) const {
+  auto MinBWs = CM.getMinimalBitwidths();
+  if (MinBWs.contains(I))
+    ValTy = IntegerType::get(ValTy->getContext(), MinBWs[I]);
+  return ValTy;
+}
+
+InstructionCost VPlanCostModel::getLegacyInstructionCost(Instruction *I,
+                                                         ElementCount VF) {
+  VectorizationCostTy Cost = CM.getInstructionCost(I, VF);
+  return Cost.first;
+}
+
 void LoopVectorizationPlanner::buildVPlansWithVPRecipes(ElementCount MinVF,
                                                         ElementCount MaxVF) {
   assert(OrigLoop->isInnermost() && "Inner loop expected.");
 
   // Add assume instructions we need to drop to DeadInstructions, to prevent
   // them from being added to the VPlan.
   // TODO: We only need to drop assumes in blocks that get flattend. If the
   // control flow is preserved, we should keep them.
for (ElementCount VF = MinVF; ElementCount::isKnownLT(VF, MaxVFTimes2);) {
     }
     VPlans.emplace_back(std::move(*Plan));
     VF = SubRange.End;
   }
 
   for (const VPlanPtr &Plan : VPlans) {
     SmallVector<VectorizationFactor> Costs;
     for (ElementCount CostVF : Plan->getVFs()) {
-      auto [VecCost, IsVec] = CM.expectedCost(CostVF, &InvalidCosts);
+      VectorizationCostTy C;
+      if (CostUsingVPlan) {
+        VPlanCostModel VPCM(*Plan, *TTI, CM);
+        C.first = VPCM.expectedCost(CostVF, C.second);
+      } else
+        C = CM.expectedCost(CostVF, &InvalidCosts);
+      auto [VecCost, IsVec] = C;
 #ifndef NDEBUG
       unsigned AssumedMinimumVscale = 1;
       if (std::optional<unsigned> VScale = getVScaleForTuning())
         AssumedMinimumVscale = *VScale;
       unsigned Width = CostVF.isScalable()
                            ? CostVF.getKnownMinValue() * AssumedMinimumVscale
                            : CostVF.getFixedValue();
       LLVM_DEBUG(dbgs() << "LV: Vector loop of width " << CostVF

llvm/lib/Transforms/Vectorize/VPlan.h

public:
   /// otherwise.
   Instruction *getUnderlyingInstr() {
     return cast<Instruction>(getVPSingleValue()->getUnderlyingValue());
   }
   const Instruction *getUnderlyingInstr() const {
     return cast<Instruction>(getVPSingleValue()->getUnderlyingValue());
   }
 
+  bool hasUnderlyingInstr() const {
+    return getNumDefinedValues() == 1 &&
+           getVPSingleValue()->getUnderlyingValue() != nullptr;
+  }
+
   /// Method to support type inquiry through isa, cast, and dyn_cast.
   static inline bool classof(const VPDef *D) {
     // All VPDefs are also VPRecipeBases.
     return true;
   }
 
   static inline bool classof(const VPUser *U) {
     return U->getVPUserID() == VPUser::VPUserID::Recipe;

llvm/lib/Transforms/Vectorize/VPlanCostModel.h

+//===- SiFive_VPlanCostModel.cpp - Vectorizer Cost Model ------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+///
+/// \file
+/// VPlan-based cost model
+///
+//===----------------------------------------------------------------------===//
+#include "llvm/Analysis/TargetTransformInfo.h"
+#include "llvm/IR/Instruction.h"
+
+#include "VPlan.h"
+#include "VPlanValue.h"
+
+namespace llvm {
+class Type;
+class TargetTransformInfo;
+class LoopVectorizationCostModel;
+
+class VPlanCostModel {
+public:
+  explicit VPlanCostModel(const VPlan &Plan, const TargetTransformInfo &TTI,
+                          LoopVectorizationCostModel &CM)
+      : Plan(Plan), TTI(TTI), CM(CM) {}
+
+  /// Return cost of the VPlan for a given \p VF
+  InstructionCost expectedCost(ElementCount VF, bool &IsVec);
+
+private:
+  /// Return individual cost of the \p VPBasicBlock for a given \p VF
+  InstructionCost getCost(const VPBlockBase *Block, ElementCount VF,
+                          bool &IsVec);
+
+  /// Return individual cost of the \p Recipe for a given \p VF
+  InstructionCost getCost(const VPRecipeBase *Recipe, ElementCount VF,
+                          bool &IsVec);
+
+  /// Return individual cost of the \p Recipe for a given \p VF
+  InstructionCost getLegacyInstructionCost(Instruction *I, ElementCount VF);
+
+  InstructionCost getMemoryOpCost(const VPWidenMemoryInstructionRecipe *VPWMIR,
+                                  ElementCount VF, Type *&VectorTy);
+
+  /// Return cost of the individual memory operation of a instruction \p I of a
+  /// given type \p Ty
+  InstructionCost getMemoryOpCost(const Instruction *I, Type *Ty,
+                                  bool IsConsecutive, bool IsMasked,
+                                  bool IsReverse);
+
+  Type *truncateToMinimalBitwidth(Type *ValTy, Instruction *I) const;
+  /// VPlan for which cost is computed
+  const VPlan &Plan;
+
+  /// Vector target information.
+  const TargetTransformInfo &TTI;
+
+  /// FIXME: Legacy model is only here during our transition to the vplan-based
+  /// model
+  LoopVectorizationCostModel &CM;
+
+  /// Use same cost kind in the cost model
+  const TargetTransformInfo::TargetCostKind CostKind = TTI::TCK_RecipThroughput;
+};
+} // namespace llvm

llvm/lib/Transforms/Vectorize/VPlanCostModel.cpp

+//===- VPlanCostModel.h - VPlan-based Vectorizer Cost Model ---------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+///
+/// \file
+/// VPlan-based cost model
+///
+//===----------------------------------------------------------------------===//
+
+#include "llvm/ADT/TypeSwitch.h"
+#include "llvm/Analysis/TargetTransformInfo.h"
+#include "llvm/Analysis/VectorUtils.h"
+#include "llvm/IR/Instruction.h"
+#include "llvm/IR/Operator.h"
+#include "llvm/Support/Debug.h"
+
+#include "VPlan.h"
+#include "VPlanCFG.h"
+#include "VPlanCostModel.h"
+#include "VPlanValue.h"
+
+using namespace llvm;
+
+#define DEBUG_TYPE "vplan-cost-model"
+
+namespace llvm {
+InstructionCost VPlanCostModel::expectedCost(ElementCount VF, bool &IsVec) {
+  InstructionCost VectorIterCost = 0;
+  for (const VPBlockBase *Block : vp_depth_first_deep(Plan.getEntry()))
+    VectorIterCost += getCost(Block, VF, IsVec);
+
+  return VectorIterCost;
+}
+
+InstructionCost VPlanCostModel::getCost(const VPBlockBase *Block,
+                                        ElementCount VF, bool &IsVec) {
+  return TypeSwitch<const VPBlockBase *, InstructionCost>(Block)
+      .Case<VPBasicBlock>([&](const VPBasicBlock *BBlock) {
+        InstructionCost Cost = 0;
+        for (const VPRecipeBase &Recipe : *BBlock)
+          Cost += getCost(Block, VF, IsVec);
+        return Cost;
+      })
+      .Default([&](const VPBlockBase *BBlock) -> InstructionCost { return 0; });
+}
+
+InstructionCost VPlanCostModel::getCost(const VPRecipeBase *Recipe,
+                                        ElementCount VF, bool &IsVec) {
+  Type *VectorTy = nullptr;
+  InstructionCost Cost =
+      TypeSwitch<const VPRecipeBase *, InstructionCost>(Recipe)
+          .Case<VPWidenMemoryInstructionRecipe>(
+              [&](const VPWidenMemoryInstructionRecipe *VPWMIR) {
+                return getMemoryOpCost(VPWMIR, VF, VectorTy);
+              })
+          .Default([&](const VPRecipeBase *R) -> InstructionCost {
+            if (!R->hasUnderlyingInstr()) {
+              LLVM_DEBUG(
+                  dbgs() << "VPlanCM: unsupported recipe ";
+                  VPSlotTracker SlotTracker((Recipe->getParent())
+                                                ? Recipe->getParent()->getPlan()
+                                                : nullptr);
+                  Recipe->print(dbgs(), Twine(), SlotTracker); dbgs() << '\n');
+              return 0;
+            }
+            Instruction *I = const_cast<Instruction *>(R->getUnderlyingInstr());
+            return getLegacyInstructionCost(I, VF);
+          });
+
+  LLVM_DEBUG(dbgs() << "VPlanCM: cost " << Cost << " for VF " << VF
+                    << " for VPInstruction: ";
+             VPSlotTracker SlotTracker((Recipe->getParent())
+                                           ? Recipe->getParent()->getPlan()
+                                           : nullptr);
+             Recipe->print(dbgs(), Twine(), SlotTracker); dbgs() << '\n');
+  if (VF.isVector() && VectorTy->isVectorTy()) {
+    if (unsigned NumParts = TTI.getNumberOfParts(VectorTy)) {
+      if (VF.isScalable())
+        // <vscale x 1 x iN> is assumed to be profitable over iN because
+        // scalable registers are a distinct register class from scalar ones.
+        // If we ever find a target which wants to lower scalable vectors
+        // back to scalars, we'll need to update this code to explicitly
+        // ask TTI about the register class uses for each part.
+        IsVec |= NumParts <= VF.getKnownMinValue();
+      else
+        IsVec |= NumParts < VF.getKnownMinValue();
+    } else
+      Cost = InstructionCost::getInvalid();
+  }
+  return Cost;
+}
+
+InstructionCost VPlanCostModel::getMemoryOpCost(const Instruction *I, Type *Ty,
+                                                bool IsConsecutive,
+                                                bool IsMasked, bool IsReverse) {
+  const Align Alignment = getLoadStoreAlignment(const_cast<Instruction *>(I));
+  const Value *Ptr = getLoadStorePointerOperand(I);
+  unsigned AS = getLoadStoreAddressSpace(const_cast<Instruction *>(I));
+  if (IsConsecutive) {
+    InstructionCost Cost = 0;
+    if (IsMasked) {
+      Cost += TTI.getMaskedMemoryOpCost(I->getOpcode(), Ty, Alignment, AS,
+                                        CostKind);
+    } else {
+      TTI::OperandValueInfo OpInfo = TTI::getOperandInfo(I->getOperand(0));
+      Cost += TTI.getMemoryOpCost(I->getOpcode(), Ty, Alignment, AS, CostKind,
+                                  OpInfo, I);
+    }
+    if (IsReverse)
+      Cost +=
+          TTI.getShuffleCost(TargetTransformInfo::SK_Reverse,
+                             cast<VectorType>(Ty), std::nullopt, CostKind, 0);
+    return Cost;
+  }
+  return TTI.getAddressComputationCost(Ty) +
+         TTI.getGatherScatterOpCost(I->getOpcode(), Ty, Ptr, IsMasked,
+                                    Alignment, CostKind, I);
+}
+
+InstructionCost
+VPlanCostModel::getMemoryOpCost(const VPWidenMemoryInstructionRecipe *VPWMIR,
+                                ElementCount VF, Type *&VectorTy) {
+  Instruction *I = &VPWMIR->getIngredient();
+  Type *ValTy = getLoadStoreType(I);
+  ValTy = truncateToMinimalBitwidth(ValTy, I);
+  const bool IsMasked = VPWMIR->getMask() != nullptr;
+  VectorTy = VectorType::get(ValTy, VF);
+
+  return getMemoryOpCost(I, VectorTy, VPWMIR->isConsecutive(), IsMasked,
+                         VPWMIR->isReverse());
+}
+
+} // namespace llvm