Preserve NVVM for both intrinsics and library functions (rust-lang#504)

wsmoses · web-flow · commit 88a5efc99a78 · 2022-02-08T19:11:20.000-05:00
* Preserve NVVM for both intrinsics and library functions

* Require preservation of NVVM intrinsics via Clang plugin

* Update CMakeLists.txt
diff --git a/.github/workflows/bcload.yml b/.github/workflows/bcload.yml
@@ -25,7 +25,7 @@ jobs:
       run: |
           wget -O - https://apt.llvm.org/llvm-snapshot.gpg.key|sudo apt-key add -
           sudo apt-add-repository "deb http://apt.llvm.org/`lsb_release -c | cut -f2`/ llvm-toolchain-`lsb_release -c | cut -f2`-${{ matrix.llvm }} main" || true
-          sudo apt-get install -y autoconf cmake gcc g++ libtool gfortran llvm-${{ matrix.llvm }}-dev libomp-${{ matrix.llvm }}-dev clang-${{ matrix.llvm }} libeigen3-dev libboost-dev
+          sudo apt-get install -y autoconf cmake gcc g++ libtool gfortran llvm-${{ matrix.llvm }}-dev libomp-${{ matrix.llvm }}-dev clang-${{ matrix.llvm }} libclang-${{ matrix.llvm }}-dev libeigen3-dev libboost-dev
           sudo python3 -m pip install --upgrade pip setuptools
           sudo python3 -m pip install lit
           sudo touch /usr/lib/llvm-${{ matrix.llvm }}/bin/yaml-bench
diff --git a/.github/workflows/ccpp.yml b/.github/workflows/ccpp.yml
@@ -26,7 +26,7 @@ jobs:
       run: |
           wget -O - https://apt.llvm.org/llvm-snapshot.gpg.key|sudo apt-key add -
           sudo apt-add-repository "deb http://apt.llvm.org/`lsb_release -c | cut -f2`/ llvm-toolchain-`lsb_release -c | cut -f2`-${{ matrix.llvm }} main" || true
-          sudo apt-get install -y autoconf cmake gcc g++ libtool gfortran llvm-${{ matrix.llvm }}-dev libomp-${{ matrix.llvm }}-dev clang-${{ matrix.llvm }} libeigen3-dev libboost-dev
+          sudo apt-get install -y autoconf cmake gcc g++ libtool gfortran llvm-${{ matrix.llvm }}-dev libomp-${{ matrix.llvm }}-dev clang-${{ matrix.llvm }} libclang-${{ matrix.llvm }}-dev libeigen3-dev libboost-dev
           sudo python3 -m pip install --upgrade pip setuptools
           sudo python3 -m pip install lit
           sudo touch /usr/lib/llvm-${{ matrix.llvm }}/bin/yaml-bench
diff --git a/enzyme/CMakeLists.txt b/enzyme/CMakeLists.txt
@@ -41,7 +41,18 @@ get_filename_component(LLVM_ABSOLUTE_DIR
 
 set(LLVM_DIR "${LLVM_ABSOLUTE_DIR}" CACHE FILEPATH "b" FORCE)
 
-message("found llvm dir " ${LLVM_DIR})
+if (NOT DEFINED LLVM_EXTERNAL_LIT)
+  if(LLVM_DIR MATCHES ".*/cmake/llvm/?$")
+      message("found llvm match ${CMAKE_MATCH_1} dir ${LLVM_DIR}")
+    if (EXISTS ${LLVM_DIR}/../../../bin/llvm-lit) 
+      set(LLVM_EXTERNAL_LIT ${LLVM_DIR}/../../../bin/llvm-lit) 
+    endif()
+  else()
+    if (EXISTS ${LLVM_DIR}/bin/llvm-lit) 
+      set(LLVM_EXTERNAL_LIT ${LLVM_DIR}/bin/llvm-lit) 
+    endif()
+  endif()
+endif()
 
 get_filename_component(LLVM_ABSOLUTE_LIT
                        "${LLVM_EXTERNAL_LIT}"
@@ -50,15 +61,55 @@ get_filename_component(LLVM_ABSOLUTE_LIT
                    set(LLVM_EXTERNAL_LIT "${LLVM_ABSOLUTE_LIT}" CACHE FILEPATH "a" FORCE)
 message("found llvm lit " ${LLVM_EXTERNAL_LIT})
 
-
-
 list(INSERT CMAKE_PREFIX_PATH 0 "${LLVM_DIR}")
+if (DEFINED Clang_DIR)
+    get_filename_component(Clang_ABSOLUTE_DIR
+      "${Clang_DIR}"
+    REALPATH BASE_DIR "${CMAKE_BINARY_DIR}")
+set(Clang_DIR "${Clang_ABSOLUTE_DIR}" CACHE FILEPATH "b" FORCE)
+  list(INSERT CMAKE_PREFIX_PATH 0 "${Clang_DIR}")
+  message("clang dir defined ${Clang_DIR}")
+else()
+  if(LLVM_DIR MATCHES ".*/cmake/llvm/?$")
+    if (EXISTS ${LLVM_DIR}/../clang/../../libclangBasic.a)
+      set(Clang_DIR ${LLVM_DIR}/../clang)
+     list(INSERT CMAKE_PREFIX_PATH 0 "${Clang_DIR}")
+    endif()
+  elseif(LLVM_DIR MATCHES ".*/llvm-([0-9]+)/?$")
+    if (EXISTS ${LLVM_DIR}/lib/libclangBasic.a)
+      set(Clang_DIR ${LLVM_DIR}/lib/cmake/clang)
+     list(INSERT CMAKE_PREFIX_PATH 0 "${Clang_DIR}")
+    endif()
+  else()
+    if (EXISTS ${LLVM_DIR}/lib/libclangBasic.a)
+      set(Clang_DIR ${LLVM_DIR})
+      list(INSERT CMAKE_PREFIX_PATH 0 "${Clang_DIR}")
+    endif()
+  endif()
+  message("clang dir from llvm ${Clang_DIR}")
+endif()
 message("CMAKE_PREFIX_PATH " ${CMAKE_PREFIX_PATH})
+
 find_package(LLVM REQUIRED CONFIG)
 
 list(APPEND CMAKE_MODULE_PATH "${LLVM_CMAKE_DIR}")
 include(AddLLVM)
 
+message("clang dir ${Clang_DIR}")
+
+if (DEFINED Clang_DIR)
+    find_package(Clang REQUIRED CONFIG PATHS ${Clang_DIR} NO_DEFAULT_PATH)
+if (${Clang_FOUND})
+    include_directories(${CLANG_INCLUDE_DIRS})
+    message("clang inc dir ${CLANG_INCLUDE_DIRS}")
+    add_definitions(${CLANG_DEFINITIONS})
+endif()
+else()
+    set(Clang_FOUND 0) 
+endif()
+message("found ${Clang_FOUND}")
+# include(AddClang)
+
 add_definitions(${LLVM_DEFINITIONS})
 include_directories(${LLVM_INCLUDE_DIRS})
 message("LLVM_INSTALL_PREFIX: ${LLVM_INSTALL_PREFIX}")
@@ -132,9 +183,13 @@ if(NOT IS_ABSOLUTE "${${var}}")
 endif()
 endforeach()
 
-export(TARGETS LLVMEnzyme-${LLVM_VERSION_MAJOR} ClangEnzyme-${LLVM_VERSION_MAJOR}
+export(TARGETS LLVMEnzyme-${LLVM_VERSION_MAJOR}
   FILE "${PROJECT_BINARY_DIR}/EnzymeTargets.cmake")
 
+if (${Clang_FOUND})
+export(TARGETS ClangEnzyme-${LLVM_VERSION_MAJOR}
+    APPEND FILE "${PROJECT_BINARY_DIR}/EnzymeTargets.cmake")
+endif()
 export(PACKAGE Enzyme)
 
 set(CONF_LLVM_VERSION_MAJOR ${LLVM_VERSION_MAJOR})
diff --git a/enzyme/Enzyme/CMakeLists.txt b/enzyme/Enzyme/CMakeLists.txt
@@ -19,13 +19,15 @@ if (${LLVM_VERSION_MAJOR} LESS 8)
         PLUGIN_TOOL
         opt
     )
+if (${Clang_FOUND})
     add_llvm_loadable_module( ClangEnzyme-${LLVM_VERSION_MAJOR}
         ${ENZYME_SRC} Clang/EnzymeClang.cpp
         DEPENDS
         intrinsics_gen
         PLUGIN_TOOL
         opt
     )
+endif()
 else()
 # on windows `PLUGIN_TOOL` doesn't link against LLVM.dll
 if ((WIN32 OR CYGWIN) AND LLVM_LINK_LLVM_DYLIB)
@@ -37,6 +39,7 @@ if ((WIN32 OR CYGWIN) AND LLVM_LINK_LLVM_DYLIB)
 	LINK_COMPONENTS
 	LLVM
     )
+if (${Clang_FOUND})
     add_llvm_library( ClangEnzyme-${LLVM_VERSION_MAJOR}
         ${ENZYME_SRC} Clang/EnzymeClang.cpp
         MODULE
@@ -45,6 +48,7 @@ if ((WIN32 OR CYGWIN) AND LLVM_LINK_LLVM_DYLIB)
 	LINK_COMPONENTS
 	LLVM
     )
+endif()
 else()
     add_llvm_library( LLVMEnzyme-${LLVM_VERSION_MAJOR}
         ${ENZYME_SRC}
@@ -54,6 +58,7 @@ else()
         PLUGIN_TOOL
         opt
     )
+if (${Clang_FOUND})
     add_llvm_library( ClangEnzyme-${LLVM_VERSION_MAJOR}
         ${ENZYME_SRC} Clang/EnzymeClang.cpp
         MODULE
@@ -64,6 +69,7 @@ else()
     )
 endif()
 endif()
+endif()
 
 if (${ENZYME_EXTERNAL_SHARED_LIB})
     add_library( Enzyme-${LLVM_VERSION_MAJOR}
@@ -82,12 +88,22 @@ if (APPLE)
 # Darwin-specific linker flags for loadable modules.
 set_target_properties(LLVMEnzyme-${LLVM_VERSION_MAJOR} PROPERTIES
     LINK_FLAGS "-Wl,-flat_namespace -Wl,-undefined -Wl,suppress")
+if (${Clang_FOUND})
 set_target_properties(ClangEnzyme-${LLVM_VERSION_MAJOR} PROPERTIES
         LINK_FLAGS "-Wl,-flat_namespace -Wl,-undefined -Wl,suppress")
 endif()
+endif()
+
+install(TARGETS LLVMEnzyme-${LLVM_VERSION_MAJOR}
+    EXPORT EnzymeTargets
+    LIBRARY DESTINATION lib COMPONENT shlib
+    PUBLIC_HEADER DESTINATION "${INSTALL_INCLUDE_DIR}/Enzyme"
+    COMPONENT dev)
 
-install(TARGETS LLVMEnzyme-${LLVM_VERSION_MAJOR} ClangEnzyme-${LLVM_VERSION_MAJOR}
+if (${Clang_FOUND})
+install(TARGETS ClangEnzyme-${LLVM_VERSION_MAJOR}
     EXPORT EnzymeTargets
     LIBRARY DESTINATION lib COMPONENT shlib
     PUBLIC_HEADER DESTINATION "${INSTALL_INCLUDE_DIR}/Enzyme"
     COMPONENT dev)
+endif()
diff --git a/enzyme/Enzyme/Clang/EnzymeClang.cpp b/enzyme/Enzyme/Clang/EnzymeClang.cpp
@@ -83,4 +83,63 @@ static void loadLTOPass(const PassManagerBuilder &Builder,
 static RegisterStandardPasses
     clangtoolLoader_LTO(PassManagerBuilder::EP_FullLinkTimeOptimizationEarly,
                         loadLTOPass);
+
+#include "clang/AST/Attr.h"
+#include "clang/AST/DeclGroup.h"
+#include "clang/Frontend/CompilerInstance.h"
+#include "clang/Frontend/FrontendAction.h"
+#include "clang/Frontend/FrontendPluginRegistry.h"
+
+template <typename ConsumerType>
+class EnzymeAction : public clang::PluginASTAction {
+protected:
+  std::unique_ptr<clang::ASTConsumer>
+  CreateASTConsumer(clang::CompilerInstance &CI, llvm::StringRef InFile) {
+    return std::unique_ptr<clang::ASTConsumer>(new ConsumerType(CI));
+  }
+
+  bool ParseArgs(const clang::CompilerInstance &CI,
+                 const std::vector<std::string> &args) {
+    return true;
+  }
+
+  PluginASTAction::ActionType getActionType() override {
+    return AddBeforeMainAction;
+  }
+};
+
+class EnzymePlugin : public clang::ASTConsumer {
+  clang::CompilerInstance &CI;
+
+public:
+  EnzymePlugin(clang::CompilerInstance &CI) : CI(CI) {}
+  ~EnzymePlugin() {}
+  bool HandleTopLevelDecl(clang::DeclGroupRef dg) override {
+    using namespace clang;
+    DeclGroupRef::iterator it;
+
+    // Forcibly require emission of all libdevice
+    for (it = dg.begin(); it != dg.end(); ++it) {
+      auto FD = dyn_cast<FunctionDecl>(*it);
+      if (!FD)
+        continue;
+
+      if (!FD->hasAttr<clang::CUDADeviceAttr>())
+        continue;
+
+      if (!FD->getIdentifier())
+        continue;
+      if (!StringRef(FD->getLocation().printToString(CI.getSourceManager()))
+               .contains("/__clang_cuda_math.h"))
+        continue;
+
+      FD->addAttr(UsedAttr::CreateImplicit(CI.getASTContext()));
+    }
+    return true;
+  }
+};
+
+// register the PluginASTAction in the registry.
+static clang::FrontendPluginRegistry::Add<EnzymeAction<EnzymePlugin>>
+    X("enzyme", "Enzyme Plugin");
 #endif
diff --git a/enzyme/Enzyme/FunctionUtils.cpp b/enzyme/Enzyme/FunctionUtils.cpp
@@ -2084,36 +2084,40 @@ void PreProcessCache::optimizeIntermediate(Function *F) {
   // EarlyCSEPass(/*memoryssa*/ true).run(*F, FAM);
 
   for (Function &Impl : *F->getParent()) {
-    if (!Impl.hasFnAttribute("implements"))
-      continue;
-    const Attribute &A = Impl.getFnAttribute("implements");
-
-    const StringRef SpecificationName = A.getValueAsString();
-    Function *Specification = F->getParent()->getFunction(SpecificationName);
-    if (!Specification) {
-      LLVM_DEBUG(dbgs() << "Found implementation '" << Impl.getName()
-                        << "' but no matching specification with name '"
-                        << SpecificationName
-                        << "', potentially inlined and/or eliminated.\n");
-      continue;
-    }
-    LLVM_DEBUG(dbgs() << "Replace specification '" << Specification->getName()
-                      << "' with implementation '" << Impl.getName() << "'\n");
-
-    for (auto I = Specification->use_begin(), UE = Specification->use_end();
-         I != UE;) {
-      auto &use = *I;
-      ++I;
-      auto cext = ConstantExpr::getBitCast(&Impl, Specification->getType());
-      use.set(cext);
-      if (auto CI = dyn_cast<CallInst>(use.getUser())) {
+    for (auto attr : {"implements", "implements2"}) {
+      if (!Impl.hasFnAttribute(attr))
+        continue;
+      const Attribute &A = Impl.getFnAttribute(attr);
+
+      const StringRef SpecificationName = A.getValueAsString();
+      Function *Specification = F->getParent()->getFunction(SpecificationName);
+      if (!Specification) {
+        LLVM_DEBUG(dbgs() << "Found implementation '" << Impl.getName()
+                          << "' but no matching specification with name '"
+                          << SpecificationName
+                          << "', potentially inlined and/or eliminated.\n");
+        continue;
+      }
+      LLVM_DEBUG(dbgs() << "Replace specification '" << Specification->getName()
+                        << "' with implementation '" << Impl.getName()
+                        << "'\n");
+
+      for (auto I = Specification->use_begin(), UE = Specification->use_end();
+           I != UE;) {
+        auto &use = *I;
+        ++I;
+        auto cext = ConstantExpr::getBitCast(&Impl, Specification->getType());
+        use.set(cext);
+        if (auto CI = dyn_cast<CallInst>(use.getUser())) {
 #if LLVM_VERSION_MAJOR >= 11
-        if (CI->getCalledOperand() == cext || CI->getCalledFunction() == &Impl)
+          if (CI->getCalledOperand() == cext ||
+              CI->getCalledFunction() == &Impl)
 #else
-        if (CI->getCalledValue() == cext || CI->getCalledFunction() == &Impl)
+          if (CI->getCalledValue() == cext || CI->getCalledFunction() == &Impl)
 #endif
-        {
-          CI->setCallingConv(Impl.getCallingConv());
+          {
+            CI->setCallingConv(Impl.getCallingConv());
+          }
         }
       }
     }
diff --git a/enzyme/Enzyme/PreserveNVVM.cpp b/enzyme/Enzyme/PreserveNVVM.cpp
@@ -55,11 +55,8 @@ class PreserveNVVM : public FunctionPass {
 
   bool runOnFunction(Function &F) override {
     bool changed = false;
-    std::map<std::pair<Type *, std::string>,
-             std::pair<std::string, std::string>>
-        Implements;
-    for (Type *T : {Type::getFloatTy(F.getContext()),
-                    Type::getDoubleTy(F.getContext())}) {
+    std::map<std::string, std::pair<std::string, std::string>> Implements;
+    for (std::string T : {"", "f"}) {
       // sincos, sinpi, cospi, sincospi, cyl_bessel_i1
       for (std::string name :
            {"sin",        "cos",     "tan",       "log2",   "exp",    "exp2",
@@ -75,19 +72,20 @@ class PreserveNVVM : public FunctionPass {
             "sqrt"}) {
         std::string nvname = "__nv_" + name;
         std::string llname = "llvm." + name + ".";
+        std::string mathname = name;
 
-        if (T->isFloatTy()) {
+        if (T == "f") {
+          mathname += "f";
           nvname += "f";
           llname += "f32";
         } else {
           llname += "f64";
         }
 
-        Implements[std::make_pair(T, nvname)] = std::make_pair(name, llname);
+        Implements[nvname] = std::make_pair(mathname, llname);
       }
     }
-    auto idx = std::make_pair(F.getReturnType(), F.getName().str());
-    auto found = Implements.find(idx);
+    auto found = Implements.find(F.getName().str());
     if (found != Implements.end()) {
       if (Begin) {
         F.removeFnAttr(Attribute::AlwaysInline);
@@ -96,6 +94,7 @@ class PreserveNVVM : public FunctionPass {
         // cannot be erased.
         F.setLinkage(Function::LinkageTypes::ExternalLinkage);
         F.addFnAttr("implements", found->second.second);
+        F.addFnAttr("implements2", found->second.first);
         F.addFnAttr("enzyme_math", found->second.first);
         changed = true;
       } else {
diff --git a/enzyme/test/CMakeLists.txt b/enzyme/test/CMakeLists.txt
@@ -10,5 +10,7 @@ set(ENZYME_TEST_DEPS LLVMEnzyme-${LLVM_VERSION_MAJOR} BCPass-${LLVM_VERSION_MAJO
 add_subdirectory(ActivityAnalysis)
 add_subdirectory(TypeAnalysis)
 add_subdirectory(Enzyme)
+if (${Clang_FOUND})
 add_subdirectory(Integration)
+endif()
 add_subdirectory(BCLoader)
