From 96af4ddc2977be3e5141d9c0e1eeb3a2428ad314 Mon Sep 17 00:00:00 2001 From: MithunR Date: Tue, 3 Feb 2026 15:50:15 -0800 Subject: [PATCH 01/20] Set max-dynamic-shared-mem with thread-safety This commit ensures that when the setup runs for the CAGRA search kernel, the calls to set `cudaFuncAttributeMaxDynamicSharedMemorySize` are thread-safe. The CUDA-context in a process is shared across all (CPU) threads. Consider 2 (CPU) threads A and B, attempting to set max-dynamic-shared-mem sizes to 3KB and 2KB respectively. Consider the following race condition: 1. Thread A successfully sets 3KB, but then gets switched before the CAGRA kernel is called. 2. Thread B sets 2KB, calls the kernel successfully, and exits. 3. Thread A calls its kernel, which attempts to use 3KB of smem. But the CUDA context's max-smem is currently set to 2KB. This causes the kernel launch to fail. The error manifests as listed in https://github.com/rapidsai/cuvs-lucene/issues/93: ``` CUDA error encountered at: file=/.../search_single_cta_kernel-inl.cuh line=2207: call='cudaPeekAtLastError()', Reason=cudaErrorInvalidValue:invalid argument Obtained 8 stack frames ... ``` This commit adds logic to track the high-water-mark for max-dynamic-shared-memory, and makes setting that value in the CUDA-context uninterruptible. This way, every thread of the same kernel-type gets the largest maximum threshold set so far. Signed-off-by: MithunR --- .../cagra/search_single_cta_kernel-inl.cuh | 26 +++++++++++++++++-- 1 file changed, 24 insertions(+), 2 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh index 0f40ff86e8..a313ae9a1d 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh @@ -2113,6 +2113,29 @@ auto get_runner(Args... args) -> std::shared_ptr return runner; } +/** + * @brief Optionally set the larger max dynamic shared memory size for the kernel. + * This is required because `cudaFuncSetAttribute` is not thread-safe. + * In the event of concurrent calls, we'd like to accommodate the largest requested size. + * @tparam KernelT The type of the kernel. + * @param smem_size The size of the dynamic shared memory to be set. + * @param kernel The kernel to be set. + */ +template // inline +void optionally_set_larger_max_smem_size(uint32_t smem_size, KernelT& kernel) +{ + static auto mutex = std::mutex{}; + static auto running_max_smem_size = uint32_t{0}; + if (smem_size > running_max_smem_size) { + auto guard = std::lock_guard{mutex}; + if (smem_size > running_max_smem_size) { + running_max_smem_size = smem_size; + RAFT_CUDA_TRY( + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, running_max_smem_size)); + } + } +} + template ; auto kernel = search_kernel_config:: choose_itopk_and_mx_candidates(ps.itopk_size, num_itopk_candidates, block_size); - RAFT_CUDA_TRY( - cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + optionally_set_larger_max_smem_size(smem_size, kernel); dim3 thread_dims(block_size, 1, 1); dim3 block_dims(1, num_queries, 1); RAFT_LOG_DEBUG( From 305cb5596756dcff0cef2b61514bbb4910a8728c Mon Sep 17 00:00:00 2001 From: MithunR Date: Wed, 4 Feb 2026 10:56:00 -0800 Subject: [PATCH 02/20] Copyright date. Formatting. Signed-off-by: MithunR --- .../detail/cagra/search_single_cta_kernel-inl.cuh | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh index a313ae9a1d..bb2e818e31 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ #pragma once @@ -2121,17 +2121,17 @@ auto get_runner(Args... args) -> std::shared_ptr * @param smem_size The size of the dynamic shared memory to be set. * @param kernel The kernel to be set. */ -template // inline +template // inline void optionally_set_larger_max_smem_size(uint32_t smem_size, KernelT& kernel) { - static auto mutex = std::mutex{}; + static auto mutex = std::mutex{}; static auto running_max_smem_size = uint32_t{0}; if (smem_size > running_max_smem_size) { auto guard = std::lock_guard{mutex}; if (smem_size > running_max_smem_size) { running_max_smem_size = smem_size; - RAFT_CUDA_TRY( - cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, running_max_smem_size)); + RAFT_CUDA_TRY(cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, running_max_smem_size)); } } } From 677f6fe9462d7c75d7e27bb1a08e8356e488cee8 Mon Sep 17 00:00:00 2001 From: MithunR Date: Wed, 4 Feb 2026 14:48:52 -0800 Subject: [PATCH 03/20] Moved other call to set cudaFuncAttributeMaxDynamicSharedMemorySize. Signed-off-by: MithunR --- .../neighbors/detail/cagra/search_single_cta_kernel-inl.cuh | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh index bb2e818e31..8628bda658 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh @@ -1880,8 +1880,7 @@ struct alignas(kCacheLineBytes) persistent_runner_t : public persistent_runner_b auto* dd_dev_ptr = dd_host.dev_ptr(stream); // set kernel attributes same as in normal kernel - RAFT_CUDA_TRY( - cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + optionally_set_larger_max_smem_size(smem_size, kernel); // set kernel launch parameters dim3 gs = calc_coop_grid_size(block_size, smem_size, persistent_device_usage); From 5e9e410f7e1aa3b7039b123ce9bd18d2fa772c2a Mon Sep 17 00:00:00 2001 From: MithunR Date: Wed, 4 Feb 2026 15:23:15 -0800 Subject: [PATCH 04/20] Moved call-sites in other files. Signed-off-by: MithunR --- .../cagra/search_multi_cta_kernel-inl.cuh | 4 +- .../cagra/search_single_cta_kernel-inl.cuh | 28 ++------------ cpp/src/neighbors/detail/smem_utils.cuh | 37 +++++++++++++++++++ 3 files changed, 42 insertions(+), 27 deletions(-) create mode 100644 cpp/src/neighbors/detail/smem_utils.cuh diff --git a/cpp/src/neighbors/detail/cagra/search_multi_cta_kernel-inl.cuh b/cpp/src/neighbors/detail/cagra/search_multi_cta_kernel-inl.cuh index 1720841c91..b332b59ce2 100644 --- a/cpp/src/neighbors/detail/cagra/search_multi_cta_kernel-inl.cuh +++ b/cpp/src/neighbors/detail/cagra/search_multi_cta_kernel-inl.cuh @@ -26,6 +26,7 @@ // TODO: This shouldn't be invoking anything from spatial/knn #include "../ann_utils.cuh" +#include "../smem_utils.cuh" #include #include // RAFT_CUDA_TRY_NOT_THROW is used TODO(tfeher): consider moving this to cuda_rt_essentials.hpp @@ -536,8 +537,7 @@ void select_and_run(const dataset_descriptor_host& dat SourceIndexT, SampleFilterT>::choose_buffer_size(result_buffer_size, block_size); - RAFT_CUDA_TRY( - cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + cuvs::neighbors::detail::optionally_set_larger_max_smem_size(smem_size, kernel); // Initialize hash table const uint32_t traversed_hash_size = hashmap::get_size(traversed_hash_bitlen); set_value_batch(traversed_hashmap_ptr, diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh index 8628bda658..dac5b227f3 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh @@ -26,6 +26,7 @@ // TODO: This shouldn't be invoking anything from spatial/knn #include "../ann_utils.cuh" +#include "../smem_utils.cuh" #include #include @@ -1880,7 +1881,7 @@ struct alignas(kCacheLineBytes) persistent_runner_t : public persistent_runner_b auto* dd_dev_ptr = dd_host.dev_ptr(stream); // set kernel attributes same as in normal kernel - optionally_set_larger_max_smem_size(smem_size, kernel); + cuvs::neighbors::detail::optionally_set_larger_max_smem_size(smem_size, kernel); // set kernel launch parameters dim3 gs = calc_coop_grid_size(block_size, smem_size, persistent_device_usage); @@ -2112,29 +2113,6 @@ auto get_runner(Args... args) -> std::shared_ptr return runner; } -/** - * @brief Optionally set the larger max dynamic shared memory size for the kernel. - * This is required because `cudaFuncSetAttribute` is not thread-safe. - * In the event of concurrent calls, we'd like to accommodate the largest requested size. - * @tparam KernelT The type of the kernel. - * @param smem_size The size of the dynamic shared memory to be set. - * @param kernel The kernel to be set. - */ -template // inline -void optionally_set_larger_max_smem_size(uint32_t smem_size, KernelT& kernel) -{ - static auto mutex = std::mutex{}; - static auto running_max_smem_size = uint32_t{0}; - if (smem_size > running_max_smem_size) { - auto guard = std::lock_guard{mutex}; - if (smem_size > running_max_smem_size) { - running_max_smem_size = smem_size; - RAFT_CUDA_TRY(cudaFuncSetAttribute( - kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, running_max_smem_size)); - } - } -} - template ; auto kernel = search_kernel_config:: choose_itopk_and_mx_candidates(ps.itopk_size, num_itopk_candidates, block_size); - optionally_set_larger_max_smem_size(smem_size, kernel); + cuvs::neighbors::detail::optionally_set_larger_max_smem_size(smem_size, kernel); dim3 thread_dims(block_size, 1, 1); dim3 block_dims(1, num_queries, 1); RAFT_LOG_DEBUG( diff --git a/cpp/src/neighbors/detail/smem_utils.cuh b/cpp/src/neighbors/detail/smem_utils.cuh new file mode 100644 index 0000000000..a2b4cb9fad --- /dev/null +++ b/cpp/src/neighbors/detail/smem_utils.cuh @@ -0,0 +1,37 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +#include + +#include +#include + +namespace cuvs::neighbors::detail { + +/** + * @brief Optionally set the larger max dynamic shared memory size for the kernel. + * This is required because `cudaFuncSetAttribute` is not thread-safe. + * In the event of concurrent calls, we'd like to accommodate the largest requested size. + * @tparam KernelT The type of the kernel. + * @param smem_size The size of the dynamic shared memory to be set. + * @param kernel The kernel to be set. + */ +template +void optionally_set_larger_max_smem_size(uint32_t smem_size, KernelT& kernel) +{ + static auto mutex = std::mutex{}; + static auto running_max_smem_size = uint32_t{0}; + if (smem_size > running_max_smem_size) { + auto guard = std::lock_guard{mutex}; + if (smem_size > running_max_smem_size) { + running_max_smem_size = smem_size; + RAFT_CUDA_TRY(cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, running_max_smem_size)); + } + } +} + +} // namespace cuvs::neighbors::detail From f55cbd39341892d7849c9b20a7db49d93a4e26b4 Mon Sep 17 00:00:00 2001 From: MithunR Date: Wed, 4 Feb 2026 21:06:07 -0800 Subject: [PATCH 05/20] Copyright date. Signed-off-by: MithunR --- cpp/src/neighbors/detail/cagra/search_multi_cta_kernel-inl.cuh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/src/neighbors/detail/cagra/search_multi_cta_kernel-inl.cuh b/cpp/src/neighbors/detail/cagra/search_multi_cta_kernel-inl.cuh index b332b59ce2..de68d2ac9d 100644 --- a/cpp/src/neighbors/detail/cagra/search_multi_cta_kernel-inl.cuh +++ b/cpp/src/neighbors/detail/cagra/search_multi_cta_kernel-inl.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ #pragma once From 2a5a00fb7cf1ca9af823f006b89f075d7f33c28b Mon Sep 17 00:00:00 2001 From: MithunR Date: Tue, 10 Feb 2026 15:18:36 -0800 Subject: [PATCH 06/20] Invoke kernel within critical section. In this version, rather than ratchet up the max smem size, we do the following in a critical section: 1. Set the max smem size in the CUDA context. 2. Register the kernel to be invoked. The kernel registration should be quite quick, and return immediately. This should avoid the loss of GPU occupancy due to an increasingly large smem size. Signed-off-by: MithunR --- .../cagra/search_multi_cta_kernel-inl.cuh | 44 ++++++----- .../cagra/search_single_cta_kernel-inl.cuh | 77 ++++++++++--------- cpp/src/neighbors/detail/smem_utils.cuh | 23 ++++++ 3 files changed, 86 insertions(+), 58 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/search_multi_cta_kernel-inl.cuh b/cpp/src/neighbors/detail/cagra/search_multi_cta_kernel-inl.cuh index 7d73ad9d4e..4ea9ad12a5 100644 --- a/cpp/src/neighbors/detail/cagra/search_multi_cta_kernel-inl.cuh +++ b/cpp/src/neighbors/detail/cagra/search_multi_cta_kernel-inl.cuh @@ -590,7 +590,6 @@ void select_and_run(const dataset_descriptor_host& dat THROW("Result buffer size %u larger than max buffer size %u", result_buffer_size, 256); } - cuvs::neighbors::detail::optionally_set_larger_max_smem_size(smem_size, kernel); // Initialize hash table const uint32_t traversed_hash_size = hashmap::get_size(traversed_hash_bitlen); set_value_batch(traversed_hashmap_ptr, @@ -608,26 +607,29 @@ void select_and_run(const dataset_descriptor_host& dat num_queries, smem_size); - kernel<<>>(topk_indices_ptr, - topk_distances_ptr, - dataset_desc.dev_ptr(stream), - queries_ptr, - graph.data_handle(), - max_elements, - graph.extent(1), - source_indices_ptr, - ps.num_random_samplings, - ps.rand_xor_mask, - dev_seed_ptr, - num_seeds, - visited_hash_bitlen, - traversed_hashmap_ptr, - traversed_hash_bitlen, - ps.itopk_size, - ps.min_iterations, - ps.max_iterations, - num_executed_iterations, - sample_filter); + auto const& invoke_kernel = [&]() { + kernel<<>>(topk_indices_ptr, + topk_distances_ptr, + dataset_desc.dev_ptr(stream), + queries_ptr, + graph.data_handle(), + max_elements, + graph.extent(1), + source_indices_ptr, + ps.num_random_samplings, + ps.rand_xor_mask, + dev_seed_ptr, + num_seeds, + visited_hash_bitlen, + traversed_hashmap_ptr, + traversed_hash_bitlen, + ps.itopk_size, + ps.min_iterations, + ps.max_iterations, + num_executed_iterations, + sample_filter); + }; + cuvs::neighbors::detail::safely_invoke_kernel_with_smem_size(kernel, smem_size, invoke_kernel); } } // namespace multi_cta_search diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh index 0b2c2a6cd0..bce3985103 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh @@ -1980,9 +1980,6 @@ struct alignas(kCacheLineBytes) persistent_runner_t : public persistent_runner_b // initialize the dataset/distance descriptor auto* dd_dev_ptr = dd_host.dev_ptr(stream); - // set kernel attributes same as in normal kernel - cuvs::neighbors::detail::optionally_set_larger_max_smem_size(smem_size, kernel); - // set kernel launch parameters dim3 gs = calc_coop_grid_size(block_size, smem_size, persistent_device_usage); dim3 bs(block_size, 1, 1); @@ -2049,15 +2046,19 @@ struct alignas(kCacheLineBytes) persistent_runner_t : public persistent_runner_b &small_hash_bitlen, &small_hash_reset_interval, &sample_filter}; - cuda::atomic_thread_fence(cuda::memory_order_seq_cst, cuda::thread_scope_system); - RAFT_CUDA_TRY(cudaLaunchCooperativeKernel>( - kernel, gs, bs, args, smem_size, stream)); - RAFT_LOG_INFO( - "Initialized the kernel %p in stream %zd; job_queue size = %u; worker_queue size = %u", - reinterpret_cast(kernel), - int64_t((cudaStream_t)stream), - job_queue.capacity(), - worker_queue.capacity()); + auto const& invoke_kernel = [&]() { + cuda::atomic_thread_fence(cuda::memory_order_seq_cst, cuda::thread_scope_system); + RAFT_CUDA_TRY(cudaLaunchCooperativeKernel>( + kernel, gs, bs, args, smem_size, stream)); + RAFT_LOG_INFO( + "Initialized the kernel %p in stream %zd; job_queue size = %u; worker_queue size = %u", + reinterpret_cast(kernel), + int64_t((cudaStream_t)stream), + job_queue.capacity(), + worker_queue.capacity()); + }; + // set kernel attributes same as in normal kernel + cuvs::neighbors::detail::safely_invoke_kernel_with_smem_size(kernel, smem_size, invoke_kernel); last_touch.store(std::chrono::system_clock::now(), std::memory_order_relaxed); } @@ -2312,35 +2313,37 @@ control is returned in this thread (in persistent_runner_t constructor), so we'r using descriptor_base_type = dataset_descriptor_base_t; auto kernel = search_kernel_config:: choose_itopk_and_mx_candidates(ps.itopk_size, num_itopk_candidates, block_size); - cuvs::neighbors::detail::optionally_set_larger_max_smem_size(smem_size, kernel); dim3 thread_dims(block_size, 1, 1); dim3 block_dims(1, num_queries, 1); RAFT_LOG_DEBUG( "Launching kernel with %u threads, %u block %u smem", block_size, num_queries, smem_size); - kernel<<>>(topk_indices_ptr, - topk_distances_ptr, - topk, - dataset_desc.dev_ptr(stream), - queries_ptr, - graph.data_handle(), - graph.extent(1), - source_indices_ptr, - ps.num_random_samplings, - ps.rand_xor_mask, - dev_seed_ptr, - num_seeds, - hashmap_ptr, - max_candidates, - max_itopk, - ps.itopk_size, - ps.search_width, - ps.min_iterations, - ps.max_iterations, - num_executed_iterations, - hash_bitlen, - small_hash_bitlen, - small_hash_reset_interval, - sample_filter); + auto const& invoke_kernel = [&]() { + kernel<<>>(topk_indices_ptr, + topk_distances_ptr, + topk, + dataset_desc.dev_ptr(stream), + queries_ptr, + graph.data_handle(), + graph.extent(1), + source_indices_ptr, + ps.num_random_samplings, + ps.rand_xor_mask, + dev_seed_ptr, + num_seeds, + hashmap_ptr, + max_candidates, + max_itopk, + ps.itopk_size, + ps.search_width, + ps.min_iterations, + ps.max_iterations, + num_executed_iterations, + hash_bitlen, + small_hash_bitlen, + small_hash_reset_interval, + sample_filter); + }; + cuvs::neighbors::detail::safely_invoke_kernel_with_smem_size(kernel, smem_size, invoke_kernel); RAFT_CUDA_TRY(cudaPeekAtLastError()); } } diff --git a/cpp/src/neighbors/detail/smem_utils.cuh b/cpp/src/neighbors/detail/smem_utils.cuh index a2b4cb9fad..b6e17b2212 100644 --- a/cpp/src/neighbors/detail/smem_utils.cuh +++ b/cpp/src/neighbors/detail/smem_utils.cuh @@ -34,4 +34,27 @@ void optionally_set_larger_max_smem_size(uint32_t smem_size, KernelT& kernel) } } +/** + * @brief (Thread-)Safely invoke a kernel with a dynamic shared memory size. + * This is required because `cudaFuncSetAttribute` is not thread-safe. + * In the event of concurrent calls, each kernel will be registered to run with + * its specified smem-size. + * @tparam KernelT The type of the kernel. + * @tparam InvocationT The type of the invocation function. + * @param kernel The kernel function address (for whom the smem-size is specified). + * @param smem_size The size of the dynamic shared memory to be set. + * @param invoke_kernel The kernel invocation function/lambda. + */ +template +void safely_invoke_kernel_with_smem_size(KernelT& kernel, + uint32_t smem_size, + InvocationT const& invoke_kernel) +{ + static auto mutex = std::mutex{}; + auto guard = std::lock_guard{mutex}; + RAFT_CUDA_TRY( + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + invoke_kernel(); +} + } // namespace cuvs::neighbors::detail From b1b97ccaabacbcd431c64292486e7ca078153ec0 Mon Sep 17 00:00:00 2001 From: MithunR Date: Tue, 10 Feb 2026 15:22:10 -0800 Subject: [PATCH 07/20] Removed old function. Signed-off-by: MithunR --- cpp/src/neighbors/detail/smem_utils.cuh | 23 ----------------------- 1 file changed, 23 deletions(-) diff --git a/cpp/src/neighbors/detail/smem_utils.cuh b/cpp/src/neighbors/detail/smem_utils.cuh index b6e17b2212..675173d615 100644 --- a/cpp/src/neighbors/detail/smem_utils.cuh +++ b/cpp/src/neighbors/detail/smem_utils.cuh @@ -11,29 +11,6 @@ namespace cuvs::neighbors::detail { -/** - * @brief Optionally set the larger max dynamic shared memory size for the kernel. - * This is required because `cudaFuncSetAttribute` is not thread-safe. - * In the event of concurrent calls, we'd like to accommodate the largest requested size. - * @tparam KernelT The type of the kernel. - * @param smem_size The size of the dynamic shared memory to be set. - * @param kernel The kernel to be set. - */ -template -void optionally_set_larger_max_smem_size(uint32_t smem_size, KernelT& kernel) -{ - static auto mutex = std::mutex{}; - static auto running_max_smem_size = uint32_t{0}; - if (smem_size > running_max_smem_size) { - auto guard = std::lock_guard{mutex}; - if (smem_size > running_max_smem_size) { - running_max_smem_size = smem_size; - RAFT_CUDA_TRY(cudaFuncSetAttribute( - kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, running_max_smem_size)); - } - } -} - /** * @brief (Thread-)Safely invoke a kernel with a dynamic shared memory size. * This is required because `cudaFuncSetAttribute` is not thread-safe. From f50d4d45490899c9ad0730a9d261d52d2761814a Mon Sep 17 00:00:00 2001 From: MithunR Date: Wed, 11 Feb 2026 13:42:57 -0800 Subject: [PATCH 08/20] Tie the kernel to its launcher. Signed-off-by: MithunR --- .../detail/cagra/search_multi_cta_kernel-inl.cuh | 4 ++-- .../detail/cagra/search_single_cta_kernel-inl.cuh | 10 ++++++---- cpp/src/neighbors/detail/smem_utils.cuh | 10 +++++----- 3 files changed, 13 insertions(+), 11 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/search_multi_cta_kernel-inl.cuh b/cpp/src/neighbors/detail/cagra/search_multi_cta_kernel-inl.cuh index 4ea9ad12a5..ff24724bdc 100644 --- a/cpp/src/neighbors/detail/cagra/search_multi_cta_kernel-inl.cuh +++ b/cpp/src/neighbors/detail/cagra/search_multi_cta_kernel-inl.cuh @@ -607,7 +607,7 @@ void select_and_run(const dataset_descriptor_host& dat num_queries, smem_size); - auto const& invoke_kernel = [&]() { + auto const& kernel_launcher = [&](auto const& kernel) -> void { kernel<<>>(topk_indices_ptr, topk_distances_ptr, dataset_desc.dev_ptr(stream), @@ -629,7 +629,7 @@ void select_and_run(const dataset_descriptor_host& dat num_executed_iterations, sample_filter); }; - cuvs::neighbors::detail::safely_invoke_kernel_with_smem_size(kernel, smem_size, invoke_kernel); + cuvs::neighbors::detail::safely_launch_kernel_with_smem_size(kernel, smem_size, kernel_launcher); } } // namespace multi_cta_search diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh index bce3985103..e2c00f36bc 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh @@ -2046,7 +2046,7 @@ struct alignas(kCacheLineBytes) persistent_runner_t : public persistent_runner_b &small_hash_bitlen, &small_hash_reset_interval, &sample_filter}; - auto const& invoke_kernel = [&]() { + auto const& kernel_launcher = [&](auto const& kernel) -> void { cuda::atomic_thread_fence(cuda::memory_order_seq_cst, cuda::thread_scope_system); RAFT_CUDA_TRY(cudaLaunchCooperativeKernel>( kernel, gs, bs, args, smem_size, stream)); @@ -2058,7 +2058,8 @@ struct alignas(kCacheLineBytes) persistent_runner_t : public persistent_runner_b worker_queue.capacity()); }; // set kernel attributes same as in normal kernel - cuvs::neighbors::detail::safely_invoke_kernel_with_smem_size(kernel, smem_size, invoke_kernel); + cuvs::neighbors::detail::safely_launch_kernel_with_smem_size( + kernel, smem_size, kernel_launcher); last_touch.store(std::chrono::system_clock::now(), std::memory_order_relaxed); } @@ -2317,7 +2318,7 @@ control is returned in this thread (in persistent_runner_t constructor), so we'r dim3 block_dims(1, num_queries, 1); RAFT_LOG_DEBUG( "Launching kernel with %u threads, %u block %u smem", block_size, num_queries, smem_size); - auto const& invoke_kernel = [&]() { + auto const& kernel_launcher = [&](auto const& kernel) -> void { kernel<<>>(topk_indices_ptr, topk_distances_ptr, topk, @@ -2343,7 +2344,8 @@ control is returned in this thread (in persistent_runner_t constructor), so we'r small_hash_reset_interval, sample_filter); }; - cuvs::neighbors::detail::safely_invoke_kernel_with_smem_size(kernel, smem_size, invoke_kernel); + cuvs::neighbors::detail::safely_launch_kernel_with_smem_size( + kernel, smem_size, kernel_launcher); RAFT_CUDA_TRY(cudaPeekAtLastError()); } } diff --git a/cpp/src/neighbors/detail/smem_utils.cuh b/cpp/src/neighbors/detail/smem_utils.cuh index 675173d615..1efb9bc805 100644 --- a/cpp/src/neighbors/detail/smem_utils.cuh +++ b/cpp/src/neighbors/detail/smem_utils.cuh @@ -20,18 +20,18 @@ namespace cuvs::neighbors::detail { * @tparam InvocationT The type of the invocation function. * @param kernel The kernel function address (for whom the smem-size is specified). * @param smem_size The size of the dynamic shared memory to be set. - * @param invoke_kernel The kernel invocation function/lambda. + * @param launch The kernel launch function/lambda. */ -template -void safely_invoke_kernel_with_smem_size(KernelT& kernel, +template +void safely_launch_kernel_with_smem_size(KernelT const& kernel, uint32_t smem_size, - InvocationT const& invoke_kernel) + KernelLauncherT const& launch) { static auto mutex = std::mutex{}; auto guard = std::lock_guard{mutex}; RAFT_CUDA_TRY( cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); - invoke_kernel(); + launch(kernel); } } // namespace cuvs::neighbors::detail From 8d66233f06d8c1de98e22a15c191ccb06c6ad353 Mon Sep 17 00:00:00 2001 From: MithunR Date: Wed, 11 Feb 2026 14:37:11 -0800 Subject: [PATCH 09/20] Better error reporting. Signed-off-by: MithunR --- cpp/src/neighbors/detail/smem_utils.cuh | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/cpp/src/neighbors/detail/smem_utils.cuh b/cpp/src/neighbors/detail/smem_utils.cuh index 1efb9bc805..4b1c003dcb 100644 --- a/cpp/src/neighbors/detail/smem_utils.cuh +++ b/cpp/src/neighbors/detail/smem_utils.cuh @@ -4,7 +4,7 @@ */ #pragma once -#include +#include #include #include @@ -29,8 +29,11 @@ void safely_launch_kernel_with_smem_size(KernelT const& kernel, { static auto mutex = std::mutex{}; auto guard = std::lock_guard{mutex}; - RAFT_CUDA_TRY( - cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + auto launch_status = + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size); + RAFT_EXPECTS(launch_status == cudaSuccess, + "Failed to set max dynamic shared memory size to %zu bytes", + smem_size); launch(kernel); } From 1d413736011fa0d467f1e1ba9622d88ad4ca2d46 Mon Sep 17 00:00:00 2001 From: achirkin Date: Thu, 12 Feb 2026 12:33:52 +0100 Subject: [PATCH 10/20] Remove the safety fix for persistent kernel (only one kernel must run at a time anyway) --- .../cagra/search_single_cta_kernel-inl.cuh | 23 ++++++++----------- 1 file changed, 9 insertions(+), 14 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh index e2c00f36bc..5e1ea20fae 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh @@ -2046,20 +2046,15 @@ struct alignas(kCacheLineBytes) persistent_runner_t : public persistent_runner_b &small_hash_bitlen, &small_hash_reset_interval, &sample_filter}; - auto const& kernel_launcher = [&](auto const& kernel) -> void { - cuda::atomic_thread_fence(cuda::memory_order_seq_cst, cuda::thread_scope_system); - RAFT_CUDA_TRY(cudaLaunchCooperativeKernel>( - kernel, gs, bs, args, smem_size, stream)); - RAFT_LOG_INFO( - "Initialized the kernel %p in stream %zd; job_queue size = %u; worker_queue size = %u", - reinterpret_cast(kernel), - int64_t((cudaStream_t)stream), - job_queue.capacity(), - worker_queue.capacity()); - }; - // set kernel attributes same as in normal kernel - cuvs::neighbors::detail::safely_launch_kernel_with_smem_size( - kernel, smem_size, kernel_launcher); + cuda::atomic_thread_fence(cuda::memory_order_seq_cst, cuda::thread_scope_system); + RAFT_CUDA_TRY(cudaLaunchCooperativeKernel>( + kernel, gs, bs, args, smem_size, stream)); + RAFT_LOG_INFO( + "Initialized the kernel %p in stream %zd; job_queue size = %u; worker_queue size = %u", + reinterpret_cast(kernel), + int64_t((cudaStream_t)stream), + job_queue.capacity(), + worker_queue.capacity()); last_touch.store(std::chrono::system_clock::now(), std::memory_order_relaxed); } From f6cf7d3d4fbbf6dbf0a3c62560c312791eca11c9 Mon Sep 17 00:00:00 2001 From: achirkin Date: Thu, 12 Feb 2026 12:34:03 +0100 Subject: [PATCH 11/20] Add a reproducer --- cpp/tests/CMakeLists.txt | 1 + .../ann_cagra/bug_issue_93_reproducer.cu | 126 ++++++++++++++++++ 2 files changed, 127 insertions(+) create mode 100644 cpp/tests/neighbors/ann_cagra/bug_issue_93_reproducer.cu diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index 9fc620b4cb..ab235a33ba 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -156,6 +156,7 @@ ConfigureTest( ConfigureTest( NAME NEIGHBORS_ANN_CAGRA_TEST_BUGS PATH neighbors/ann_cagra/bug_extreme_inputs_oob.cu neighbors/ann_cagra/bug_multi_cta_crash.cu + neighbors/ann_cagra/bug_issue_93_reproducer.cu GPUS 1 PERCENT 100 ) diff --git a/cpp/tests/neighbors/ann_cagra/bug_issue_93_reproducer.cu b/cpp/tests/neighbors/ann_cagra/bug_issue_93_reproducer.cu new file mode 100644 index 0000000000..c13fd76bdc --- /dev/null +++ b/cpp/tests/neighbors/ann_cagra/bug_issue_93_reproducer.cu @@ -0,0 +1,126 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + * + * Reproducer for https://github.com/rapidsai/cuvs-lucene/issues/93 + * cuvsCagraSearch returned 0 (Reason=cudaErrorInvalidValue:invalid argument) + * + * ROOT CAUSE: + * `cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)` + * is not thread-safe. It sets a CUDA-context-wide attribute. When two threads call it + * concurrently with different smem_size values, the following race occurs: + * 1. Thread A sets max-dynamic-shared-mem to SIZE_A (larger). + * 2. Thread B overwrites it with SIZE_B (smaller). + * 3. Thread A launches its kernel requesting SIZE_A of shared memory, + * but the CUDA context now only allows SIZE_B → cudaErrorInvalidValue. + * + * HOW IT MANIFESTS IN cuvs-lucene: + * Lucene's TaskExecutor dispatches per-segment CAGRA searches to a thread pool. + * Each segment has a different number of vectors (e.g. 25, 26, 47), leading to + * different graph degrees after reduction, and therefore different smem_size values + * in the single-CTA search kernel. The concurrent cudaFuncSetAttribute calls race. + * + * REPRODUCTION STRATEGY: + * Build CAGRA indices with different dataset sizes (different graph degrees), + * then search them concurrently from separate threads, each with its own raft::resources. + * This mirrors the cuvs-lucene setup where each thread gets a ThreadLocal CuVSResources. + */ + +#include + +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include + +namespace cuvs::neighbors::cagra { + +// NOLINTNEXTLINE(readability-identifier-naming) +TEST(Issue93Reproducer, ConcurrentSearchDifferentGraphDegrees) +{ + raft::resources handle; + raft::random::RngState rng(6181234567890123459ULL); + + // Dataset sizes from REPRODUCER.md warnings (different sizes → different graph degrees). + // NOLINTNEXTLINE(cppcoreguidelines-avoid-magic-numbers) + std::vector dataset_sizes = {25, 26, 47, 25}; + constexpr int dim = 64; + constexpr int top_k = 10; + + // Build indices on the main thread. + std::vector> indices; + for (int n_rows : dataset_sizes) { + auto database = raft::make_device_matrix(handle, n_rows, dim); + raft::random::uniform( + handle, rng, database.data_handle(), n_rows * dim, -1.0F, 1.0F); // NOLINT + + cagra::index_params ip; + ip.metric = cuvs::distance::DistanceType::L2Expanded; + ip.intermediate_graph_degree = 128; // NOLINT + ip.graph_degree = 64; // NOLINT + ip.graph_build_params = + graph_build_params::nn_descent_params(ip.intermediate_graph_degree, ip.metric); + + indices.push_back(cagra::build(handle, ip, raft::make_const_mdspan(database.view()))); + } + raft::resource::sync_stream(handle); + + // Search concurrently from multiple threads until the first failure. + const int num_threads = static_cast(indices.size()); + std::mutex error_mutex; + std::string first_error; + + // NOLINTNEXTLINE(cppcoreguidelines-avoid-magic-numbers) + for (int iter = 0; iter < 50 && first_error.empty(); ++iter) { + std::vector threads; + for (int t = 0; t < num_threads; ++t) { + threads.emplace_back([&, t, iter]() { + raft::resources thread_handle; + raft::random::RngState thread_rng(42ULL + static_cast(t) + + static_cast(iter) * 1000ULL); + try { + auto query = raft::make_device_matrix(thread_handle, 1, dim); + raft::random::uniform(thread_handle, thread_rng, query.data_handle(), dim, -1.0F, 1.0F); + + // Match cuvs-lucene params: Java's Panama zero-initializes the struct, + // and SINGLE_CTA = 0 in the enum, so algo is SINGLE_CTA. + cagra::search_params sp; + sp.itopk_size = top_k; + sp.search_width = 1; + sp.algo = search_algo::SINGLE_CTA; + + auto neighbors = raft::make_device_matrix(thread_handle, 1, top_k); + auto distances = raft::make_device_matrix(thread_handle, 1, top_k); + + cagra::search(thread_handle, + sp, + indices[static_cast(t)], + raft::make_const_mdspan(query.view()), + neighbors.view(), + distances.view()); + + raft::resource::sync_stream(thread_handle); + } catch (const std::exception& e) { + std::lock_guard lock(error_mutex); + if (first_error.empty()) { first_error = e.what(); } + } + }); + } + for (auto& th : threads) { + th.join(); + } + } + + ASSERT_TRUE(first_error.empty()) << first_error; +} + +} // namespace cuvs::neighbors::cagra From 9b135004811c90753b2525310ef6d72cce9dde40 Mon Sep 17 00:00:00 2001 From: achirkin Date: Thu, 12 Feb 2026 12:57:15 +0100 Subject: [PATCH 12/20] Fix style --- cpp/tests/neighbors/ann_cagra/bug_issue_93_reproducer.cu | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/tests/neighbors/ann_cagra/bug_issue_93_reproducer.cu b/cpp/tests/neighbors/ann_cagra/bug_issue_93_reproducer.cu index c13fd76bdc..6b4b037167 100644 --- a/cpp/tests/neighbors/ann_cagra/bug_issue_93_reproducer.cu +++ b/cpp/tests/neighbors/ann_cagra/bug_issue_93_reproducer.cu @@ -64,7 +64,7 @@ TEST(Issue93Reproducer, ConcurrentSearchDifferentGraphDegrees) handle, rng, database.data_handle(), n_rows * dim, -1.0F, 1.0F); // NOLINT cagra::index_params ip; - ip.metric = cuvs::distance::DistanceType::L2Expanded; + ip.metric = cuvs::distance::DistanceType::L2Expanded; ip.intermediate_graph_degree = 128; // NOLINT ip.graph_degree = 64; // NOLINT ip.graph_build_params = From 8d6fb1a143f068e68cb59c895ccd2d49b2d24c5a Mon Sep 17 00:00:00 2001 From: MithunR Date: Thu, 12 Feb 2026 23:37:55 -0800 Subject: [PATCH 13/20] Apply suggestion from @achirkin 1. Fixed documentation to indicate the right critical section: Setting the smem_size _and_ the kernel launch. 2. Switched back to ratcheting up the max smem-size. 3. Switched tracking the max smem-size via an atomic. Co-authored-by: Artem M. Chirkin <9253178+achirkin@users.noreply.github.com> --- cpp/src/neighbors/detail/smem_utils.cuh | 47 ++++++++++++++++++------- 1 file changed, 35 insertions(+), 12 deletions(-) diff --git a/cpp/src/neighbors/detail/smem_utils.cuh b/cpp/src/neighbors/detail/smem_utils.cuh index 4b1c003dcb..eccf23eec8 100644 --- a/cpp/src/neighbors/detail/smem_utils.cuh +++ b/cpp/src/neighbors/detail/smem_utils.cuh @@ -6,16 +6,20 @@ #include +#include #include #include namespace cuvs::neighbors::detail { /** - * @brief (Thread-)Safely invoke a kernel with a dynamic shared memory size. - * This is required because `cudaFuncSetAttribute` is not thread-safe. - * In the event of concurrent calls, each kernel will be registered to run with - * its specified smem-size. + * @brief (Thread-)Safely invoke a kernel with a maximum dynamic shared memory size. + * This is required because the sequence `cudaFuncSetAttribute` + kernel launch is not executed + * atomically. + * + * Used this way, the cudaFuncAttributeMaxDynamicSharedMemorySize can only grow and thus + * guarantees that the kernel is safe to launch. + * * @tparam KernelT The type of the kernel. * @tparam InvocationT The type of the invocation function. * @param kernel The kernel function address (for whom the smem-size is specified). @@ -27,14 +31,33 @@ void safely_launch_kernel_with_smem_size(KernelT const& kernel, uint32_t smem_size, KernelLauncherT const& launch) { - static auto mutex = std::mutex{}; - auto guard = std::lock_guard{mutex}; - auto launch_status = - cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size); - RAFT_EXPECTS(launch_status == cudaSuccess, - "Failed to set max dynamic shared memory size to %zu bytes", - smem_size); - launch(kernel); + // the last smem size is parameterized by the kernel thanks to the template parameter. + static std::atomic current_smem_size{0}; + auto last_smem_size = current_smem_size.load(std::memory_order_relaxed); + if (smem_size > last_smem_size) { + // We still need a mutex for the critical section: actualize last_smem_size and set the + // attribute. + static auto mutex = std::mutex{}; + auto guard = std::lock_guard{mutex}; + if (!current_smem_size.compare_exchange_strong( + last_smem_size, smem_size, std::memory_order_relaxed, std::memory_order_relaxed)) { + // The value has been updated by another thread between the load and the mutex acquisition. + if (smem_size > last_smem_size) { + current_smem_size.store(smem_size, std::memory_order_relaxed); + } + } + // Only update if the last seen value is smaller than the new one. + if (smem_size > last_smem_size) { + auto launch_status = + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size); + RAFT_EXPECTS(launch_status == cudaSuccess, + "Failed to set max dynamic shared memory size to %zu bytes", + smem_size); + } + } + // We don't need to guard the kernel launch because the smem_size can only grow. + return launch(kernel); } } // namespace cuvs::neighbors::detail + From 1955dceee24a48f91d195359ce50b757adb4e476 Mon Sep 17 00:00:00 2001 From: MithunR Date: Fri, 13 Feb 2026 11:14:32 -0800 Subject: [PATCH 14/20] Fixed formatting again. Signed-off-by: MithunR --- cpp/src/neighbors/detail/smem_utils.cuh | 1 - 1 file changed, 1 deletion(-) diff --git a/cpp/src/neighbors/detail/smem_utils.cuh b/cpp/src/neighbors/detail/smem_utils.cuh index eccf23eec8..dabb2a2711 100644 --- a/cpp/src/neighbors/detail/smem_utils.cuh +++ b/cpp/src/neighbors/detail/smem_utils.cuh @@ -60,4 +60,3 @@ void safely_launch_kernel_with_smem_size(KernelT const& kernel, } } // namespace cuvs::neighbors::detail - From 45dff116fd8a9acf427d9e996bfd2f550ac7ed3b Mon Sep 17 00:00:00 2001 From: "Artem M. Chirkin" <9253178+achirkin@users.noreply.github.com> Date: Wed, 18 Feb 2026 13:03:44 +0100 Subject: [PATCH 15/20] Apply suggestion from @achirkin --- .../neighbors/detail/cagra/search_single_cta_kernel-inl.cuh | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh index 5e1ea20fae..a42322b958 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh @@ -1980,6 +1980,10 @@ struct alignas(kCacheLineBytes) persistent_runner_t : public persistent_runner_b // initialize the dataset/distance descriptor auto* dd_dev_ptr = dd_host.dev_ptr(stream); + // set kernel attributes same as in normal kernel + RAFT_CUDA_TRY( + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + // set kernel launch parameters dim3 gs = calc_coop_grid_size(block_size, smem_size, persistent_device_usage); dim3 bs(block_size, 1, 1); From d3f6ecd6a7ffaa137cbb2ff568a6e56977809ae7 Mon Sep 17 00:00:00 2001 From: MithunR Date: Thu, 19 Feb 2026 14:16:08 -0800 Subject: [PATCH 16/20] Fixed format string. Signed-off-by: MithunR --- cpp/src/neighbors/detail/smem_utils.cuh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/src/neighbors/detail/smem_utils.cuh b/cpp/src/neighbors/detail/smem_utils.cuh index dabb2a2711..41c95c0ccd 100644 --- a/cpp/src/neighbors/detail/smem_utils.cuh +++ b/cpp/src/neighbors/detail/smem_utils.cuh @@ -51,7 +51,7 @@ void safely_launch_kernel_with_smem_size(KernelT const& kernel, auto launch_status = cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size); RAFT_EXPECTS(launch_status == cudaSuccess, - "Failed to set max dynamic shared memory size to %zu bytes", + "Failed to set max dynamic shared memory size to %u bytes", smem_size); } } From 088fcc8b711e024b1795390869269c195698f2c8 Mon Sep 17 00:00:00 2001 From: MithunR Date: Thu, 19 Feb 2026 14:46:00 -0800 Subject: [PATCH 17/20] Fixed thrust header. Signed-off-by: MithunR --- cpp/src/neighbors/detail/knn_brute_force.cuh | 1 + 1 file changed, 1 insertion(+) diff --git a/cpp/src/neighbors/detail/knn_brute_force.cuh b/cpp/src/neighbors/detail/knn_brute_force.cuh index da9a2e78d7..aa7c9ce61b 100644 --- a/cpp/src/neighbors/detail/knn_brute_force.cuh +++ b/cpp/src/neighbors/detail/knn_brute_force.cuh @@ -40,6 +40,7 @@ #include #include #include +#include #include #include From 856e695e6af109ceac275b63d63c6cfee38547fa Mon Sep 17 00:00:00 2001 From: MithunR Date: Thu, 19 Feb 2026 15:23:37 -0800 Subject: [PATCH 18/20] Fix compile error for thrust make_counting_iterator MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This commit fixes the compile error one sees because of the use of `thrust::make_counting_iterator`: ``` │ │ $SRC_DIR/cpp/tests/neighbors/all_neighbors/../../../src/neighbors/detail/knn_brute_force.cuh(230): error: namespace "thrust" has no member "make_counting_iterator" │ │ auto count = thrust::make_counting_iterator(0); │ │ ^ │ │ $SRC_DIR/cpp/tests/neighbors/all_neighbors/../../../src/neighbors/detail/knn_brute_force.cuh(230): error: type name is not allowed │ │ auto count = thrust::make_counting_iterator(0); ``` The RAPIDS projects seem to be moving away from using the thrust function, and switching to the cuda version of the same. This change moves cuVS in that direction. Signed-off-by: MithunR --- cpp/src/neighbors/detail/knn_brute_force.cuh | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/cpp/src/neighbors/detail/knn_brute_force.cuh b/cpp/src/neighbors/detail/knn_brute_force.cuh index aa7c9ce61b..ada7fab978 100644 --- a/cpp/src/neighbors/detail/knn_brute_force.cuh +++ b/cpp/src/neighbors/detail/knn_brute_force.cuh @@ -37,11 +37,11 @@ #include #include +#include #include #include #include #include -#include #include #include @@ -227,7 +227,7 @@ void tiled_brute_force_knn(const raft::resources& handle, } auto distances_ptr = temp_distances.data(); - auto count = thrust::make_counting_iterator(0); + auto count = cuda::make_counting_iterator(0); DistanceT masked_distance = select_min ? std::numeric_limits::infinity() : std::numeric_limits::lowest(); @@ -277,7 +277,7 @@ void tiled_brute_force_knn(const raft::resources& handle, DistanceT* out_distances = temp_out_distances.data(); IndexType* out_indices = temp_out_indices.data(); - auto count = thrust::make_counting_iterator(0); + auto count = cuda::make_counting_iterator(0); thrust::for_each(raft::resource::get_thrust_policy(handle), count, count + current_query_size * current_k, From eff47c8a50cbf9c4d63cb36cef5100ba7741716f Mon Sep 17 00:00:00 2001 From: MithunR Date: Thu, 19 Feb 2026 16:00:21 -0800 Subject: [PATCH 19/20] Revert "Fix compile error for thrust make_counting_iterator" This reverts commit 856e695e6af109ceac275b63d63c6cfee38547fa. It'll be simpler to use thrust/counting_iterator for now, per 1825. --- cpp/src/neighbors/detail/knn_brute_force.cuh | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/cpp/src/neighbors/detail/knn_brute_force.cuh b/cpp/src/neighbors/detail/knn_brute_force.cuh index ada7fab978..aa7c9ce61b 100644 --- a/cpp/src/neighbors/detail/knn_brute_force.cuh +++ b/cpp/src/neighbors/detail/knn_brute_force.cuh @@ -37,11 +37,11 @@ #include #include -#include #include #include #include #include +#include #include #include @@ -227,7 +227,7 @@ void tiled_brute_force_knn(const raft::resources& handle, } auto distances_ptr = temp_distances.data(); - auto count = cuda::make_counting_iterator(0); + auto count = thrust::make_counting_iterator(0); DistanceT masked_distance = select_min ? std::numeric_limits::infinity() : std::numeric_limits::lowest(); @@ -277,7 +277,7 @@ void tiled_brute_force_knn(const raft::resources& handle, DistanceT* out_distances = temp_out_distances.data(); IndexType* out_indices = temp_out_indices.data(); - auto count = cuda::make_counting_iterator(0); + auto count = thrust::make_counting_iterator(0); thrust::for_each(raft::resource::get_thrust_policy(handle), count, count + current_query_size * current_k, From b59b8c3fb084f18e66314d826ebec03e82c5658e Mon Sep 17 00:00:00 2001 From: MithunR Date: Thu, 19 Feb 2026 16:09:24 -0800 Subject: [PATCH 20/20] Cherrypick fix from #1825. Signed-off-by: MithunR --- cpp/src/neighbors/detail/knn_brute_force.cuh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/src/neighbors/detail/knn_brute_force.cuh b/cpp/src/neighbors/detail/knn_brute_force.cuh index aa7c9ce61b..a226e1d0cf 100644 --- a/cpp/src/neighbors/detail/knn_brute_force.cuh +++ b/cpp/src/neighbors/detail/knn_brute_force.cuh @@ -41,7 +41,7 @@ #include #include #include -#include +#include #include #include