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 42c8597ad..9e8f80129 100644 --- a/cmake/ProteusFunctions.cmake +++ b/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 "$") + 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 +") + + 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() 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..aea5b4c42 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,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"; + 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, + const std::vector &Plugins = {}) { PipelineTuningOptions PTO; std::optional 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,10 +284,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 +295,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..d97818216 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..c8288c875 --- /dev/null +++ b/src/runtime/JITPassPluginRegistry.cpp @@ -0,0 +1,99 @@ +#include "proteus/impl/JITPassPluginRegistry.h" + +#include "proteus/Error.h" +#include "proteus/impl/Hashing.h" + +#include +#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)); + 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; + } + +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) { + auto BufOrErr = llvm::MemoryBuffer::getFile(Config.Path); + if (!BufOrErr) + return Config.Path + "|" + Config.Pipeline; + + return hashValue(BufOrErr.get()->getBuffer()).toString(); + } + + std::mutex Mutex; + std::atomic HasPlugins{false}; + 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..967c7adc2 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -28,6 +28,29 @@ else() message(STATUS "Found Filecheck at ${FILECHECK}") endif() +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}) +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) +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..d0a9ff71c --- /dev/null +++ b/tests/JITTestPass.cpp @@ -0,0 +1,38 @@ +#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/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: *; +}; 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..024e5c000 --- /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-DAG: [JITTestPass] +// CHECK-DAG: [CustomPipeline] default,jit-test-pass +// CHECK-DAG: KernelPassPlugin +// CHECK: [proteus][JitEngineDevice] MemoryCache rank 0 hits 0 accesses 1