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

[Graham Hunter 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 =
----------------
huntergr-arm 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.

If we use `-force-target-supports-masked-memory-ops` and `-force-target-supports-gather-scatter-ops`, we don't reach this code (and therefore don't crash), so it's only in the case that we need to scalarize that we have problems.

I guess the problem here is that we're requesting the cost here before we have the opportunity to create an appropriate recipe in makeCallWideningDecisions.

@fhahn, any objections?

(Side note, if a masked variant is available, we could hoist the call out of the conditional block and only scalarize the scatter. But that's separate from fixing the crash...)

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