[llvm] [LoopVectorize] Don't assert in getVectorCallCost for vector library variants (PR #202085)

[Florian Hahn 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 =
----------------
fhahn wrote:

> This will fix the crash, but I'm not sure we need to return a cost for the vector variant here. The output IR ends up scalarizing the call, so it would be better to return the cost for that as below. I think we may just need to remove the assert.

I think the code-path we reach `getVectorCallCost` from `computePredInstDiscount` use that to get the cost of the wide variant, and compares this to the cost of VF x scalar call cost, to decide if forced scalarization is profitable for a whole tree of operations. I think for that to work as expected, we should return the most profitable cost for the call at the given VF, which is either the cost of the vector call (or intrinsic) or scalarization (in which case `computePredInstDiscount`  would not need to force the call to scalar).

Given that the latest change did not impact the test, it may be good to try to add a test where the cost returned actually impacts the decision. Not sure if that's possible though.

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