From 9f43ebebba00435552d1c756f8f29f25cb2cee11 Mon Sep 17 00:00:00 2001 From: john bowen Date: Fri, 22 May 2026 13:14:16 -0700 Subject: [PATCH 01/12] Add interface for plugin pass --- include/proteus/Init.h | 6 ++ src/include/proteus/impl/CoreLLVM.h | 84 ++++++++++++++-- src/include/proteus/impl/CoreLLVMHIP.h | 8 +- src/include/proteus/impl/Hashing.h | 6 ++ .../proteus/impl/JITPassPluginRegistry.h | 22 +++++ src/runtime/CMakeLists.txt | 1 + src/runtime/Frontend/CppJitModule.cpp | 3 + src/runtime/Frontend/LLVMIRJitModule.cpp | 3 + src/runtime/Frontend/MLIRJitModule.cpp | 3 + src/runtime/Init.cpp | 8 ++ src/runtime/JITPassPluginRegistry.cpp | 95 +++++++++++++++++++ tests/CMakeLists.txt | 19 ++++ tests/JITTestPass.cpp | 39 ++++++++ tests/cpu/CMakeLists.txt | 9 ++ tests/cpu/jit_pass_plugin.cpp | 27 ++++++ tests/cpu/jit_pass_plugin_cmake.cpp | 24 +++++ tests/gpu/CMakeLists.txt | 4 + tests/gpu/kernel_pass_plugin.cpp | 28 ++++++ 18 files changed, 380 insertions(+), 9 deletions(-) create mode 100644 src/include/proteus/impl/JITPassPluginRegistry.h create mode 100644 src/runtime/JITPassPluginRegistry.cpp create mode 100644 tests/JITTestPass.cpp create mode 100644 tests/cpu/jit_pass_plugin.cpp create mode 100644 tests/cpu/jit_pass_plugin_cmake.cpp create mode 100644 tests/gpu/kernel_pass_plugin.cpp diff --git a/include/proteus/Init.h b/include/proteus/Init.h index 3a28dcde4..0f6284226 100644 --- a/include/proteus/Init.h +++ b/include/proteus/Init.h @@ -9,6 +9,8 @@ #ifndef PROTEUS_INIT_H #define PROTEUS_INIT_H +#include + namespace proteus { [[deprecated("it is a no-op and will be removed in a future version.")]] @@ -16,6 +18,10 @@ 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(); diff --git a/src/include/proteus/impl/CoreLLVM.h b/src/include/proteus/impl/CoreLLVM.h index 87d49ba5a..643d59c17 100644 --- a/src/include/proteus/impl/CoreLLVM.h +++ b/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 @@ -17,6 +18,13 @@ static_assert(__cplusplus >= 201703L, #include #include #include +#if __has_include() +#include +#elif __has_include() +#include +#else +#error "Cannot find LLVM PassPlugin.h" +#endif #include #include #include @@ -47,6 +55,7 @@ static_assert(__cplusplus >= 201703L, #include #include #include +#include namespace proteus { using namespace llvm; @@ -88,9 +97,59 @@ createTargetMachine(Module &M, StringRef Arch, unsigned OptLevel = 3) { return TM; } +inline std::string getDefaultOptimizationPipeline(char OptLevel) { + switch (OptLevel) { + case '0': + return "default"; + case '1': + return "default"; + case '2': + return "default"; + case '3': + return "default"; + case 's': + return "default"; + case 'z': + return "default"; + default: + reportFatalError(std::string("Unsupported optimization level ") + OptLevel); + } + return ""; +} + +inline std::string +composeOptimizationPassPipeline(std::optional PassPipeline, + char OptLevel, + const std::vector &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 +loadJITPassPlugins(const std::vector &Plugins) { + std::vector 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 = 3) { + unsigned CodegenOptLevel, + const std::vector + &Plugins = {}) { PipelineTuningOptions PTO; std::optional PGOOpt; @@ -99,7 +158,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 +269,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,10 +286,10 @@ 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); @@ -227,8 +297,8 @@ inline void optimizeIR(Module &M, StringRef Arch, 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"); diff --git a/src/include/proteus/impl/CoreLLVMHIP.h b/src/include/proteus/impl/CoreLLVMHIP.h index 38c0fa2a6..ea21030bc 100644 --- a/src/include/proteus/impl/CoreLLVMHIP.h +++ b/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(OptConfig.CodegenOptLevel); unsigned ParallelCodeGenParallelismLevel = diff --git a/src/include/proteus/impl/Hashing.h b/src/include/proteus/impl/Hashing.h index e103599e4..cd7bb8ff4 100644 --- a/src/include/proteus/impl/Hashing.h +++ b/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 @@ -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; } diff --git a/src/include/proteus/impl/JITPassPluginRegistry.h b/src/include/proteus/impl/JITPassPluginRegistry.h new file mode 100644 index 000000000..46b86614c --- /dev/null +++ b/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 +#include + +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 getJITPassPluginConfigs(); + +} // namespace proteus + +#endif diff --git a/src/runtime/CMakeLists.txt b/src/runtime/CMakeLists.txt index 12921aaf6..555064617 100644 --- a/src/runtime/CMakeLists.txt +++ b/src/runtime/CMakeLists.txt @@ -5,6 +5,7 @@ set(SOURCES Error.cpp KernelMetadata.cpp JitEngine.cpp + JITPassPluginRegistry.cpp JitEngineHost.cpp Init.cpp TimeTracing.cpp diff --git a/src/runtime/Frontend/CppJitModule.cpp b/src/runtime/Frontend/CppJitModule.cpp index 4215eb2ef..3bb76c46e 100644 --- a/src/runtime/Frontend/CppJitModule.cpp +++ b/src/runtime/Frontend/CppJitModule.cpp @@ -36,6 +36,9 @@ CppJitModule::~CppJitModule() = default; void CppJitModule::compile() { TIMESCOPE(CppJitModule, compile); + ModuleHash = std::make_unique(computeCppJitModuleHash( + TargetModel, CompilerBackend, Code, ExtraArgs)); + if ((Library = Dispatch.lookupCompiledLibrary(*ModuleHash))) { IsCompiled = true; return; diff --git a/src/runtime/Frontend/LLVMIRJitModule.cpp b/src/runtime/Frontend/LLVMIRJitModule.cpp index ded4f97f2..016994a36 100644 --- a/src/runtime/Frontend/LLVMIRJitModule.cpp +++ b/src/runtime/Frontend/LLVMIRJitModule.cpp @@ -84,6 +84,9 @@ void LLVMIRJitModule::compile(bool Verify) { if (IsCompiled) return; + ModuleHash = std::make_unique( + hash(static_cast(TargetModel), Code, Config::get().getCGConfig())); + if ((Library = Dispatch.lookupCompiledLibrary(*ModuleHash))) { IsCompiled = true; return; diff --git a/src/runtime/Frontend/MLIRJitModule.cpp b/src/runtime/Frontend/MLIRJitModule.cpp index 6a46707c3..bd191962d 100644 --- a/src/runtime/Frontend/MLIRJitModule.cpp +++ b/src/runtime/Frontend/MLIRJitModule.cpp @@ -37,6 +37,9 @@ void MLIRJitModule::compile(bool Verify) { if (IsCompiled) return; + ModuleHash = std::make_unique( + hash(static_cast(TargetModel), Code, Config::get().getCGConfig())); + if ((Library = Dispatch.lookupCompiledLibrary(*ModuleHash))) { IsCompiled = true; return; diff --git a/src/runtime/Init.cpp b/src/runtime/Init.cpp index 3b1934057..6d9587634 100644 --- a/src/runtime/Init.cpp +++ b/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 diff --git a/src/runtime/JITPassPluginRegistry.cpp b/src/runtime/JITPassPluginRegistry.cpp new file mode 100644 index 000000000..0b4430687 --- /dev/null +++ b/src/runtime/JITPassPluginRegistry.cpp @@ -0,0 +1,95 @@ +#include "proteus/impl/JITPassPluginRegistry.h" + +#include "proteus/Error.h" + +#include +#include +#include + +#include +#include + +namespace proteus { +namespace { + +class JITPassPluginRegistry { +public: + static JITPassPluginRegistry &instance() { + static JITPassPluginRegistry Registry; + return Registry; + } + + void registerPlugin(const std::string &PluginPath, + const std::string &PassPipeline) { + if (PluginPath.empty()) + reportFatalError("JIT pass plugin path must be non-empty"); + if (PassPipeline.empty()) + reportFatalError("JIT pass plugin pipeline must be non-empty"); + + JITPassPluginConfig Config{normalizePath(PluginPath), PassPipeline, {}}; + 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; + }); + if (It != Plugins.end()) { + It->Fingerprint = std::move(Config.Fingerprint); + return; + } + + Plugins.push_back(std::move(Config)); + } + + void clear() { + std::lock_guard Lock(Mutex); + Plugins.clear(); + } + + std::vector snapshot() { + std::lock_guard Lock(Mutex); + return Plugins; + } + +private: + static std::string normalizePath(const std::string &PluginPath) { + llvm::SmallString<256> RealPath; + if (!llvm::sys::fs::real_path(PluginPath, RealPath)) + return std::string(RealPath.str()); + + return PluginPath; + } + + static std::string computeFingerprint(const JITPassPluginConfig &Config) { + llvm::sys::fs::file_status Status; + if (llvm::sys::fs::status(Config.Path, Status)) + return Config.Path + "|" + Config.Pipeline; + + llvm::SmallString<128> Storage; + llvm::raw_svector_ostream OS(Storage); + const llvm::sys::fs::UniqueID ID = Status.getUniqueID(); + OS << Config.Path << "|" << Config.Pipeline << "|" << ID.getDevice() << ":" + << ID.getFile() << "|" << Status.getSize() << "|" + << Status.getLastModificationTime().time_since_epoch().count(); + return std::string(OS.str()); + } + + std::mutex Mutex; + std::vector Plugins; +}; + +} // namespace + +void registerJITPassPluginImpl(const std::string &PluginPath, + const std::string &PassPipeline) { + JITPassPluginRegistry::instance().registerPlugin(PluginPath, PassPipeline); +} + +void clearJITPassPluginsImpl() { JITPassPluginRegistry::instance().clear(); } + +std::vector getJITPassPluginConfigs() { + return JITPassPluginRegistry::instance().snapshot(); +} + +} // namespace proteus diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 866c6e6dc..aaa625d69 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -28,6 +28,25 @@ else() message(STATUS "Found Filecheck at ${FILECHECK}") endif() +add_library(JITTestPass SHARED JITTestPass.cpp) +target_compile_definitions(JITTestPass PRIVATE ${LLVM_DEFINITIONS}) +target_include_directories(JITTestPass SYSTEM PRIVATE "${LLVM_INSTALL_DIR}/include") + +if(LLVM_LINK_LLVM_DYLIB) + set(jit_test_pass_llvm_libs LLVM) +else() + llvm_map_components_to_libnames(jit_test_pass_llvm_libs + Core + Passes + Support + ) +endif() +target_link_libraries(JITTestPass PRIVATE ${jit_test_pass_llvm_libs}) + +if(NOT LLVM_ENABLE_RTTI) + target_compile_options(JITTestPass PRIVATE -fno-rtti) +endif() + function(proteus_attach_pass_plugin_rebuild_dep target) set(proteus_pass_stamp "${CMAKE_CURRENT_BINARY_DIR}/${target}.proteuspass.stamp") diff --git a/tests/JITTestPass.cpp b/tests/JITTestPass.cpp new file mode 100644 index 000000000..a740cbe7d --- /dev/null +++ b/tests/JITTestPass.cpp @@ -0,0 +1,39 @@ +#include +#include +#include +#if __has_include() +#include +#elif __has_include() +#include +#else +#error "Cannot find LLVM PassPlugin.h" +#endif +#include + +namespace { + +class JITTestPass : public llvm::PassInfoMixin { +public: + llvm::PreservedAnalyses run(llvm::Module &M, + llvm::ModuleAnalysisManager &) { + llvm::outs() << "[JITTestPass] " << M.getName() << "\n"; + return llvm::PreservedAnalyses::all(); + } +}; + +} // namespace + +extern "C" LLVM_ATTRIBUTE_WEAK ::llvm::PassPluginLibraryInfo +llvmGetPassPluginInfo() { + 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") + return false; + MPM.addPass(JITTestPass()); + return true; + }); + }}; +} diff --git a/tests/cpu/CMakeLists.txt b/tests/cpu/CMakeLists.txt index 2fe7efe58..0bffd4cb6 100644 --- a/tests/cpu/CMakeLists.txt +++ b/tests/cpu/CMakeLists.txt @@ -65,5 +65,14 @@ CREATE_CPU_TEST(types_jit_array types_jit_array.cpp) CREATE_CPU_TEST(dynamic_jit_array dynamic_jit_array.cpp) CREATE_CPU_TEST(jit_struct jit_struct.cpp) CREATE_CPU_TEST(custom_pipeline custom_pipeline.cpp) +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_cmake jit_pass_plugin_cmake.cpp) +proteus_register_jit_pass_plugin( + jit_pass_plugin_cmake + PLUGIN_TARGET JITTestPass + PIPELINE jit-test-pass) 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 new file mode 100644 index 000000000..f99ef8335 --- /dev/null +++ b/tests/cpu/jit_pass_plugin.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" %build/jit_pass_plugin | %FILECHECK %s +// RUN: rm -rf "%t.$$.proteus" +// clang-format on + +#include + +#include +#include + +__attribute__((annotate("jit"))) int add_one(int x) { + proteus::jit_arg(x); + return x + 1; +} + +int main() { + proteus::registerJITPassPlugin(PROTEUS_TEST_JIT_PASS_PLUGIN_PATH, + "jit-test-pass"); + std::cout << add_one(4) << "\n"; + return 0; +} + +// CHECK: [JITTestPass] +// 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 new file mode 100644 index 000000000..e1f9710af --- /dev/null +++ b/tests/cpu/jit_pass_plugin_cmake.cpp @@ -0,0 +1,24 @@ +// 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_cmake | %FILECHECK %s +// RUN: rm -rf "%t.$$.proteus" +// clang-format on + +#include + +#include + +__attribute__((annotate("jit"))) int add_two(int x) { + proteus::jit_arg(x); + return x + 2; +} + +int main() { + std::cout << add_two(5) << "\n"; + return 0; +} + +// CHECK: [JITTestPass] +// CHECK: [CustomPipeline] default,jit-test-pass +// CHECK: 7 +// CHECK: [proteus][JitEngineHost] MemoryCache rank 0 hits 0 accesses 1 diff --git a/tests/gpu/CMakeLists.txt b/tests/gpu/CMakeLists.txt index ea01374e3..24158a915 100644 --- a/tests/gpu/CMakeLists.txt +++ b/tests/gpu/CMakeLists.txt @@ -193,6 +193,10 @@ CREATE_GPU_TEST(dim_spec_cache_test dim_spec_cache_test.cpp) CREATE_GPU_TEST(kernel kernel.cpp) CREATE_GPU_TEST(kernel_metadata kernel_metadata.cpp) CREATE_GPU_TEST(kernel_pass_pipeline kernel_pass_pipeline.cpp) +CREATE_GPU_TEST(kernel_pass_plugin kernel_pass_plugin.cpp) +target_compile_definitions(kernel_pass_plugin.${lang} PRIVATE + PROTEUS_TEST_JIT_PASS_PLUGIN_PATH="$") +add_dependencies(kernel_pass_plugin.${lang} JITTestPass) CREATE_GPU_TEST(kernel_cache kernel_cache.cpp) CREATE_GPU_TEST(kernel_args kernel_args.cpp) CREATE_GPU_TEST(kernel_args_api kernel_args_api.cpp) diff --git a/tests/gpu/kernel_pass_plugin.cpp b/tests/gpu/kernel_pass_plugin.cpp new file mode 100644 index 000000000..254dafeee --- /dev/null +++ b/tests/gpu/kernel_pass_plugin.cpp @@ -0,0 +1,28 @@ +// 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: rm -rf "%t.$$.proteus" +// clang-format on + +#include + +#include "gpu_common.h" +#include +#include + +__global__ __attribute__((annotate("jit"))) void kernel_pass_plugin() { + printf("KernelPassPlugin\n"); +} + +int main() { + proteus::registerJITPassPlugin(PROTEUS_TEST_JIT_PASS_PLUGIN_PATH, + "jit-test-pass"); + kernel_pass_plugin<<<1, 1>>>(); + gpuErrCheck(gpuDeviceSynchronize()); + return 0; +} + +// CHECK: [JITTestPass] +// CHECK: [CustomPipeline] default,jit-test-pass +// CHECK: KernelPassPlugin +// CHECK: [proteus][JitEngineDevice] MemoryCache rank 0 hits 0 accesses 1 From 0d3b2773b253b84b0b2d7c3da67d0b529dfefb1b Mon Sep 17 00:00:00 2001 From: john bowen Date: Wed, 27 May 2026 10:21:46 -0700 Subject: [PATCH 02/12] Add CMake method for connecting plugin pass --- cmake/ProteusFunctions.cmake | 57 ++++++++++++++++++++++++++++++++++++ 1 file changed, 57 insertions(+) diff --git a/cmake/ProteusFunctions.cmake b/cmake/ProteusFunctions.cmake index 42c8597ad..64fbbd55c 100644 --- a/cmake/ProteusFunctions.cmake +++ b/cmake/ProteusFunctions.cmake @@ -28,3 +28,60 @@ 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 "$") + 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 + +namespace { +struct AutoRegisterProteusJITPassPlugin { + AutoRegisterProteusJITPassPlugin() { + proteus::registerJITPassPlugin( + R\"(${_proteus_jit_pass_plugin_path})\", + R\"(${PROTEUS_JIT_PASS_PIPELINE})\"); + } +}; + +AutoRegisterProteusJITPassPlugin AutoRegisterProteusJITPassPluginInstance; +} // namespace +") + + target_sources(${target} PRIVATE "${_proteus_jit_pass_source}") +endfunction() From f168c6c4c17033a74c490e07fe295b38680f6f75 Mon Sep 17 00:00:00 2001 From: john bowen Date: Thu, 28 May 2026 10:53:26 -0700 Subject: [PATCH 03/12] fix linkage errors in the plugin tests --- tests/CMakeLists.txt | 17 ++++++----------- tests/gpu/CMakeLists.txt | 6 ++++++ 2 files changed, 12 insertions(+), 11 deletions(-) diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index aaa625d69..4b87c0f55 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -28,20 +28,15 @@ else() message(STATUS "Found Filecheck at ${FILECHECK}") endif() -add_library(JITTestPass SHARED JITTestPass.cpp) +add_library(JITTestPass MODULE JITTestPass.cpp) target_compile_definitions(JITTestPass PRIVATE ${LLVM_DEFINITIONS}) target_include_directories(JITTestPass SYSTEM PRIVATE "${LLVM_INSTALL_DIR}/include") -if(LLVM_LINK_LLVM_DYLIB) - set(jit_test_pass_llvm_libs LLVM) -else() - llvm_map_components_to_libnames(jit_test_pass_llvm_libs - Core - Passes - Support - ) -endif() -target_link_libraries(JITTestPass PRIVATE ${jit_test_pass_llvm_libs}) +# Build the test pass like a normal LLVM pass plugin: it should resolve LLVM +# symbols from the hosting process rather than embedding a second copy of LLVM +# into the plugin itself. +target_link_options(JITTestPass + PRIVATE "$<$:LINKER:SHELL:-undefined dynamic_lookup>") if(NOT LLVM_ENABLE_RTTI) target_compile_options(JITTestPass PRIVATE -fno-rtti) diff --git a/tests/gpu/CMakeLists.txt b/tests/gpu/CMakeLists.txt index 24158a915..1339c23ec 100644 --- a/tests/gpu/CMakeLists.txt +++ b/tests/gpu/CMakeLists.txt @@ -14,6 +14,9 @@ function(CREATE_GPU_TEST exe check_source) set_source_files_properties(${check_source} ${ARGN} PROPERTIES LANGUAGE ${lang}) target_link_libraries(${exe}.${lang} PUBLIC proteus) proteus_attach_pass_plugin_rebuild_dep(${exe}.${lang} ${check_source} ${ARGN}) + target_link_options(${exe}.${lang} PRIVATE + $ + ) add_test(NAME ${exe}.${lang} COMMAND ${LIT} -vv -D EXT=${lang} -DFILECHECK=${FILECHECK} ${check_source}) set_tests_properties(${exe}.${lang} PROPERTIES LABELS "gpu;gpu-basic") @@ -24,6 +27,9 @@ function(CREATE_GPU_TEST_RDC exe check_source) set_source_files_properties(${check_source} ${ARGN} PROPERTIES LANGUAGE ${lang}) target_link_libraries(${exe}.${lang}.rdc PUBLIC proteus) proteus_attach_pass_plugin_rebuild_dep(${exe}.${lang}.rdc ${check_source} ${ARGN}) + # target_link_options(${exe}.${lang}.rdc PRIVATE + # $ + # ) if(PROTEUS_ENABLE_HIP) # This is unsupported see: https://gitlab.kitware.com/cmake/cmake/-/issues/23210 From f1b9a55d1b58fdaf97bea2a4433cf38eed2d0c75 Mon Sep 17 00:00:00 2001 From: john bowen Date: Thu, 28 May 2026 15:34:59 -0700 Subject: [PATCH 04/12] attempt to fix linkage errors --- tests/CMakeLists.txt | 19 ++++++++++++++----- tests/JITTestPass.exports.map | 5 +++++ 2 files changed, 19 insertions(+), 5 deletions(-) create mode 100644 tests/JITTestPass.exports.map diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 4b87c0f55..967c7adc2 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -32,11 +32,20 @@ add_library(JITTestPass MODULE JITTestPass.cpp) target_compile_definitions(JITTestPass PRIVATE ${LLVM_DEFINITIONS}) target_include_directories(JITTestPass SYSTEM PRIVATE "${LLVM_INSTALL_DIR}/include") -# Build the test pass like a normal LLVM pass plugin: it should resolve LLVM -# symbols from the hosting process rather than embedding a second copy of LLVM -# into the plugin itself. -target_link_options(JITTestPass - PRIVATE "$<$:LINKER:SHELL:-undefined dynamic_lookup>") +if(LLVM_LINK_LLVM_DYLIB) + set(jit_test_pass_llvm_libs LLVM) +else() + llvm_map_components_to_libnames(jit_test_pass_llvm_libs + Core + Passes + Support + ) +endif() +target_link_libraries(JITTestPass PRIVATE ${jit_test_pass_llvm_libs}) +target_link_options(JITTestPass PRIVATE + "$<$:LINKER:SHELL:-undefined dynamic_lookup>" + "$<$>,$>:-Wl,--version-script=${CMAKE_CURRENT_SOURCE_DIR}/JITTestPass.exports.map>" +) if(NOT LLVM_ENABLE_RTTI) target_compile_options(JITTestPass PRIVATE -fno-rtti) diff --git a/tests/JITTestPass.exports.map b/tests/JITTestPass.exports.map new file mode 100644 index 000000000..d352de5c2 --- /dev/null +++ b/tests/JITTestPass.exports.map @@ -0,0 +1,5 @@ +{ + global: + llvmGetPassPluginInfo; + local: *; +}; From 42f491b1a85fda7f9dce0242b8c5564201088695 Mon Sep 17 00:00:00 2001 From: john bowen Date: Thu, 28 May 2026 15:41:28 -0700 Subject: [PATCH 05/12] clang format --- src/include/proteus/impl/CoreLLVM.h | 22 ++++++++++------------ src/runtime/Frontend/CppJitModule.cpp | 4 ++-- 2 files changed, 12 insertions(+), 14 deletions(-) diff --git a/src/include/proteus/impl/CoreLLVM.h b/src/include/proteus/impl/CoreLLVM.h index 643d59c17..aea5b4c42 100644 --- a/src/include/proteus/impl/CoreLLVM.h +++ b/src/include/proteus/impl/CoreLLVM.h @@ -117,13 +117,12 @@ inline std::string getDefaultOptimizationPipeline(char OptLevel) { return ""; } -inline std::string -composeOptimizationPassPipeline(std::optional PassPipeline, - char OptLevel, - const std::vector &Plugins) { - std::string Pipeline = - PassPipeline ? std::move(PassPipeline.value()) - : getDefaultOptimizationPipeline(OptLevel); +inline std::string composeOptimizationPassPipeline( + std::optional PassPipeline, char OptLevel, + const std::vector &Plugins) { + std::string Pipeline = PassPipeline + ? std::move(PassPipeline.value()) + : getDefaultOptimizationPipeline(OptLevel); for (const auto &Plugin : Plugins) { Pipeline += ","; Pipeline += Plugin.Pipeline; @@ -145,11 +144,10 @@ loadJITPassPlugins(const std::vector &Plugins) { return LoadedPlugins; } -inline void runOptimizationPassPipeline(Module &M, StringRef Arch, - const std::string &PassPipeline, - unsigned CodegenOptLevel, - const std::vector - &Plugins = {}) { +inline void runOptimizationPassPipeline( + Module &M, StringRef Arch, const std::string &PassPipeline, + unsigned CodegenOptLevel, + const std::vector &Plugins = {}) { PipelineTuningOptions PTO; std::optional PGOOpt; diff --git a/src/runtime/Frontend/CppJitModule.cpp b/src/runtime/Frontend/CppJitModule.cpp index 3bb76c46e..d97818216 100644 --- a/src/runtime/Frontend/CppJitModule.cpp +++ b/src/runtime/Frontend/CppJitModule.cpp @@ -36,8 +36,8 @@ CppJitModule::~CppJitModule() = default; void CppJitModule::compile() { TIMESCOPE(CppJitModule, compile); - ModuleHash = std::make_unique(computeCppJitModuleHash( - TargetModel, CompilerBackend, Code, ExtraArgs)); + ModuleHash = std::make_unique( + computeCppJitModuleHash(TargetModel, CompilerBackend, Code, ExtraArgs)); if ((Library = Dispatch.lookupCompiledLibrary(*ModuleHash))) { IsCompiled = true; From 3b0c3dddd597bdf0923f5d774377ab111dbeb13b Mon Sep 17 00:00:00 2001 From: john bowen Date: Tue, 2 Jun 2026 13:38:16 -0700 Subject: [PATCH 06/12] Remove rdynamic flag --- tests/gpu/CMakeLists.txt | 6 ------ 1 file changed, 6 deletions(-) diff --git a/tests/gpu/CMakeLists.txt b/tests/gpu/CMakeLists.txt index 1339c23ec..24158a915 100644 --- a/tests/gpu/CMakeLists.txt +++ b/tests/gpu/CMakeLists.txt @@ -14,9 +14,6 @@ function(CREATE_GPU_TEST exe check_source) set_source_files_properties(${check_source} ${ARGN} PROPERTIES LANGUAGE ${lang}) target_link_libraries(${exe}.${lang} PUBLIC proteus) proteus_attach_pass_plugin_rebuild_dep(${exe}.${lang} ${check_source} ${ARGN}) - target_link_options(${exe}.${lang} PRIVATE - $ - ) add_test(NAME ${exe}.${lang} COMMAND ${LIT} -vv -D EXT=${lang} -DFILECHECK=${FILECHECK} ${check_source}) set_tests_properties(${exe}.${lang} PROPERTIES LABELS "gpu;gpu-basic") @@ -27,9 +24,6 @@ function(CREATE_GPU_TEST_RDC exe check_source) set_source_files_properties(${check_source} ${ARGN} PROPERTIES LANGUAGE ${lang}) target_link_libraries(${exe}.${lang}.rdc PUBLIC proteus) proteus_attach_pass_plugin_rebuild_dep(${exe}.${lang}.rdc ${check_source} ${ARGN}) - # target_link_options(${exe}.${lang}.rdc PRIVATE - # $ - # ) if(PROTEUS_ENABLE_HIP) # This is unsupported see: https://gitlab.kitware.com/cmake/cmake/-/issues/23210 From 88e5edb0cf56ed77af5b8515a08caf1e9fcfc34f Mon Sep 17 00:00:00 2001 From: john bowen Date: Tue, 2 Jun 2026 17:02:36 -0700 Subject: [PATCH 07/12] Loosen ordering restrictions on plugin unit test --- tests/gpu/kernel_pass_plugin.cpp | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/gpu/kernel_pass_plugin.cpp b/tests/gpu/kernel_pass_plugin.cpp index 254dafeee..024e5c000 100644 --- a/tests/gpu/kernel_pass_plugin.cpp +++ b/tests/gpu/kernel_pass_plugin.cpp @@ -22,7 +22,7 @@ int main() { return 0; } -// CHECK: [JITTestPass] -// CHECK: [CustomPipeline] default,jit-test-pass -// CHECK: KernelPassPlugin +// CHECK-DAG: [JITTestPass] +// CHECK-DAG: [CustomPipeline] default,jit-test-pass +// CHECK-DAG: KernelPassPlugin // CHECK: [proteus][JitEngineDevice] MemoryCache rank 0 hits 0 accesses 1 From b0152430989ec33507b968fb1c6ae528ce86a45b Mon Sep 17 00:00:00 2001 From: john bowen Date: Wed, 3 Jun 2026 13:41:24 -0700 Subject: [PATCH 08/12] remove proteus plugin registration file from gcov --- cmake/ProteusFunctions.cmake | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/cmake/ProteusFunctions.cmake b/cmake/ProteusFunctions.cmake index 64fbbd55c..f71a963fa 100644 --- a/cmake/ProteusFunctions.cmake +++ b/cmake/ProteusFunctions.cmake @@ -83,5 +83,13 @@ 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. + 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() From cb090c434ea27523bf17991401af8ebb16c58eb6 Mon Sep 17 00:00:00 2001 From: Zane Fink Date: Wed, 1 Jul 2026 14:56:01 -0700 Subject: [PATCH 09/12] Fingerprint use file contents instead of node/rank-specific info --- src/runtime/JITPassPluginRegistry.cpp | 15 +++++---------- 1 file changed, 5 insertions(+), 10 deletions(-) diff --git a/src/runtime/JITPassPluginRegistry.cpp b/src/runtime/JITPassPluginRegistry.cpp index 0b4430687..80a17e835 100644 --- a/src/runtime/JITPassPluginRegistry.cpp +++ b/src/runtime/JITPassPluginRegistry.cpp @@ -1,10 +1,11 @@ #include "proteus/impl/JITPassPluginRegistry.h" #include "proteus/Error.h" +#include "proteus/impl/Hashing.h" #include #include -#include +#include #include #include @@ -62,17 +63,11 @@ class JITPassPluginRegistry { } static std::string computeFingerprint(const JITPassPluginConfig &Config) { - llvm::sys::fs::file_status Status; - if (llvm::sys::fs::status(Config.Path, Status)) + auto BufOrErr = llvm::MemoryBuffer::getFile(Config.Path); + if (!BufOrErr) return Config.Path + "|" + Config.Pipeline; - llvm::SmallString<128> Storage; - llvm::raw_svector_ostream OS(Storage); - const llvm::sys::fs::UniqueID ID = Status.getUniqueID(); - OS << Config.Path << "|" << Config.Pipeline << "|" << ID.getDevice() << ":" - << ID.getFile() << "|" << Status.getSize() << "|" - << Status.getLastModificationTime().time_since_epoch().count(); - return std::string(OS.str()); + return hashValue(BufOrErr.get()->getBuffer()).toString(); } std::mutex Mutex; From e454fe24b91738b6a7c088dcadd803b2a8407684 Mon Sep 17 00:00:00 2001 From: Zane Fink Date: Wed, 1 Jul 2026 14:58:14 -0700 Subject: [PATCH 10/12] Atomic hasplugins --- src/runtime/JITPassPluginRegistry.cpp | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/src/runtime/JITPassPluginRegistry.cpp b/src/runtime/JITPassPluginRegistry.cpp index 80a17e835..c8288c875 100644 --- a/src/runtime/JITPassPluginRegistry.cpp +++ b/src/runtime/JITPassPluginRegistry.cpp @@ -8,6 +8,7 @@ #include #include +#include #include namespace proteus { @@ -41,14 +42,21 @@ class JITPassPluginRegistry { } Plugins.push_back(std::move(Config)); + HasPlugins.store(true, std::memory_order_release); } void clear() { std::lock_guard Lock(Mutex); Plugins.clear(); + HasPlugins.store(false, std::memory_order_release); } std::vector snapshot() { + // Fast path: avoid taking the lock on the common case where no plugins are + // registered, which happens on every JIT kernel launch. + if (!HasPlugins.load(std::memory_order_acquire)) + return {}; + std::lock_guard Lock(Mutex); return Plugins; } @@ -71,6 +79,7 @@ class JITPassPluginRegistry { } std::mutex Mutex; + std::atomic HasPlugins{false}; std::vector Plugins; }; From 5e56e3cc93db292e6701ed1dd2eab721ae4ba754 Mon Sep 17 00:00:00 2001 From: Zane Fink Date: Wed, 1 Jul 2026 15:05:39 -0700 Subject: [PATCH 11/12] fix formatting --- tests/JITTestPass.cpp | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/JITTestPass.cpp b/tests/JITTestPass.cpp index a740cbe7d..d0a9ff71c 100644 --- a/tests/JITTestPass.cpp +++ b/tests/JITTestPass.cpp @@ -14,8 +14,7 @@ namespace { class JITTestPass : public llvm::PassInfoMixin { public: - llvm::PreservedAnalyses run(llvm::Module &M, - llvm::ModuleAnalysisManager &) { + llvm::PreservedAnalyses run(llvm::Module &M, llvm::ModuleAnalysisManager &) { llvm::outs() << "[JITTestPass] " << M.getName() << "\n"; return llvm::PreservedAnalyses::all(); } From 0c893a096eeb0513fa0a0c7e97713afafab24a72 Mon Sep 17 00:00:00 2001 From: Zane Fink Date: Thu, 2 Jul 2026 11:32:47 -0700 Subject: [PATCH 12/12] fix unused flag --- CMakeLists.txt | 7 +++++-- cmake/ProteusFunctions.cmake | 3 +++ 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 162122c9c..807e0c90e 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -67,8 +67,11 @@ if(ENABLE_COVERAGE) add_compile_options( "$<$:-g>" "$<$:-g>" - "$<$:--coverage>" - "$<$:--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). + "$<$:SHELL:-fprofile-arcs -ftest-coverage>" + "$<$:SHELL:-fprofile-arcs -ftest-coverage>" "$<$:SHELL:-Xarch_host -g>" "$<$:SHELL:-Xarch_host --coverage>" "$<$:SHELL:-Xarch_host -g>" diff --git a/cmake/ProteusFunctions.cmake b/cmake/ProteusFunctions.cmake index f71a963fa..9e8f80129 100644 --- a/cmake/ProteusFunctions.cmake +++ b/cmake/ProteusFunctions.cmake @@ -87,6 +87,9 @@ AutoRegisterProteusJITPassPlugin AutoRegisterProteusJITPassPluginInstance; # 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()