1 files changed, 16 insertions, 49 deletions

diff --git a/llvm/lib/Transforms/Utils/InjectTLIMappings.cpp b/llvm/lib/Transforms/Utils/InjectTLIMappings.cpp
index 55bcb6f3b121..dab0be3a9fde 100644
--- a/llvm/lib/Transforms/Utils/InjectTLIMappings.cpp
+++ b/llvm/lib/Transforms/Utils/InjectTLIMappings.cpp
@@ -19,7 +19,6 @@
 #include "llvm/Analysis/TargetLibraryInfo.h"
 #include "llvm/Analysis/VectorUtils.h"
 #include "llvm/IR/InstIterator.h"
-#include "llvm/Transforms/Utils.h"
 #include "llvm/Transforms/Utils/ModuleUtils.h"
 
 using namespace llvm;
@@ -40,7 +39,7 @@ STATISTIC(NumCompUsedAdded,
 /// CI (other than void) need to be widened to a VectorType of VF
 /// lanes.
 static void addVariantDeclaration(CallInst &CI, const ElementCount &VF,
-                                  const StringRef VFName) {
+                                  bool Predicate, const StringRef VFName) {
   Module *M = CI.getModule();
 
   // Add function declaration.
@@ -50,6 +49,8 @@ static void addVariantDeclaration(CallInst &CI, const ElementCount &VF,
     Tys.push_back(ToVectorTy(ArgOperand->getType(), VF));
   assert(!CI.getFunctionType()->isVarArg() &&
          "VarArg functions are not supported.");
+  if (Predicate)
+    Tys.push_back(ToVectorTy(Type::getInt1Ty(RetTy->getContext()), VF));
   FunctionType *FTy = FunctionType::get(RetTy, Tys, /*isVarArg=*/false);
   Function *VectorF =
       Function::Create(FTy, Function::ExternalLinkage, VFName, M);
@@ -89,19 +90,19 @@ static void addMappingsFromTLI(const TargetLibraryInfo &TLI, CallInst &CI) {
   const SetVector<StringRef> OriginalSetOfMappings(Mappings.begin(),
                                                    Mappings.end());
 
-  auto AddVariantDecl = [&](const ElementCount &VF) {
+  auto AddVariantDecl = [&](const ElementCount &VF, bool Predicate) {
     const std::string TLIName =
-        std::string(TLI.getVectorizedFunction(ScalarName, VF));
+        std::string(TLI.getVectorizedFunction(ScalarName, VF, Predicate));
     if (!TLIName.empty()) {
-      std::string MangledName =
-          VFABI::mangleTLIVectorName(TLIName, ScalarName, CI.arg_size(), VF);
+      std::string MangledName = VFABI::mangleTLIVectorName(
+          TLIName, ScalarName, CI.arg_size(), VF, Predicate);
       if (!OriginalSetOfMappings.count(MangledName)) {
         Mappings.push_back(MangledName);
         ++NumCallInjected;
       }
       Function *VariantF = M->getFunction(TLIName);
       if (!VariantF)
-        addVariantDeclaration(CI, VF, TLIName);
+        addVariantDeclaration(CI, VF, Predicate, TLIName);
     }
   };
 
@@ -109,13 +110,15 @@ static void addMappingsFromTLI(const TargetLibraryInfo &TLI, CallInst &CI) {
   ElementCount WidestFixedVF, WidestScalableVF;
   TLI.getWidestVF(ScalarName, WidestFixedVF, WidestScalableVF);
 
-  for (ElementCount VF = ElementCount::getFixed(2);
-       ElementCount::isKnownLE(VF, WidestFixedVF); VF *= 2)
-    AddVariantDecl(VF);
+  for (bool Predicated : {false, true}) {
+    for (ElementCount VF = ElementCount::getFixed(2);
+         ElementCount::isKnownLE(VF, WidestFixedVF); VF *= 2)
+      AddVariantDecl(VF, Predicated);
 
-  // TODO: Add scalable variants once we're able to test them.
-  assert(WidestScalableVF.isZero() &&
-         "Scalable vector mappings not yet supported");
+    for (ElementCount VF = ElementCount::getScalable(2);
+         ElementCount::isKnownLE(VF, WidestScalableVF); VF *= 2)
+      AddVariantDecl(VF, Predicated);
+  }
 
   VFABI::setVectorVariantNames(&CI, Mappings);
 }
@@ -138,39 +141,3 @@ PreservedAnalyses InjectTLIMappings::run(Function &F,
   // Even if the pass adds IR attributes, the analyses are preserved.
   return PreservedAnalyses::all();
 }
-
-////////////////////////////////////////////////////////////////////////////////
-// Legacy PM Implementation.
-////////////////////////////////////////////////////////////////////////////////
-bool InjectTLIMappingsLegacy::runOnFunction(Function &F) {
-  const TargetLibraryInfo &TLI =
-      getAnalysis<TargetLibraryInfoWrapperPass>().getTLI(F);
-  return runImpl(TLI, F);
-}
-
-void InjectTLIMappingsLegacy::getAnalysisUsage(AnalysisUsage &AU) const {
-  AU.setPreservesCFG();
-  AU.addRequired<TargetLibraryInfoWrapperPass>();
-  AU.addPreserved<TargetLibraryInfoWrapperPass>();
-  AU.addPreserved<ScalarEvolutionWrapperPass>();
-  AU.addPreserved<AAResultsWrapperPass>();
-  AU.addPreserved<LoopAccessLegacyAnalysis>();
-  AU.addPreserved<DemandedBitsWrapperPass>();
-  AU.addPreserved<OptimizationRemarkEmitterWrapperPass>();
-  AU.addPreserved<GlobalsAAWrapperPass>();
-}
-
-////////////////////////////////////////////////////////////////////////////////
-// Legacy Pass manager initialization
-////////////////////////////////////////////////////////////////////////////////
-char InjectTLIMappingsLegacy::ID = 0;
-
-INITIALIZE_PASS_BEGIN(InjectTLIMappingsLegacy, DEBUG_TYPE,
-                      "Inject TLI Mappings", false, false)
-INITIALIZE_PASS_DEPENDENCY(TargetLibraryInfoWrapperPass)
-INITIALIZE_PASS_END(InjectTLIMappingsLegacy, DEBUG_TYPE, "Inject TLI Mappings",
-                    false, false)
-
-FunctionPass *llvm::createInjectTLIMappingsLegacyPass() {
-  return new InjectTLIMappingsLegacy();
-}