From f100087227ca487fc6ac7244ea1099ad7df5e61b Mon Sep 17 00:00:00 2001 From: Vyas Ramasubramani Date: Thu, 28 May 2026 15:36:04 -0700 Subject: [PATCH 1/4] feat(rtcx): add LTO-IR fragment compilation and runtime link-launch API - fragments.cmake: add_fragment(), add_composite_fragment(), get_jit_fragment_architectures() - compute_matrix_product.cmake/.py: Cartesian product of JSON matrices - embed.cmake: add OUTPUT_DIRECTORY parameter - rtcx.hpp: lto_kernel, get_lto_linked_kernel, AlgorithmPlanner, AlgorithmLauncher - rtcx.cpp: implementation of link-launch API --- cpp/librtcx/compute_matrix_product.cmake | 55 +++++ cpp/librtcx/compute_matrix_product.py | 243 +++++++++++++++++++++++ cpp/librtcx/embed.cmake | 8 +- cpp/librtcx/fragments.cmake | 219 ++++++++++++++++++++ cpp/librtcx/rtcx.cpp | 182 +++++++++++++++++ cpp/librtcx/rtcx.hpp | 90 +++++++++ 6 files changed, 793 insertions(+), 4 deletions(-) create mode 100644 cpp/librtcx/compute_matrix_product.cmake create mode 100644 cpp/librtcx/compute_matrix_product.py create mode 100644 cpp/librtcx/fragments.cmake diff --git a/cpp/librtcx/compute_matrix_product.cmake b/cpp/librtcx/compute_matrix_product.cmake new file mode 100644 index 000000000000..d643b1d26d5c --- /dev/null +++ b/cpp/librtcx/compute_matrix_product.cmake @@ -0,0 +1,55 @@ +# ============================================================================= +# cmake-format: off +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 +# cmake-format: on +# ============================================================================= + +include_guard(GLOBAL) + +# Compute the Cartesian product of a JSON matrix file or string and return it as a JSON array. +function(compute_matrix_product output_var) + set(options) + set(one_value MATRIX_JSON_FILE MATRIX_JSON_STRING) + set(multi_value) + + cmake_parse_arguments(_CMP "${options}" "${one_value}" "${multi_value}" ${ARGN}) + + find_package(Python3 REQUIRED COMPONENTS Interpreter) + + if(_CMP_MATRIX_JSON_FILE) + execute_process( + COMMAND "${Python3_EXECUTABLE}" "${CMAKE_CURRENT_FUNCTION_LIST_DIR}/compute_matrix_product.py" + "${_CMP_MATRIX_JSON_FILE}" # + OUTPUT_VARIABLE output COMMAND_ERROR_IS_FATAL ANY + ) + else() + execute_process( + COMMAND ${CMAKE_COMMAND} -E echo "${_CMP_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() + +# Populate parent-scope CMake variables from one JSON matrix entry object. +function(populate_matrix_variables matrix_json_entry) + string(JSON len LENGTH "${matrix_json_entry}") + 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/cpp/librtcx/compute_matrix_product.py b/cpp/librtcx/compute_matrix_product.py new file mode 100644 index 000000000000..76262f6597ae --- /dev/null +++ b/cpp/librtcx/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/cpp/librtcx/embed.cmake b/cpp/librtcx/embed.cmake index a2b60635fa4b..ee3f2251271f 100644 --- a/cpp/librtcx/embed.cmake +++ b/cpp/librtcx/embed.cmake @@ -187,7 +187,7 @@ endfunction() # for JIT compilation. function(embed TARGET) set(OPTIONS "") - set(ONE_VALUE_ARGS "COMPRESSION") + set(ONE_VALUE_ARGS "COMPRESSION" "OUTPUT_DIRECTORY") set(MULTI_VALUE_ARGS "") cmake_parse_arguments(ARG "${OPTIONS}" "${ONE_VALUE_ARGS}" "${MULTI_VALUE_ARGS}" ${ARGN}) @@ -195,8 +195,8 @@ function(embed TARGET) message(FATAL_ERROR "embed target '${TARGET}' has not been initialized with add_embed()") endif() - if(NOT DEFINED ARG_COMPRESSION) - message(FATAL_ERROR "COMPRESSION argument is required") + if(NOT DEFINED ARG_COMPRESSION OR NOT DEFINED ARG_OUTPUT_DIRECTORY) + message(FATAL_ERROR "COMPRESSION and OUTPUT_DIRECTORY arguments are required") endif() if(NOT ARG_COMPRESSION STREQUAL "none" AND NOT ARG_COMPRESSION STREQUAL "zstd") @@ -243,7 +243,7 @@ function(embed TARGET) PROPERTY EMBED_INCLUDE_DIRECTORIES ) - set(OUTPUT_DIR "${CUDF_GENERATED_INCLUDE_DIR}/rtcx_embed") + set(OUTPUT_DIR "${ARG_OUTPUT_DIRECTORY}") set(EMBED_SCRIPT_TEMPLATE "${CMAKE_CURRENT_FUNCTION_LIST_DIR}/embed.in.cpp") set(CONFIGURED_EMBED_SCRIPT "${CMAKE_CURRENT_BINARY_DIR}/${TARGET}__embed_cfg.cpp") set(EMBED_SCRIPT "${CMAKE_CURRENT_BINARY_DIR}/${TARGET}__embed.cpp") diff --git a/cpp/librtcx/fragments.cmake b/cpp/librtcx/fragments.cmake new file mode 100644 index 000000000000..42c30caa2c74 --- /dev/null +++ b/cpp/librtcx/fragments.cmake @@ -0,0 +1,219 @@ +# cmake-format: off +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# cmake-format: on +# cmake-lint: disable=C0113,E1120 + +# ============================================================================= +# librtcx JIT Fragment Compilation Infrastructure +# ============================================================================= +# Provides: add_fragment(), add_composite_fragment(), get_jit_fragment_architectures() These macros +# compile CUDA sources into per-architecture LTO-IR fatbins and embed them into a target via +# embed_blob(). Any CUDA project using librtcx can use these to register precompiled fragments. +# +# Usage: include(librtcx/fragments.cmake) add_embed(my_fragments) add_fragment(my_fragments FRAGMENT +# my_kernel SOURCE my_kernel.cu LINK_LIBRARIES dep1::dep1 dep2::dep2 INCLUDE_DIRECTORIES +# "${CMAKE_CURRENT_SOURCE_DIR}/include" COMPILE_DEFINITIONS MY_DISABLE_EXPORTS) embed(my_fragments +# COMPRESSION zstd OUTPUT_DIRECTORY ...) +# ============================================================================= + +# Fragment name convention: fragments must end with _sm for the runtime arch detection +# heuristic. +function(get_jit_fragment_architectures OUT_VAR) + get_property(_RAPIDS_CUDA_ARCHITECTURES GLOBAL PROPERTY rapids_cuda_architectures) + if(_RAPIDS_CUDA_ARCHITECTURES) + set(_ARCHITECTURES ${_RAPIDS_CUDA_ARCHITECTURES}) + else() + set(_ARCHITECTURES ${CMAKE_CUDA_ARCHITECTURES}) + endif() + + set(_FRAGMENT_ARCHITECTURES) + foreach(_ARCH IN LISTS _ARCHITECTURES) + if(_ARCH STREQUAL "RAPIDS" OR _ARCH STREQUAL "ALL") + if(CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 13.0.0) + list(APPEND _FRAGMENT_ARCHITECTURES 75-real 80-real 86-real 90a-real 100f-real 120a-real + 120 + ) + elseif(CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 12.9.0) + list(APPEND _FRAGMENT_ARCHITECTURES 70-real 75-real 80-real 86-real 90a-real 100f-real + 120a-real 120 + ) + else() + list(APPEND _FRAGMENT_ARCHITECTURES 70-real 75-real 80-real 86-real 90a-real 90-virtual) + endif() + elseif(_ARCH MATCHES "^[0-9]+[a-z]*(-real)?$") + list(APPEND _FRAGMENT_ARCHITECTURES ${_ARCH}) + endif() + endforeach() + + list(REMOVE_DUPLICATES _FRAGMENT_ARCHITECTURES) + set(${OUT_VAR} + ${_FRAGMENT_ARCHITECTURES} + PARENT_SCOPE + ) +endfunction() + +# This macro is used to create object libraries for JIT compilation fragments, and embed them as +# per-architecture fatbins in the final library. It compiles the specified source file with the +# appropriate flags to generate fatbins containing the specified kernel, and then embeds those +# fatbins in the final library with metadata that allows them to be looked up at runtime. +# +# For architectures below sm80, nvJitLink fails to accept LTO fatbins in CUDA 12.9. Disable IPO to +# emit standard device-code fatbins instead. +macro(add_fragment) + set(TARGET ${ARGV0}) + set(OPTIONS KERNEL_ONLY) + set(ONE_VALUE_ARGS FRAGMENT SOURCE ENTRY_NAME) + set(MULTI_VALUE_ARGS DEFINITIONS ARRAY_IDS ARRAY_VALUES LINK_LIBRARIES INCLUDE_DIRECTORIES + COMPILE_DEFINITIONS + ) + cmake_parse_arguments(ARG "${OPTIONS}" "${ONE_VALUE_ARGS}" "${MULTI_VALUE_ARGS}" ${ARGN}) + + if(NOT ARG_FRAGMENT) + message(FATAL_ERROR "add_fragment requires FRAGMENT argument") + endif() + + if(NOT ARG_SOURCE) + message(FATAL_ERROR "add_fragment requires SOURCE argument") + endif() + + if(NOT ARG_ENTRY_NAME) + set(ARG_ENTRY_NAME "kernel_entry") + endif() + + get_jit_fragment_architectures(FRAGMENT_ARCHITECTURES) + foreach(FRAGMENT_ARCH IN LISTS FRAGMENT_ARCHITECTURES) + string(REGEX REPLACE "-.*$" "" FRAGMENT_ARCH_ID "${FRAGMENT_ARCH}") + string(REGEX REPLACE "[^0-9]" "" FRAGMENT_ARCH_NUM "${FRAGMENT_ARCH_ID}") + if(FRAGMENT_ARCH_NUM AND FRAGMENT_ARCH_NUM LESS 80) + set(FRAGMENT_LTO OFF) + else() + set(FRAGMENT_LTO ON) + endif() + set(OBJECT_ID ${TARGET}_${ARG_FRAGMENT}_sm${FRAGMENT_ARCH_ID}) + add_library(${OBJECT_ID} OBJECT ${ARG_SOURCE}) + target_compile_options( + ${OBJECT_ID} PRIVATE --compress-mode=size --expt-relaxed-constexpr --extended-lambda + -Xfatbin=--compress-all + ) + + if(ARG_KERNEL_ONLY) + # ensure that the FATBIN symbols only contain the specified kernel + target_compile_options(${OBJECT_ID} PRIVATE -Xnvlink=--kernels-used=${ARG_ENTRY_NAME}) + endif() + + target_compile_definitions(${OBJECT_ID} PRIVATE ${ARG_COMPILE_DEFINITIONS} ${ARG_DEFINITIONS}) + set_target_properties( + ${OBJECT_ID} + PROPERTIES CUDA_ARCHITECTURES ${FRAGMENT_ARCH} + CUDA_SEPARABLE_COMPILATION ON + CUDA_FATBIN_COMPILATION ON + POSITION_INDEPENDENT_CODE ON + INTERPROCEDURAL_OPTIMIZATION ${FRAGMENT_LTO} + CXX_STANDARD 20 + CXX_STANDARD_REQUIRED ON + CXX_EXTENSIONS ON + CXX_VISIBILITY_PRESET hidden + CUDA_STANDARD 20 + CUDA_STANDARD_REQUIRED ON + CUDA_VISIBILITY_PRESET hidden + ) + if(ARG_LINK_LIBRARIES) + target_link_libraries(${OBJECT_ID} PRIVATE ${ARG_LINK_LIBRARIES}) + endif() + if(ARG_INCLUDE_DIRECTORIES) + target_include_directories(${OBJECT_ID} PRIVATE ${ARG_INCLUDE_DIRECTORIES}) + endif() + + embed_blob( + ${TARGET} + FILE + $ + DEST + fragments/${ARG_FRAGMENT}_sm${FRAGMENT_ARCH_ID}.fatbin + ID + ${ARG_FRAGMENT}_sm${FRAGMENT_ARCH_ID} + ARRAY_IDS + ${ARG_ARRAY_IDS} + fragment_arch + ARRAY_VALUES + ${ARG_ARRAY_VALUES} + sm${FRAGMENT_ARCH_ID} + ) + endforeach() +endmacro() + +# ============================================================================= +# This macro pre-links individual LTO-IR fragment fatbins at CMake build time into a single +# composite fatbin using nvJitLink, then embeds it via embed_blob(). If nvJitLink is not found, the +# composite is skipped with a warning. +# ============================================================================= +# NOTE: add_composite_fragment has no call sites as of 2026-05-27. It is aspirational +# infrastructure. +macro(add_composite_fragment) + set(_ACF_TARGET ${ARGV0}) + set(_ACF_COMPOSITE_NAME ${ARGV1}) + set(_ACF_OPTIONS "") + set(_ACF_ONE_VALUE_ARGS ENTRY_NAME) + set(_ACF_MULTI_VALUE_ARGS FRAGMENTS) + cmake_parse_arguments( + _ACF "${_ACF_OPTIONS}" "${_ACF_ONE_VALUE_ARGS}" "${_ACF_MULTI_VALUE_ARGS}" ${ARGN} + ) + + if(NOT _ACF_COMPOSITE_NAME) + message(FATAL_ERROR "add_composite_fragment requires COMPOSITE_NAME argument") + endif() + + if(NOT _ACF_FRAGMENTS) + message(FATAL_ERROR "add_composite_fragment requires FRAGMENTS argument") + endif() + + if(NOT _ACF_ENTRY_NAME) + set(_ACF_ENTRY_NAME "kernel_entry") + endif() + + # nvJitLink is a library API, not a CLI tool. Look for a helper binary that wraps it. If not + # found, the composite fragment is skipped gracefully. + find_program( + NVJITLINK_EXECUTABLE nvJitLink HINTS "${CUDAToolkit_BIN_DIR}" + "${CMAKE_CUDA_COMPILER_TOOLKIT_ROOT}/bin" ENV PATH + ) + + if(NOT NVJITLINK_EXECUTABLE) + message(STATUS "nvJitLink tool not found — composite fragment '${_ACF_COMPOSITE_NAME}' skipped") + else() + # Collect input fatbin files from OBJECT targets of registered fragments + set(_ACF_INPUT_FILES "") + set(_ACF_INPUT_DEPS "") + foreach(_ACF_FRAG_ID IN LISTS _ACF_FRAGMENTS) + # Each fragment was registered as OBJECT library named ${TARGET}_${FRAGMENT} + set(_ACF_OBJ_TARGET "${_ACF_TARGET}_${_ACF_FRAG_ID}") + list(APPEND _ACF_INPUT_FILES "$") + list(APPEND _ACF_INPUT_DEPS "${_ACF_OBJ_TARGET}") + endforeach() + + # Output composite fatbin path + set(_ACF_COMPOSITE_OUTPUT "${CMAKE_CURRENT_BINARY_DIR}/fragments/${_ACF_COMPOSITE_NAME}.fatbin") + + # Custom command to run nvJitLink to pre-link fragments at build time + add_custom_command( + OUTPUT "${_ACF_COMPOSITE_OUTPUT}" + COMMAND ${CMAKE_COMMAND} -E make_directory "${CMAKE_CURRENT_BINARY_DIR}/fragments" + COMMAND ${NVJITLINK_EXECUTABLE} -o "${_ACF_COMPOSITE_OUTPUT}" ${_ACF_INPUT_FILES} + DEPENDS ${_ACF_INPUT_FILES} + COMMENT "Pre-linking composite fragment: ${_ACF_COMPOSITE_NAME}" + VERBATIM + ) + + # Custom target to ensure the composite is built before embedding + add_custom_target( + "${_ACF_TARGET}_composite_${_ACF_COMPOSITE_NAME}" ALL DEPENDS "${_ACF_COMPOSITE_OUTPUT}" + ) + + # Embed the composite fatbin via the existing embed_blob infrastructure + embed_blob( + ${_ACF_TARGET} FILE "${_ACF_COMPOSITE_OUTPUT}" DEST "fragments/${_ACF_COMPOSITE_NAME}.fatbin" + ID "${_ACF_COMPOSITE_NAME}" + ) + endif() +endmacro() diff --git a/cpp/librtcx/rtcx.cpp b/cpp/librtcx/rtcx.cpp index ae2a970c5ec4..0090e6663337 100644 --- a/cpp/librtcx/rtcx.cpp +++ b/cpp/librtcx/rtcx.cpp @@ -21,11 +21,18 @@ #include #include +#include +#include +#include #include #include #include +#include #include +#include #include +#include +#include #define RTCX_EXPECTS(condition_, reason_, exception_type_) \ do { \ @@ -1336,4 +1343,179 @@ rtcx::byte_buffer decompress_blob(std::span compressed_binar return decompressed; } +namespace { + +void check_cuda(cudaError_t status) +{ + if (status != cudaSuccess) { throw std::runtime_error{cudaGetErrorString(status)}; } +} + +std::string current_device_arch_option() +{ + int device = 0; + int cc_major = 0; + int cc_minor = 0; + check_cuda(cudaGetDevice(&device)); + check_cuda(cudaDeviceGetAttribute(&cc_major, cudaDevAttrComputeCapabilityMajor, device)); + check_cuda(cudaDeviceGetAttribute(&cc_minor, cudaDevAttrComputeCapabilityMinor, device)); + return std::format("-arch=sm_{}{}", cc_major, cc_minor); +} + +std::optional arch_option_from_fragment_name(char const* name) +{ + if (name == nullptr) { return std::nullopt; } + auto const fragment_name = std::string_view{name}; + auto const sm_pos = fragment_name.rfind("_sm"); + if (sm_pos == std::string_view::npos) { return std::nullopt; } + + auto arch = std::string{fragment_name.substr(sm_pos + 1)}; + if (arch.size() <= 2) { return std::nullopt; } + arch.insert(2, "_"); + return std::format("-arch={}", arch); +} + +bool lto_supported_by_arch_option(std::string_view arch_option) +{ + constexpr auto prefix = std::string_view{"-arch=sm_"}; + if (!arch_option.starts_with(prefix)) { return true; } + + auto const arch = arch_option.substr(prefix.size()); + auto value = 0; + for (auto const c : arch) { + if (c < '0' || c > '9') { break; } + value = value * 10 + (c - '0'); + } + return value >= 80; +} + +std::vector env_nvjitlink_options() +{ + auto const* env = std::getenv("RTCX_NVJITLINK_OPTIONS"); + if (env == nullptr) { return {}; } + + std::istringstream stream{env}; + return {std::istream_iterator{stream}, std::istream_iterator{}}; +} + +bool link_timing_enabled() +{ + auto const* env = std::getenv("RTCX_LINK_TIMING"); + return env != nullptr && std::string_view{env} == "1"; +} + +char const* link_timing_meta() +{ + if (auto const* env = std::getenv("RTCX_LINK_TIMING_META")) { return env; } + return ""; +} + +void emit_link_timing(std::chrono::steady_clock::time_point start) +{ + auto const elapsed = + std::chrono::duration(std::chrono::steady_clock::now() - start).count(); + auto const* meta = link_timing_meta(); + if (std::string_view{meta}.empty()) { + std::fprintf(stderr, "RTCX_LINK_TIMING build_ms=%.3f\n", elapsed); + } else { + std::fprintf(stderr, "RTCX_LINK_TIMING build_ms=%.3f %s\n", elapsed, meta); + } +} + +} // namespace + +lto_kernel get_lto_linked_kernel(cache_t& cache, + std::string const& name, + std::string const& entry_name, + std::span memory_fragments, + std::span link_options) +{ + auto arch_option = current_device_arch_option(); + if (!memory_fragments.empty()) { + if (auto fragment_arch_option = arch_option_from_fragment_name(memory_fragments.front().name)) { + arch_option = *fragment_arch_option; + } + } + auto env_options = env_nvjitlink_options(); + + std::vector resolved_link_options; + resolved_link_options.reserve(link_options.size() + env_options.size() + 2); + if (lto_supported_by_arch_option(arch_option)) { resolved_link_options.push_back("-lto"); } + resolved_link_options.push_back(arch_option.c_str()); +#if defined(RTCX_NVJITLINK_PROFILE_O3) + resolved_link_options.push_back("-O3"); +#elif defined(RTCX_NVJITLINK_PROFILE_SPLIT_COMPILE) + resolved_link_options.push_back("-split-compile=0"); +#elif defined(RTCX_NVJITLINK_PROFILE_SPLIT_COMPILE_EXTENDED) + resolved_link_options.push_back("-split-compile-extended=0"); +#endif + for (auto const& option : env_options) { + resolved_link_options.push_back(option.c_str()); + } + resolved_link_options.insert( + resolved_link_options.end(), link_options.begin(), link_options.end()); + + auto update_string = [](sha256_context& hasher, std::string_view value) { + hasher.update(std::span{reinterpret_cast(value.data()), + value.size()}); + }; + auto update_scalar = [](sha256_context& hasher, auto const& value) { + hasher.update( + std::span{reinterpret_cast(&value), sizeof(value)}); + }; + + sha256_context hasher; + update_string(hasher, "rtcx_lto_kernel\n"); + update_string(hasher, name); + update_string(hasher, "\nentry="); + update_string(hasher, entry_name); + for (auto const option : resolved_link_options) { + update_string(hasher, "\noption="); + update_string(hasher, option); + } + for (auto const& f : memory_fragments) { + update_string(hasher, "\nfragment="); + update_string(hasher, f.name != nullptr ? std::string_view{f.name} : std::string_view{}); + update_string(hasher, "\ntype="); + auto const type = static_cast(f.type); + update_scalar(hasher, type); + update_string(hasher, "\nsize="); + auto const size = f.data.size(); + update_scalar(hasher, size); + update_string(hasher, "\ndata="); + hasher.update(f.data); + } + auto cache_key_sha = hasher.finalize(); + + auto compile = [&]() -> std::tuple { + link_params params; + params.name = name.c_str(); + params.output_type = binary_type::CUBIN; + params.memory_fragments = memory_fragments; + params.link_options = resolved_link_options; + auto buf = link_library(params); + auto raw_blob = blob_t::from_buffer(std::move(buf)); + auto shared = std::make_shared(std::move(raw_blob)); + auto lib = load_library(shared->view()); + return {lib, shared}; + }; + + auto const emit_timing = link_timing_enabled(); + auto const start = std::chrono::steady_clock::now(); + auto lib_future = + cache.get_or_add_library(cache_key_sha, library_compile_func::from_functor(compile)); + auto lib = lib_future.get(); + if (emit_timing) { emit_link_timing(start); } + return lto_kernel{lib, lib->get_kernel(entry_name.c_str())}; +} + +void launch_lto_linked_kernel(lto_kernel const& k, + cuda_dim3 grid_dim, + cuda_dim3 block_dim, + std::uint32_t shared_mem_bytes, + CUstream stream, + void** kernel_params) +{ + k.ref.launch(grid_dim, block_dim, shared_mem_bytes, stream, kernel_params); +} + } // namespace rtcx diff --git a/cpp/librtcx/rtcx.hpp b/cpp/librtcx/rtcx.hpp index d0ea4bff235c..f561ee9fdfe7 100644 --- a/cpp/librtcx/rtcx.hpp +++ b/cpp/librtcx/rtcx.hpp @@ -8,6 +8,7 @@ #include "sha256.hpp" #include +#include #include #include #include @@ -16,6 +17,7 @@ #include #include #include +#include #include #include #include @@ -1019,4 +1021,92 @@ rtcx::byte_buffer decompress_blob(std::span compressed_binar std::size_t uncompressed_size, std::string_view compression); +/// @brief Holds a linked library and a kernel reference. +/// Library MUST outlive kernel_ref. This struct owns both. +struct [[nodiscard]] lto_kernel { + library lib; ///< Shared ownership of the linked library + kernel_ref ref; ///< Reference to the specific kernel in lib +}; + +/// @brief Get a JIT-LTO linked kernel by linking memory fragments together. +/// Declaration: rtcx::get_lto_linked_kernel(cache_t& cache, ...) +/// @param cache The rtcx cache to use for caching linked libraries +/// @param name Debug name (used for cache key) +/// @param entry_name Name of the entry kernel in the linked library +/// @param memory_fragments Spans of binary fragment data to link +/// @param link_options Additional nvJitLink options (optional) +/// @return lto_kernel struct holding library lifetime + kernel reference +[[nodiscard]] __attribute__((visibility("default"))) lto_kernel +get_lto_linked_kernel(cache_t& cache, + std::string const& name, + std::string const& entry_name, + std::span memory_fragments, + std::span link_options = {}); + +/// @brief Launch a previously linked LTO kernel +__attribute__((visibility("default"))) void launch_lto_linked_kernel(lto_kernel const& k, + cuda_dim3 grid_dim, + cuda_dim3 block_dim, + std::uint32_t shared_mem_bytes, + CUstream stream, + void** kernel_params); + +/// @brief Type-safe launcher for a JIT-LTO linked kernel. +/// +/// Library MUST outlive kernel_ref; this struct guarantees that by storing the lto_kernel which +/// bundles both. +struct AlgorithmLauncher { + explicit AlgorithmLauncher(lto_kernel kernel) : kernel_(std::move(kernel)) {} + + AlgorithmLauncher(AlgorithmLauncher const&) = delete; + AlgorithmLauncher& operator=(AlgorithmLauncher const&) = delete; + AlgorithmLauncher(AlgorithmLauncher&&) = default; + AlgorithmLauncher& operator=(AlgorithmLauncher&&) = default; + + template + void dispatch( + CUstream stream, cuda_dim3 grid, cuda_dim3 block, std::uint32_t shared_mem, Args&&... args) + { + static_assert(std::is_same_v...)>, + "dispatch() argument types do not match the kernel function signature FuncT"); + std::array kernel_params = { + const_cast(static_cast(&args))...}; + kernel_.ref.launch(grid, block, shared_mem, stream, kernel_params.data()); + } + + private: + lto_kernel kernel_; +}; + +/// @brief Accumulates LTO-IR fragments for a single kernel entry point. +/// +/// Call get_launcher() to link and load the assembled fragments. Caching is handled +/// by get_lto_linked_kernel() internally via rtcx::cache_t; this struct itself holds +/// no cache. +/// +/// @note This struct is intended to be stack-allocated per call. Do NOT store as a +/// static or global — the underlying rtcx cache is process-global already. +/// +/// @note Runtime kernel arguments (e.g. a `nullable` bool, seed values, output pointers) +/// are kernel PARAMETERS passed via AlgorithmLauncher::dispatch(), NOT +/// planner concerns. The planner only decides WHICH code fragments to link. +struct AlgorithmPlanner { + explicit AlgorithmPlanner(std::string entrypoint) : entrypoint_(std::move(entrypoint)) {} + + void add_fragment(memory_fragment fragment) { fragments_.push_back(std::move(fragment)); } + + void add_link_option(char const* option) { link_options_.push_back(option); } + + [[nodiscard]] AlgorithmLauncher get_launcher(cache_t& cache) const + { + auto result = get_lto_linked_kernel(cache, entrypoint_, entrypoint_, fragments_, link_options_); + return AlgorithmLauncher{std::move(result)}; + } + + private: + std::string entrypoint_; + std::vector fragments_; + std::vector link_options_; +}; + } // namespace rtcx From d98b05774eb119c4e26e34698144e3033062dd92 Mon Sep 17 00:00:00 2001 From: Vyas Ramasubramani Date: Thu, 28 May 2026 15:36:25 -0700 Subject: [PATCH 2/4] Align rtcx fragment SDK documentation style --- cpp/librtcx/compute_matrix_product.cmake | 43 ++++++- cpp/librtcx/fragments.cmake | 147 +++++++++++++++++------ cpp/librtcx/rtcx.hpp | 110 ++++++++++++----- 3 files changed, 232 insertions(+), 68 deletions(-) diff --git a/cpp/librtcx/compute_matrix_product.cmake b/cpp/librtcx/compute_matrix_product.cmake index d643b1d26d5c..8cb22d2e38f0 100644 --- a/cpp/librtcx/compute_matrix_product.cmake +++ b/cpp/librtcx/compute_matrix_product.cmake @@ -4,10 +4,36 @@ # SPDX-License-Identifier: Apache-2.0 # cmake-format: on # ============================================================================= +# cmake-lint: disable=C0103,C0112,E1120 include_guard(GLOBAL) -# Compute the Cartesian product of a JSON matrix file or string and return it as a JSON array. +#[=======================================================================[.rst: +compute_matrix_product +---------------------- + +Computes the Cartesian product of a JSON matrix file or string. + +.. code-block:: cmake + + compute_matrix_product( + (MATRIX_JSON_FILE | MATRIX_JSON_STRING )) + +The generated matrix product is returned as a JSON array string in ````. + +``MATRIX_JSON_FILE`` + Path to a JSON matrix file. + +``MATRIX_JSON_STRING`` + JSON matrix content. + +Result Variables +^^^^^^^^^^^^^^^^ + +```` + JSON array containing the Cartesian product of the input matrix. + +#]=======================================================================] function(compute_matrix_product output_var) set(options) set(one_value MATRIX_JSON_FILE MATRIX_JSON_STRING) @@ -38,12 +64,23 @@ function(compute_matrix_product output_var) ) endfunction() -# Populate parent-scope CMake variables from one JSON matrix entry object. +#[=======================================================================[.rst: +populate_matrix_variables +------------------------- + +Populates CMake variables from a JSON matrix entry object. + +.. code-block:: cmake + + populate_matrix_variables() + +Each member in ```` is written as a variable in the caller's scope. + +#]=======================================================================] function(populate_matrix_variables matrix_json_entry) string(JSON len LENGTH "${matrix_json_entry}") 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}") diff --git a/cpp/librtcx/fragments.cmake b/cpp/librtcx/fragments.cmake index 42c30caa2c74..50e832754d06 100644 --- a/cpp/librtcx/fragments.cmake +++ b/cpp/librtcx/fragments.cmake @@ -2,23 +2,29 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # cmake-format: on -# cmake-lint: disable=C0113,E1120 - -# ============================================================================= -# librtcx JIT Fragment Compilation Infrastructure -# ============================================================================= -# Provides: add_fragment(), add_composite_fragment(), get_jit_fragment_architectures() These macros -# compile CUDA sources into per-architecture LTO-IR fatbins and embed them into a target via -# embed_blob(). Any CUDA project using librtcx can use these to register precompiled fragments. -# -# Usage: include(librtcx/fragments.cmake) add_embed(my_fragments) add_fragment(my_fragments FRAGMENT -# my_kernel SOURCE my_kernel.cu LINK_LIBRARIES dep1::dep1 dep2::dep2 INCLUDE_DIRECTORIES -# "${CMAKE_CURRENT_SOURCE_DIR}/include" COMPILE_DEFINITIONS MY_DISABLE_EXPORTS) embed(my_fragments -# COMPRESSION zstd OUTPUT_DIRECTORY ...) -# ============================================================================= - -# Fragment name convention: fragments must end with _sm for the runtime arch detection -# heuristic. +# cmake-lint: disable=C0112,C0113,E1120 + +#[=======================================================================[.rst: +get_jit_fragment_architectures +------------------------------ + +Returns the CUDA architectures used for JIT fragment compilation. + +.. code-block:: cmake + + get_jit_fragment_architectures() + +This function reads the ``rapids_cuda_architectures`` global property when it is available and falls +back to :cmake:variable:`CMAKE_CUDA_ARCHITECTURES` otherwise. Only real architectures and the +virtual architectures required for fragment compatibility are returned. + +Result Variables +^^^^^^^^^^^^^^^^ + +```` + List of CUDA architectures for JIT fragment targets. + +#]=======================================================================] function(get_jit_fragment_architectures OUT_VAR) get_property(_RAPIDS_CUDA_ARCHITECTURES GLOBAL PROPERTY rapids_cuda_architectures) if(_RAPIDS_CUDA_ARCHITECTURES) @@ -53,13 +59,59 @@ function(get_jit_fragment_architectures OUT_VAR) ) endfunction() -# This macro is used to create object libraries for JIT compilation fragments, and embed them as -# per-architecture fatbins in the final library. It compiles the specified source file with the -# appropriate flags to generate fatbins containing the specified kernel, and then embeds those -# fatbins in the final library with metadata that allows them to be looked up at runtime. -# -# For architectures below sm80, nvJitLink fails to accept LTO fatbins in CUDA 12.9. Disable IPO to -# emit standard device-code fatbins instead. +#[=======================================================================[.rst: +add_fragment +------------ + +Compiles and embeds a CUDA source file as per-architecture JIT fragments. + +.. code-block:: cmake + + add_fragment( + FRAGMENT + SOURCE + [ENTRY_NAME ] + [KERNEL_ONLY] + [DEFINITIONS ...] + [ARRAY_IDS ...] + [ARRAY_VALUES ...] + [LINK_LIBRARIES ...] + [INCLUDE_DIRECTORIES ...] + [COMPILE_DEFINITIONS ...]) + +The generated fragment files are embedded in ```` with names of the form +``fragments/_sm.fatbin``. Fragment identifiers must use the ``_sm`` suffix so the +runtime can select fragments for the current device architecture. + +``FRAGMENT`` + Fragment name used for the generated object targets and embedded blob identifiers. + +``SOURCE`` + CUDA source file compiled for each fragment architecture. + +``ENTRY_NAME`` + Kernel entry symbol. Defaults to ``kernel_entry``. + +``KERNEL_ONLY`` + Passes ``--kernels-used=`` to nvlink so each fragment fatbin only retains ``ENTRY_NAME`` + and device code reachable from that kernel. + +``DEFINITIONS`` + Additional values passed to :cmake:command:`target_compile_definitions`. + +``ARRAY_IDS`` and ``ARRAY_VALUES`` + Additional metadata arrays passed to :cmake:command:`embed_blob`. + +``LINK_LIBRARIES`` + Targets linked by each fragment object target. + +``INCLUDE_DIRECTORIES`` + Include directories used by each fragment object target. + +``COMPILE_DEFINITIONS`` + Compile definitions used by each fragment object target. + +#]=======================================================================] macro(add_fragment) set(TARGET ${ARGV0}) set(OPTIONS KERNEL_ONLY) @@ -98,7 +150,7 @@ macro(add_fragment) ) if(ARG_KERNEL_ONLY) - # ensure that the FATBIN symbols only contain the specified kernel + # Ensure that the FATBIN symbols only contain the specified kernel. target_compile_options(${OBJECT_ID} PRIVATE -Xnvlink=--kernels-used=${ARG_ENTRY_NAME}) endif() @@ -143,13 +195,32 @@ macro(add_fragment) endforeach() endmacro() -# ============================================================================= -# This macro pre-links individual LTO-IR fragment fatbins at CMake build time into a single -# composite fatbin using nvJitLink, then embeds it via embed_blob(). If nvJitLink is not found, the -# composite is skipped with a warning. -# ============================================================================= -# NOTE: add_composite_fragment has no call sites as of 2026-05-27. It is aspirational -# infrastructure. +#[=======================================================================[.rst: +add_composite_fragment +---------------------- + +Pre-links LTO-IR fragments into a composite fatbin and embeds it in a target. + +.. code-block:: cmake + + add_composite_fragment( + + FRAGMENTS ... + [ENTRY_NAME ]) + +If a compatible ``nvJitLink`` executable is not available, the composite fragment is skipped. + +.. note:: + This macro currently has no call sites. It is retained as supporting infrastructure for future + pre-linked fragments. + +``FRAGMENTS`` + Fragment object targets to pre-link into the composite fatbin. + +``ENTRY_NAME`` + Kernel entry symbol. Defaults to ``kernel_entry``. + +#]=======================================================================] macro(add_composite_fragment) set(_ACF_TARGET ${ARGV0}) set(_ACF_COMPOSITE_NAME ${ARGV1}) @@ -182,20 +253,20 @@ macro(add_composite_fragment) if(NOT NVJITLINK_EXECUTABLE) message(STATUS "nvJitLink tool not found — composite fragment '${_ACF_COMPOSITE_NAME}' skipped") else() - # Collect input fatbin files from OBJECT targets of registered fragments + # Collect input fatbin files from OBJECT targets of registered fragments. set(_ACF_INPUT_FILES "") set(_ACF_INPUT_DEPS "") foreach(_ACF_FRAG_ID IN LISTS _ACF_FRAGMENTS) - # Each fragment was registered as OBJECT library named ${TARGET}_${FRAGMENT} + # Each fragment was registered as an OBJECT library named ${TARGET}_${FRAGMENT}. set(_ACF_OBJ_TARGET "${_ACF_TARGET}_${_ACF_FRAG_ID}") list(APPEND _ACF_INPUT_FILES "$") list(APPEND _ACF_INPUT_DEPS "${_ACF_OBJ_TARGET}") endforeach() - # Output composite fatbin path + # Output composite fatbin path. set(_ACF_COMPOSITE_OUTPUT "${CMAKE_CURRENT_BINARY_DIR}/fragments/${_ACF_COMPOSITE_NAME}.fatbin") - # Custom command to run nvJitLink to pre-link fragments at build time + # Custom command to run nvJitLink to pre-link fragments at build time. add_custom_command( OUTPUT "${_ACF_COMPOSITE_OUTPUT}" COMMAND ${CMAKE_COMMAND} -E make_directory "${CMAKE_CURRENT_BINARY_DIR}/fragments" @@ -205,12 +276,12 @@ macro(add_composite_fragment) VERBATIM ) - # Custom target to ensure the composite is built before embedding + # Custom target to ensure the composite is built before embedding. add_custom_target( "${_ACF_TARGET}_composite_${_ACF_COMPOSITE_NAME}" ALL DEPENDS "${_ACF_COMPOSITE_OUTPUT}" ) - # Embed the composite fatbin via the existing embed_blob infrastructure + # Embed the composite fatbin via the existing embed_blob infrastructure. embed_blob( ${_ACF_TARGET} FILE "${_ACF_COMPOSITE_OUTPUT}" DEST "fragments/${_ACF_COMPOSITE_NAME}.fatbin" ID "${_ACF_COMPOSITE_NAME}" diff --git a/cpp/librtcx/rtcx.hpp b/cpp/librtcx/rtcx.hpp index f561ee9fdfe7..50212607b1b6 100644 --- a/cpp/librtcx/rtcx.hpp +++ b/cpp/librtcx/rtcx.hpp @@ -1021,21 +1021,30 @@ rtcx::byte_buffer decompress_blob(std::span compressed_binar std::size_t uncompressed_size, std::string_view compression); -/// @brief Holds a linked library and a kernel reference. -/// Library MUST outlive kernel_ref. This struct owns both. +/** + * @brief Linked library and kernel reference. + * + * The library must outlive the kernel reference. This type owns both objects to preserve that + * lifetime relationship. + */ struct [[nodiscard]] lto_kernel { library lib; ///< Shared ownership of the linked library kernel_ref ref; ///< Reference to the specific kernel in lib }; -/// @brief Get a JIT-LTO linked kernel by linking memory fragments together. -/// Declaration: rtcx::get_lto_linked_kernel(cache_t& cache, ...) -/// @param cache The rtcx cache to use for caching linked libraries -/// @param name Debug name (used for cache key) -/// @param entry_name Name of the entry kernel in the linked library -/// @param memory_fragments Spans of binary fragment data to link -/// @param link_options Additional nvJitLink options (optional) -/// @return lto_kernel struct holding library lifetime + kernel reference +/** + * @brief Links memory fragments and returns a loaded kernel. + * + * The linked library is cached using the supplied `cache` and a key derived from the kernel name, + * entry point, link options, and fragment contents. + * + * @param cache Cache used for linked libraries + * @param name Name used for diagnostics and cache key generation + * @param entry_name Name of the kernel entry point in the linked library + * @param memory_fragments Binary fragments to link + * @param link_options Additional nvJitLink options + * @return Linked library and kernel reference + */ [[nodiscard]] __attribute__((visibility("default"))) lto_kernel get_lto_linked_kernel(cache_t& cache, std::string const& name, @@ -1043,7 +1052,16 @@ get_lto_linked_kernel(cache_t& cache, std::span memory_fragments, std::span link_options = {}); -/// @brief Launch a previously linked LTO kernel +/** + * @brief Launches a previously linked LTO kernel. + * + * @param k Linked kernel to launch + * @param grid_dim CUDA grid dimensions + * @param block_dim CUDA block dimensions + * @param shared_mem_bytes Dynamic shared memory size in bytes + * @param stream CUDA stream used for the kernel launch + * @param kernel_params Kernel argument pointer array + */ __attribute__((visibility("default"))) void launch_lto_linked_kernel(lto_kernel const& k, cuda_dim3 grid_dim, cuda_dim3 block_dim, @@ -1051,11 +1069,17 @@ __attribute__((visibility("default"))) void launch_lto_linked_kernel(lto_kernel CUstream stream, void** kernel_params); -/// @brief Type-safe launcher for a JIT-LTO linked kernel. -/// -/// Library MUST outlive kernel_ref; this struct guarantees that by storing the lto_kernel which -/// bundles both. +/** + * @brief Type-safe launcher for a JIT-LTO linked kernel. + * + * This type stores the linked kernel object, which keeps the linked library alive for all launches. + */ struct AlgorithmLauncher { + /** + * @brief Constructs a launcher from a linked kernel. + * + * @param kernel Linked kernel to launch + */ explicit AlgorithmLauncher(lto_kernel kernel) : kernel_(std::move(kernel)) {} AlgorithmLauncher(AlgorithmLauncher const&) = delete; @@ -1063,6 +1087,17 @@ struct AlgorithmLauncher { AlgorithmLauncher(AlgorithmLauncher&&) = default; AlgorithmLauncher& operator=(AlgorithmLauncher&&) = default; + /** + * @brief Launches the linked kernel with type-checked arguments. + * + * @tparam FuncT Function signature expected by the linked kernel + * @tparam Args Kernel argument types + * @param stream CUDA stream used for the kernel launch + * @param grid CUDA grid dimensions + * @param block CUDA block dimensions + * @param shared_mem Dynamic shared memory size in bytes + * @param args Kernel arguments + */ template void dispatch( CUstream stream, cuda_dim3 grid, cuda_dim3 block, std::uint32_t shared_mem, Args&&... args) @@ -1078,25 +1113,46 @@ struct AlgorithmLauncher { lto_kernel kernel_; }; -/// @brief Accumulates LTO-IR fragments for a single kernel entry point. -/// -/// Call get_launcher() to link and load the assembled fragments. Caching is handled -/// by get_lto_linked_kernel() internally via rtcx::cache_t; this struct itself holds -/// no cache. -/// -/// @note This struct is intended to be stack-allocated per call. Do NOT store as a -/// static or global — the underlying rtcx cache is process-global already. -/// -/// @note Runtime kernel arguments (e.g. a `nullable` bool, seed values, output pointers) -/// are kernel PARAMETERS passed via AlgorithmLauncher::dispatch(), NOT -/// planner concerns. The planner only decides WHICH code fragments to link. +/** + * @brief Accumulates LTO-IR fragments for a single kernel entry point. + * + * Call `get_launcher()` to link and load the assembled fragments. Caching is handled by + * `get_lto_linked_kernel()` using the supplied `rtcx::cache_t`; this object does not own a cache. + * + * Runtime kernel arguments, such as nullable flags, seed values, and output pointers, are passed to + * `AlgorithmLauncher::dispatch()`. This type only selects which code fragments to link. + * + * @note This type is intended to be stack-allocated per call. Do not store it as static or global + * state. + */ struct AlgorithmPlanner { + /** + * @brief Constructs a planner for a kernel entry point. + * + * @param entrypoint Kernel entry point name + */ explicit AlgorithmPlanner(std::string entrypoint) : entrypoint_(std::move(entrypoint)) {} + /** + * @brief Adds a fragment to the link plan. + * + * @param fragment Fragment to link + */ void add_fragment(memory_fragment fragment) { fragments_.push_back(std::move(fragment)); } + /** + * @brief Adds an nvJitLink option to the link plan. + * + * @param option Null-terminated nvJitLink option string + */ void add_link_option(char const* option) { link_options_.push_back(option); } + /** + * @brief Links the planned fragments and returns a launcher. + * + * @param cache Cache used for linked libraries + * @return Type-safe launcher for the linked kernel + */ [[nodiscard]] AlgorithmLauncher get_launcher(cache_t& cache) const { auto result = get_lto_linked_kernel(cache, entrypoint_, entrypoint_, fragments_, link_options_); From 729ed610dcb805961903856cb5247d2e0b9dc2cf Mon Sep 17 00:00:00 2001 From: Vyas Ramasubramani Date: Thu, 28 May 2026 15:36:27 -0700 Subject: [PATCH 3/4] Remove unused composite fragment helper --- cpp/librtcx/fragments.cmake | 94 ------------------------------------- 1 file changed, 94 deletions(-) diff --git a/cpp/librtcx/fragments.cmake b/cpp/librtcx/fragments.cmake index 50e832754d06..dc76386af779 100644 --- a/cpp/librtcx/fragments.cmake +++ b/cpp/librtcx/fragments.cmake @@ -194,97 +194,3 @@ macro(add_fragment) ) endforeach() endmacro() - -#[=======================================================================[.rst: -add_composite_fragment ----------------------- - -Pre-links LTO-IR fragments into a composite fatbin and embeds it in a target. - -.. code-block:: cmake - - add_composite_fragment( - - FRAGMENTS ... - [ENTRY_NAME ]) - -If a compatible ``nvJitLink`` executable is not available, the composite fragment is skipped. - -.. note:: - This macro currently has no call sites. It is retained as supporting infrastructure for future - pre-linked fragments. - -``FRAGMENTS`` - Fragment object targets to pre-link into the composite fatbin. - -``ENTRY_NAME`` - Kernel entry symbol. Defaults to ``kernel_entry``. - -#]=======================================================================] -macro(add_composite_fragment) - set(_ACF_TARGET ${ARGV0}) - set(_ACF_COMPOSITE_NAME ${ARGV1}) - set(_ACF_OPTIONS "") - set(_ACF_ONE_VALUE_ARGS ENTRY_NAME) - set(_ACF_MULTI_VALUE_ARGS FRAGMENTS) - cmake_parse_arguments( - _ACF "${_ACF_OPTIONS}" "${_ACF_ONE_VALUE_ARGS}" "${_ACF_MULTI_VALUE_ARGS}" ${ARGN} - ) - - if(NOT _ACF_COMPOSITE_NAME) - message(FATAL_ERROR "add_composite_fragment requires COMPOSITE_NAME argument") - endif() - - if(NOT _ACF_FRAGMENTS) - message(FATAL_ERROR "add_composite_fragment requires FRAGMENTS argument") - endif() - - if(NOT _ACF_ENTRY_NAME) - set(_ACF_ENTRY_NAME "kernel_entry") - endif() - - # nvJitLink is a library API, not a CLI tool. Look for a helper binary that wraps it. If not - # found, the composite fragment is skipped gracefully. - find_program( - NVJITLINK_EXECUTABLE nvJitLink HINTS "${CUDAToolkit_BIN_DIR}" - "${CMAKE_CUDA_COMPILER_TOOLKIT_ROOT}/bin" ENV PATH - ) - - if(NOT NVJITLINK_EXECUTABLE) - message(STATUS "nvJitLink tool not found — composite fragment '${_ACF_COMPOSITE_NAME}' skipped") - else() - # Collect input fatbin files from OBJECT targets of registered fragments. - set(_ACF_INPUT_FILES "") - set(_ACF_INPUT_DEPS "") - foreach(_ACF_FRAG_ID IN LISTS _ACF_FRAGMENTS) - # Each fragment was registered as an OBJECT library named ${TARGET}_${FRAGMENT}. - set(_ACF_OBJ_TARGET "${_ACF_TARGET}_${_ACF_FRAG_ID}") - list(APPEND _ACF_INPUT_FILES "$") - list(APPEND _ACF_INPUT_DEPS "${_ACF_OBJ_TARGET}") - endforeach() - - # Output composite fatbin path. - set(_ACF_COMPOSITE_OUTPUT "${CMAKE_CURRENT_BINARY_DIR}/fragments/${_ACF_COMPOSITE_NAME}.fatbin") - - # Custom command to run nvJitLink to pre-link fragments at build time. - add_custom_command( - OUTPUT "${_ACF_COMPOSITE_OUTPUT}" - COMMAND ${CMAKE_COMMAND} -E make_directory "${CMAKE_CURRENT_BINARY_DIR}/fragments" - COMMAND ${NVJITLINK_EXECUTABLE} -o "${_ACF_COMPOSITE_OUTPUT}" ${_ACF_INPUT_FILES} - DEPENDS ${_ACF_INPUT_FILES} - COMMENT "Pre-linking composite fragment: ${_ACF_COMPOSITE_NAME}" - VERBATIM - ) - - # Custom target to ensure the composite is built before embedding. - add_custom_target( - "${_ACF_TARGET}_composite_${_ACF_COMPOSITE_NAME}" ALL DEPENDS "${_ACF_COMPOSITE_OUTPUT}" - ) - - # Embed the composite fatbin via the existing embed_blob infrastructure. - embed_blob( - ${_ACF_TARGET} FILE "${_ACF_COMPOSITE_OUTPUT}" DEST "fragments/${_ACF_COMPOSITE_NAME}.fatbin" - ID "${_ACF_COMPOSITE_NAME}" - ) - endif() -endmacro() From a01ba9b592dd214953e2ceccf5e8b0afe2764a17 Mon Sep 17 00:00:00 2001 From: Vyas Ramasubramani Date: Thu, 28 May 2026 15:36:29 -0700 Subject: [PATCH 4/4] Remove unnecessary LTO arch < 80 guard MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit nvJitLink 12.9 successfully compiles, links, and executes LTO fatbins for sm_70. The guard was speculative — tested empirically on V100 (sm_70) and confirmed working. Arch deprecation is handled by get_jit_fragment_architectures dropping old archs from default lists in future toolkit versions. --- cpp/librtcx/fragments.cmake | 68 +++---------------------------------- cpp/librtcx/rtcx.cpp | 16 +-------- 2 files changed, 6 insertions(+), 78 deletions(-) diff --git a/cpp/librtcx/fragments.cmake b/cpp/librtcx/fragments.cmake index dc76386af779..0801f57c718c 100644 --- a/cpp/librtcx/fragments.cmake +++ b/cpp/librtcx/fragments.cmake @@ -4,61 +4,6 @@ # cmake-format: on # cmake-lint: disable=C0112,C0113,E1120 -#[=======================================================================[.rst: -get_jit_fragment_architectures ------------------------------- - -Returns the CUDA architectures used for JIT fragment compilation. - -.. code-block:: cmake - - get_jit_fragment_architectures() - -This function reads the ``rapids_cuda_architectures`` global property when it is available and falls -back to :cmake:variable:`CMAKE_CUDA_ARCHITECTURES` otherwise. Only real architectures and the -virtual architectures required for fragment compatibility are returned. - -Result Variables -^^^^^^^^^^^^^^^^ - -```` - List of CUDA architectures for JIT fragment targets. - -#]=======================================================================] -function(get_jit_fragment_architectures OUT_VAR) - get_property(_RAPIDS_CUDA_ARCHITECTURES GLOBAL PROPERTY rapids_cuda_architectures) - if(_RAPIDS_CUDA_ARCHITECTURES) - set(_ARCHITECTURES ${_RAPIDS_CUDA_ARCHITECTURES}) - else() - set(_ARCHITECTURES ${CMAKE_CUDA_ARCHITECTURES}) - endif() - - set(_FRAGMENT_ARCHITECTURES) - foreach(_ARCH IN LISTS _ARCHITECTURES) - if(_ARCH STREQUAL "RAPIDS" OR _ARCH STREQUAL "ALL") - if(CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 13.0.0) - list(APPEND _FRAGMENT_ARCHITECTURES 75-real 80-real 86-real 90a-real 100f-real 120a-real - 120 - ) - elseif(CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 12.9.0) - list(APPEND _FRAGMENT_ARCHITECTURES 70-real 75-real 80-real 86-real 90a-real 100f-real - 120a-real 120 - ) - else() - list(APPEND _FRAGMENT_ARCHITECTURES 70-real 75-real 80-real 86-real 90a-real 90-virtual) - endif() - elseif(_ARCH MATCHES "^[0-9]+[a-z]*(-real)?$") - list(APPEND _FRAGMENT_ARCHITECTURES ${_ARCH}) - endif() - endforeach() - - list(REMOVE_DUPLICATES _FRAGMENT_ARCHITECTURES) - set(${OUT_VAR} - ${_FRAGMENT_ARCHITECTURES} - PARENT_SCOPE - ) -endfunction() - #[=======================================================================[.rst: add_fragment ------------ @@ -133,15 +78,12 @@ macro(add_fragment) set(ARG_ENTRY_NAME "kernel_entry") endif() - get_jit_fragment_architectures(FRAGMENT_ARCHITECTURES) + get_property(FRAGMENT_ARCHITECTURES GLOBAL PROPERTY rapids_cuda_architectures) + if(NOT FRAGMENT_ARCHITECTURES) + set(FRAGMENT_ARCHITECTURES ${CMAKE_CUDA_ARCHITECTURES}) + endif() foreach(FRAGMENT_ARCH IN LISTS FRAGMENT_ARCHITECTURES) string(REGEX REPLACE "-.*$" "" FRAGMENT_ARCH_ID "${FRAGMENT_ARCH}") - string(REGEX REPLACE "[^0-9]" "" FRAGMENT_ARCH_NUM "${FRAGMENT_ARCH_ID}") - if(FRAGMENT_ARCH_NUM AND FRAGMENT_ARCH_NUM LESS 80) - set(FRAGMENT_LTO OFF) - else() - set(FRAGMENT_LTO ON) - endif() set(OBJECT_ID ${TARGET}_${ARG_FRAGMENT}_sm${FRAGMENT_ARCH_ID}) add_library(${OBJECT_ID} OBJECT ${ARG_SOURCE}) target_compile_options( @@ -161,7 +103,7 @@ macro(add_fragment) CUDA_SEPARABLE_COMPILATION ON CUDA_FATBIN_COMPILATION ON POSITION_INDEPENDENT_CODE ON - INTERPROCEDURAL_OPTIMIZATION ${FRAGMENT_LTO} + INTERPROCEDURAL_OPTIMIZATION ON CXX_STANDARD 20 CXX_STANDARD_REQUIRED ON CXX_EXTENSIONS ON diff --git a/cpp/librtcx/rtcx.cpp b/cpp/librtcx/rtcx.cpp index 0090e6663337..b824b21f61ce 100644 --- a/cpp/librtcx/rtcx.cpp +++ b/cpp/librtcx/rtcx.cpp @@ -1374,20 +1374,6 @@ std::optional arch_option_from_fragment_name(char const* name) return std::format("-arch={}", arch); } -bool lto_supported_by_arch_option(std::string_view arch_option) -{ - constexpr auto prefix = std::string_view{"-arch=sm_"}; - if (!arch_option.starts_with(prefix)) { return true; } - - auto const arch = arch_option.substr(prefix.size()); - auto value = 0; - for (auto const c : arch) { - if (c < '0' || c > '9') { break; } - value = value * 10 + (c - '0'); - } - return value >= 80; -} - std::vector env_nvjitlink_options() { auto const* env = std::getenv("RTCX_NVJITLINK_OPTIONS"); @@ -1439,7 +1425,7 @@ lto_kernel get_lto_linked_kernel(cache_t& cache, std::vector resolved_link_options; resolved_link_options.reserve(link_options.size() + env_options.size() + 2); - if (lto_supported_by_arch_option(arch_option)) { resolved_link_options.push_back("-lto"); } + resolved_link_options.push_back("-lto"); resolved_link_options.push_back(arch_option.c_str()); #if defined(RTCX_NVJITLINK_PROFILE_O3) resolved_link_options.push_back("-O3");