From 91a9e522d2ec91a6c64efb42e8b1343e1a608e38 Mon Sep 17 00:00:00 2001 From: Arha Gatram Date: Thu, 9 Jul 2026 09:10:14 -0700 Subject: [PATCH 1/2] Copied cuVS JIT-LTO logic wholesale --- AlgorithmLauncher.cpp | 66 +++++++++ AlgorithmLauncher.hpp | 60 ++++++++ AlgorithmPlanner.cpp | 116 ++++++++++++++++ AlgorithmPlanner.hpp | 60 ++++++++ CMakeLists.txt | 6 +- FragmentEntry.cpp | 14 ++ FragmentEntry.hpp | 64 +++++++++ NVRTCLTOFragmentCompiler.cpp | 108 +++++++++++++++ NVRTCLTOFragmentCompiler.hpp | 29 ++++ compute_matrix_product.cmake | 56 ++++++++ compute_matrix_product.py | 243 +++++++++++++++++++++++++++++++++ generate_jit_lto_kernels.cmake | 136 ++++++++++++++++++ nvjitlink_checker.cpp | 27 ++++ nvjitlink_checker.hpp | 11 ++ register_fatbin.cpp.in | 22 +++ 15 files changed, 1017 insertions(+), 1 deletion(-) create mode 100644 AlgorithmLauncher.cpp create mode 100644 AlgorithmLauncher.hpp create mode 100644 AlgorithmPlanner.cpp create mode 100644 AlgorithmPlanner.hpp create mode 100644 FragmentEntry.cpp create mode 100644 FragmentEntry.hpp create mode 100644 NVRTCLTOFragmentCompiler.cpp create mode 100644 NVRTCLTOFragmentCompiler.hpp create mode 100644 compute_matrix_product.cmake create mode 100644 compute_matrix_product.py create mode 100644 generate_jit_lto_kernels.cmake create mode 100644 nvjitlink_checker.cpp create mode 100644 nvjitlink_checker.hpp create mode 100644 register_fatbin.cpp.in diff --git a/AlgorithmLauncher.cpp b/AlgorithmLauncher.cpp new file mode 100644 index 0000000..7ee1e01 --- /dev/null +++ b/AlgorithmLauncher.cpp @@ -0,0 +1,66 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * 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..10418c4 --- /dev/null +++ b/AlgorithmLauncher.hpp @@ -0,0 +1,60 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * 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(const AlgorithmLauncher&) = delete; + AlgorithmLauncher& operator=(const AlgorithmLauncher&) = 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..ecf7a3c --- /dev/null +++ b/AlgorithmPlanner.cpp @@ -0,0 +1,116 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +#include "cuda_runtime.h" +#include "nvJitLink.h" + +#include +#include + +std::string AlgorithmPlanner::get_fragments_key() const +{ + std::string key = ""; + for (const auto& 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 (const auto& 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 (const auto& 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..7f275b1 --- /dev/null +++ b/AlgorithmPlanner.hpp @@ -0,0 +1,60 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +#include "AlgorithmLauncher.hpp" +#include "FragmentEntry.hpp" + +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..219faec 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -51,7 +51,7 @@ if(NOT TARGET nvtx3::nvtx3-cpp) rapids_cpm_nvtx3() endif() -add_library(rtcx STATIC hash.cpp rtcx.cpp) +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 @@ -95,6 +95,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 +105,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,6 +128,8 @@ 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 diff --git a/FragmentEntry.cpp b/FragmentEntry.cpp new file mode 100644 index 0000000..b3270f1 --- /dev/null +++ b/FragmentEntry.cpp @@ -0,0 +1,14 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * 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..35aa466 --- /dev/null +++ b/FragmentEntry.hpp @@ -0,0 +1,64 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include +#include +#include +#include + +#include + +#include "nvjitlink_checker.hpp" + +struct FragmentEntry { + virtual ~FragmentEntry() = default; + + virtual bool add_to(nvJitLinkHandle& handle) const = 0; + + virtual const char* get_key() const = 0; +}; + +struct FatbinFragmentEntry : FragmentEntry { + virtual const uint8_t* get_data() const = 0; + + virtual size_t get_length() const = 0; + + bool add_to(nvJitLinkHandle& handle) const override final; +}; + +template +struct StaticFatbinFragmentEntry final : FatbinFragmentEntry { + const uint8_t* get_data() const override { return StaticFatbinFragmentEntry::data; } + + size_t get_length() const override { return StaticFatbinFragmentEntry::length; } + + const char* get_key() const override + { + return typeid(StaticFatbinFragmentEntry).name(); + } + + static const uint8_t* const data; + static const size_t length; +}; + +struct UDFFatbinFragment final : FatbinFragmentEntry { + UDFFatbinFragment(std::string key, std::vector bytes) + : key_(std::move(key)), bytes_(std::move(bytes)) + { + } + + const uint8_t* get_data() const override { return bytes_.data(); } + + size_t get_length() const override { return bytes_.size(); } + + const char* 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..1d4f054 --- /dev/null +++ b/NVRTCLTOFragmentCompiler.cpp @@ -0,0 +1,108 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include + +#include + +#include +#include + +#include "cuda.h" +#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 (const auto& 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..b4796e6 --- /dev/null +++ b/NVRTCLTOFragmentCompiler.hpp @@ -0,0 +1,29 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * 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..82a34f9 --- /dev/null +++ b/compute_matrix_product.cmake @@ -0,0 +1,56 @@ +# ============================================================================= +# cmake-format: off +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# 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..f427fc4 --- /dev/null +++ b/generate_jit_lto_kernels.cmake @@ -0,0 +1,136 @@ +# ============================================================================= +# cmake-format: off +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# 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} + ) +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..f0d1213 --- /dev/null +++ b/nvjitlink_checker.cpp @@ -0,0 +1,27 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * 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..12b062d --- /dev/null +++ b/nvjitlink_checker.hpp @@ -0,0 +1,11 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * 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); From 87adf46a3b87317294537aeeef3deb0369605589 Mon Sep 17 00:00:00 2001 From: Arha Gatram Date: Thu, 9 Jul 2026 15:41:34 -0700 Subject: [PATCH 2/2] add raft as dependency --- AlgorithmLauncher.cpp | 3 +- AlgorithmLauncher.hpp | 15 +++--- AlgorithmPlanner.cpp | 27 +++++----- AlgorithmPlanner.hpp | 8 +-- CMakeLists.txt | 53 ++++++++++++-------- FragmentEntry.cpp | 2 +- FragmentEntry.hpp | 26 +++++----- NVRTCLTOFragmentCompiler.cpp | 15 +++--- NVRTCLTOFragmentCompiler.hpp | 2 +- compute_matrix_product.cmake | 31 ++++-------- generate_jit_lto_kernels.cmake | 91 ++++++++++++++-------------------- nvjitlink_checker.cpp | 7 ++- nvjitlink_checker.hpp | 2 +- 13 files changed, 131 insertions(+), 151 deletions(-) diff --git a/AlgorithmLauncher.cpp b/AlgorithmLauncher.cpp index 7ee1e01..981575f 100644 --- a/AlgorithmLauncher.cpp +++ b/AlgorithmLauncher.cpp @@ -1,10 +1,9 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * 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} {} diff --git a/AlgorithmLauncher.hpp b/AlgorithmLauncher.hpp index 10418c4..821c3db 100644 --- a/AlgorithmLauncher.hpp +++ b/AlgorithmLauncher.hpp @@ -1,18 +1,19 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * 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 +#include #include +#include +#include struct AlgorithmLauncher { AlgorithmLauncher() : kernel{nullptr}, library{nullptr} {} @@ -21,8 +22,8 @@ struct AlgorithmLauncher { ~AlgorithmLauncher(); - AlgorithmLauncher(const AlgorithmLauncher&) = delete; - AlgorithmLauncher& operator=(const AlgorithmLauncher&) = delete; + AlgorithmLauncher(AlgorithmLauncher const&) = delete; + AlgorithmLauncher& operator=(AlgorithmLauncher const&) = delete; AlgorithmLauncher(AlgorithmLauncher&& other) noexcept; AlgorithmLauncher& operator=(AlgorithmLauncher&& other) noexcept; diff --git a/AlgorithmPlanner.cpp b/AlgorithmPlanner.cpp index ecf7a3c..9d9da43 100644 --- a/AlgorithmPlanner.cpp +++ b/AlgorithmPlanner.cpp @@ -1,8 +1,16 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * 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 @@ -12,19 +20,10 @@ #include #include -#include -#include - -#include "cuda_runtime.h" -#include "nvJitLink.h" - -#include -#include - std::string AlgorithmPlanner::get_fragments_key() const { std::string key = ""; - for (const auto& fragment : this->fragments) { + for (auto const& fragment : this->fragments) { key += fragment->get_key(); } return key; @@ -50,7 +49,7 @@ std::shared_ptr AlgorithmPlanner::get_launcher() std::string log_message = "JIT compiling launcher for kernel: " + this->entrypoint + " and device functions: "; - for (const auto& fragment : this->fragments) { + for (auto const& fragment : this->fragments) { log_message += std::string{fragment->get_key()} + ","; } log_message.pop_back(); @@ -73,7 +72,7 @@ std::shared_ptr AlgorithmPlanner::build() // Load the generated LTO IR and link them together nvJitLinkHandle handle; - std::vector lopts; + std::vector lopts; lopts.reserve(2 + linktime_extra_options.size()); lopts.push_back("-lto"); lopts.push_back(archs.c_str()); @@ -83,7 +82,7 @@ std::shared_ptr AlgorithmPlanner::build() auto result = nvJitLinkCreate(&handle, static_cast(lopts.size()), lopts.data()); check_nvjitlink_result(handle, result); - for (const auto& frag : this->fragments) { + for (auto const& frag : this->fragments) { frag->add_to(handle); } diff --git a/AlgorithmPlanner.hpp b/AlgorithmPlanner.hpp index 7f275b1..4fdb81e 100644 --- a/AlgorithmPlanner.hpp +++ b/AlgorithmPlanner.hpp @@ -1,10 +1,13 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * 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 @@ -13,9 +16,6 @@ #include #include -#include "AlgorithmLauncher.hpp" -#include "FragmentEntry.hpp" - struct LauncherJitCache { std::shared_mutex mutex; std::unordered_map> launchers; diff --git a/CMakeLists.txt b/CMakeLists.txt index 219faec..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 AlgorithmLauncher.cpp AlgorithmPlanner.cpp FragmentEntry.cpp nvjitlink_checker.cpp NVRTCLTOFragmentCompiler.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) @@ -128,20 +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(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 index b3270f1..48a5eb7 100644 --- a/FragmentEntry.cpp +++ b/FragmentEntry.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/FragmentEntry.hpp b/FragmentEntry.hpp index 35aa466..c2f7b0a 100644 --- a/FragmentEntry.hpp +++ b/FragmentEntry.hpp @@ -1,30 +1,30 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * 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 -#include - -#include "nvjitlink_checker.hpp" - struct FragmentEntry { virtual ~FragmentEntry() = default; virtual bool add_to(nvJitLinkHandle& handle) const = 0; - virtual const char* get_key() const = 0; + virtual char const* get_key() const = 0; }; struct FatbinFragmentEntry : FragmentEntry { - virtual const uint8_t* get_data() const = 0; + virtual uint8_t const* get_data() const = 0; virtual size_t get_length() const = 0; @@ -33,17 +33,17 @@ struct FatbinFragmentEntry : FragmentEntry { template struct StaticFatbinFragmentEntry final : FatbinFragmentEntry { - const uint8_t* get_data() const override { return StaticFatbinFragmentEntry::data; } + uint8_t const* get_data() const override { return StaticFatbinFragmentEntry::data; } size_t get_length() const override { return StaticFatbinFragmentEntry::length; } - const char* get_key() const override + char const* get_key() const override { return typeid(StaticFatbinFragmentEntry).name(); } - static const uint8_t* const data; - static const size_t length; + static uint8_t const* const data; + static size_t const length; }; struct UDFFatbinFragment final : FatbinFragmentEntry { @@ -52,11 +52,11 @@ struct UDFFatbinFragment final : FatbinFragmentEntry { { } - const uint8_t* get_data() const override { return bytes_.data(); } + uint8_t const* get_data() const override { return bytes_.data(); } size_t get_length() const override { return bytes_.size(); } - const char* get_key() const override { return key_.c_str(); } + char const* get_key() const override { return key_.c_str(); } private: std::string key_; diff --git a/NVRTCLTOFragmentCompiler.cpp b/NVRTCLTOFragmentCompiler.cpp index 1d4f054..e64f64b 100644 --- a/NVRTCLTOFragmentCompiler.cpp +++ b/NVRTCLTOFragmentCompiler.cpp @@ -1,17 +1,16 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ -#include - -#include +#include "cuda.h" +#include +#include #include #include -#include "cuda.h" -#include +#include #define NVRTC_SAFE_CALL(_call) \ { \ @@ -64,9 +63,9 @@ std::unique_ptr NVRTCLTOFragmentCompiler::compile(std::string nvrtcCreateProgram(&prog, code.c_str(), "nvrtc_lto_fragment", 0, nullptr, nullptr)); // Convert std::vector to std::vector for nvrtc API - std::vector opts; + std::vector opts; opts.reserve(this->standard_compile_opts.size()); - for (const auto& opt : this->standard_compile_opts) { + for (auto const& opt : this->standard_compile_opts) { opts.push_back(opt.c_str()); } diff --git a/NVRTCLTOFragmentCompiler.hpp b/NVRTCLTOFragmentCompiler.hpp index b4796e6..56011c3 100644 --- a/NVRTCLTOFragmentCompiler.hpp +++ b/NVRTCLTOFragmentCompiler.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/compute_matrix_product.cmake b/compute_matrix_product.cmake index 82a34f9..fc36d8d 100644 --- a/compute_matrix_product.cmake +++ b/compute_matrix_product.cmake @@ -1,6 +1,6 @@ # ============================================================================= # cmake-format: off -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # cmake-format: on # ============================================================================= @@ -17,24 +17,18 @@ function(compute_matrix_product output_var) 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 - ) + 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 - ) + 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 - ) + set(${output_var} "${output}" PARENT_SCOPE) endfunction() function(populate_matrix_variables matrix_json_entry) @@ -48,9 +42,6 @@ function(populate_matrix_variables matrix_json_entry) 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 - ) + set(${key} "${value}" PARENT_SCOPE) endforeach() endfunction() diff --git a/generate_jit_lto_kernels.cmake b/generate_jit_lto_kernels.cmake index f427fc4..9c90438 100644 --- a/generate_jit_lto_kernels.cmake +++ b/generate_jit_lto_kernels.cmake @@ -1,6 +1,6 @@ # ============================================================================= # cmake-format: off -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # cmake-format: on # ============================================================================= @@ -25,27 +25,21 @@ function(add_jit_lto_kernel kernel_target) 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} - ) + 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 - ) + 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}) @@ -69,42 +63,26 @@ function(process_jit_lto_matrix_entry source_list_var) 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} - ) + 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 - ) + 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 - ) + 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} - ) + 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 - ) + 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}") @@ -116,21 +94,24 @@ function(generate_jit_lto_kernels source_list_var) # 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} - ) + 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 - ) + set(${source_list_var} "${${source_list_var}}" PARENT_SCOPE) endfunction() diff --git a/nvjitlink_checker.cpp b/nvjitlink_checker.cpp index f0d1213..da53108 100644 --- a/nvjitlink_checker.cpp +++ b/nvjitlink_checker.cpp @@ -1,16 +1,15 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ +#include #include +#include #include -#include #include -#include - void check_nvjitlink_result(nvJitLinkHandle handle, nvJitLinkResult result) { if (result != NVJITLINK_SUCCESS) { diff --git a/nvjitlink_checker.hpp b/nvjitlink_checker.hpp index 12b062d..4329884 100644 --- a/nvjitlink_checker.hpp +++ b/nvjitlink_checker.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */