[Codegen][ReplaceWithVecLib] add pass to replace vector intrinsics wi…

…th calls to vector library

This patch adds a pass to replace calls to vector intrinsics
(i.e., LLVM intrinsics operating on vector operands) with
calls to a vector library.

Currently, calls to LLVM intrinsics are only replaced with
calls to vector libraries when scalar calls to intrinsics are
vectorized by the Loop- or SLP-Vectorizer.

With this pass, it is now possible to replace calls to LLVM
intrinsics already operating on vector operands, e.g., if
such code was generated by MLIR. For the replacement,
information from the TargetLibraryInfo, e.g., as specified
via -vector-library is used.

Differential Revision: https://reviews.llvm.org/D95373

llvm/include/llvm/CodeGen/CodeGenPassBuilder.h

@@ -29,6 +29,7 @@
 #include "llvm/CodeGen/MachineModuleInfo.h"
 #include "llvm/CodeGen/MachinePassManager.h"
 #include "llvm/CodeGen/PreISelIntrinsicLowering.h"
+#include "llvm/CodeGen/ReplaceWithVeclib.h"
 #include "llvm/CodeGen/UnreachableBlockElim.h"
 #include "llvm/IR/IRPrintingPasses.h"
 #include "llvm/IR/PassManager.h"
@@ -650,6 +651,12 @@ void CodeGenPassBuilder<Derived>::addIRPasses(AddIRPass &addPass) const {
   if (getOptLevel() != CodeGenOpt::None && !Opt.DisableConstantHoisting)
     addPass(ConstantHoistingPass());
 
+  if (getOptLevel() != CodeGenOpt::None) {
+    // Replace calls to LLVM intrinsics (e.g., exp, log) operating on vector
+    // operands with calls to the corresponding functions in a vector library.
+    addPass(ReplaceWithVeclib());
+  }
+
   if (getOptLevel() != CodeGenOpt::None && !Opt.DisablePartialLibcallInlining)
     addPass(PartiallyInlineLibCallsPass());
 

llvm/include/llvm/CodeGen/MachinePassRegistry.def

@@ -39,6 +39,7 @@ FUNCTION_PASS("mergeicmps", MergeICmpsPass, ())
 FUNCTION_PASS("lower-constant-intrinsics", LowerConstantIntrinsicsPass, ())
 FUNCTION_PASS("unreachableblockelim", UnreachableBlockElimPass, ())
 FUNCTION_PASS("consthoist", ConstantHoistingPass, ())
+FUNCTION_PASS("replace-with-veclib", ReplaceWithVeclib, ())
 FUNCTION_PASS("partially-inline-libcalls", PartiallyInlineLibCallsPass, ())
 FUNCTION_PASS("ee-instrument", EntryExitInstrumenterPass, (false))
 FUNCTION_PASS("post-inline-ee-instrument", EntryExitInstrumenterPass, (true))

llvm/include/llvm/CodeGen/Passes.h

@@ -448,6 +448,10 @@ namespace llvm {
   /// shuffles.
   FunctionPass *createExpandReductionsPass();
 
+  // This pass replaces intrinsics operating on vector operands with calls to
+  // the corresponding function in a vector library (e.g., SVML, libmvec).
+  FunctionPass *createReplaceWithVeclibLegacyPass();
+
   // This pass expands memcmp() to load/stores.
   FunctionPass *createExpandMemCmpPass();
 

llvm/include/llvm/CodeGen/ReplaceWithVeclib.h

@@ -0,0 +1,38 @@
+//===- ReplaceWithVeclib.h - Replace vector instrinsics with veclib calls -===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Replaces calls to LLVM vector intrinsics (i.e., calls to LLVM intrinsics
+// with vector operands) with matching calls to functions from a vector
+// library (e.g., libmvec, SVML) according to TargetLibraryInfo.
+//
+//===----------------------------------------------------------------------===//
+#ifndef LLVM_TRANSFORMS_UTILS_REPLACEWITHVECLIB_H
+#define LLVM_TRANSFORMS_UTILS_REPLACEWITHVECLIB_H
+
+#include "llvm/IR/PassManager.h"
+#include "llvm/InitializePasses.h"
+
+namespace llvm {
+class ReplaceWithVeclib : public PassInfoMixin<ReplaceWithVeclib> {
+public:
+  PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
+};
+
+// Legacy pass
+class ReplaceWithVeclibLegacy : public FunctionPass {
+public:
+  static char ID;
+  ReplaceWithVeclibLegacy() : FunctionPass(ID) {
+    initializeReplaceWithVeclibLegacyPass(*PassRegistry::getPassRegistry());
+  }
+  void getAnalysisUsage(AnalysisUsage &AU) const override;
+  bool runOnFunction(Function &F) override;
+};
+
+} // End namespace llvm
+#endif // LLVM_TRANSFORMS_UTILS_REPLACEWITHVECLIB_H

llvm/include/llvm/InitializePasses.h

@@ -380,6 +380,7 @@ void initializeRegionPrinterPass(PassRegistry&);
 void initializeRegionViewerPass(PassRegistry&);
 void initializeRegisterCoalescerPass(PassRegistry&);
 void initializeRenameIndependentSubregsPass(PassRegistry&);
+void initializeReplaceWithVeclibLegacyPass(PassRegistry &);
 void initializeResetMachineFunctionPass(PassRegistry&);
 void initializeReversePostOrderFunctionAttrsLegacyPassPass(PassRegistry&);
 void initializeRewriteStatepointsForGCLegacyPassPass(PassRegistry &);

llvm/lib/CodeGen/CMakeLists.txt

@@ -147,6 +147,7 @@ add_llvm_component_library(LLVMCodeGen
   RegisterUsageInfo.cpp
   RegUsageInfoCollector.cpp
   RegUsageInfoPropagate.cpp
+  ReplaceWithVeclib.cpp
   ResetMachineFunctionPass.cpp
   SafeStack.cpp
   SafeStackLayout.cpp

llvm/lib/CodeGen/ReplaceWithVeclib.cpp

@@ -0,0 +1,256 @@
+//=== ReplaceWithVeclib.cpp - Replace vector instrinsics with veclib calls ===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Replaces calls to LLVM vector intrinsics (i.e., calls to LLVM intrinsics
+// with vector operands) with matching calls to functions from a vector
+// library (e.g., libmvec, SVML) according to TargetLibraryInfo.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/CodeGen/ReplaceWithVeclib.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/Statistic.h"
+#include "llvm/Analysis/DemandedBits.h"
+#include "llvm/Analysis/GlobalsModRef.h"
+#include "llvm/Analysis/OptimizationRemarkEmitter.h"
+#include "llvm/Analysis/TargetLibraryInfo.h"
+#include "llvm/Analysis/VectorUtils.h"
+#include "llvm/CodeGen/Passes.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/InstIterator.h"
+#include "llvm/IR/IntrinsicInst.h"
+#include "llvm/Transforms/Utils/ModuleUtils.h"
+
+using namespace llvm;
+
+#define DEBUG_TYPE "replace-with-veclib"
+
+STATISTIC(NumCallsReplaced,
+          "Number of calls to intrinsics that have been replaced.");
+
+STATISTIC(NumTLIFuncDeclAdded,
+          "Number of vector library function declarations added.");
+
+STATISTIC(NumFuncUsedAdded,
+          "Number of functions added to `llvm.compiler.used`");
+
+static bool replaceWithTLIFunction(CallInst &CI, const StringRef TLIName) {
+  Module *M = CI.getModule();
+
+  Function *OldFunc = CI.getCalledFunction();
+
+  // Check if the vector library function is already declared in this module,
+  // otherwise insert it.
+  Function *TLIFunc = M->getFunction(TLIName);
+  if (!TLIFunc) {
+    TLIFunc = Function::Create(OldFunc->getFunctionType(),
+                               Function::ExternalLinkage, TLIName, *M);
+    TLIFunc->copyAttributesFrom(OldFunc);
+
+    LLVM_DEBUG(dbgs() << DEBUG_TYPE << ": Added vector library function `"
+                      << TLIName << "` of type `" << *(TLIFunc->getType())
+                      << "` to module.\n");
+
+    ++NumTLIFuncDeclAdded;
+
+    // Add the freshly created function to llvm.compiler.used,
+    // similar to as it is done in InjectTLIMappings
+    appendToCompilerUsed(*M, {TLIFunc});
+
+    LLVM_DEBUG(dbgs() << DEBUG_TYPE << ": Adding `" << TLIName
+                      << "` to `@llvm.compiler.used`.\n");
+    ++NumFuncUsedAdded;
+  }
+
+  // Replace the call to the vector intrinsic with a call
+  // to the corresponding function from the vector library.
+  IRBuilder<> IRBuilder{&CI};
+  SmallVector<Value *> Args(CI.arg_operands());
+  // Preserve the operand bundles.
+  SmallVector<OperandBundleDef, 1> OpBundles;
+  CI.getOperandBundlesAsDefs(OpBundles);
+  CallInst *Replacement = IRBuilder.CreateCall(TLIFunc, Args, OpBundles);
+  assert(OldFunc->getFunctionType() == TLIFunc->getFunctionType() &&
+         "Expecting function types to be identical");
+  CI.replaceAllUsesWith(Replacement);
+  if (isa<FPMathOperator>(Replacement)) {
+    // Preserve fast math flags for FP math.
+    Replacement->copyFastMathFlags(&CI);
+  }
+
+  LLVM_DEBUG(dbgs() << DEBUG_TYPE << ": Replaced call to `"
+                    << OldFunc->getName() << "` with call to `" << TLIName
+                    << "`.\n");
+  ++NumCallsReplaced;
+  return true;
+}
+
+static bool replaceWithCallToVeclib(const TargetLibraryInfo &TLI,
+                                    CallInst &CI) {
+  if (!CI.getCalledFunction()) {
+    return false;
+  }
+
+  auto IntrinsicID = CI.getCalledFunction()->getIntrinsicID();
+  if (IntrinsicID == Intrinsic::not_intrinsic) {
+    // Replacement is only performed for intrinsic functions
+    return false;
+  }
+
+  // Convert vector arguments to scalar type and check that
+  // all vector operands have identical vector width.
+  unsigned VF = 0;
+  SmallVector<Type *> ScalarTypes;
+  for (auto Arg : enumerate(CI.arg_operands())) {
+    auto *ArgType = Arg.value()->getType();
+    // Vector calls to intrinsics can still have
+    // scalar operands for specific arguments.
+    if (hasVectorInstrinsicScalarOpd(IntrinsicID, Arg.index())) {
+      ScalarTypes.push_back(ArgType);
+    } else {
+      // The argument in this place should be a vector if
+      // this is a call to a vector intrinsic.
+      auto *VectorArgTy = dyn_cast<VectorType>(ArgType);
+      if (!VectorArgTy) {
+        // The argument is not a vector, do not perform
+        // the replacement.
+        return false;
+      }
+      auto NumElements = VectorArgTy->getElementCount();
+      if (NumElements.isScalable()) {
+        // The current implementation does not support
+        // scalable vectors.
+        return false;
+      }
+      if (VF && VF != NumElements.getFixedValue()) {
+        // The different arguments differ in vector size.
+        return false;
+      } else {
+        VF = NumElements.getFixedValue();
+      }
+      ScalarTypes.push_back(VectorArgTy->getElementType());
+    }
+  }
+
+  // Try to reconstruct the name for the scalar version of this
+  // intrinsic using the intrinsic ID and the argument types
+  // converted to scalar above.
+  std::string ScalarName;
+  if (Intrinsic::isOverloaded(IntrinsicID)) {
+    ScalarName = Intrinsic::getName(IntrinsicID, ScalarTypes);
+  } else {
+    ScalarName = Intrinsic::getName(IntrinsicID).str();
+  }
+
+  if (!TLI.isFunctionVectorizable(ScalarName)) {
+    // The TargetLibraryInfo does not contain a vectorized version of
+    // the scalar function.
+    return false;
+  }
+
+  // Try to find the mapping for the scalar version of this intrinsic
+  // and the exact vector width of the call operands in the
+  // TargetLibraryInfo.
+  const std::string TLIName =
+      std::string(TLI.getVectorizedFunction(ScalarName, VF));
+
+  LLVM_DEBUG(dbgs() << DEBUG_TYPE << ": Looking up TLI mapping for `"
+                    << ScalarName << "` and vector width " << VF << ".\n");
+
+  if (!TLIName.empty()) {
+    // Found the correct mapping in the TargetLibraryInfo,
+    // replace the call to the intrinsic with a call to
+    // the vector library function.
+    LLVM_DEBUG(dbgs() << DEBUG_TYPE << ": Found TLI function `" << TLIName
+                      << "`.\n");
+    return replaceWithTLIFunction(CI, TLIName);
+  }
+
+  return false;
+}
+
+static bool runImpl(const TargetLibraryInfo &TLI, Function &F) {
+  bool Changed = false;
+  SmallVector<CallInst *> ReplacedCalls;
+  for (auto &I : instructions(F)) {
+    if (auto *CI = dyn_cast<CallInst>(&I)) {
+      if (replaceWithCallToVeclib(TLI, *CI)) {
+        ReplacedCalls.push_back(CI);
+        Changed = true;
+      }
+    }
+  }
+  // Erase the calls to the intrinsics that have been replaced
+  // with calls to the vector library.
+  for (auto *CI : ReplacedCalls) {
+    CI->eraseFromParent();
+  }
+  return Changed;
+}
+
+////////////////////////////////////////////////////////////////////////////////
+// New pass manager implementation.
+////////////////////////////////////////////////////////////////////////////////
+PreservedAnalyses ReplaceWithVeclib::run(Function &F,
+                                         FunctionAnalysisManager &AM) {
+  const TargetLibraryInfo &TLI = AM.getResult<TargetLibraryAnalysis>(F);
+  auto Changed = runImpl(TLI, F);
+  if (Changed) {
+    PreservedAnalyses PA;
+    PA.preserveSet<CFGAnalyses>();
+    PA.preserve<TargetLibraryAnalysis>();
+    PA.preserve<ScalarEvolutionAnalysis>();
+    PA.preserve<AAManager>();
+    PA.preserve<LoopAccessAnalysis>();
+    PA.preserve<DemandedBitsAnalysis>();
+    PA.preserve<OptimizationRemarkEmitterAnalysis>();
+    PA.preserve<GlobalsAA>();
+    return PA;
+  } else {
+    // The pass did not replace any calls, hence it preserves all analyses.
+    return PreservedAnalyses::all();
+  }
+}
+
+////////////////////////////////////////////////////////////////////////////////
+// Legacy PM Implementation.
+////////////////////////////////////////////////////////////////////////////////
+bool ReplaceWithVeclibLegacy::runOnFunction(Function &F) {
+  const TargetLibraryInfo &TLI =
+      getAnalysis<TargetLibraryInfoWrapperPass>().getTLI(F);
+  return runImpl(TLI, F);
+}
+
+void ReplaceWithVeclibLegacy::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 ReplaceWithVeclibLegacy::ID = 0;
+
+INITIALIZE_PASS_BEGIN(ReplaceWithVeclibLegacy, DEBUG_TYPE,
+                      "Replace intrinsics with calls to vector library", false,
+                      false)
+INITIALIZE_PASS_DEPENDENCY(TargetLibraryInfoWrapperPass)
+INITIALIZE_PASS_END(ReplaceWithVeclibLegacy, DEBUG_TYPE,
+                    "Replace intrinsics with calls to vector library", false,
+                    false)
+
+FunctionPass *llvm::createReplaceWithVeclibLegacyPass() {
+  return new ReplaceWithVeclibLegacy();
+}

llvm/lib/CodeGen/TargetPassConfig.cpp

@@ -858,6 +858,9 @@ void TargetPassConfig::addIRPasses() {
   if (getOptLevel() != CodeGenOpt::None && !DisableConstantHoisting)
     addPass(createConstantHoistingPass());
 
+  if (getOptLevel() != CodeGenOpt::None)
+    addPass(createReplaceWithVeclibLegacyPass());
+
   if (getOptLevel() != CodeGenOpt::None && !DisablePartialLibcallInlining)
     addPass(createPartiallyInlineLibCallsPass());
 

llvm/test/CodeGen/AArch64/O3-pipeline.ll

@@ -54,6 +54,7 @@
 ; CHECK-NEXT:       Branch Probability Analysis
 ; CHECK-NEXT:       Block Frequency Analysis
 ; CHECK-NEXT:       Constant Hoisting
+; CHECK-NEXT:       Replace intrinsics with calls to vector library
 ; CHECK-NEXT:       Partially inline calls to library functions
 ; CHECK-NEXT:       Instrument function entry/exit with calls to e.g. mcount() (post inlining)
 ; CHECK-NEXT:       Scalarize Masked Memory Intrinsics

llvm/test/CodeGen/ARM/O3-pipeline.ll

@@ -34,6 +34,7 @@
 ; CHECK-NEXT:      Branch Probability Analysis
 ; CHECK-NEXT:      Block Frequency Analysis
 ; CHECK-NEXT:      Constant Hoisting
+; CHECK-NEXT:      Replace intrinsics with calls to vector library
 ; CHECK-NEXT:      Partially inline calls to library functions
 ; CHECK-NEXT:      Instrument function entry/exit with calls to e.g. mcount() (post inlining)
 ; CHECK-NEXT:      Scalarize Masked Memory Intrinsics