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

[via llvm-commits]

================
@@ -2098,29 +2098,43 @@ 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
 /// VPWidenCallRecipe::computeCost rather than by scalarization.
 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 =
----------------
seantalts wrote:

Gotcha, sounds good to me. Thank you! 

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