From b651d59235275c4111f4edbf8647a0a80cff7ccd Mon Sep 17 00:00:00 2001 From: Arha Gatram Date: Thu, 9 Jul 2026 10:07:11 -0700 Subject: [PATCH 1/3] use librtcx for jit-lto --- cpp/CMakeLists.txt | 11 +- .../modules/generate_jit_lto_kernels.cmake | 136 ------------------ cpp/cmake/modules/register_fatbin.cpp.in | 22 --- cpp/cmake/thirdparty/get_rtcx.cmake | 34 +++++ .../cuvs/detail/jit_lto/AlgorithmLauncher.hpp | 60 -------- .../cuvs/detail/jit_lto/AlgorithmPlanner.hpp | 60 -------- .../cuvs/detail/jit_lto/FragmentEntry.hpp | 64 --------- .../jit_lto/NVRTCLTOFragmentCompiler.hpp | 29 ---- .../cuvs/detail/jit_lto/nvjitlink_checker.hpp | 11 -- cpp/src/detail/jit_lto/AlgorithmLauncher.cpp | 66 --------- cpp/src/detail/jit_lto/AlgorithmPlanner.cpp | 116 --------------- cpp/src/detail/jit_lto/FragmentEntry.cpp | 14 -- .../jit_lto/NVRTCLTOFragmentCompiler.cpp | 108 -------------- cpp/src/detail/jit_lto/nvjitlink_checker.cpp | 27 ---- .../pairwise_matrix_planner.hpp | 4 +- .../cagra_jit_launcher_factory.hpp | 2 +- .../jit_lto_kernels/cagra_planner_base.hpp | 2 +- .../jit_lto_kernels/sample_filter_udf.cuh | 2 +- .../search_multi_cta_kernel_launcher_jit.cuh | 4 +- .../search_multi_kernel_launcher_jit.cuh | 4 +- .../search_single_cta_kernel_launcher_jit.cuh | 4 +- .../interleaved_scan_planner.hpp | 6 +- .../ivf_flat_interleaved_scan_jit.cuh | 4 +- .../compute_similarity_planner.hpp | 6 +- .../ivf_pq/ivf_pq_compute_similarity.hpp | 4 +- ...oducts_with_bitwise_block_sort_planner.hpp | 6 +- ...te_inner_products_with_bitwise_planner.hpp | 6 +- ...ucts_with_lut16_opt_block_sort_planner.hpp | 6 +- ..._inner_products_with_lut16_opt_planner.hpp | 6 +- ...r_products_with_lut_block_sort_planner.hpp | 6 +- ...ompute_inner_products_with_lut_planner.hpp | 6 +- .../jit_lto_kernels/launcher_factory.hpp | 4 +- .../detail/jit_lto_kernels/scan_planner.hpp | 4 +- fern/pages/jit_lto_guide.md | 4 +- 34 files changed, 83 insertions(+), 765 deletions(-) delete mode 100644 cpp/cmake/modules/generate_jit_lto_kernels.cmake delete mode 100644 cpp/cmake/modules/register_fatbin.cpp.in create mode 100644 cpp/cmake/thirdparty/get_rtcx.cmake delete mode 100644 cpp/include/cuvs/detail/jit_lto/AlgorithmLauncher.hpp delete mode 100644 cpp/include/cuvs/detail/jit_lto/AlgorithmPlanner.hpp delete mode 100644 cpp/include/cuvs/detail/jit_lto/FragmentEntry.hpp delete mode 100644 cpp/include/cuvs/detail/jit_lto/NVRTCLTOFragmentCompiler.hpp delete mode 100644 cpp/include/cuvs/detail/jit_lto/nvjitlink_checker.hpp delete mode 100644 cpp/src/detail/jit_lto/AlgorithmLauncher.cpp delete mode 100644 cpp/src/detail/jit_lto/AlgorithmPlanner.cpp delete mode 100644 cpp/src/detail/jit_lto/FragmentEntry.cpp delete mode 100644 cpp/src/detail/jit_lto/NVRTCLTOFragmentCompiler.cpp delete mode 100644 cpp/src/detail/jit_lto/nvjitlink_checker.cpp diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 96c4e70cfd..32dfb898ef 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -353,8 +353,9 @@ if(NOT BUILD_CPU_ONLY) set(JIT_LTO_TARGET_ARCHITECTURE "75-real") endif() - # Generate interleaved scan kernel files at build time - include(cmake/modules/generate_jit_lto_kernels.cmake) + # Generate interleaved scan kernel files at build time add_subdirectory(../../librtcx + # src/detail/librtcx) + include(cmake/thirdparty/get_rtcx.cmake) add_library(jit_lto_kernel_usage_requirements INTERFACE) target_include_directories( @@ -1331,11 +1332,6 @@ if(NOT BUILD_CPU_ONLY) src/core/omp_wrapper.cpp src/util/file_io.cpp src/util/host_memory.cpp - src/detail/jit_lto/AlgorithmLauncher.cpp - src/detail/jit_lto/AlgorithmPlanner.cpp - src/detail/jit_lto/FragmentEntry.cpp - src/detail/jit_lto/nvjitlink_checker.cpp - src/detail/jit_lto/NVRTCLTOFragmentCompiler.cpp src/distance/detail/kernels/gram_matrix.cu src/distance/detail/kernels/kernel_factory.cu src/distance/detail/kernels/kernel_matrices.cu @@ -1460,6 +1456,7 @@ if(NOT BUILD_CPU_ONLY) PRIVATE cuvs::cuvs_cpp_headers cuco::cuco nvidia::cutlass::cutlass + rtcx::rtcx ${CUVS_CTK_MATH_DEPENDENCIES} $ $ diff --git a/cpp/cmake/modules/generate_jit_lto_kernels.cmake b/cpp/cmake/modules/generate_jit_lto_kernels.cmake deleted file mode 100644 index f427fc4e5e..0000000000 --- a/cpp/cmake/modules/generate_jit_lto_kernels.cmake +++ /dev/null @@ -1,136 +0,0 @@ -# ============================================================================= -# 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/cpp/cmake/modules/register_fatbin.cpp.in b/cpp/cmake/modules/register_fatbin.cpp.in deleted file mode 100644 index b154132358..0000000000 --- a/cpp/cmake/modules/register_fatbin.cpp.in +++ /dev/null @@ -1,22 +0,0 @@ -/* - * 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); diff --git a/cpp/cmake/thirdparty/get_rtcx.cmake b/cpp/cmake/thirdparty/get_rtcx.cmake new file mode 100644 index 0000000000..7a4476c903 --- /dev/null +++ b/cpp/cmake/thirdparty/get_rtcx.cmake @@ -0,0 +1,34 @@ +#============================================================================== +# cmake-format: off +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# cmake-format: on +#=============================================================================- + +# This function finds rtcx +function(find_and_configure_rtcx VERSION) + + # Ensure rtcx installs its targets unconditionally. In shared builds the static library is + # absorbed into libcuvs.so, but CMake still requires the target to be in an export set for + # install(EXPORT) validation. In static builds consumers need to link librtcx.a directly. + set(RTCX_INSTALL ON) + + rapids_cpm_find( + rtcx ${VERSION} + GLOBAL_TARGETS rtcx::rtcx + CPM_ARGS + GIT_REPOSITORY https://github.com/arhag23/librtcx.git + GIT_TAG dc819febd2539c086e144b2632c3f5857c5d74b8 + GIT_SHALLOW FALSE + EXCLUDE_FROM_ALL TRUE + ) + + # When CPM fetches from source (add_subdirectory), embed.cmake is not auto-included. Include it + # explicitly so add_embed/embed_includes/embed functions are available. + if(rtcx_ADDED OR DEFINED CPM_rtcx_SOURCE) + include("${rtcx_SOURCE_DIR}/generate_jit_lto_kernels.cmake") + endif() +endfunction() + +set(RTCX_MIN_VERSION_cuvs "0.1") +find_and_configure_rtcx(${RTCX_MIN_VERSION_cuvs}) diff --git a/cpp/include/cuvs/detail/jit_lto/AlgorithmLauncher.hpp b/cpp/include/cuvs/detail/jit_lto/AlgorithmLauncher.hpp deleted file mode 100644 index 10418c430a..0000000000 --- a/cpp/include/cuvs/detail/jit_lto/AlgorithmLauncher.hpp +++ /dev/null @@ -1,60 +0,0 @@ -/* - * 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/cpp/include/cuvs/detail/jit_lto/AlgorithmPlanner.hpp b/cpp/include/cuvs/detail/jit_lto/AlgorithmPlanner.hpp deleted file mode 100644 index 7f275b1285..0000000000 --- a/cpp/include/cuvs/detail/jit_lto/AlgorithmPlanner.hpp +++ /dev/null @@ -1,60 +0,0 @@ -/* - * 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/cpp/include/cuvs/detail/jit_lto/FragmentEntry.hpp b/cpp/include/cuvs/detail/jit_lto/FragmentEntry.hpp deleted file mode 100644 index 35aa46633c..0000000000 --- a/cpp/include/cuvs/detail/jit_lto/FragmentEntry.hpp +++ /dev/null @@ -1,64 +0,0 @@ -/* - * 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/cpp/include/cuvs/detail/jit_lto/NVRTCLTOFragmentCompiler.hpp b/cpp/include/cuvs/detail/jit_lto/NVRTCLTOFragmentCompiler.hpp deleted file mode 100644 index 6fb3d4a80c..0000000000 --- a/cpp/include/cuvs/detail/jit_lto/NVRTCLTOFragmentCompiler.hpp +++ /dev/null @@ -1,29 +0,0 @@ -/* - * 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/cpp/include/cuvs/detail/jit_lto/nvjitlink_checker.hpp b/cpp/include/cuvs/detail/jit_lto/nvjitlink_checker.hpp deleted file mode 100644 index 12b062d795..0000000000 --- a/cpp/include/cuvs/detail/jit_lto/nvjitlink_checker.hpp +++ /dev/null @@ -1,11 +0,0 @@ -/* - * 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/cpp/src/detail/jit_lto/AlgorithmLauncher.cpp b/cpp/src/detail/jit_lto/AlgorithmLauncher.cpp deleted file mode 100644 index b300f4d367..0000000000 --- a/cpp/src/detail/jit_lto/AlgorithmLauncher.cpp +++ /dev/null @@ -1,66 +0,0 @@ -/* - * 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/cpp/src/detail/jit_lto/AlgorithmPlanner.cpp b/cpp/src/detail/jit_lto/AlgorithmPlanner.cpp deleted file mode 100644 index 7416ea396d..0000000000 --- a/cpp/src/detail/jit_lto/AlgorithmPlanner.cpp +++ /dev/null @@ -1,116 +0,0 @@ -/* - * 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/cpp/src/detail/jit_lto/FragmentEntry.cpp b/cpp/src/detail/jit_lto/FragmentEntry.cpp deleted file mode 100644 index bf0893c8a6..0000000000 --- a/cpp/src/detail/jit_lto/FragmentEntry.cpp +++ /dev/null @@ -1,14 +0,0 @@ -/* - * 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/cpp/src/detail/jit_lto/NVRTCLTOFragmentCompiler.cpp b/cpp/src/detail/jit_lto/NVRTCLTOFragmentCompiler.cpp deleted file mode 100644 index 72b56ace2b..0000000000 --- a/cpp/src/detail/jit_lto/NVRTCLTOFragmentCompiler.cpp +++ /dev/null @@ -1,108 +0,0 @@ -/* - * 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/cpp/src/detail/jit_lto/nvjitlink_checker.cpp b/cpp/src/detail/jit_lto/nvjitlink_checker.cpp deleted file mode 100644 index 95f4725362..0000000000 --- a/cpp/src/detail/jit_lto/nvjitlink_checker.cpp +++ /dev/null @@ -1,27 +0,0 @@ -/* - * 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/cpp/src/distance/detail/pairwise_matrix/jit_lto_kernels/pairwise_matrix_planner.hpp b/cpp/src/distance/detail/pairwise_matrix/jit_lto_kernels/pairwise_matrix_planner.hpp index 0d00b3eca6..bdc1942cbd 100644 --- a/cpp/src/distance/detail/pairwise_matrix/jit_lto_kernels/pairwise_matrix_planner.hpp +++ b/cpp/src/distance/detail/pairwise_matrix/jit_lto_kernels/pairwise_matrix_planner.hpp @@ -1,11 +1,11 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once -#include +#include #include namespace cuvs::distance::detail { diff --git a/cpp/src/neighbors/detail/cagra/jit_lto_kernels/cagra_jit_launcher_factory.hpp b/cpp/src/neighbors/detail/cagra/jit_lto_kernels/cagra_jit_launcher_factory.hpp index 9b7e44bd8e..bd96e64503 100644 --- a/cpp/src/neighbors/detail/cagra/jit_lto_kernels/cagra_jit_launcher_factory.hpp +++ b/cpp/src/neighbors/detail/cagra/jit_lto_kernels/cagra_jit_launcher_factory.hpp @@ -12,7 +12,7 @@ #include "search_multi_kernel_planner.hpp" #include "search_single_cta_planner.hpp" -#include +#include #include #include diff --git a/cpp/src/neighbors/detail/cagra/jit_lto_kernels/cagra_planner_base.hpp b/cpp/src/neighbors/detail/cagra/jit_lto_kernels/cagra_planner_base.hpp index 07b432aa1f..0e78bb7a67 100644 --- a/cpp/src/neighbors/detail/cagra/jit_lto_kernels/cagra_planner_base.hpp +++ b/cpp/src/neighbors/detail/cagra/jit_lto_kernels/cagra_planner_base.hpp @@ -5,7 +5,7 @@ #pragma once -#include +#include #include #include #include diff --git a/cpp/src/neighbors/detail/cagra/jit_lto_kernels/sample_filter_udf.cuh b/cpp/src/neighbors/detail/cagra/jit_lto_kernels/sample_filter_udf.cuh index 3e16cf2bcc..efbc2d5a4f 100644 --- a/cpp/src/neighbors/detail/cagra/jit_lto_kernels/sample_filter_udf.cuh +++ b/cpp/src/neighbors/detail/cagra/jit_lto_kernels/sample_filter_udf.cuh @@ -7,7 +7,7 @@ #include "../cagra_filter_payload.hpp" -#include +#include #include #include diff --git a/cpp/src/neighbors/detail/cagra/search_multi_cta_kernel_launcher_jit.cuh b/cpp/src/neighbors/detail/cagra/search_multi_cta_kernel_launcher_jit.cuh index 8a673405b7..2efcbef0dc 100644 --- a/cpp/src/neighbors/detail/cagra/search_multi_cta_kernel_launcher_jit.cuh +++ b/cpp/src/neighbors/detail/cagra/search_multi_cta_kernel_launcher_jit.cuh @@ -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 */ @@ -17,7 +17,7 @@ #include "search_plan.cuh" // For search_params #include "set_value_batch.cuh" // For set_value_batch #include "shared_launcher_jit.hpp" // For shared JIT helper functions -#include +#include #include #include #include diff --git a/cpp/src/neighbors/detail/cagra/search_multi_kernel_launcher_jit.cuh b/cpp/src/neighbors/detail/cagra/search_multi_kernel_launcher_jit.cuh index bc341b9082..9500450adc 100644 --- a/cpp/src/neighbors/detail/cagra/search_multi_kernel_launcher_jit.cuh +++ b/cpp/src/neighbors/detail/cagra/search_multi_kernel_launcher_jit.cuh @@ -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 */ @@ -14,7 +14,7 @@ #include "jit_lto_kernels/search_multi_kernel_planner.hpp" #include "search_plan.cuh" // For search_params #include "shared_launcher_jit.hpp" // sample-filter payload helpers and JIT tags -#include +#include #include #include #include diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel_launcher_jit.cuh b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel_launcher_jit.cuh index a8731d06b5..14346a4656 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel_launcher_jit.cuh +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel_launcher_jit.cuh @@ -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 */ @@ -27,7 +27,7 @@ #include #include -#include +#include #include #include #include diff --git a/cpp/src/neighbors/ivf_flat/detail/jit_lto_kernels/interleaved_scan_planner.hpp b/cpp/src/neighbors/ivf_flat/detail/jit_lto_kernels/interleaved_scan_planner.hpp index ed8191016b..3de1f2f8b3 100644 --- a/cpp/src/neighbors/ivf_flat/detail/jit_lto_kernels/interleaved_scan_planner.hpp +++ b/cpp/src/neighbors/ivf_flat/detail/jit_lto_kernels/interleaved_scan_planner.hpp @@ -1,12 +1,12 @@ /* - * 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 diff --git a/cpp/src/neighbors/ivf_flat/ivf_flat_interleaved_scan_jit.cuh b/cpp/src/neighbors/ivf_flat/ivf_flat_interleaved_scan_jit.cuh index a9748ef836..c9e3d18abb 100644 --- a/cpp/src/neighbors/ivf_flat/ivf_flat_interleaved_scan_jit.cuh +++ b/cpp/src/neighbors/ivf_flat/ivf_flat_interleaved_scan_jit.cuh @@ -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 */ @@ -8,8 +8,8 @@ #include "../ivf_common.cuh" #include "detail/jit_lto_kernels/interleaved_scan_planner.hpp" #include "detail/jit_lto_kernels/kernel_def.hpp" +#include #include -#include #include #include #include diff --git a/cpp/src/neighbors/ivf_pq/detail/jit_lto_kernels/compute_similarity_planner.hpp b/cpp/src/neighbors/ivf_pq/detail/jit_lto_kernels/compute_similarity_planner.hpp index 0621966cad..9b8d8732ed 100644 --- a/cpp/src/neighbors/ivf_pq/detail/jit_lto_kernels/compute_similarity_planner.hpp +++ b/cpp/src/neighbors/ivf_pq/detail/jit_lto_kernels/compute_similarity_planner.hpp @@ -1,12 +1,12 @@ /* - * 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 diff --git a/cpp/src/neighbors/ivf_pq/ivf_pq_compute_similarity.hpp b/cpp/src/neighbors/ivf_pq/ivf_pq_compute_similarity.hpp index 93c7dbfe3d..4670e84572 100644 --- a/cpp/src/neighbors/ivf_pq/ivf_pq_compute_similarity.hpp +++ b/cpp/src/neighbors/ivf_pq/ivf_pq_compute_similarity.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -9,7 +9,7 @@ #include -#include +#include namespace cuvs::neighbors::ivf_pq::detail { diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_block_sort_planner.hpp b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_block_sort_planner.hpp index dbe6ec5fd4..bc04f78f6d 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_block_sort_planner.hpp +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_block_sort_planner.hpp @@ -1,12 +1,12 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once -#include -#include +#include +#include #include namespace cuvs::neighbors::ivf_rabitq::detail { diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_planner.hpp b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_planner.hpp index 3ec4809395..6f1e5da064 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_planner.hpp +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_planner.hpp @@ -1,12 +1,12 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once -#include -#include +#include +#include #include namespace cuvs::neighbors::ivf_rabitq::detail { diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_block_sort_planner.hpp b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_block_sort_planner.hpp index acf4e27e3a..3f0cf47ae9 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_block_sort_planner.hpp +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_block_sort_planner.hpp @@ -1,12 +1,12 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once -#include -#include +#include +#include #include namespace cuvs::neighbors::ivf_rabitq::detail { diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_planner.hpp b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_planner.hpp index 95e0694c9c..786fdd227a 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_planner.hpp +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_planner.hpp @@ -1,12 +1,12 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once -#include -#include +#include +#include #include namespace cuvs::neighbors::ivf_rabitq::detail { diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_block_sort_planner.hpp b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_block_sort_planner.hpp index 9db3839880..658971117c 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_block_sort_planner.hpp +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_block_sort_planner.hpp @@ -1,12 +1,12 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once -#include -#include +#include +#include #include namespace cuvs::neighbors::ivf_rabitq::detail { diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_planner.hpp b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_planner.hpp index 3559a9bee1..71fda33a44 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_planner.hpp +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_planner.hpp @@ -1,12 +1,12 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once -#include -#include +#include +#include #include namespace cuvs::neighbors::ivf_rabitq::detail { diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/launcher_factory.hpp b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/launcher_factory.hpp index d9c94821f1..3dd9088d6f 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/launcher_factory.hpp +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/launcher_factory.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 */ @@ -12,7 +12,7 @@ #include "compute_inner_products_with_lut_block_sort_planner.hpp" #include "compute_inner_products_with_lut_planner.hpp" -#include +#include #include #include diff --git a/cpp/src/neighbors/ivf_sq/detail/jit_lto_kernels/scan_planner.hpp b/cpp/src/neighbors/ivf_sq/detail/jit_lto_kernels/scan_planner.hpp index 05ea34532e..80921a941c 100644 --- a/cpp/src/neighbors/ivf_sq/detail/jit_lto_kernels/scan_planner.hpp +++ b/cpp/src/neighbors/ivf_sq/detail/jit_lto_kernels/scan_planner.hpp @@ -1,11 +1,11 @@ /* - * 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 diff --git a/fern/pages/jit_lto_guide.md b/fern/pages/jit_lto_guide.md index c6c4e279ef..2504fb316d 100644 --- a/fern/pages/jit_lto_guide.md +++ b/fern/pages/jit_lto_guide.md @@ -523,8 +523,8 @@ The planner is responsible for: ```cpp #pragma once -#include -#include +#include +#include #include #include #include From 930827ceaf7750081a22f7c09a994a142d5fbe05 Mon Sep 17 00:00:00 2001 From: Arha Gatram Date: Fri, 10 Jul 2026 15:22:28 -0700 Subject: [PATCH 2/3] link tests against rtcx --- cpp/cmake/thirdparty/get_rtcx.cmake | 2 +- cpp/tests/CMakeLists.txt | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/cpp/cmake/thirdparty/get_rtcx.cmake b/cpp/cmake/thirdparty/get_rtcx.cmake index 7a4476c903..7291af0828 100644 --- a/cpp/cmake/thirdparty/get_rtcx.cmake +++ b/cpp/cmake/thirdparty/get_rtcx.cmake @@ -18,7 +18,7 @@ function(find_and_configure_rtcx VERSION) GLOBAL_TARGETS rtcx::rtcx CPM_ARGS GIT_REPOSITORY https://github.com/arhag23/librtcx.git - GIT_TAG dc819febd2539c086e144b2632c3f5857c5d74b8 + GIT_TAG 75f90a8503766ce37d62f4418449315874bf2d87 GIT_SHALLOW FALSE EXCLUDE_FROM_ALL TRUE ) diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index 13a07b10b5..dae1a5da83 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -1,6 +1,6 @@ # ============================================================================= # cmake-format: off -# SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # cmake-format: on # ============================================================================= @@ -42,6 +42,7 @@ function(ConfigureTest) PRIVATE cuvs cuvs::cuvs raft::raft + rtcx::rtcx GTest::gtest GTest::gtest_main Threads::Threads From 8a0f2a108dc3244244d8afabd6cd353226578b8d Mon Sep 17 00:00:00 2001 From: Arha Gatram Date: Fri, 10 Jul 2026 17:07:22 -0700 Subject: [PATCH 3/3] Add missing rtcx dependencies --- cpp/CMakeLists.txt | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 32dfb898ef..45a2443a48 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -1537,7 +1537,7 @@ if(NOT BUILD_CPU_ONLY) $> $<$:CUDA::nvtx3> PRIVATE $ $ - $ CUDA::nvJitLink CUDA::nvrtc + $ CUDA::nvJitLink CUDA::nvrtc rtcx::rtcx ) # ensure CUDA symbols aren't relocated to the middle of the debug build binaries @@ -1594,6 +1594,7 @@ SECTIONS PRIVATE $ CUDA::nvJitLink CUDA::nvrtc + rtcx::rtcx $<$:CUDA::nvtx3> $ $