feature: Add interface for plugin pass

CMakeLists.txt

@@ -67,8 +67,11 @@ if(ENABLE_COVERAGE)
     add_compile_options(
       "$<$<COMPILE_LANGUAGE:C>:-g>"
       "$<$<COMPILE_LANGUAGE:CXX>:-g>"
-      "$<$<COMPILE_LANGUAGE:C>:--coverage>"
-      "$<$<COMPILE_LANGUAGE:CXX>:--coverage>"
+      # Explicit form, not the --coverage alias: only these let a later
+      # per-source -fno-profile-arcs cancel instrumentation (see
+      # proteus_register_jit_pass_plugin in cmake/ProteusFunctions.cmake).
+      "$<$<COMPILE_LANGUAGE:C>:SHELL:-fprofile-arcs -ftest-coverage>"
+      "$<$<COMPILE_LANGUAGE:CXX>:SHELL:-fprofile-arcs -ftest-coverage>"
       "$<$<COMPILE_LANGUAGE:CUDA>:SHELL:-Xarch_host -g>"
       "$<$<COMPILE_LANGUAGE:CUDA>:SHELL:-Xarch_host --coverage>"
       "$<$<COMPILE_LANGUAGE:HIP>:SHELL:-Xarch_host -g>"

cmake/ProteusFunctions.cmake

@@ -28,3 +28,71 @@ function(add_proteus target)
 
     message(STATUS "Linked target '${target}' with libproteus")
 endfunction()
+
+function(proteus_register_jit_pass_plugin target)
+    if(NOT TARGET ${target})
+        message(FATAL_ERROR "Target '${target}' does not exist")
+    endif()
+
+    set(one_value_args PLUGIN_TARGET PLUGIN_PATH PIPELINE)
+    cmake_parse_arguments(PROTEUS_JIT_PASS "" "${one_value_args}" "" ${ARGN})
+
+    if(NOT PROTEUS_JIT_PASS_PIPELINE)
+        message(FATAL_ERROR "proteus_register_jit_pass_plugin requires PIPELINE")
+    endif()
+
+    if(PROTEUS_JIT_PASS_PLUGIN_TARGET AND PROTEUS_JIT_PASS_PLUGIN_PATH)
+        message(FATAL_ERROR
+            "proteus_register_jit_pass_plugin accepts either PLUGIN_TARGET or PLUGIN_PATH, not both")
+    endif()
+
+    if(NOT PROTEUS_JIT_PASS_PLUGIN_TARGET AND NOT PROTEUS_JIT_PASS_PLUGIN_PATH)
+        message(FATAL_ERROR
+            "proteus_register_jit_pass_plugin requires PLUGIN_TARGET or PLUGIN_PATH")
+    endif()
+
+    if(PROTEUS_JIT_PASS_PLUGIN_TARGET)
+        if(NOT TARGET ${PROTEUS_JIT_PASS_PLUGIN_TARGET})
+            message(FATAL_ERROR
+                "Plugin target '${PROTEUS_JIT_PASS_PLUGIN_TARGET}' does not exist")
+        endif()
+        set(_proteus_jit_pass_plugin_path "$<TARGET_FILE:${PROTEUS_JIT_PASS_PLUGIN_TARGET}>")
+        add_dependencies(${target} ${PROTEUS_JIT_PASS_PLUGIN_TARGET})
+    else()
+        set(_proteus_jit_pass_plugin_path "${PROTEUS_JIT_PASS_PLUGIN_PATH}")
+    endif()
+
+    string(MD5 _proteus_jit_pass_key
+        "${target};${_proteus_jit_pass_plugin_path};${PROTEUS_JIT_PASS_PIPELINE}")
+    set(_proteus_jit_pass_source
+        "${CMAKE_CURRENT_BINARY_DIR}/${target}.proteus_jit_pass_${_proteus_jit_pass_key}.cpp")
+
+    file(GENERATE OUTPUT "${_proteus_jit_pass_source}" CONTENT
+"#include <proteus/Init.h>
+
+namespace {
+struct AutoRegisterProteusJITPassPlugin {
+  AutoRegisterProteusJITPassPlugin() {
+    proteus::registerJITPassPlugin(
+        R\"(${_proteus_jit_pass_plugin_path})\",
+        R\"(${PROTEUS_JIT_PASS_PIPELINE})\");
+  }
+};
+
+AutoRegisterProteusJITPassPlugin AutoRegisterProteusJITPassPluginInstance;
+} // namespace
+")
+
+    if(ENABLE_COVERAGE)
+        # This generated TU only registers the JIT pass plugin and does not
+        # represent product code. Exclude it from gcov instrumentation so
+        # gcovr does not try to resolve a synthetic build-tree-only source.
+        # Relies on coverage being enabled with the explicit
+        # -fprofile-arcs/-ftest-coverage flags (see CMakeLists.txt), not the
+        # --coverage alias, which these -fno- flags cannot cancel.
+        set_source_files_properties("${_proteus_jit_pass_source}" PROPERTIES
+            COMPILE_OPTIONS "-fno-profile-arcs;-fno-test-coverage")
+    endif()
+
+    target_sources(${target} PRIVATE "${_proteus_jit_pass_source}")
+endfunction()

include/proteus/Init.h

@@ -9,13 +9,19 @@
 #ifndef PROTEUS_INIT_H
 #define PROTEUS_INIT_H
 
+#include <string>
+
 namespace proteus {
 
 [[deprecated("it is a no-op and will be removed in a future version.")]]
 void init();
 [[deprecated("it is a no-op and will be removed in a future version.")]]
 void finalize();
 
+void registerJITPassPlugin(const std::string &PluginPath,
+                           const std::string &PassPipeline);
+void clearJITPassPlugins();
+
 void enable();
 void disable();
 

src/include/proteus/impl/CoreLLVM.h

@@ -8,6 +8,7 @@ static_assert(__cplusplus >= 201703L,
 #include "proteus/TimeTracing.h"
 #include "proteus/impl/Config.h"
 #include "proteus/impl/Debug.h"
+#include "proteus/impl/JITPassPluginRegistry.h"
 #include "proteus/impl/Logger.h"
 
 #include <llvm/CodeGen/CommandFlags.h>
@@ -17,6 +18,13 @@ static_assert(__cplusplus >= 201703L,
 #include <llvm/Linker/Linker.h>
 #include <llvm/MC/TargetRegistry.h>
 #include <llvm/Passes/PassBuilder.h>
+#if __has_include(<llvm/Plugins/PassPlugin.h>)
+#include <llvm/Plugins/PassPlugin.h>
+#elif __has_include(<llvm/Passes/PassPlugin.h>)
+#include <llvm/Passes/PassPlugin.h>
+#else
+#error "Cannot find LLVM PassPlugin.h"
+#endif
 #include <llvm/Support/TargetSelect.h>
 #include <llvm/Target/TargetMachine.h>
 #include <llvm/Transforms/IPO/MergeFunctions.h>
@@ -47,6 +55,7 @@ static_assert(__cplusplus >= 201703L,
 #include <optional>
 #include <string>
 #include <utility>
+#include <vector>
 
 namespace proteus {
 using namespace llvm;
@@ -88,9 +97,57 @@ createTargetMachine(Module &M, StringRef Arch, unsigned OptLevel = 3) {
   return TM;
 }
 
-inline void runOptimizationPassPipeline(Module &M, StringRef Arch,
-                                        const std::string &PassPipeline,
-                                        unsigned CodegenOptLevel = 3) {
+inline std::string getDefaultOptimizationPipeline(char OptLevel) {
+  switch (OptLevel) {
+  case '0':
+    return "default<O0>";
+  case '1':
+    return "default<O1>";
+  case '2':
+    return "default<O2>";
+  case '3':
+    return "default<O3>";
+  case 's':
+    return "default<Os>";
+  case 'z':
+    return "default<Oz>";
+  default:
+    reportFatalError(std::string("Unsupported optimization level ") + OptLevel);
+  }
+  return "";
+}
+
+inline std::string composeOptimizationPassPipeline(
+    std::optional<std::string> PassPipeline, char OptLevel,
+    const std::vector<JITPassPluginConfig> &Plugins) {
+  std::string Pipeline = PassPipeline
+                             ? std::move(PassPipeline.value())
+                             : getDefaultOptimizationPipeline(OptLevel);
+  for (const auto &Plugin : Plugins) {
+    Pipeline += ",";
+    Pipeline += Plugin.Pipeline;
+  }
+  return Pipeline;
+}
+
+inline std::vector<PassPlugin>
+loadJITPassPlugins(const std::vector<JITPassPluginConfig> &Plugins) {
+  std::vector<PassPlugin> LoadedPlugins;
+  LoadedPlugins.reserve(Plugins.size());
+  for (const auto &Plugin : Plugins) {
+    auto LoadedPlugin = PassPlugin::Load(Plugin.Path);
+    if (!LoadedPlugin)
+      reportFatalError("Failed to load JIT pass plugin '" + Plugin.Path +
+                       "': " + toString(LoadedPlugin.takeError()));
+    LoadedPlugins.push_back(std::move(*LoadedPlugin));
+  }
+  return LoadedPlugins;
+}
+
+inline void runOptimizationPassPipeline(
+    Module &M, StringRef Arch, const std::string &PassPipeline,
+    unsigned CodegenOptLevel,
+    const std::vector<JITPassPluginConfig> &Plugins = {}) {
   PipelineTuningOptions PTO;
 
   std::optional<PGOOptions> PGOOpt;
@@ -99,7 +156,10 @@ inline void runOptimizationPassPipeline(Module &M, StringRef Arch,
     report_fatal_error(std::move(Err));
   TargetLibraryInfoImpl TLII(Triple(M.getTargetTriple()));
 
+  auto LoadedPlugins = loadJITPassPlugins(Plugins);
   PassBuilder PB(TM->get(), PTO, PGOOpt, nullptr);
+  for (const auto &Plugin : LoadedPlugins)
+    Plugin.registerPassBuilderCallbacks(PB);
   LoopAnalysisManager LAM;
   FunctionAnalysisManager FAM;
   CGSCCAnalysisManager CGAM;
@@ -207,7 +267,15 @@ inline void optimizeIR(Module &M, StringRef Arch,
   TIMESCOPE("proteus::optimizeIR");
   Timer T(Config::get().ProteusEnableTimers);
 
-  if (OptConfig.PassPipeline) {
+  const auto Plugins = getJITPassPluginConfigs();
+  const bool UseTextualPipeline = OptConfig.PassPipeline || !Plugins.empty();
+  const std::string FinalPipeline =
+      UseTextualPipeline
+          ? detail::composeOptimizationPassPipeline(OptConfig.PassPipeline,
+                                                    OptConfig.OptLevel, Plugins)
+          : std::string();
+
+  if (UseTextualPipeline) {
     auto TraceOut = [](const std::string &PassPipeline) {
       SmallString<128> S;
       raw_svector_ostream OS(S);
@@ -216,19 +284,19 @@ inline void optimizeIR(Module &M, StringRef Arch,
     };
 
     if (Config::get().traceSpecializations())
-      Logger::trace(TraceOut(OptConfig.PassPipeline.value()));
+      Logger::trace(TraceOut(FinalPipeline));
 
-    detail::runOptimizationPassPipeline(M, Arch, OptConfig.PassPipeline.value(),
-                                        OptConfig.CodegenOptLevel);
+    detail::runOptimizationPassPipeline(M, Arch, FinalPipeline,
+                                        OptConfig.CodegenOptLevel, Plugins);
   } else {
     detail::runOptimizationPassPipeline(M, Arch, OptConfig.OptLevel,
                                         OptConfig.CodegenOptLevel);
   }
 
   PROTEUS_TIMER_OUTPUT(Logger::outs("proteus")
                        << "optimizeIR optlevel "
-                       << (OptConfig.PassPipeline
-                               ? StringRef(OptConfig.PassPipeline.value())
+                       << (UseTextualPipeline
+                               ? StringRef(FinalPipeline)
                                : StringRef(&OptConfig.OptLevel, 1))
                        << " codegenopt " << OptConfig.CodegenOptLevel << " "
                        << T.elapsed() << " ms\n");

src/include/proteus/impl/CoreLLVMHIP.h

@@ -249,10 +249,14 @@ codegenParallel(Module &M, StringRef DeviceArch,
   Conf.VerifyEach = false;
   Conf.DiagHandler = DiagnosticHandler;
   Conf.OptLevel = OptConfig.OptLevel;
+  const auto Plugins = getJITPassPluginConfigs();
   // Parallel codegen lets LTO own optimization, so custom textual pipelines
   // must be forwarded to the LTO configuration instead of run beforehand.
-  if (OptConfig.PassPipeline)
-    Conf.OptPipeline = OptConfig.PassPipeline.value();
+  if (OptConfig.PassPipeline || !Plugins.empty())
+    Conf.OptPipeline = proteus::detail::composeOptimizationPassPipeline(
+        OptConfig.PassPipeline, OptConfig.OptLevel, Plugins);
+  for (const auto &Plugin : Plugins)
+    Conf.PassPlugins.push_back(Plugin.Path);
   Conf.CGOptLevel = static_cast<CodeGenOptLevel>(OptConfig.CodegenOptLevel);
 
   unsigned ParallelCodeGenParallelismLevel =

src/include/proteus/impl/Hashing.h

@@ -4,6 +4,7 @@
 #include "proteus/CompilerInterfaceTypes.h"
 #include "proteus/TimeTracing.h"
 #include "proteus/impl/Config.h"
+#include "proteus/impl/JITPassPluginRegistry.h"
 #include "proteus/impl/RuntimeConstantTypeHelpers.h"
 
 #include <llvm/ADT/ArrayRef.h>
@@ -145,6 +146,11 @@ inline HashT hashCodeGenConfig(const CodeGenerationConfig &CGConfig) {
   H = hashCombine(H, hashValue(CGConfig.codeGenOptLevel()));
   if (auto Pipeline = CGConfig.optPipeline())
     H = hashCombine(H, hashValue(Pipeline.value()));
+  for (const auto &Plugin : getJITPassPluginConfigs()) {
+    H = hashCombine(H, hashValue(Plugin.Path));
+    H = hashCombine(H, hashValue(Plugin.Pipeline));
+    H = hashCombine(H, hashValue(Plugin.Fingerprint));
+  }
   return H;
 }
 

src/include/proteus/impl/JITPassPluginRegistry.h

@@ -0,0 +1,22 @@
+#ifndef PROTEUS_JIT_PASS_PLUGIN_REGISTRY_H
+#define PROTEUS_JIT_PASS_PLUGIN_REGISTRY_H
+
+#include <string>
+#include <vector>
+
+namespace proteus {
+
+struct JITPassPluginConfig {
+  std::string Path;
+  std::string Pipeline;
+  std::string Fingerprint;
+};
+
+void registerJITPassPluginImpl(const std::string &PluginPath,
+                               const std::string &PassPipeline);
+void clearJITPassPluginsImpl();
+std::vector<JITPassPluginConfig> getJITPassPluginConfigs();
+
+} // namespace proteus
+
+#endif

src/runtime/CMakeLists.txt

@@ -5,6 +5,7 @@ set(SOURCES
   Error.cpp
   KernelMetadata.cpp
   JitEngine.cpp
+  JITPassPluginRegistry.cpp
   JitEngineHost.cpp
   Init.cpp
   TimeTracing.cpp

src/runtime/Frontend/CppJitModule.cpp

@@ -36,6 +36,9 @@ CppJitModule::~CppJitModule() = default;
 void CppJitModule::compile() {
   TIMESCOPE(CppJitModule, compile);
 
+  ModuleHash = std::make_unique<HashT>(
+      computeCppJitModuleHash(TargetModel, CompilerBackend, Code, ExtraArgs));
+
   if ((Library = Dispatch.lookupCompiledLibrary(*ModuleHash))) {
     IsCompiled = true;
     return;

src/runtime/Frontend/LLVMIRJitModule.cpp

@@ -84,6 +84,9 @@ void LLVMIRJitModule::compile(bool Verify) {
   if (IsCompiled)
     return;
 
+  ModuleHash = std::make_unique<HashT>(
+      hash(static_cast<int>(TargetModel), Code, Config::get().getCGConfig()));
+
   if ((Library = Dispatch.lookupCompiledLibrary(*ModuleHash))) {
     IsCompiled = true;
     return;

src/runtime/Frontend/MLIRJitModule.cpp

@@ -37,6 +37,9 @@ void MLIRJitModule::compile(bool Verify) {
   if (IsCompiled)
     return;
 
+  ModuleHash = std::make_unique<HashT>(
+      hash(static_cast<int>(TargetModel), Code, Config::get().getCGConfig()));
+
   if ((Library = Dispatch.lookupCompiledLibrary(*ModuleHash))) {
     IsCompiled = true;
     return;

src/runtime/Init.cpp

@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------===//
 
 #include "proteus/Init.h"
+#include "proteus/impl/JITPassPluginRegistry.h"
 
 // NOLINTBEGIN(readability-identifier-naming)
 extern "C" void __proteus_enable_host();
@@ -20,6 +21,13 @@ namespace proteus {
 void init() {}
 void finalize() {}
 
+void registerJITPassPlugin(const std::string &PluginPath,
+                           const std::string &PassPipeline) {
+  registerJITPassPluginImpl(PluginPath, PassPipeline);
+}
+
+void clearJITPassPlugins() { clearJITPassPluginsImpl(); }
+
 void enable() {
   __proteus_enable_host();
 #if PROTEUS_ENABLE_HIP || PROTEUS_ENABLE_CUDA