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

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: seantalts

<details>
<summary>Changes</summary>

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.

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


1 Files Affected:

- (modified) llvm/lib/Transforms/Vectorize/LoopVectorize.cpp (+25-11) 


``````````diff
diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
index 7ad454d4a1797..fbfd3dc9a11a2 100644
--- a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
+++ b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
@@ -2098,6 +2098,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
@@ -2105,22 +2121,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;

``````````

</details>


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