Enable specific plugin pass placement

cmake/ProteusFunctions.cmake

@@ -34,11 +34,24 @@ function(proteus_register_jit_pass_plugin target)
         message(FATAL_ERROR "Target '${target}' does not exist")
     endif()
 
-    set(one_value_args PLUGIN_TARGET PLUGIN_PATH PIPELINE)
+    set(one_value_args PLUGIN_TARGET PLUGIN_PATH PIPELINE POSITION)
     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")
+    if(PROTEUS_JIT_PASS_UNPARSED_ARGUMENTS)
+        message(FATAL_ERROR
+            "proteus_register_jit_pass_plugin received unknown arguments: ${PROTEUS_JIT_PASS_UNPARSED_ARGUMENTS}")
+    endif()
+
+    if(PROTEUS_JIT_PASS_POSITION AND NOT PROTEUS_JIT_PASS_PIPELINE)
+        message(FATAL_ERROR
+            "proteus_register_jit_pass_plugin POSITION requires a nonempty PIPELINE")
+    endif()
+
+    if(PROTEUS_JIT_PASS_POSITION AND
+       NOT PROTEUS_JIT_PASS_POSITION STREQUAL "PREPEND" AND
+       NOT PROTEUS_JIT_PASS_POSITION STREQUAL "APPEND")
+        message(FATAL_ERROR
+            "proteus_register_jit_pass_plugin POSITION must be PREPEND or APPEND")
     endif()
 
     if(PROTEUS_JIT_PASS_PLUGIN_TARGET AND PROTEUS_JIT_PASS_PLUGIN_PATH)
@@ -62,8 +75,26 @@ function(proteus_register_jit_pass_plugin target)
         set(_proteus_jit_pass_plugin_path "${PROTEUS_JIT_PASS_PLUGIN_PATH}")
     endif()
 
+    if(PROTEUS_JIT_PASS_PIPELINE)
+        if(PROTEUS_JIT_PASS_POSITION STREQUAL "PREPEND")
+            set(_proteus_jit_pass_position "Prepend")
+        else()
+            set(_proteus_jit_pass_position "Append")
+        endif()
+        set(_proteus_jit_pass_registration
+"    proteus::registerJITPassPlugin(
+        R\"(${_proteus_jit_pass_plugin_path})\",
+        R\"(${PROTEUS_JIT_PASS_PIPELINE})\",
+        proteus::JITPassPluginPosition::${_proteus_jit_pass_position});")
+    else()
+        set(_proteus_jit_pass_position "LoadOnly")
+        set(_proteus_jit_pass_registration
+"    proteus::registerJITPassPlugin(
+        R\"(${_proteus_jit_pass_plugin_path})\");")
+    endif()
+
     string(MD5 _proteus_jit_pass_key
-        "${target};${_proteus_jit_pass_plugin_path};${PROTEUS_JIT_PASS_PIPELINE}")
+        "${target};${_proteus_jit_pass_plugin_path};${_proteus_jit_pass_position};${PROTEUS_JIT_PASS_PIPELINE}")
     set(_proteus_jit_pass_source
         "${CMAKE_CURRENT_BINARY_DIR}/${target}.proteus_jit_pass_${_proteus_jit_pass_key}.cpp")
 
