1 files changed, 16 insertions, 5 deletions
diff --git a/llvm/lib/Transforms/Utils/InjectTLIMappings.cpp b/llvm/lib/Transforms/Utils/InjectTLIMappings.cpp
index a2b72e4e7f03..a1e160d144dc 100644
--- a/llvm/lib/Transforms/Utils/InjectTLIMappings.cpp
+++ b/llvm/lib/Transforms/Utils/InjectTLIMappings.cpp
@@ -40,7 +40,7 @@ STATISTIC(NumCompUsedAdded,
 /// lanes. The TLI assumes that all parameters and the return type of
 /// CI (other than void) need to be widened to a VectorType of VF
 /// lanes.
-static void addVariantDeclaration(CallInst &CI, const unsigned VF,
+static void addVariantDeclaration(CallInst &CI, const ElementCount &VF,
                                   const StringRef VFName) {
   Module *M = CI.getModule();
 
@@ -89,9 +89,8 @@ static void addMappingsFromTLI(const TargetLibraryInfo &TLI, CallInst &CI) {
   Module *M = CI.getModule();
   const SetVector<StringRef> OriginalSetOfMappings(Mappings.begin(),
                                                    Mappings.end());
-  //  All VFs in the TLI are powers of 2.
-  for (unsigned VF = 2, WidestVF = TLI.getWidestVF(ScalarName); VF <= WidestVF;
-       VF *= 2) {
+
+  auto AddVariantDecl = [&](const ElementCount &VF) {
     const std::string TLIName =
         std::string(TLI.getVectorizedFunction(ScalarName, VF));
     if (!TLIName.empty()) {
@@ -105,7 +104,19 @@ static void addMappingsFromTLI(const TargetLibraryInfo &TLI, CallInst &CI) {
       if (!VariantF)
         addVariantDeclaration(CI, VF, TLIName);
     }
-  }
+  };
+
+  //  All VFs in the TLI are powers of 2.
+  ElementCount WidestFixedVF, WidestScalableVF;
+  TLI.getWidestVF(ScalarName, WidestFixedVF, WidestScalableVF);
+
+  for (ElementCount VF = ElementCount::getFixed(2);
+       ElementCount::isKnownLE(VF, WidestFixedVF); VF *= 2)
+    AddVariantDecl(VF);
+
+  // TODO: Add scalable variants once we're able to test them.
+  assert(WidestScalableVF.isZero() &&
+         "Scalable vector mappings not yet supported");
 
   VFABI::setVectorVariantNames(&CI, Mappings);
 }