[llvm] [LoopVectorize] Price vector library variants in getVectorCallCost (PR #202085)

[via llvm-commits]

https://github.com/seantalts updated https://github.com/llvm/llvm-project/pull/202085

>From 7344a61d09712e0bce85977897d9b2932a7a9c5a Mon Sep 17 00:00:00 2001
From: Sean Talts <sean.talts at gmail.com>
Date: Fri, 5 Jun 2026 18:32:33 +0000
Subject: [PATCH] [LoopVectorize] Price vector library variants in
 getVectorCallCost instead of asserting

During loop vectorization, computePredInstDiscount queries the cost of instructions
at vector VF using getInstructionCost. If an instruction is a CallInst that has a
vector library variant, getInstructionCost delegates to getVectorCallCost, which
previously asserted that vector library variants should not reach it.

Instead of asserting, query the cost of the vector library variant using
TTI.getCallInstrCost and return it.
---
 .../Transforms/Vectorize/LoopVectorize.cpp    | 36 +++++++++++++------
 1 file changed, 25 insertions(+), 11 deletions(-)

diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
index fcfe11b6e9de6..d7335cbbda912 100644
--- a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
+++ b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
@@ -2102,6 +2102,22 @@ static unsigned estimateElementCount(ElementCount VF,
   return EstimatedVF;
 }
 
+/// Returns the vector variant function of \p CI usable at \p VF, respecting \p
+/// MaskRequired. Returns nullptr if none is found.
+static Function *getVectorLibraryVariantFor(const CallInst &CI, ElementCount VF,
+                                            bool MaskRequired,
+                                            const TargetLibraryInfo *TLI) {
+  if (!TLI || CI.isNoBuiltin())
+    return nullptr;
+  for (const auto &Info : VFDatabase::getMappings(CI)) {
+    if (Info.Shape.VF == VF && (!MaskRequired || Info.isMasked())) {
+      if (auto *F = CI.getModule()->getFunction(Info.VectorName))
+        return F;
+    }
+  }
+  return nullptr;
+}
+
 /// Returns true iff \p CI has a library vector variant usable at \p VF: a
 /// mapping with matching VF, masked if required, whose vector function is
 /// declared in the module. Such variants are priced by
@@ -2109,22 +2125,20 @@ static unsigned estimateElementCount(ElementCount VF,
 static bool hasVectorLibraryVariantFor(const CallInst &CI, ElementCount VF,
                                        bool MaskRequired,
                                        const TargetLibraryInfo *TLI) {
-  if (!TLI || CI.isNoBuiltin())
-    return false;
-  return any_of(VFDatabase::getMappings(CI), [&](const VFInfo &Info) {
-    return Info.Shape.VF == VF && (!MaskRequired || Info.isMasked()) &&
-           CI.getModule()->getFunction(Info.VectorName);
-  });
+  return getVectorLibraryVariantFor(CI, VF, MaskRequired, TLI) != nullptr;
 }
 
 InstructionCost
 LoopVectorizationCostModel::getVectorCallCost(CallInst *CI,
                                               ElementCount VF) const {
-  // Vector library variants are priced by VPWidenCallRecipe::computeCost and
-  // should not reach this function.
-  assert((VF.isScalar() ||
-          !hasVectorLibraryVariantFor(*CI, VF, isMaskRequired(CI), TLI)) &&
-         "getVectorCallCost does not price vector library variants");
+  if (VF.isVector()) {
+    if (Function *Variant =
+            getVectorLibraryVariantFor(*CI, VF, isMaskRequired(CI), TLI)) {
+      return TTI.getCallInstrCost(nullptr, Variant->getReturnType(),
+                                  Variant->getFunctionType()->params(),
+                                  Config.CostKind);
+    }
+  }
 
   Type *RetTy = CI->getType();
   SmallVector<Type *, 4> Tys;