diff --git a/cmake/ProteusFunctions.cmake b/cmake/ProteusFunctions.cmake index 9e8f80129..fd35e2928 100644 --- a/cmake/ProteusFunctions.cmake +++ b/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} } }; diff --git a/docs/dev/optimization-pipeline.md b/docs/dev/optimization-pipeline.md index 9f06722d7..bbed4b9d8 100644 --- a/docs/dev/optimization-pipeline.md +++ b/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,6 +166,7 @@ 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 @@ -108,9 +174,11 @@ 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. diff --git a/include/proteus/Init.h b/include/proteus/Init.h index 0f6284226..df6a9ca32 100644 --- a/include/proteus/Init.h +++ b/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(); diff --git a/src/include/proteus/impl/CoreLLVM.h b/src/include/proteus/impl/CoreLLVM.h index aea5b4c42..efbc5619d 100644 --- a/src/include/proteus/impl/CoreLLVM.h +++ b/src/include/proteus/impl/CoreLLVM.h @@ -52,6 +52,7 @@ static_assert(__cplusplus >= 201703L, #include #include +#include #include #include #include @@ -120,24 +121,60 @@ inline std::string getDefaultOptimizationPipeline(char OptLevel) { inline std::string composeOptimizationPassPipeline( std::optional PassPipeline, char OptLevel, const std::vector &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 &Plugins) { + return std::any_of(Plugins.begin(), Plugins.end(), + [](const JITPassPluginConfig &Plugin) { + return Plugin.Insertion.has_value(); + }); +} + +inline std::vector +getUniqueJITPassPluginPaths(const std::vector &Plugins) { + std::vector 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 loadJITPassPlugins(const std::vector &Plugins) { std::vector 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 &Plugins = {}) { PipelineTuningOptions PTO; std::optional 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") diff --git a/src/include/proteus/impl/CoreLLVMHIP.h b/src/include/proteus/impl/CoreLLVMHIP.h index ea21030bc..3e724a61a 100644 --- a/src/include/proteus/impl/CoreLLVMHIP.h +++ b/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(OptConfig.CodegenOptLevel); unsigned ParallelCodeGenParallelismLevel = diff --git a/src/include/proteus/impl/Hashing.h b/src/include/proteus/impl/Hashing.h index cd7bb8ff4..da3e12474 100644 --- a/src/include/proteus/impl/Hashing.h +++ b/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(Plugin.Insertion->Position))); + } H = hashCombine(H, hashValue(Plugin.Fingerprint)); } return H; diff --git a/src/include/proteus/impl/JITPassPluginRegistry.h b/src/include/proteus/impl/JITPassPluginRegistry.h index 46b86614c..af6956b48 100644 --- a/src/include/proteus/impl/JITPassPluginRegistry.h +++ b/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 #include #include 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 Insertion; std::string Fingerprint; }; void registerJITPassPluginImpl(const std::string &PluginPath, - const std::string &PassPipeline); + std::optional Insertion); void clearJITPassPluginsImpl(); std::vector getJITPassPluginConfigs(); diff --git a/src/runtime/Init.cpp b/src/runtime/Init.cpp index 6d9587634..f58f76d85 100644 --- a/src/runtime/Init.cpp +++ b/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(); } diff --git a/src/runtime/JITPassPluginRegistry.cpp b/src/runtime/JITPassPluginRegistry.cpp index c8288c875..28b81fb3e 100644 --- a/src/runtime/JITPassPluginRegistry.cpp +++ b/src/runtime/JITPassPluginRegistry.cpp @@ -10,6 +10,8 @@ #include #include #include +#include +#include namespace proteus { namespace { @@ -22,20 +24,22 @@ class JITPassPluginRegistry { } void registerPlugin(const std::string &PluginPath, - const std::string &PassPipeline) { + std::optional Insertion) { if (PluginPath.empty()) reportFatalError("JIT pass plugin path must be non-empty"); - if (PassPipeline.empty()) + if (Insertion && Insertion->Pipeline.empty()) reportFatalError("JIT pass plugin pipeline must be non-empty"); - JITPassPluginConfig Config{normalizePath(PluginPath), PassPipeline, {}}; + JITPassPluginConfig Config{ + normalizePath(PluginPath), std::move(Insertion), {}}; Config.Fingerprint = computeFingerprint(Config); std::lock_guard Lock(Mutex); - auto It = std::find_if( - Plugins.begin(), Plugins.end(), [&](const JITPassPluginConfig &Entry) { - return Entry.Path == Config.Path && Entry.Pipeline == Config.Pipeline; - }); + auto It = std::find_if(Plugins.begin(), Plugins.end(), + [&](const JITPassPluginConfig &Entry) { + return Entry.Path == Config.Path && + Entry.Insertion == Config.Insertion; + }); if (It != Plugins.end()) { It->Fingerprint = std::move(Config.Fingerprint); return; @@ -72,8 +76,12 @@ class JITPassPluginRegistry { static std::string computeFingerprint(const JITPassPluginConfig &Config) { auto BufOrErr = llvm::MemoryBuffer::getFile(Config.Path); - if (!BufOrErr) - return Config.Path + "|" + Config.Pipeline; + if (!BufOrErr) { + std::string Fingerprint = Config.Path + "|"; + if (Config.Insertion) + Fingerprint += Config.Insertion->Pipeline; + return Fingerprint; + } return hashValue(BufOrErr.get()->getBuffer()).toString(); } @@ -85,9 +93,11 @@ class JITPassPluginRegistry { } // namespace -void registerJITPassPluginImpl(const std::string &PluginPath, - const std::string &PassPipeline) { - JITPassPluginRegistry::instance().registerPlugin(PluginPath, PassPipeline); +void registerJITPassPluginImpl( + const std::string &PluginPath, + std::optional Insertion) { + JITPassPluginRegistry::instance().registerPlugin(PluginPath, + std::move(Insertion)); } void clearJITPassPluginsImpl() { JITPassPluginRegistry::instance().clear(); } diff --git a/tests/JITTestPass.cpp b/tests/JITTestPass.cpp index d0a9ff71c..461a8dd99 100644 --- a/tests/JITTestPass.cpp +++ b/tests/JITTestPass.cpp @@ -10,28 +10,39 @@ #endif #include +#include +#include + namespace { class JITTestPass : public llvm::PassInfoMixin { public: + explicit JITTestPass(std::string PipelineName) + : PipelineName(std::move(PipelineName)) {} + llvm::PreservedAnalyses run(llvm::Module &M, llvm::ModuleAnalysisManager &) { - llvm::outs() << "[JITTestPass] " << M.getName() << "\n"; + llvm::outs() << "[JITTestPass] " << PipelineName << " " << M.getName() + << "\n"; return llvm::PreservedAnalyses::all(); } + +private: + std::string PipelineName; }; } // namespace extern "C" LLVM_ATTRIBUTE_WEAK ::llvm::PassPluginLibraryInfo llvmGetPassPluginInfo() { + llvm::outs() << "[JITTestPluginInfo]\n"; return {LLVM_PLUGIN_API_VERSION, "JITTestPass", "0.1", [](llvm::PassBuilder &PB) { PB.registerPipelineParsingCallback( [](llvm::StringRef Name, llvm::ModulePassManager &MPM, llvm::ArrayRef) { - if (Name != "jit-test-pass") + if (!Name.starts_with("jit-test-pass")) return false; - MPM.addPass(JITTestPass()); + MPM.addPass(JITTestPass(Name.str())); return true; }); }}; diff --git a/tests/cpu/CMakeLists.txt b/tests/cpu/CMakeLists.txt index 0bffd4cb6..21860f2e5 100644 --- a/tests/cpu/CMakeLists.txt +++ b/tests/cpu/CMakeLists.txt @@ -69,10 +69,20 @@ CREATE_CPU_TEST(jit_pass_plugin jit_pass_plugin.cpp) target_compile_definitions(jit_pass_plugin PRIVATE PROTEUS_TEST_JIT_PASS_PLUGIN_PATH="$") add_dependencies(jit_pass_plugin JITTestPass) +CREATE_CPU_TEST(jit_pass_plugin_position jit_pass_plugin_position.cpp) +target_compile_definitions(jit_pass_plugin_position PRIVATE + PROTEUS_TEST_JIT_PASS_PLUGIN_PATH="$") +add_dependencies(jit_pass_plugin_position JITTestPass) CREATE_CPU_TEST(jit_pass_plugin_cmake jit_pass_plugin_cmake.cpp) proteus_register_jit_pass_plugin( jit_pass_plugin_cmake PLUGIN_TARGET JITTestPass - PIPELINE jit-test-pass) + PIPELINE jit-test-pass + POSITION PREPEND) +CREATE_CPU_TEST(jit_pass_plugin_cmake_load_only + jit_pass_plugin_cmake_load_only.cpp) +proteus_register_jit_pass_plugin( + jit_pass_plugin_cmake_load_only + PLUGIN_TARGET JITTestPass) CREATE_CPU_TEST(modify_gvar modify_gvar.cpp) CREATE_CPU_TEST(jit_eh jit_eh.cpp) diff --git a/tests/cpu/jit_pass_plugin.cpp b/tests/cpu/jit_pass_plugin.cpp index f99ef8335..5db9641a4 100644 --- a/tests/cpu/jit_pass_plugin.cpp +++ b/tests/cpu/jit_pass_plugin.cpp @@ -21,7 +21,10 @@ int main() { return 0; } -// CHECK: [JITTestPass] +// CHECK: [JITTestPluginInfo] +// CHECK-NOT: [JITTestPluginInfo] +// CHECK: [JITTestPass] jit-test-pass +// CHECK-NOT: [JITTestPass] jit-test-pass // CHECK: [CustomPipeline] default,jit-test-pass // CHECK: 5 // CHECK: [proteus][JitEngineHost] MemoryCache rank 0 hits 0 accesses 1 diff --git a/tests/cpu/jit_pass_plugin_cmake.cpp b/tests/cpu/jit_pass_plugin_cmake.cpp index e1f9710af..008f7675b 100644 --- a/tests/cpu/jit_pass_plugin_cmake.cpp +++ b/tests/cpu/jit_pass_plugin_cmake.cpp @@ -18,7 +18,10 @@ int main() { return 0; } -// CHECK: [JITTestPass] -// CHECK: [CustomPipeline] default,jit-test-pass +// CHECK: [JITTestPluginInfo] +// CHECK-NOT: [JITTestPluginInfo] +// CHECK: [JITTestPass] jit-test-pass +// CHECK-NOT: [JITTestPass] jit-test-pass +// CHECK: [CustomPipeline] jit-test-pass,default // CHECK: 7 // CHECK: [proteus][JitEngineHost] MemoryCache rank 0 hits 0 accesses 1 diff --git a/tests/cpu/jit_pass_plugin_cmake_load_only.cpp b/tests/cpu/jit_pass_plugin_cmake_load_only.cpp new file mode 100644 index 000000000..3325f0815 --- /dev/null +++ b/tests/cpu/jit_pass_plugin_cmake_load_only.cpp @@ -0,0 +1,27 @@ +// clang-format off +// RUN: rm -rf "%t.$$.proteus" +// RUN: PROTEUS_CACHE_DIR="%t.$$.proteus" PROTEUS_CODEGEN=serial PROTEUS_TRACE_OUTPUT="specialization;cache-stats" PROTEUS_OPT_PIPELINE="default,jit-test-pass" %build/jit_pass_plugin_cmake_load_only | %FILECHECK %s +// RUN: rm -rf "%t.$$.proteus" +// clang-format on + +#include + +#include + +__attribute__((annotate("jit"))) int add_four(int x) { + proteus::jit_arg(x); + return x + 4; +} + +int main() { + std::cout << add_four(7) << "\n"; + return 0; +} + +// CHECK: [JITTestPluginInfo] +// CHECK-NOT: [JITTestPluginInfo] +// CHECK: [JITTestPass] jit-test-pass +// CHECK-NOT: [JITTestPass] jit-test-pass +// CHECK: [CustomPipeline] default,jit-test-pass +// CHECK: 11 +// CHECK: [proteus][JitEngineHost] MemoryCache rank 0 hits 0 accesses 1 diff --git a/tests/cpu/jit_pass_plugin_position.cpp b/tests/cpu/jit_pass_plugin_position.cpp new file mode 100644 index 000000000..3cfb91481 --- /dev/null +++ b/tests/cpu/jit_pass_plugin_position.cpp @@ -0,0 +1,93 @@ +// clang-format off +// RUN: rm -rf "%t.$$.proteus" +// RUN: PROTEUS_CACHE_DIR="%t.$$.proteus" PROTEUS_CODEGEN=serial PROTEUS_TRACE_OUTPUT="specialization;cache-stats" %build/jit_pass_plugin_position append | %FILECHECK %s --check-prefix=APPEND +// RUN: PROTEUS_CACHE_DIR="%t.$$.proteus" PROTEUS_CODEGEN=serial PROTEUS_TRACE_OUTPUT="specialization;cache-stats" %build/jit_pass_plugin_position prepend | %FILECHECK %s --check-prefix=PREPEND +// RUN: PROTEUS_CACHE_DIR="%t.$$.proteus" PROTEUS_CODEGEN=serial PROTEUS_TRACE_OUTPUT="specialization;cache-stats" %build/jit_pass_plugin_position load-only | %FILECHECK %s --check-prefix=LOAD-DEFAULT +// RUN: rm -rf "%t.$$.proteus" +// RUN: PROTEUS_CACHE_DIR="%t.$$.proteus" PROTEUS_CODEGEN=serial PROTEUS_TRACE_OUTPUT="specialization;cache-stats" PROTEUS_OPT_PIPELINE="default,jit-test-pass" %build/jit_pass_plugin_position load-only | %FILECHECK %s --check-prefix=LOAD-ONLY +// RUN: rm -rf "%t.$$.proteus" +// RUN: PROTEUS_CACHE_DIR="%t.$$.proteus" PROTEUS_CODEGEN=serial PROTEUS_TRACE_OUTPUT="specialization;cache-stats" %build/jit_pass_plugin_position order | %FILECHECK %s --check-prefix=ORDER +// RUN: rm -rf "%t.$$.proteus" +// clang-format on + +#include +#include + +#include +#include + +__attribute__((annotate("jit"))) int add_three(int x) { + proteus::jit_arg(x); + return x + 3; +} + +int main(int argc, char **argv) { + const std::string Mode = argc > 1 ? argv[1] : ""; + if (Mode == "append") { + proteus::registerJITPassPlugin(PROTEUS_TEST_JIT_PASS_PLUGIN_PATH, + "jit-test-pass", + proteus::JITPassPluginPosition::Append); + } else if (Mode == "prepend") { + proteus::registerJITPassPlugin(PROTEUS_TEST_JIT_PASS_PLUGIN_PATH, + "jit-test-pass", + proteus::JITPassPluginPosition::Prepend); + } else if (Mode == "load-only") { + proteus::registerJITPassPlugin(PROTEUS_TEST_JIT_PASS_PLUGIN_PATH); + } else if (Mode == "order") { + proteus::registerJITPassPlugin(PROTEUS_TEST_JIT_PASS_PLUGIN_PATH, + "jit-test-pass-prepend-first", + proteus::JITPassPluginPosition::Prepend); + proteus::registerJITPassPlugin(PROTEUS_TEST_JIT_PASS_PLUGIN_PATH, + "jit-test-pass-append-first", + proteus::JITPassPluginPosition::Append); + proteus::registerJITPassPlugin(PROTEUS_TEST_JIT_PASS_PLUGIN_PATH); + proteus::registerJITPassPlugin(PROTEUS_TEST_JIT_PASS_PLUGIN_PATH, + "jit-test-pass-prepend-second", + proteus::JITPassPluginPosition::Prepend); + proteus::registerJITPassPlugin(PROTEUS_TEST_JIT_PASS_PLUGIN_PATH, + "jit-test-pass-append-second", + proteus::JITPassPluginPosition::Append); + } else { + return 1; + } + + std::cout << add_three(6) << "\n"; + return 0; +} + +// APPEND: [JITTestPluginInfo] +// APPEND-NOT: [JITTestPluginInfo] +// APPEND: [JITTestPass] jit-test-pass +// APPEND-NOT: [JITTestPass] jit-test-pass +// APPEND: [CustomPipeline] default,jit-test-pass +// APPEND: 9 + +// PREPEND: [JITTestPluginInfo] +// PREPEND-NOT: [JITTestPluginInfo] +// PREPEND: [JITTestPass] jit-test-pass +// PREPEND-NOT: [JITTestPass] jit-test-pass +// PREPEND: [CustomPipeline] jit-test-pass,default +// PREPEND: 9 + +// LOAD-DEFAULT: [JITTestPluginInfo] +// LOAD-DEFAULT-NOT: [JITTestPass] +// LOAD-DEFAULT-NOT: [CustomPipeline] +// LOAD-DEFAULT: 9 + +// LOAD-ONLY: [JITTestPluginInfo] +// LOAD-ONLY-NOT: [JITTestPluginInfo] +// LOAD-ONLY: [JITTestPass] jit-test-pass +// LOAD-ONLY-NOT: [JITTestPass] jit-test-pass +// LOAD-ONLY: [CustomPipeline] default,jit-test-pass +// LOAD-ONLY: 9 + +// ORDER: [JITTestPluginInfo] +// ORDER-NOT: [JITTestPluginInfo] +// ORDER: [JITTestPass] jit-test-pass-prepend-first +// ORDER: [JITTestPass] jit-test-pass-prepend-second +// ORDER: [JITTestPass] jit-test-pass-append-first +// ORDER: [JITTestPass] jit-test-pass-append-second +// ORDER: [CustomPipeline] jit-test-pass-prepend-first, +// ORDER-SAME: jit-test-pass-prepend-second,default, +// ORDER-SAME: jit-test-pass-append-first,jit-test-pass-append-second +// ORDER: 9 diff --git a/tests/gpu/kernel_pass_plugin.cpp b/tests/gpu/kernel_pass_plugin.cpp index 024e5c000..81642f3f8 100644 --- a/tests/gpu/kernel_pass_plugin.cpp +++ b/tests/gpu/kernel_pass_plugin.cpp @@ -1,10 +1,17 @@ // clang-format off // RUN: rm -rf "%t.$$.proteus" -// RUN: PROTEUS_CACHE_DIR="%t.$$.proteus" PROTEUS_CODEGEN=serial PROTEUS_TRACE_OUTPUT="specialization;cache-stats" %build/kernel_pass_plugin.%ext | %FILECHECK %s +// RUN: PROTEUS_CACHE_DIR="%t.$$.proteus" PROTEUS_CODEGEN=serial PROTEUS_TRACE_OUTPUT="specialization;cache-stats" %build/kernel_pass_plugin.%ext append | %FILECHECK %s --check-prefixes=CHECK,APPEND +// RUN: rm -rf "%t.$$.proteus" +// RUN: PROTEUS_CACHE_DIR="%t.$$.proteus" PROTEUS_CODEGEN=serial PROTEUS_TRACE_OUTPUT="specialization;cache-stats" %build/kernel_pass_plugin.%ext prepend | %FILECHECK %s --check-prefixes=CHECK,PREPEND +// RUN: rm -rf "%t.$$.proteus" +// RUN: PROTEUS_CACHE_DIR="%t.$$.proteus" PROTEUS_CODEGEN=serial PROTEUS_TRACE_OUTPUT="specialization;cache-stats" PROTEUS_OPT_PIPELINE="default,jit-test-pass" %build/kernel_pass_plugin.%ext load-only | %FILECHECK %s --check-prefixes=CHECK,LOAD-ONLY +// RUN: rm -rf "%t.$$.proteus" +// RUN: if [ "%device_lang" = "HIP" ]; then PROTEUS_CACHE_DIR="%t.$$.proteus" PROTEUS_CODEGEN=parallel PROTEUS_TRACE_OUTPUT="specialization;cache-stats" PROTEUS_OPT_PIPELINE="default,jit-test-pass" %build/kernel_pass_plugin.%ext load-only | %FILECHECK %s --check-prefixes=CHECK,PARALLEL; fi // RUN: rm -rf "%t.$$.proteus" // clang-format on #include +#include #include "gpu_common.h" #include @@ -14,15 +21,41 @@ __global__ __attribute__((annotate("jit"))) void kernel_pass_plugin() { printf("KernelPassPlugin\n"); } -int main() { - proteus::registerJITPassPlugin(PROTEUS_TEST_JIT_PASS_PLUGIN_PATH, - "jit-test-pass"); +int main(int argc, char **argv) { + const std::string Mode = argc > 1 ? argv[1] : ""; + if (Mode == "append") { + proteus::registerJITPassPlugin(PROTEUS_TEST_JIT_PASS_PLUGIN_PATH, + "jit-test-pass"); + } else if (Mode == "prepend") { + proteus::registerJITPassPlugin(PROTEUS_TEST_JIT_PASS_PLUGIN_PATH, + "jit-test-pass", + proteus::JITPassPluginPosition::Prepend); + } else if (Mode == "load-only") { + proteus::registerJITPassPlugin(PROTEUS_TEST_JIT_PASS_PLUGIN_PATH); + } else { + return 1; + } + kernel_pass_plugin<<<1, 1>>>(); gpuErrCheck(gpuDeviceSynchronize()); return 0; } -// CHECK-DAG: [JITTestPass] -// CHECK-DAG: [CustomPipeline] default,jit-test-pass +// APPEND-DAG: [JITTestPluginInfo] +// APPEND-DAG: [JITTestPass] jit-test-pass +// APPEND-DAG: [CustomPipeline] default,jit-test-pass + +// PREPEND-DAG: [JITTestPluginInfo] +// PREPEND-DAG: [JITTestPass] jit-test-pass +// PREPEND-DAG: [CustomPipeline] jit-test-pass,default + +// LOAD-ONLY-DAG: [JITTestPluginInfo] +// LOAD-ONLY-DAG: [JITTestPass] jit-test-pass +// LOAD-ONLY-DAG: [CustomPipeline] default,jit-test-pass + +// PARALLEL-DAG: [JITTestPluginInfo] +// PARALLEL-DAG: [JITTestPass] jit-test-pass + // CHECK-DAG: KernelPassPlugin +// LOAD-ONLY-NOT: [JITTestPass] jit-test-pass // CHECK: [proteus][JitEngineDevice] MemoryCache rank 0 hits 0 accesses 1