llvm/lib/Transforms/Utils/InjectTLIMappings.cpp

   1 //===- InjectTLIMAppings.cpp - TLI to VFABI attribute injection  ----------===//
   2 //
   3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
   4 // See https://llvm.org/LICENSE.txt for license information.
   5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
   6 //
   7 //===----------------------------------------------------------------------===//
   8 //
   9 // Populates the VFABI attribute with the scalar-to-vector mappings
  10 // from the TargetLibraryInfo.
  11 //
  12 //===----------------------------------------------------------------------===//
  13
  14 #include "llvm/Transforms/Utils/InjectTLIMappings.h"
  15 #include "llvm/ADT/Statistic.h"
  16 #include "llvm/Analysis/DemandedBits.h"
  17 #include "llvm/Analysis/GlobalsModRef.h"
  18 #include "llvm/Analysis/OptimizationRemarkEmitter.h"
  19 #include "llvm/Analysis/TargetLibraryInfo.h"
  20 #include "llvm/Analysis/VectorUtils.h"
  21 #include "llvm/IR/InstIterator.h"
  22 #include "llvm/IR/VFABIDemangler.h"
  23 #include "llvm/Transforms/Utils/ModuleUtils.h"
  24
  25 using namespace llvm;
  26
  27 #define DEBUG_TYPE "inject-tli-mappings"
  28
  29 STATISTIC(NumCallInjected,
  30           "Number of calls in which the mappings have been injected.");
  31
  32 STATISTIC(NumVFDeclAdded,
  33           "Number of function declarations that have been added.");
  34 STATISTIC(NumCompUsedAdded,
  35           "Number of `@llvm.compiler.used` operands that have been added.");
  36
  37 /// A helper function that adds the vector variant declaration for vectorizing
  38 /// the CallInst \p CI with a vectorization factor of \p VF lanes. For each
  39 /// mapping, TLI provides a VABI prefix, which contains all information required
  40 /// to create vector function declaration.
  41 static void addVariantDeclaration(CallInst &CI, const ElementCount &VF,
  42                                   const VecDesc *VD) {
  43   Module *M = CI.getModule();
  44   FunctionType *ScalarFTy = CI.getFunctionType();
  45
  46   assert(!ScalarFTy->isVarArg() && "VarArg functions are not supported.");
  47
  48   const std::optional<VFInfo> Info = VFABI::tryDemangleForVFABI(
  49       VD->getVectorFunctionABIVariantString(), ScalarFTy);
  50
  51   assert(Info && "Failed to demangle vector variant");
  52   assert(Info->Shape.VF == VF && "Mangled name does not match VF");
  53
  54   const StringRef VFName = VD->getVectorFnName();
  55   FunctionType *VectorFTy = VFABI::createFunctionType(*Info, ScalarFTy);
  56   Function *VecFunc =
  57       Function::Create(VectorFTy, Function::ExternalLinkage, VFName, M);
  58   VecFunc->copyAttributesFrom(CI.getCalledFunction());
  59   ++NumVFDeclAdded;
  60   LLVM_DEBUG(dbgs() << DEBUG_TYPE << ": Added to the module: `" << VFName
  61                     << "` of type " << *VectorFTy << "\n");
  62
  63   // Make function declaration (without a body) "sticky" in the IR by
  64   // listing it in the @llvm.compiler.used intrinsic.
  65   assert(!VecFunc->size() && "VFABI attribute requires `@llvm.compiler.used` "
  66                              "only on declarations.");
  67   appendToCompilerUsed(*M, {VecFunc});
  68   LLVM_DEBUG(dbgs() << DEBUG_TYPE << ": Adding `" << VFName
  69                     << "` to `@llvm.compiler.used`.\n");
  70   ++NumCompUsedAdded;
  71 }
  72
  73 static void addMappingsFromTLI(const TargetLibraryInfo &TLI, CallInst &CI) {
  74   // This is needed to make sure we don't query the TLI for calls to
  75   // bitcast of function pointers, like `%call = call i32 (i32*, ...)
  76   // bitcast (i32 (...)* @goo to i32 (i32*, ...)*)(i32* nonnull %i)`,
  77   // as such calls make the `isFunctionVectorizable` raise an
  78   // exception.
  79   if (CI.isNoBuiltin() || !CI.getCalledFunction())
  80     return;
  81
  82   StringRef ScalarName = CI.getCalledFunction()->getName();
  83
  84   // Nothing to be done if the TLI thinks the function is not
  85   // vectorizable.
  86   if (!TLI.isFunctionVectorizable(ScalarName))
  87     return;
  88   SmallVector<std::string, 8> Mappings;
  89   VFABI::getVectorVariantNames(CI, Mappings);
  90   Module *M = CI.getModule();
  91   const SetVector<StringRef> OriginalSetOfMappings(Mappings.begin(),
  92                                                    Mappings.end());
  93
  94   auto AddVariantDecl = [&](const ElementCount &VF, bool Predicate) {
  95     const VecDesc *VD = TLI.getVectorMappingInfo(ScalarName, VF, Predicate);
  96     if (VD && !VD->getVectorFnName().empty()) {
  97       std::string MangledName = VD->getVectorFunctionABIVariantString();
  98       if (!OriginalSetOfMappings.count(MangledName)) {
  99         Mappings.push_back(MangledName);
 100         ++NumCallInjected;
 101       }
 102       Function *VariantF = M->getFunction(VD->getVectorFnName());
 103       if (!VariantF)
 104         addVariantDeclaration(CI, VF, VD);
 105     }
 106   };
 107
 108   //  All VFs in the TLI are powers of 2.
 109   ElementCount WidestFixedVF, WidestScalableVF;
 110   TLI.getWidestVF(ScalarName, WidestFixedVF, WidestScalableVF);
 111
 112   for (bool Predicated : {false, true}) {
 113     for (ElementCount VF = ElementCount::getFixed(2);
 114          ElementCount::isKnownLE(VF, WidestFixedVF); VF *= 2)
 115       AddVariantDecl(VF, Predicated);
 116
 117     for (ElementCount VF = ElementCount::getScalable(2);
 118          ElementCount::isKnownLE(VF, WidestScalableVF); VF *= 2)
 119       AddVariantDecl(VF, Predicated);
 120   }
 121
 122   VFABI::setVectorVariantNames(&CI, Mappings);
 123 }
 124
 125 static bool runImpl(const TargetLibraryInfo &TLI, Function &F) {
 126   for (auto &I : instructions(F))
 127     if (auto CI = dyn_cast<CallInst>(&I))
 128       addMappingsFromTLI(TLI, *CI);
 129   // Even if the pass adds IR attributes, the analyses are preserved.
 130   return false;
 131 }
 132
 133 ////////////////////////////////////////////////////////////////////////////////
 134 // New pass manager implementation.
 135 ////////////////////////////////////////////////////////////////////////////////
 136 PreservedAnalyses InjectTLIMappings::run(Function &F,
 137                                          FunctionAnalysisManager &AM) {
 138   const TargetLibraryInfo &TLI = AM.getResult<TargetLibraryAnalysis>(F);
 139   runImpl(TLI, F);
 140   // Even if the pass adds IR attributes, the analyses are preserved.
 141   return PreservedAnalyses::all();
 142 }