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

@@ -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 <SOURCE_DIR> &&
+      git checkout <SOURCE_DIR> && git apply <SOURCE_DIR>/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 $<TARGET_FILE:LLVMEnzyme> -Wl,-force_load cladDifferentiator)
+    elseif(MSVC)
+      set(_enzyme_link_flags "-WHOLEARCHIVE:" $<TARGET_FILE:LLVMEnzyme>)
+    else()
+      set(_enzyme_link_flags -Wl,--whole-archive $<TARGET_FILE:LLVMEnzyme> -Wl,--no-whole-archive)
+    endif()
+    target_link_libraries(clad PUBLIC ${_enzyme_link_flags})
+
+    add_dependencies(clad LLVMEnzyme)
+  endif(ENABLE_ENZYME_BACKEND)
 endif()

tools/ClangBackendPlugin.cpp

@@ -0,0 +1,30 @@
+//--------------------------------------------------------------------*- C++ -*-
+// clad - the C++ Clang-based Automatic Differentiator
+// version: $Id$
+// author:  Vassil Vassilev <vvasilev-at-cern.ch>
+//------------------------------------------------------------------------------
+
+#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<PassBuilder::PipelineElement> InnerPipeline) {
+        if (Name == "plugin-pass") {
+          PM.addPass(ClangBackendPluginPass());
+          return true;
+        }
+        return false;
+      });
+}
+} // namespace clad
+
+#endif // CLANG_VERSION_MAJOR > 8

tools/ClangBackendPlugin.h

@@ -0,0 +1,31 @@
+//--------------------------------------------------------------------*- C++ -*-
+// clad - the C++ Clang-based Automatic Differentiator
+// version: $Id$
+// author:  Vassil Vassilev <vvasilev-at-cern.ch>
+//------------------------------------------------------------------------------
+
+#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<ClangBackendPluginPass> {
+  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

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,10 +363,30 @@ namespace clad {
   }
 } // end namespace clad
 
+// Attach the frontend plugin.
+
 using namespace clad::plugin;
 // register the PluginASTAction in the registry.
 static clang::FrontendPluginRegistry::Add<Action<CladPlugin> >
 X("clad", "Produces derivatives or arbitrary functions");
 
 static PragmaHandlerRegistry::Add<CladPragmaHandler>
     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