diff --git a/AlgorithmLauncher.cpp b/AlgorithmLauncher.cpp new file mode 100644 index 0000000..981575f --- /dev/null +++ b/AlgorithmLauncher.cpp @@ -0,0 +1,65 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include + +AlgorithmLauncher::AlgorithmLauncher(cudaKernel_t k, cudaLibrary_t lib) : kernel{k}, library{lib} {} + +AlgorithmLauncher::~AlgorithmLauncher() +{ + if (library != nullptr) { (void)cudaLibraryUnload(library); } +} + +AlgorithmLauncher::AlgorithmLauncher(AlgorithmLauncher&& other) noexcept + : kernel{other.kernel}, library{other.library} +{ + other.kernel = nullptr; + other.library = nullptr; +} + +AlgorithmLauncher& AlgorithmLauncher::operator=(AlgorithmLauncher&& other) noexcept +{ + if (this != &other) { + if (library != nullptr) { cudaLibraryUnload(library); } + kernel = other.kernel; + library = other.library; + other.kernel = nullptr; + other.library = nullptr; + } + return *this; +} + +void AlgorithmLauncher::call( + cudaStream_t stream, dim3 grid, dim3 block, std::size_t shared_mem, void** kernel_args) +{ + cudaLaunchConfig_t config{}; + config.gridDim = grid; + config.blockDim = block; + config.stream = stream; + config.dynamicSmemBytes = shared_mem; + config.numAttrs = 0; + config.attrs = NULL; + + RAFT_CUDA_TRY(cudaLaunchKernelExC(&config, kernel, kernel_args)); +} + +void AlgorithmLauncher::call_cooperative( + cudaStream_t stream, dim3 grid, dim3 block, std::size_t shared_mem, void** kernel_args) +{ + cudaLaunchAttribute attribute[1]; + attribute[0].id = cudaLaunchAttributeCooperative; + attribute[0].val.cooperative = 1; + + cudaLaunchConfig_t config{}; + config.gridDim = grid; + config.blockDim = block; + config.stream = stream; + config.dynamicSmemBytes = shared_mem; + config.numAttrs = 1; + config.attrs = attribute; + + RAFT_CUDA_TRY(cudaLaunchKernelExC(&config, kernel, kernel_args)); +} diff --git a/AlgorithmLauncher.hpp b/AlgorithmLauncher.hpp new file mode 100644 index 0000000..821c3db --- /dev/null +++ b/AlgorithmLauncher.hpp @@ -0,0 +1,61 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include + +#include +#include + +#include +#include +#include +#include + +struct AlgorithmLauncher { + AlgorithmLauncher() : kernel{nullptr}, library{nullptr} {} + + AlgorithmLauncher(cudaKernel_t k, cudaLibrary_t lib); + + ~AlgorithmLauncher(); + + AlgorithmLauncher(AlgorithmLauncher const&) = delete; + AlgorithmLauncher& operator=(AlgorithmLauncher const&) = delete; + + AlgorithmLauncher(AlgorithmLauncher&& other) noexcept; + AlgorithmLauncher& operator=(AlgorithmLauncher&& other) noexcept; + + template + void dispatch(cudaStream_t stream, dim3 grid, dim3 block, std::size_t shared_mem, Args&&... args) + { + static_assert(std::is_same_v...)>, + "dispatch() argument types do not match the kernel function signature FuncT"); + + void* kernel_args[] = {const_cast(static_cast(&args))...}; + this->call(stream, grid, block, shared_mem, kernel_args); + } + + template + void dispatch_cooperative( + cudaStream_t stream, dim3 grid, dim3 block, std::size_t shared_mem, Args&&... args) + { + static_assert( + std::is_same_v...)>, + "dispatch_cooperative() argument types do not match the kernel function signature FuncT"); + + void* kernel_args[] = {const_cast(static_cast(&args))...}; + this->call_cooperative(stream, grid, block, shared_mem, kernel_args); + } + + cudaKernel_t get_kernel() { return this->kernel; } + + private: + void call(cudaStream_t stream, dim3 grid, dim3 block, std::size_t shared_mem, void** args); + void call_cooperative( + cudaStream_t stream, dim3 grid, dim3 block, std::size_t shared_mem, void** args); + cudaKernel_t kernel; + cudaLibrary_t library; +}; diff --git a/AlgorithmPlanner.cpp b/AlgorithmPlanner.cpp new file mode 100644 index 0000000..9d9da43 --- /dev/null +++ b/AlgorithmPlanner.cpp @@ -0,0 +1,115 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "cuda_runtime.h" +#include "nvJitLink.h" + +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +std::string AlgorithmPlanner::get_fragments_key() const +{ + std::string key = ""; + for (auto const& fragment : this->fragments) { + key += fragment->get_key(); + } + return key; +} + +std::shared_ptr AlgorithmPlanner::read_cache(std::string const& launch_key) const +{ + auto& launchers = jit_cache_.launchers; + std::shared_lock read_lock(jit_cache_.mutex); + if (auto it = launchers.find(launch_key); it != launchers.end()) { return it->second; } + return nullptr; +} + +std::shared_ptr AlgorithmPlanner::get_launcher() +{ + auto& launchers = jit_cache_.launchers; + auto launch_key = this->get_fragments_key(); + + if (auto hit = read_cache(launch_key)) { return hit; } + + std::unique_lock write_lock(jit_cache_.mutex); + if (auto it = launchers.find(launch_key); it != launchers.end()) { return it->second; } + + std::string log_message = + "JIT compiling launcher for kernel: " + this->entrypoint + " and device functions: "; + for (auto const& fragment : this->fragments) { + log_message += std::string{fragment->get_key()} + ","; + } + log_message.pop_back(); + RAFT_LOG_DEBUG("%s", log_message.c_str()); + auto launcher = this->build(); + launchers[launch_key] = launcher; + return launcher; +} + +std::shared_ptr AlgorithmPlanner::build() +{ + int device = 0; + int major = 0; + int minor = 0; + RAFT_CUDA_TRY(cudaGetDevice(&device)); + RAFT_CUDA_TRY(cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, device)); + RAFT_CUDA_TRY(cudaDeviceGetAttribute(&minor, cudaDevAttrComputeCapabilityMinor, device)); + + std::string archs = "-arch=sm_" + std::to_string((major * 10 + minor)); + + // Load the generated LTO IR and link them together + nvJitLinkHandle handle; + std::vector lopts; + lopts.reserve(2 + linktime_extra_options.size()); + lopts.push_back("-lto"); + lopts.push_back(archs.c_str()); + for (auto const& opt : linktime_extra_options) { + lopts.push_back(opt.c_str()); + } + auto result = nvJitLinkCreate(&handle, static_cast(lopts.size()), lopts.data()); + check_nvjitlink_result(handle, result); + + for (auto const& frag : this->fragments) { + frag->add_to(handle); + } + + // Call to nvJitLinkComplete causes linker to link together all the LTO-IR + // modules perform any optimizations and generate cubin from it. + result = nvJitLinkComplete(handle); + check_nvjitlink_result(handle, result); + + // get cubin from nvJitLink + size_t cubin_size; + result = nvJitLinkGetLinkedCubinSize(handle, &cubin_size); + check_nvjitlink_result(handle, result); + + std::unique_ptr cubin{new char[cubin_size]}; + result = nvJitLinkGetLinkedCubin(handle, cubin.get()); + check_nvjitlink_result(handle, result); + + result = nvJitLinkDestroy(&handle); + RAFT_EXPECTS(result == NVJITLINK_SUCCESS, "nvJitLinkDestroy failed"); + + // cubin is linked, so now load it + cudaLibrary_t library; + RAFT_CUDA_TRY( + cudaLibraryLoadData(&library, cubin.get(), nullptr, nullptr, 0, nullptr, nullptr, 0)); + + cudaKernel_t kernel; + RAFT_CUDA_TRY(cudaLibraryGetKernel(&kernel, library, this->entrypoint.c_str())); + + return std::make_shared(kernel, library); +} diff --git a/AlgorithmPlanner.hpp b/AlgorithmPlanner.hpp new file mode 100644 index 0000000..4fdb81e --- /dev/null +++ b/AlgorithmPlanner.hpp @@ -0,0 +1,60 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include "AlgorithmLauncher.hpp" +#include "FragmentEntry.hpp" + +#include +#include +#include +#include +#include +#include +#include + +struct LauncherJitCache { + std::shared_mutex mutex; + std::unordered_map> launchers; +}; + +struct AlgorithmPlanner { + AlgorithmPlanner(std::string entrypoint, LauncherJitCache& jit_cache) + : entrypoint(std::move(entrypoint)), jit_cache_(jit_cache) + { + } + + std::shared_ptr get_launcher(); + + std::string entrypoint; + std::vector> fragments; + + template >> + void add_fragment(std::unique_ptr fragment) + { + fragments.push_back(std::unique_ptr(std::move(fragment))); + } + + template + void add_static_fragment() + { + add_fragment(std::make_unique>()); + } + + protected: + /** Extra link-time option strings passed to nvJitLink. Base build() + * always passes "-lto" and "-arch=sm_XX" first; derived planners may append here in their + * constructor body. */ + std::vector linktime_extra_options; + + private: + std::string get_fragments_key() const; + std::shared_ptr build(); + + std::shared_ptr read_cache(std::string const& launch_key) const; + + LauncherJitCache& jit_cache_; +}; diff --git a/CMakeLists.txt b/CMakeLists.txt index db7fcd7..ccdfe93 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -21,19 +21,30 @@ project(rtcx VERSION 0.1.0 LANGUAGES CXX) option(RTCX_STATIC_LINK_NVRTC "Use static linking for NVRTC" OFF) option(RTCX_STATIC_LINK_NVJITLINK "Use static linking for nvJitLink" OFF) -rapids_find_package(CUDAToolkit REQUIRED BUILD_EXPORT_SET rtcx-exports - INSTALL_EXPORT_SET rtcx-exports) +rapids_find_package(CUDAToolkit REQUIRED BUILD_EXPORT_SET rtcx-exports INSTALL_EXPORT_SET + rtcx-exports) if(NOT TARGET zstd) set(CPM_DOWNLOAD_zstd ON) - rapids_cpm_find(zstd 1.5.7 - GLOBAL_TARGETS zstd + rapids_cpm_find(zstd + 1.5.7 + GLOBAL_TARGETS + zstd CPM_ARGS - GIT_REPOSITORY https://github.com/facebook/zstd.git - GIT_TAG v1.5.7 - GIT_SHALLOW FALSE SOURCE_SUBDIR build/cmake - OPTIONS "ZSTD_BUILD_STATIC ON" "ZSTD_BUILD_SHARED OFF" "ZSTD_BUILD_TESTS OFF" - "ZSTD_BUILD_PROGRAMS OFF" "BUILD_SHARED_LIBS OFF") + GIT_REPOSITORY + https://github.com/facebook/zstd.git + GIT_TAG + v1.5.7 + GIT_SHALLOW + FALSE + SOURCE_SUBDIR + build/cmake + OPTIONS + "ZSTD_BUILD_STATIC ON" + "ZSTD_BUILD_SHARED OFF" + "ZSTD_BUILD_TESTS OFF" + "ZSTD_BUILD_PROGRAMS OFF" + "BUILD_SHARED_LIBS OFF") if(zstd_ADDED) # disable weak symbols support to hide tracing APIs as well @@ -51,7 +62,11 @@ if(NOT TARGET nvtx3::nvtx3-cpp) rapids_cpm_nvtx3() endif() -add_library(rtcx STATIC hash.cpp rtcx.cpp) +rapids_cpm_find(raft "${RAPIDS_VERSION_MAJOR_MINOR}" GLOBAL_TARGETS raft::raft GIT_REPOSITORY + https://github.com/rapidsai/raft.git) + +add_library(rtcx STATIC hash.cpp rtcx.cpp AlgorithmLauncher.cpp AlgorithmPlanner.cpp + FragmentEntry.cpp nvjitlink_checker.cpp NVRTCLTOFragmentCompiler.cpp) add_library(rtcx::rtcx ALIAS rtcx) set_target_properties(rtcx @@ -64,7 +79,7 @@ set_target_properties(rtcx target_include_directories(rtcx PUBLIC $ PRIVATE ${CUDAToolkit_INCLUDE_DIRS}) -target_link_libraries(rtcx PRIVATE zstd ${CMAKE_DL_LIBS} nvtx3::nvtx3-cpp) +target_link_libraries(rtcx PRIVATE zstd ${CMAKE_DL_LIBS} nvtx3::nvtx3-cpp raft::raft) if(RTCX_STATIC_LINK_NVRTC) target_link_libraries(rtcx PRIVATE CUDA::nvrtc_static CUDA::nvrtc_builtins_static) @@ -95,6 +110,7 @@ set(rtcx_install_code_string # Embed functions (add_embed, embed_includes, embed_blob, embed) # are included automatically so consumers don't need explicit include(). include("${CMAKE_CURRENT_LIST_DIR}/embed.cmake") +include("${CMAKE_CURRENT_LIST_DIR}/generate_jit_lto_kernels.cmake") # Set rtcx_LIBCXX_DIR for consumers using embed_includes with libcxx headers. set(rtcx_LIBCXX_DIR "${PACKAGE_PREFIX_DIR}/share/rtcx/libcxx") @@ -104,6 +120,7 @@ string(CONFIGURE [=[ # Embed functions (add_embed, embed_includes, embed_blob, embed) # are included automatically so consumers don't need explicit include(). include("@CMAKE_CURRENT_SOURCE_DIR@/embed.cmake") +include("@CMAKE_CURRENT_SOURCE_DIR@/generate_jit_lto_kernels.cmake") # Set rtcx_LIBCXX_DIR for consumers using embed_includes with libcxx headers. set(rtcx_LIBCXX_DIR "@CMAKE_CURRENT_SOURCE_DIR@/libcxx") @@ -126,18 +143,16 @@ if(RTCX_INSTALL) install(FILES rtcx.hpp hash.hpp embed.hpp DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/rtcx) install(FILES embed.cmake DESTINATION ${lib_dir}/cmake/rtcx) install(DIRECTORY embed/ DESTINATION ${lib_dir}/cmake/rtcx/embed) + install(FILES AlgorithmLauncher.hpp AlgorithmPlanner.hpp FragmentEntry.hpp nvjitlink_checker.hpp + NVRTCLTOFRagmentCompiler.hpp DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/rtcx) + install(FILES generate_jit_lto_kernels.cmake register_fatbin.cpp.in + DESTINATION ${lib_dir}/cmake/rtcx) install(DIRECTORY libcxx/ DESTINATION ${CMAKE_INSTALL_DATADIR}/rtcx/libcxx) - rapids_export(INSTALL rtcx - EXPORT_SET rtcx-exports - GLOBAL_TARGETS rtcx - NAMESPACE rtcx:: + rapids_export(INSTALL rtcx EXPORT_SET rtcx-exports GLOBAL_TARGETS rtcx NAMESPACE rtcx:: FINAL_CODE_BLOCK rtcx_install_code_string) endif() # Build-tree export (always, so CPM consumers can find_package from the build tree) -rapids_export(BUILD rtcx - EXPORT_SET rtcx-exports - GLOBAL_TARGETS rtcx - NAMESPACE rtcx:: +rapids_export(BUILD rtcx EXPORT_SET rtcx-exports GLOBAL_TARGETS rtcx NAMESPACE rtcx:: FINAL_CODE_BLOCK rtcx_build_code_string) diff --git a/FragmentEntry.cpp b/FragmentEntry.cpp new file mode 100644 index 0000000..48a5eb7 --- /dev/null +++ b/FragmentEntry.cpp @@ -0,0 +1,14 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include + +bool FatbinFragmentEntry::add_to(nvJitLinkHandle& handle) const +{ + auto result = nvJitLinkAddData(handle, NVJITLINK_INPUT_ANY, get_data(), get_length(), get_key()); + + check_nvjitlink_result(handle, result); + return true; +} diff --git a/FragmentEntry.hpp b/FragmentEntry.hpp new file mode 100644 index 0000000..c2f7b0a --- /dev/null +++ b/FragmentEntry.hpp @@ -0,0 +1,64 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include "nvjitlink_checker.hpp" + +#include + +#include +#include +#include +#include +#include + +struct FragmentEntry { + virtual ~FragmentEntry() = default; + + virtual bool add_to(nvJitLinkHandle& handle) const = 0; + + virtual char const* get_key() const = 0; +}; + +struct FatbinFragmentEntry : FragmentEntry { + virtual uint8_t const* get_data() const = 0; + + virtual size_t get_length() const = 0; + + bool add_to(nvJitLinkHandle& handle) const override final; +}; + +template +struct StaticFatbinFragmentEntry final : FatbinFragmentEntry { + uint8_t const* get_data() const override { return StaticFatbinFragmentEntry::data; } + + size_t get_length() const override { return StaticFatbinFragmentEntry::length; } + + char const* get_key() const override + { + return typeid(StaticFatbinFragmentEntry).name(); + } + + static uint8_t const* const data; + static size_t const length; +}; + +struct UDFFatbinFragment final : FatbinFragmentEntry { + UDFFatbinFragment(std::string key, std::vector bytes) + : key_(std::move(key)), bytes_(std::move(bytes)) + { + } + + uint8_t const* get_data() const override { return bytes_.data(); } + + size_t get_length() const override { return bytes_.size(); } + + char const* get_key() const override { return key_.c_str(); } + + private: + std::string key_; + std::vector bytes_; +}; diff --git a/NVRTCLTOFragmentCompiler.cpp b/NVRTCLTOFragmentCompiler.cpp new file mode 100644 index 0000000..e64f64b --- /dev/null +++ b/NVRTCLTOFragmentCompiler.cpp @@ -0,0 +1,107 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "cuda.h" + +#include +#include +#include +#include + +#include + +#define NVRTC_SAFE_CALL(_call) \ + { \ + nvrtcResult result = _call; \ + std::string error_string = \ + std::string("nvrtc error: ") + std::string(nvrtcGetErrorString(result)); \ + RAFT_EXPECTS(result == NVRTC_SUCCESS, "%s", error_string.c_str()); \ + } + +NVRTCLTOFragmentCompiler::NVRTCLTOFragmentCompiler() +{ + int device = 0; + int major = 0; + int minor = 0; + RAFT_CUDA_TRY(cudaGetDevice(&device)); + RAFT_CUDA_TRY(cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, device)); + RAFT_CUDA_TRY(cudaDeviceGetAttribute(&minor, cudaDevAttrComputeCapabilityMinor, device)); + + this->standard_compile_opts = { + std::string{"-arch=sm_" + std::to_string((major * 10 + minor))}, + std::string{"-dlto"}, + std::string{"-rdc=true"}, + std::string{"--std=c++20"}, + std::string{"-default-device"}, + }; +} + +std::unique_ptr NVRTCLTOFragmentCompiler::read_cache( + std::string const& key) const +{ + std::shared_lock read_lock(cache_mutex_); + if (auto it = cache.find(key); it != cache.end()) { + return std::make_unique(key, it->second); + } + return nullptr; +} + +std::unique_ptr NVRTCLTOFragmentCompiler::compile(std::string const& key, + std::string const& code) +{ + if (auto hit = read_cache(key)) { return hit; } + + std::unique_lock write_lock(cache_mutex_); + if (auto it = cache.find(key); it != cache.end()) { + return std::make_unique(key, it->second); + } + + nvrtcProgram prog; + NVRTC_SAFE_CALL( + nvrtcCreateProgram(&prog, code.c_str(), "nvrtc_lto_fragment", 0, nullptr, nullptr)); + + // Convert std::vector to std::vector for nvrtc API + std::vector opts; + opts.reserve(this->standard_compile_opts.size()); + for (auto const& opt : this->standard_compile_opts) { + opts.push_back(opt.c_str()); + } + + nvrtcResult compileResult = nvrtcCompileProgram(prog, // prog + opts.size(), // numOptions + opts.data()); // options + + try { + if (compileResult != NVRTC_SUCCESS) { + // Obtain compilation log from the program. + size_t log_size; + NVRTC_SAFE_CALL(nvrtcGetProgramLogSize(prog, &log_size)); + std::unique_ptr log{new char[log_size]}; + NVRTC_SAFE_CALL(nvrtcGetProgramLog(prog, log.get())); + RAFT_FAIL("nvrtc compile error log: \n%s", log.get()); + } + } catch (...) { + NVRTC_SAFE_CALL(nvrtcDestroyProgram(&prog)); + throw; + } + + // Obtain generated LTO IR from the program. + std::size_t ltoIRSize; + NVRTC_SAFE_CALL(nvrtcGetLTOIRSize(prog, <oIRSize)); + + std::vector lto_ir(ltoIRSize); + NVRTC_SAFE_CALL(nvrtcGetLTOIR(prog, reinterpret_cast(lto_ir.data()))); + + NVRTC_SAFE_CALL(nvrtcDestroyProgram(&prog)); + + cache[key] = std::move(lto_ir); + return std::make_unique(key, cache[key]); +} + +NVRTCLTOFragmentCompiler& nvrtc_compiler() +{ + static NVRTCLTOFragmentCompiler compiler; + return compiler; +} diff --git a/NVRTCLTOFragmentCompiler.hpp b/NVRTCLTOFragmentCompiler.hpp new file mode 100644 index 0000000..56011c3 --- /dev/null +++ b/NVRTCLTOFragmentCompiler.hpp @@ -0,0 +1,29 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include + +#include +#include +#include +#include +#include + +struct NVRTCLTOFragmentCompiler { + NVRTCLTOFragmentCompiler(); + + std::vector standard_compile_opts; + std::unordered_map> cache; + mutable std::shared_mutex cache_mutex_; + + std::unique_ptr compile(std::string const& key, std::string const& code); + + private: + std::unique_ptr read_cache(std::string const& key) const; +}; + +NVRTCLTOFragmentCompiler& nvrtc_compiler(); diff --git a/compute_matrix_product.cmake b/compute_matrix_product.cmake new file mode 100644 index 0000000..fc36d8d --- /dev/null +++ b/compute_matrix_product.cmake @@ -0,0 +1,47 @@ +# ============================================================================= +# cmake-format: off +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# cmake-format: on +# ============================================================================= + +include_guard(GLOBAL) + +function(compute_matrix_product output_var) + set(options) + set(one_value MATRIX_JSON_FILE MATRIX_JSON_STRING) + set(multi_value) + + cmake_parse_arguments(_JIT_LTO "${options}" "${one_value}" "${multi_value}" ${ARGN}) + + find_package(Python3 REQUIRED COMPONENTS Interpreter) + + if(_JIT_LTO_MATRIX_JSON_FILE) + execute_process(COMMAND "${Python3_EXECUTABLE}" + "${CMAKE_CURRENT_FUNCTION_LIST_DIR}/compute_matrix_product.py" + "${_JIT_LTO_MATRIX_JSON_FILE}" # + OUTPUT_VARIABLE output COMMAND_ERROR_IS_FATAL ANY) + else() + execute_process(COMMAND "${CMAKE_COMMAND}" -E echo "${_JIT_LTO_MATRIX_JSON_STRING}" + COMMAND "${Python3_EXECUTABLE}" + "${CMAKE_CURRENT_FUNCTION_LIST_DIR}/compute_matrix_product.py" - + OUTPUT_VARIABLE output COMMAND_ERROR_IS_FATAL ANY) + endif() + + set(${output_var} "${output}" PARENT_SCOPE) +endfunction() + +function(populate_matrix_variables matrix_json_entry) + string(JSON len LENGTH "${matrix_json_entry}") + if(len EQUAL 0) + return() + endif() + math(EXPR last "${len} - 1") + + # cmake-lint: disable=C0103,E1120 + foreach(i RANGE "${last}") + string(JSON key MEMBER "${matrix_json_entry}" "${i}") + string(JSON value GET "${matrix_json_entry}" "${key}") + set(${key} "${value}" PARENT_SCOPE) + endforeach() +endfunction() diff --git a/compute_matrix_product.py b/compute_matrix_product.py new file mode 100644 index 0000000..76262f6 --- /dev/null +++ b/compute_matrix_product.py @@ -0,0 +1,243 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# This algorithm takes a JSON dictionary and computes a matrix product of all +# of its arrays. We use this to compute all matrix combinations for kernel +# generation. We *could* write this in CMake with `string(JSON)`, but writing +# it in Python is much easier. Once we have a version of CMake that has +# https://gitlab.kitware.com/cmake/cmake/-/merge_requests/11516, we may be able +# to port the algorithm to CMake script and use it in other RAPIDS projects. + +import argparse +import json +import re +import sys +import warnings +from itertools import chain +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from collections.abc import Generator, Iterator + + MatrixValue = None | bool | int | float | str + Matrix = MatrixValue | list["Matrix"] | dict[str, "Matrix"] + + +class NoKeyError(ValueError): + pass + + +class UnusedKeyWarning(UserWarning): + pass + + +class UsedKeyWarning(UserWarning): + pass + + +IDENTIFIER_RE: re.Pattern = re.compile(r"^(?P_*)(?P.*)$") + + +def iterate_matrix_product( + *, + matrix: "Matrix", + warn_unused=True, + warn_used=True, +) -> "Generator[dict[str, MatrixValue]]": + """Computes a matrix product of a JSON document + + This algorithm computes the product of a matrix in a more sophisticated + way than can be done with itertools.product(). Multiple related values + can be grouped together, and a dimension can have sub-dimensions. Given + the following JSON document: + + .. code-block:: json + + { + "value": ["one", "two"], + "_group": [ + { + "subgroup_value": "three" + "subgroup_subdim": ["four", "five"] + }, + { + "subgroup_value": "six", + "subgroup_subdim": ["seven", "eight"] + } + ] + } + + The following matrix product entries will be produced: + + .. code-block:: json + + {"subgroup_subdim": "four", "subgroup_value": "three", "value": "one"} + {"subgroup_subdim": "five", "subgroup_value": "three", "value": "one"} + {"subgroup_subdim": "seven", "subgroup_value": "six", "value": "one"} + {"subgroup_subdim": "eight", "subgroup_value": "six", "value": "one"} + {"subgroup_subdim": "four", "subgroup_value": "three", "value": "two"} + {"subgroup_subdim": "five", "subgroup_value": "three", "value": "two"} + {"subgroup_subdim": "seven", "subgroup_value": "six", "value": "two"} + {"subgroup_subdim": "eight", "subgroup_value": "six", "value": "two"} + + Notice that the name ``_group`` does not appear in any of the matrix + product entries. This is because all leaf nodes beneath it are under + a different key that's closer to them in the hierarchy, which is used in + the final product. If a key is used only for grouping and does not appear + in the final product, it should be prefixed with an underscore (``_``) to + indicate that they are hidden. Likewise, keys that appear in the final + product should not be prefixed with an underscore. Failure to follow this + convention will not affect the proper functioning of the algorithm, but a + warning will be emitted unless the respective ``warn_unused`` and/or + ``warn_used`` parameters are set to ``False``. + + Every leaf node in the input document must have at least one dictionary key + in its path, or else a ``NoKeyError`` will be thrown. For example, a + document consisting only of an array of strings is invalid. + + Parameters + ---------- + matrix : Matrix + JSON document on which to compute the product. + warn_unused : bool + Whether or not to warn if unused keys are not prefixed with underscores + (true by default). + warn_used : bool + Whether or not to warn if used keys are prefixed with underscores (true + by default). + + Returns + ------- + Generator[dict[str, MatrixValue]] + Iterator of matrix product entries. + + Raises + ------ + NoKeyError + If a leaf node in the document does not have any key leading to it. + + Warns + ----- + UnusedKeyWarning + If an unused key is not prefixed with an underscore. + UsedKeyWarning + If a used key is prefixed with an underscore. + """ + + def iterate_next_dimension( + queue: "Iterator[tuple[tuple[str | int, ...], str | None, Matrix]]", + entry: "dict[str, MatrixValue]", + ) -> "Generator[tuple[dict[str, MatrixValue], bool]]": + try: + path, key, matrix = next(queue) + except StopIteration: + yield (entry, False) + else: + used = False + for e, u in iterate_impl(path, key, matrix, queue, entry): + if u: + used = True + yield e, u + + try: + last = path[-1] + except IndexError: + pass + else: + if isinstance(last, str): + match = IDENTIFIER_RE.search(last) + assert match + underscores = match.group("underscores") + rest = match.group("rest") + path_repr = "".join( + f"[{json.dumps(i)}]" for i in path[:-1] + ) + + if warn_used and used and underscores: + warnings.warn( + f"Key {json.dumps(last)} at root{path_repr} " + f"is used in a matrix product entry even though it " + f"begins with {json.dumps(underscores)}. Consider " + f"renaming it to {json.dumps(rest)} to indicate this.", + category=UsedKeyWarning, + ) + elif warn_unused and not used and not underscores: + warnings.warn( + f"Key {json.dumps(last)} at root{path_repr} " + f"is never used in a matrix product entry and is used " + f"only for grouping. Consider renaming it to " + f"{json.dumps(f'_{last}')} to indicate this.", + category=UnusedKeyWarning, + ) + + def iterate_impl( + path: tuple[str | int, ...], + key: str | None, + matrix: "Matrix", + queue: "Iterator[tuple[tuple[str | int, ...], str | None, Matrix]]", + entry: "dict[str, MatrixValue]", + ) -> "Generator[tuple[dict[str, MatrixValue], bool]]": + if isinstance(matrix, dict): + yield from ( + (e, False) + for e, _ in iterate_next_dimension( + chain( + ( + ((*path, k), k, v) + for (k, v) in sorted(matrix.items()) + ), + queue, + ), + entry, + ) + ) + elif isinstance(matrix, list): + queue_list = list(queue) + for i, v in enumerate(matrix): + yield from iterate_next_dimension( + chain([((*path, i), key, v)], queue_list), entry + ) + else: + if key is None: + raise NoKeyError + yield from ( + (e, True) + for e, _ in iterate_next_dimension( + queue, {**entry, key: matrix} + ) + ) + + yield from ( + entry for entry, _ in iterate_impl((), None, matrix, chain(), {}) + ) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--warn-unused", action=argparse.BooleanOptionalAction, default=True + ) + parser.add_argument( + "--warn-used", action=argparse.BooleanOptionalAction, default=True + ) + parser.add_argument("filename", nargs="?", default="-") + namespace = parser.parse_args() + with ( + sys.stdin + if namespace.filename == "-" + else open(namespace.filename) as f + ): + matrix: "Matrix" = json.load(f) + + json.dump( + list( + iterate_matrix_product( + matrix=matrix, + warn_unused=namespace.warn_unused, + warn_used=namespace.warn_used, + ) + ), + sys.stdout, + indent=2, + sort_keys=True, + ) diff --git a/generate_jit_lto_kernels.cmake b/generate_jit_lto_kernels.cmake new file mode 100644 index 0000000..9c90438 --- /dev/null +++ b/generate_jit_lto_kernels.cmake @@ -0,0 +1,117 @@ +# ============================================================================= +# cmake-format: off +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# cmake-format: on +# ============================================================================= + +include_guard(GLOBAL) + +include(${CMAKE_CURRENT_LIST_DIR}/compute_matrix_product.cmake) + +function(add_jit_lto_kernel kernel_target) + set(options) + set(one_value KERNEL_FILE FATBIN_HEADER_FILE) + set(multi_value LINK_LIBRARIES EXTRA_COMPILE_OPTIONS) + + cmake_parse_arguments(_JIT_LTO "${options}" "${one_value}" "${multi_value}" ${ARGN}) + + add_library(${kernel_target} OBJECT EXCLUDE_FROM_ALL "${_JIT_LTO_KERNEL_FILE}") + # Do not modify these properties, options, and libraries. Usage requirements (including CUDA + # version, etc.) should be propagated to the kernel targets via INTERFACE libraries passed in + # through the LINK_LIBRARIES argument. + target_link_libraries(${kernel_target} PRIVATE ${_JIT_LTO_LINK_LIBRARIES}) + target_compile_options(${kernel_target} PRIVATE -Xfatbin=--compress-all --compress-mode=size) + if(_JIT_LTO_EXTRA_COMPILE_OPTIONS) + target_compile_options(${kernel_target} PRIVATE ${_JIT_LTO_EXTRA_COMPILE_OPTIONS}) + endif() + set_target_properties(${kernel_target} + PROPERTIES CUDA_SEPARABLE_COMPILATION ON CUDA_FATBIN_COMPILATION ON + POSITION_INDEPENDENT_CODE ON INTERPROCEDURAL_OPTIMIZATION ON) + + add_custom_command(OUTPUT "${_JIT_LTO_FATBIN_HEADER_FILE}" + COMMAND "${bin_to_c}" --const --name embedded_fatbin --static + $ > "${_JIT_LTO_FATBIN_HEADER_FILE}" + DEPENDS $ ${kernel_target} + COMMENT "Generating fatbin header file for ${kernel_target}") +endfunction() + +function(process_jit_lto_matrix_entry source_list_var) + set(options) + set(one_value NAME_FORMAT KERNEL_INPUT_FILE OUTPUT_DIRECTORY FRAGMENT_TAG_FORMAT + MATRIX_JSON_ENTRY) + set(multi_value KERNEL_LINK_LIBRARIES FRAGMENT_TAG_HEADER_FILES KERNEL_EXTRA_COMPILE_OPTIONS) + + cmake_parse_arguments(_JIT_LTO "${options}" "${one_value}" "${multi_value}" ${ARGN}) + + populate_matrix_variables("${_JIT_LTO_MATRIX_JSON_ENTRY}") + string(CONFIGURE "${_JIT_LTO_NAME_FORMAT}" kernel_name @ONLY) + string(CONFIGURE "${_JIT_LTO_FRAGMENT_TAG_FORMAT}" fragment_tag @ONLY) + + set(fragment_tag_header_files "") + foreach(header_file IN LISTS _JIT_LTO_FRAGMENT_TAG_HEADER_FILES) + if(NOT header_file MATCHES "^(\".*\"|<.*>)$") + set(header_file "\"${header_file}\"") + endif() + string(APPEND fragment_tag_header_files "#include ${header_file}\n") + endforeach() + + set(kernel_file "${_JIT_LTO_OUTPUT_DIRECTORY}/${kernel_name}_kernel.cu") + set(kernel_target "${kernel_name}_kernel") + set(fatbin_header_file "${_JIT_LTO_OUTPUT_DIRECTORY}/${kernel_name}_fatbin.h") + set(fatbin_file "${_JIT_LTO_OUTPUT_DIRECTORY}/${kernel_name}_fatbin.cpp") + configure_file("${_JIT_LTO_KERNEL_INPUT_FILE}" "${kernel_file}" @ONLY) + configure_file("${CMAKE_CURRENT_FUNCTION_LIST_DIR}/register_fatbin.cpp.in" "${fatbin_file}" @ONLY) + + add_jit_lto_kernel(${kernel_target} KERNEL_FILE "${kernel_file}" FATBIN_HEADER_FILE + "${fatbin_header_file}" LINK_LIBRARIES ${_JIT_LTO_KERNEL_LINK_LIBRARIES} + EXTRA_COMPILE_OPTIONS ${_JIT_LTO_KERNEL_EXTRA_COMPILE_OPTIONS}) + list(APPEND ${source_list_var} "${fatbin_header_file}" "${fatbin_file}") + set(${source_list_var} "${${source_list_var}}" PARENT_SCOPE) +endfunction() + +function(generate_jit_lto_kernels source_list_var) + set(options) + set(one_value NAME_FORMAT MATRIX_JSON_FILE MATRIX_JSON_STRING KERNEL_INPUT_FILE + FRAGMENT_TAG_FORMAT OUTPUT_DIRECTORY) + set(multi_value KERNEL_LINK_LIBRARIES FRAGMENT_TAG_HEADER_FILES KERNEL_EXTRA_COMPILE_OPTIONS) + + cmake_parse_arguments(_JIT_LTO "${options}" "${one_value}" "${multi_value}" ${ARGN}) + + find_package(CUDAToolkit REQUIRED) + find_program(bin_to_c NAMES bin2c PATHS ${CUDAToolkit_BIN_DIR}) + + if(_JIT_LTO_MATRIX_JSON_FILE) + set_property(DIRECTORY PROPERTY CMAKE_CONFIGURE_DEPENDS "${_JIT_LTO_MATRIX_JSON_FILE}" APPEND) + compute_matrix_product(matrix_product MATRIX_JSON_FILE "${_JIT_LTO_MATRIX_JSON_FILE}") + else() + compute_matrix_product(matrix_product MATRIX_JSON_STRING "${_JIT_LTO_MATRIX_JSON_STRING}") + endif() + + string(JSON len LENGTH "${matrix_product}") + math(EXPR last "${len} - 1") + + # cmake-lint: disable=C0103,E1120 + foreach(i RANGE "${last}") + string(JSON matrix_json_entry GET "${matrix_product}" "${i}") + process_jit_lto_matrix_entry("${source_list_var}" + NAME_FORMAT + "${_JIT_LTO_NAME_FORMAT}" + KERNEL_INPUT_FILE + "${_JIT_LTO_KERNEL_INPUT_FILE}" + FRAGMENT_TAG_FORMAT + "${_JIT_LTO_FRAGMENT_TAG_FORMAT}" + FRAGMENT_TAG_HEADER_FILES + ${_JIT_LTO_FRAGMENT_TAG_HEADER_FILES} + OUTPUT_DIRECTORY + "${_JIT_LTO_OUTPUT_DIRECTORY}" + MATRIX_JSON_ENTRY + "${matrix_json_entry}" + KERNEL_LINK_LIBRARIES + ${_JIT_LTO_KERNEL_LINK_LIBRARIES} + KERNEL_EXTRA_COMPILE_OPTIONS + ${_JIT_LTO_KERNEL_EXTRA_COMPILE_OPTIONS}) + endforeach() + + set(${source_list_var} "${${source_list_var}}" PARENT_SCOPE) +endfunction() diff --git a/nvjitlink_checker.cpp b/nvjitlink_checker.cpp new file mode 100644 index 0000000..da53108 --- /dev/null +++ b/nvjitlink_checker.cpp @@ -0,0 +1,26 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include +#include + +#include +#include + +void check_nvjitlink_result(nvJitLinkHandle handle, nvJitLinkResult result) +{ + if (result != NVJITLINK_SUCCESS) { + std::string error_msg = "nvJITLink failed with error " + std::to_string(result); + size_t log_size = 0; + result = nvJitLinkGetErrorLogSize(handle, &log_size); + if (result == NVJITLINK_SUCCESS && log_size > 0) { + std::unique_ptr log{new char[log_size]}; + result = nvJitLinkGetErrorLog(handle, log.get()); + if (result == NVJITLINK_SUCCESS) { error_msg += "\n" + std::string(log.get()); } + } + RAFT_FAIL("AlgorithmPlanner nvJITLink error log: %s", error_msg.c_str()); + } +} diff --git a/nvjitlink_checker.hpp b/nvjitlink_checker.hpp new file mode 100644 index 0000000..4329884 --- /dev/null +++ b/nvjitlink_checker.hpp @@ -0,0 +1,11 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include + +// We can make a better RAII wrapper around nvjitlinkhandle +void check_nvjitlink_result(nvJitLinkHandle handle, nvJitLinkResult result); diff --git a/register_fatbin.cpp.in b/register_fatbin.cpp.in new file mode 100644 index 0000000..ffe6580 --- /dev/null +++ b/register_fatbin.cpp.in @@ -0,0 +1,22 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "@fatbin_header_file@" +#include + +@fragment_tag_header_files@ + +namespace { + +using fragment_tag = @fragment_tag@; +using fragment_entry = StaticFatbinFragmentEntry; + +} // namespace + +template <> +const uint8_t* const fragment_entry::data = embedded_fatbin; + +template <> +const size_t fragment_entry::length = sizeof(embedded_fatbin);