@@ -73,9 +104,7 @@ function(proteus_register_jit_pass_plugin target)
 namespace {
 struct AutoRegisterProteusJITPassPlugin {
   AutoRegisterProteusJITPassPlugin() {
-    proteus::registerJITPassPlugin(
-        R\"(${_proteus_jit_pass_plugin_path})\",
-        R\"(${PROTEUS_JIT_PASS_PIPELINE})\");
+${_proteus_jit_pass_registration}
   }
 };
 

docs/dev/optimization-pipeline.md

@@ -33,6 +33,69 @@ Config::get().getCGConfig()
 so per-kernel JSON `"Pipeline"` does not naturally apply to DSL, MLIR, or C++
 frontend modules.
 
+## JIT Pass Plugins
+
+Registering a JIT pass plugin always loads its normalized library path
+and registers its PassBuilder callbacks
+before Proteus builds or parses the optimization pipeline.
+The C++ API supports three registration modes:
+
+```cpp
+// Load callbacks and append a pipeline fragment (the default).
+proteus::registerJITPassPlugin(path, "my-pass");
+
+// Load callbacks and explicitly place a pipeline fragment.
+proteus::registerJITPassPlugin(
+    path, "my-pass", proteus::JITPassPluginPosition::Prepend);
+proteus::registerJITPassPlugin(
+    path, "my-pass", proteus::JITPassPluginPosition::Append);
+
+// Load callbacks without automatically inserting a fragment.
+proteus::registerJITPassPlugin(path);
+```
+
+For automatic insertion,
+Proteus composes the effective textual pipeline in this order:
+
+1. Prepended fragments in registration order.
+2. The configured `PROTEUS_OPT_PIPELINE` or default optimization pipeline.
+3. Appended fragments in registration order.
+
+Load-only registrations do not add text to the pipeline.
+They can still make pass names available
+to a user-authored `PROTEUS_OPT_PIPELINE`,
+and their callbacks can affect default pipelines through LLVM extension points.
+If every registration is load-only and no custom pipeline is configured,
+Proteus builds the normal LLVM default pipeline.
+
+Multiple registrations can use the same plugin library.
+Proteus preserves every distinct registration for ordering and cache identity,
+while loading each normalized library path only once per compilation.
+The registration mode, position, pipeline fragment, plugin fingerprint,
+and registration order contribute to JIT cache identity.
+
+The CMake helper exposes the same modes:
+
+```cmake
+# Append by default.
+proteus_register_jit_pass_plugin(
+    app PLUGIN_TARGET MyPassPlugin PIPELINE my-pass)
+
+# Explicit placement.
+proteus_register_jit_pass_plugin(
+    app PLUGIN_TARGET MyPassPlugin PIPELINE my-pass POSITION PREPEND)
+
+# Load only.
+proteus_register_jit_pass_plugin(
+    app PLUGIN_TARGET MyPassPlugin)
+```
+
+`POSITION` accepts `PREPEND` or `APPEND`.
+It requires a nonempty `PIPELINE`;
+when `PIPELINE` is present without `POSITION`,
+the helper appends the pipeline fragment by default.
+`PLUGIN_PATH` can be used instead of `PLUGIN_TARGET` in all three modes.
+
 ## Support Matrix
 
 | Frontend / API path | Target type | Main compile path | Uses `PROTEUS_OPT_PIPELINE`? | Notes |
@@ -75,6 +138,8 @@ This includes:
 - `PROTEUS_OPT_LEVEL`
 - `PROTEUS_CODEGEN_OPT_LEVEL`
 - `PROTEUS_OPT_PIPELINE`
+- Ordered JIT pass-plugin registrations,
+  including load-only versus inserted mode and prepend versus append position
 
 Runtime annotated JIT cache keys also include runtime specialization policy:
 
@@ -101,16 +166,19 @@ way annotated runtime JIT paths do.
 interface accepts some compiler options, but it does not expose a documented
 LLVM textual pass pipeline interface equivalent to `opt`/PassBuilder or LLVM
 LTO's `OptPipeline`.
+JIT pass-plugin registration does not change this exclusion.
 
 ### CppJit Host+CUDA / Host+HIP
 
 The Clang backend currently compiles mixed host/device C++ offload source
 directly into a shared library for `HOST_CUDA` and `HOST_HIP`. Proteus receives
 the final `.so`, so it cannot run its LLVM pass pipeline over the host and
 device modules.
+JIT pass-plugin registration does not change this exclusion.
 
 ### NVCC Backend
 
 The NVCC backend is an external compiler path. Proteus does not own LLVM IR
 optimization there, so `PROTEUS_OPT_PIPELINE` does not apply and is not included
 in NVCC CppJit cache keys.
+JIT pass-plugin registration does not change this exclusion.

include/proteus/Init.h

@@ -18,8 +18,14 @@ void init();
 [[deprecated("it is a no-op and will be removed in a future version.")]]
 void finalize();
 
+enum class JITPassPluginPosition { Prepend, Append };
+
+void registerJITPassPlugin(const std::string &PluginPath);
 void registerJITPassPlugin(const std::string &PluginPath,
                            const std::string &PassPipeline);
+void registerJITPassPlugin(const std::string &PluginPath,
+                           const std::string &PassPipeline,
+                           JITPassPluginPosition Position);
 void clearJITPassPlugins();
 
 void enable();

src/include/proteus/impl/CoreLLVM.h

@@ -52,6 +52,7 @@ static_assert(__cplusplus >= 201703L,
 #include <llvm/Transforms/IPO/StripSymbols.h>
 #include <llvm/Transforms/Utils/ModuleUtils.h>
 
+#include <algorithm>
 #include <optional>
 #include <string>
 #include <utility>
@@ -120,24 +121,60 @@ inline std::string getDefaultOptimizationPipeline(char OptLevel) {
 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);
+  std::string Pipeline;
   for (const auto &Plugin : Plugins) {
+    if (!Plugin.Insertion ||
+        Plugin.Insertion->Position != JITPassPluginPosition::Prepend)
+      continue;
+    if (!Pipeline.empty())
+      Pipeline += ",";
+    Pipeline += Plugin.Insertion->Pipeline;
+  }
+
+  if (!Pipeline.empty())
+    Pipeline += ",";
+  Pipeline += PassPipeline ? std::move(*PassPipeline)
+                           : getDefaultOptimizationPipeline(OptLevel);
+
+  for (const auto &Plugin : Plugins) {
+    if (!Plugin.Insertion ||
+        Plugin.Insertion->Position != JITPassPluginPosition::Append)
+      continue;
     Pipeline += ",";
-    Pipeline += Plugin.Pipeline;
+    Pipeline += Plugin.Insertion->Pipeline;
   }
+
   return Pipeline;
 }
 
+inline bool
+hasJITPassPluginInsertion(const std::vector<JITPassPluginConfig> &Plugins) {
+  return std::any_of(Plugins.begin(), Plugins.end(),
+                     [](const JITPassPluginConfig &Plugin) {
+                       return Plugin.Insertion.has_value();
+                     });
+}
+
+inline std::vector<std::string>
+getUniqueJITPassPluginPaths(const std::vector<JITPassPluginConfig> &Plugins) {
+  std::vector<std::string> Paths;
+  Paths.reserve(Plugins.size());
+  for (const auto &Plugin : Plugins) {
+    if (std::find(Paths.begin(), Paths.end(), Plugin.Path) == Paths.end())
+      Paths.push_back(Plugin.Path);
+  }
+  return Paths;
+}
+
 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);
+  const auto PluginPaths = getUniqueJITPassPluginPaths(Plugins);
+  LoadedPlugins.reserve(PluginPaths.size());
+  for (const auto &PluginPath : PluginPaths) {
+    auto LoadedPlugin = PassPlugin::Load(PluginPath);
     if (!LoadedPlugin)
-      reportFatalError("Failed to load JIT pass plugin '" + Plugin.Path +
+      reportFatalError("Failed to load JIT pass plugin '" + PluginPath +
                        "': " + toString(LoadedPlugin.takeError()));
     LoadedPlugins.push_back(std::move(*LoadedPlugin));
   }
@@ -179,9 +216,10 @@ inline void runOptimizationPassPipeline(
   Passes.run(M, MAM);
 }
 
-inline void runOptimizationPassPipeline(Module &M, StringRef Arch,
-                                        char OptLevel = '3',
-                                        unsigned CodegenOptLevel = 3) {
+inline void runOptimizationPassPipeline(
+    Module &M, StringRef Arch, char OptLevel = '3',
+    unsigned CodegenOptLevel = 3,
+    const std::vector<JITPassPluginConfig> &Plugins = {}) {
   PipelineTuningOptions PTO;
 
   std::optional<PGOOptions> PGOOpt;
@@ -190,7 +228,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;
@@ -268,7 +309,8 @@ inline void optimizeIR(Module &M, StringRef Arch,
   Timer T(Config::get().ProteusEnableTimers);
 
   const auto Plugins = getJITPassPluginConfigs();
-  const bool UseTextualPipeline = OptConfig.PassPipeline || !Plugins.empty();
+  const bool UseTextualPipeline =
+      OptConfig.PassPipeline || detail::hasJITPassPluginInsertion(Plugins);
   const std::string FinalPipeline =
       UseTextualPipeline
           ? detail::composeOptimizationPassPipeline(OptConfig.PassPipeline,
@@ -290,7 +332,7 @@ inline void optimizeIR(Module &M, StringRef Arch,
                                         OptConfig.CodegenOptLevel, Plugins);
   } else {
     detail::runOptimizationPassPipeline(M, Arch, OptConfig.OptLevel,
-                                        OptConfig.CodegenOptLevel);
+                                        OptConfig.CodegenOptLevel, Plugins);
   }
 
   PROTEUS_TIMER_OUTPUT(Logger::outs("proteus")

src/include/proteus/impl/CoreLLVMHIP.h

@@ -252,11 +252,13 @@ codegenParallel(Module &M, StringRef DeviceArch,
   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 || !Plugins.empty())
+  if (OptConfig.PassPipeline ||
+      proteus::detail::hasJITPassPluginInsertion(Plugins))
     Conf.OptPipeline = proteus::detail::composeOptimizationPassPipeline(
         OptConfig.PassPipeline, OptConfig.OptLevel, Plugins);
-  for (const auto &Plugin : Plugins)
-    Conf.PassPlugins.push_back(Plugin.Path);
+  for (const auto &PluginPath :
+       proteus::detail::getUniqueJITPassPluginPaths(Plugins))
+    Conf.PassPlugins.push_back(PluginPath);
   Conf.CGOptLevel = static_cast<CodeGenOptLevel>(OptConfig.CodegenOptLevel);
 
   unsigned ParallelCodeGenParallelismLevel =

src/include/proteus/impl/Hashing.h

@@ -148,7 +148,12 @@ inline HashT hashCodeGenConfig(const CodeGenerationConfig &CGConfig) {
     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.Insertion.has_value()));
+    if (Plugin.Insertion) {
+      H = hashCombine(H, hashValue(Plugin.Insertion->Pipeline));
+      H = hashCombine(H,
+                      hashValue(static_cast<int>(Plugin.Insertion->Position)));
+    }
     H = hashCombine(H, hashValue(Plugin.Fingerprint));
   }
   return H;

src/include/proteus/impl/JITPassPluginRegistry.h

@@ -1,19 +1,31 @@
 #ifndef PROTEUS_JIT_PASS_PLUGIN_REGISTRY_H
 #define PROTEUS_JIT_PASS_PLUGIN_REGISTRY_H
 
+#include "proteus/Init.h"
+
+#include <optional>
 #include <string>
 #include <vector>
 
 namespace proteus {
 
+struct JITPassPluginInsertion {
+  std::string Pipeline;
+  JITPassPluginPosition Position;
+
+  bool operator==(const JITPassPluginInsertion &Other) const {
+    return Pipeline == Other.Pipeline && Position == Other.Position;
+  }
+};
+
 struct JITPassPluginConfig {
   std::string Path;
-  std::string Pipeline;
+  std::optional<JITPassPluginInsertion> Insertion;
   std::string Fingerprint;
 };
 
 void registerJITPassPluginImpl(const std::string &PluginPath,
-                               const std::string &PassPipeline);
+                               std::optional<JITPassPluginInsertion> Insertion);
 void clearJITPassPluginsImpl();
 std::vector<JITPassPluginConfig> getJITPassPluginConfigs();
 

src/runtime/Init.cpp

@@ -21,9 +21,21 @@ namespace proteus {
 void init() {}
 void finalize() {}
 
+void registerJITPassPlugin(const std::string &PluginPath) {
+  registerJITPassPluginImpl(PluginPath, std::nullopt);
+}
+
 void registerJITPassPlugin(const std::string &PluginPath,
                            const std::string &PassPipeline) {
-  registerJITPassPluginImpl(PluginPath, PassPipeline);
+  registerJITPassPlugin(PluginPath, PassPipeline,
+                        JITPassPluginPosition::Append);
+}
+
+void registerJITPassPlugin(const std::string &PluginPath,
+                           const std::string &PassPipeline,
+                           JITPassPluginPosition Position) {
+  registerJITPassPluginImpl(
+      PluginPath, JITPassPluginInsertion{std::string(PassPipeline), Position});
 }
 
 void clearJITPassPlugins() { clearJITPassPluginsImpl(); }