From b5bfdb98fa87fd238c4d09d352ad1971ae81859e Mon Sep 17 00:00:00 2001 From: Vassil Vassilev Date: Fri, 29 Jul 2022 09:34:55 +0000 Subject: [PATCH] Implement basic clang backend plugin support. This patch teaches Clad how to attach itself as a plugin to the compiler backend. This allows us to implement our own AD-related backend passes but also to integrate tools such as Enzyme easily. --- tools/CMakeLists.txt | 96 +++++++++++++++++++++++++++++++----- tools/ClangBackendPlugin.cpp | 30 +++++++++++ tools/ClangBackendPlugin.h | 31 ++++++++++++ tools/ClangPlugin.cpp | 39 ++++++++++++++- 4 files changed, 182 insertions(+), 14 deletions(-) create mode 100644 tools/ClangBackendPlugin.cpp create mode 100644 tools/ClangBackendPlugin.h diff --git a/tools/CMakeLists.txt b/tools/CMakeLists.txt index 92a540b20..4a635a51d 100644 --- a/tools/CMakeLists.txt +++ b/tools/CMakeLists.txt @@ -2,21 +2,91 @@ if( NOT CLAD_BUILT_STANDALONE AND NOT LLVM_BUILD_TOOLS ) set(EXCLUDE_FROM_ALL ON) endif() -add_llvm_library(cladPlugin +set(CLAD_PLUGIN_SRC ClangPlugin.cpp + ClangBackendPlugin.cpp DerivedFnInfo.cpp RequiredSymbols.cpp - ) +) if (NOT CLAD_BUILD_STATIC_ONLY) - add_llvm_loadable_module(clad - ClangPlugin.cpp - DerivedFnInfo.cpp - RequiredSymbols.cpp - PLUGIN_TOOL - clang - ) - - target_link_libraries(clad PRIVATE - cladDifferentiator - ) + add_llvm_loadable_module(clad ${CLAD_PLUGIN_SRC} PLUGIN_TOOL clang) + + target_link_libraries(clad PRIVATE cladDifferentiator) + + # Add Enzyme as a backend. + if (ENABLE_ENZYME_BACKEND) + + if (NOT UNIX) + message(FATAL_ERROR "The enzyme backend is only supported on Unix platforms") + endif() + + set(_enzyme_build_type ${CMAKE_CFG_INTDIR}) + set(_enzyme_cmake_logging_settings + LOG_DOWNLOAD ON + LOG_CONFIGURE ON + LOG_BUILD ON + LOG_INSTALL ON + LOG_OUTPUT_ON_FAILURE ON + ) + + include(ExternalProject) + + set(_enzyme_patch_command + ${CMAKE_COMMAND} -E copy_if_different + ${CMAKE_SOURCE_DIR}/patches/enzyme.patch && + git checkout && git apply /enzyme.patch + ) + + # Set build byproducts only needed by Ninja. + set(_enzyme_static_archive_name ${CMAKE_LIBRARY_OUTPUT_DIRECTORY}${CMAKE_STATIC_LIBRARY_PREFIX}LLVMEnzyme-${CLANG_VERSION_MAJOR}${CMAKE_STATIC_LIBRARY_SUFFIX}) + if("${CMAKE_GENERATOR}" STREQUAL "Ninja") + set(ENZYME_BYPRODUCTS ${_enzyme_static_archive_name}) + endif() + + ExternalProject_Add( + Enzyme + GIT_REPOSITORY https://github.com/wsmoses/Enzyme + GIT_TAG v0.0.33 + GIT_SHALLOW 1 # Do not clone the history + PATCH_COMMAND ${_enzyme_patch_command} + UPDATE_COMMAND "" + CMAKE_ARGS -G ${CMAKE_GENERATOR} + -DCMAKE_BUILD_TYPE=${CMAKE_BUILD_TYPE} + -DCMAKE_C_COMPILER=${CMAKE_C_COMPILER} + -DCMAKE_C_FLAGS=${CMAKE_C_FLAGS} + -DCMAKE_CXX_COMPILER=${CMAKE_CXX_COMPILER} + -DCMAKE_CXX_FLAGS=${CMAKE_CXX_FLAGS} + -DCMAKE_INSTALL_PREFIX=${CMAKE_BINARY_DIR} + -DLLVM_DIR=${LLVM_DIR} + -DENZYME_CLANG=OFF + -DENZYME_BUILD_STATIC_ONLY=On + -DENZYME_ENABLE_BENCHMARKS=Off + -DLLVM_EXTERNAL_LIT=${LLVM_EXTERNAL_LIT} + SOURCE_SUBDIR enzyme + BUILD_COMMAND ${CMAKE_COMMAND} --build . --config ${_enzyme_build_type} -- -j8 + #TEST_COMMAND ${CMAKE_COMMAND} --build . --config ${_enzyme_build_type} --target check-enzyme + #TEST_BEFORE_INSTALL 1 + INSTALL_COMMAND ${CMAKE_COMMAND} --install . + BUILD_BYPRODUCTS ${ENZYME_BYPRODUCTS} + ${_enzyme_cmake_logging_settings} + DEPENDS LLVMSupport + ) + + # Register LLVMEnzyme for our build system. + add_library(LLVMEnzyme IMPORTED STATIC GLOBAL) + add_dependencies(LLVMEnzyme Enzyme) # Emzyme is the ExternalProject_Add target. + set_property(TARGET LLVMEnzyme PROPERTY IMPORTED_LOCATION ${_enzyme_static_archive_name}) + + # Paste the contents of libLLVMEnzyme-X.so into clad.so. + if (APPLE) + set(_enzyme_link_flags -Wl,-force_load $ -Wl,-force_load cladDifferentiator) + elseif(MSVC) + set(_enzyme_link_flags "-WHOLEARCHIVE:" $) + else() + set(_enzyme_link_flags -Wl,--whole-archive $ -Wl,--no-whole-archive) + endif() + target_link_libraries(clad PUBLIC ${_enzyme_link_flags}) + + add_dependencies(clad LLVMEnzyme) + endif(ENABLE_ENZYME_BACKEND) endif() diff --git a/tools/ClangBackendPlugin.cpp b/tools/ClangBackendPlugin.cpp new file mode 100644 index 000000000..495d1f5ee --- /dev/null +++ b/tools/ClangBackendPlugin.cpp @@ -0,0 +1,30 @@ +//--------------------------------------------------------------------*- C++ -*- +// clad - the C++ Clang-based Automatic Differentiator +// version: $Id$ +// author: Vassil Vassilev +//------------------------------------------------------------------------------ + +#include "clang/Basic/Version.h" // for CLANG_VERSION_MAJOR + +#if CLANG_VERSION_MAJOR > 8 + +#include "ClangBackendPlugin.h" + +#include "llvm/Passes/PassBuilder.h" + +namespace clad { +using namespace llvm; +void ClangBackendPluginPass::registerCallbacks(PassBuilder& PB) { + PB.registerPipelineParsingCallback( + [](StringRef Name, ModulePassManager& PM, + ArrayRef InnerPipeline) { + if (Name == "plugin-pass") { + PM.addPass(ClangBackendPluginPass()); + return true; + } + return false; + }); +} +} // namespace clad + +#endif // CLANG_VERSION_MAJOR > 8 diff --git a/tools/ClangBackendPlugin.h b/tools/ClangBackendPlugin.h new file mode 100644 index 000000000..d89b18631 --- /dev/null +++ b/tools/ClangBackendPlugin.h @@ -0,0 +1,31 @@ +//--------------------------------------------------------------------*- C++ -*- +// clad - the C++ Clang-based Automatic Differentiator +// version: $Id$ +// author: Vassil Vassilev +//------------------------------------------------------------------------------ + +#ifndef CLANG_BACKEND_PLUGIN_H +#define CLANG_BACKEND_PLUGIN_H + +#include "llvm/Config/llvm-config.h" // for CLANG_VERSION_MAJOR + +#if CLANG_VERSION_MAJOR > 8 + +#include "llvm/IR/PassManager.h" +#include "llvm/Passes/PassPlugin.h" + +namespace clad { + +struct ClangBackendPluginPass + : public llvm::PassInfoMixin { + llvm::PreservedAnalyses run(llvm::Module& M, + llvm::ModuleAnalysisManager& MAM) { + return llvm::PreservedAnalyses::all(); + } + + static void registerCallbacks(llvm::PassBuilder& PB); +}; +} // end namespace clad + +#endif // CLANG_VERSION_MAJOR > 8 +#endif // CLANG_BACKEND_PLUGIN_H diff --git a/tools/ClangPlugin.cpp b/tools/ClangPlugin.cpp index a1ef05eb7..3dac61618 100644 --- a/tools/ClangPlugin.cpp +++ b/tools/ClangPlugin.cpp @@ -99,7 +99,24 @@ namespace clad { }; CladPlugin::CladPlugin(CompilerInstance& CI, DifferentiationOptions& DO) - : m_CI(CI), m_DO(DO), m_HasRuntime(false) { } + : m_CI(CI), m_DO(DO), m_HasRuntime(false) { +#if CLANG_VERSION_MAJOR > 8 + + FrontendOptions& Opts = CI.getFrontendOpts(); + // Find the path to clad. + llvm::StringRef CladSoPath; + for (llvm::StringRef P : Opts.Plugins) + if (llvm::sys::path::stem(P).endswith("clad")) { + CladSoPath = P; + break; + } + + // Register clad as a backend pass. + CodeGenOptions& CGOpts = CI.getCodeGenOpts(); + CGOpts.PassPlugins.push_back(CladSoPath.str()); +#endif // CLANG_VERSION_MAJOR > 8 + } + CladPlugin::~CladPlugin() {} // We cannot use HandleTranslationUnit because codegen already emits code on @@ -346,6 +363,8 @@ namespace clad { } } // end namespace clad +// Attach the frontend plugin. + using namespace clad::plugin; // register the PluginASTAction in the registry. static clang::FrontendPluginRegistry::Add > @@ -353,3 +372,21 @@ X("clad", "Produces derivatives or arbitrary functions"); static PragmaHandlerRegistry::Add Y("clad", "Clad pragma directives handler."); + +#include "clang/Basic/Version.h" // for CLANG_VERSION_MAJOR +#if CLANG_VERSION_MAJOR > 8 + +// Attach the backend plugin. + +#include "ClangBackendPlugin.h" + +#define BACKEND_PLUGIN_NAME "CladBackendPlugin" +// FIXME: Add a proper versioning that's based on CLANG_VERSION_STRING and +// a similar approach for clad (see Version.cpp and VERSION). +#define BACKEND_PLUGIN_VERSION "FIXME" +extern "C" ::llvm::PassPluginLibraryInfo LLVM_ATTRIBUTE_WEAK +llvmGetPassPluginInfo() { + return {LLVM_PLUGIN_API_VERSION, BACKEND_PLUGIN_NAME, BACKEND_PLUGIN_VERSION, + clad::ClangBackendPluginPass::registerCallbacks}; +} +#endif // CLANG_VERSION_MAJOR