diff --git a/cpp/src/neighbors/detail/cagra/device_common.hpp b/cpp/src/neighbors/detail/cagra/device_common.hpp index 1cc7772476..0b75de6bab 100644 --- a/cpp/src/neighbors/detail/cagra/device_common.hpp +++ b/cpp/src/neighbors/detail/cagra/device_common.hpp @@ -102,10 +102,13 @@ RAFT_DEVICE_INLINE_FUNCTION void compute_distance_to_random_nodes( IndexT* __restrict__ traversed_hash_ptr, const uint32_t traversed_hash_bitlen, const uint32_t block_id = 0, - const uint32_t num_blocks = 1) + const uint32_t num_blocks = 1, + const IndexT graph_size = 0) { const auto team_size_bits = dataset_desc.team_size_bitshift_from_smem(); const auto max_i = raft::round_up_safe(num_pickup, warp_size >> team_size_bits); + const auto compute_distance = dataset_desc.compute_distance_impl; + const IndexT seed_index_limit = graph_size > 0 ? graph_size : dataset_desc.size; for (uint32_t i = threadIdx.x >> team_size_bits; i < max_i; i += (blockDim.x >> team_size_bits)) { const bool valid_i = (i < num_pickup); @@ -121,7 +124,7 @@ RAFT_DEVICE_INLINE_FUNCTION void compute_distance_to_random_nodes( if (seed_ptr && (gid < num_seeds)) { seed_index = seed_ptr[gid]; } else { - seed_index = device::xorshift64(gid ^ rand_xor_mask) % dataset_desc.size; + seed_index = device::xorshift64(gid ^ rand_xor_mask) % seed_index_limit; } } 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 ff24724bdc..111bd97254 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 @@ -194,7 +194,8 @@ RAFT_KERNEL __launch_bounds__(1024, 1) search_kernel( const uint32_t min_iteration, const uint32_t max_iteration, uint32_t* const num_executed_iterations, /* stats */ - SAMPLE_FILTER_T sample_filter) + SAMPLE_FILTER_T sample_filter, + const typename DATASET_DESCRIPTOR_T::INDEX_T graph_size = 0) { using DATA_T = typename DATASET_DESCRIPTOR_T::DATA_T; using INDEX_T = typename DATASET_DESCRIPTOR_T::INDEX_T; @@ -282,7 +283,8 @@ RAFT_KERNEL __launch_bounds__(1024, 1) search_kernel( local_traversed_hashmap_ptr, traversed_hash_bitlen, block_id, - num_blocks); + num_blocks, + graph_size); __syncthreads(); _CLK_REC(clk_compute_1st_distance); @@ -607,29 +609,27 @@ void select_and_run(const dataset_descriptor_host& dat num_queries, smem_size); - auto const& kernel_launcher = [&](auto const& kernel) -> void { - 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_launch_kernel_with_smem_size(kernel, smem_size, kernel_launcher); + 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, + static_cast(graph.extent(0))); } } // namespace multi_cta_search diff --git a/cpp/src/neighbors/detail/cagra/search_multi_kernel.cuh b/cpp/src/neighbors/detail/cagra/search_multi_kernel.cuh index f7d353d864..045c63fe59 100644 --- a/cpp/src/neighbors/detail/cagra/search_multi_kernel.cuh +++ b/cpp/src/neighbors/detail/cagra/search_multi_kernel.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 @@ -104,7 +104,8 @@ RAFT_KERNEL random_pickup_kernel( typename DATASET_DESCRIPTOR_T::DISTANCE_T* const result_distances_ptr, // [num_queries, ldr] const std::uint32_t ldr, // (*) ldr >= num_pickup typename DATASET_DESCRIPTOR_T::INDEX_T* const visited_hashmap_ptr, // [num_queries, 1 << bitlen] - const std::uint32_t hash_bitlen) + const std::uint32_t hash_bitlen, + const typename DATASET_DESCRIPTOR_T::INDEX_T graph_size = 0) { using DATA_T = typename DATASET_DESCRIPTOR_T::DATA_T; using INDEX_T = typename DATASET_DESCRIPTOR_T::INDEX_T; @@ -119,6 +120,8 @@ RAFT_KERNEL random_pickup_kernel( dataset_desc = dataset_desc->setup_workspace(smem, queries_ptr, query_id); __syncthreads(); + const INDEX_T seed_index_limit = graph_size > 0 ? graph_size : dataset_desc->size; + INDEX_T best_index_team_local; DISTANCE_T best_norm2_team_local = utils::get_max_value(); for (unsigned i = 0; i < num_distilation; i++) { @@ -128,7 +131,7 @@ RAFT_KERNEL random_pickup_kernel( } else { // Chose a seed node randomly seed_index = - device::xorshift64((global_team_index ^ rand_xor_mask) * (i + 1)) % dataset_desc->size; + device::xorshift64((global_team_index ^ rand_xor_mask) * (i + 1)) % seed_index_limit; } DISTANCE_T norm2 = dataset_desc->compute_distance(seed_index, true); @@ -166,7 +169,8 @@ void random_pickup(const dataset_descriptor_host& data std::size_t ldr, // (*) ldr >= num_pickup IndexT* visited_hashmap_ptr, // [num_queries, 1 << bitlen] std::uint32_t hash_bitlen, - cudaStream_t cuda_stream) + cudaStream_t cuda_stream, + IndexT graph_size = 0) { const auto block_size = 256u; const auto num_teams_per_threadblock = block_size / dataset_desc.team_size; @@ -185,7 +189,8 @@ void random_pickup(const dataset_descriptor_host& data result_distances_ptr, ldr, visited_hashmap_ptr, - hash_bitlen); + hash_bitlen, + graph_size); } template @@ -826,7 +831,8 @@ struct search result_buffer_allocation_size, hashmap.data(), hash_bitlen, - stream); + stream, + static_cast(this->dataset_size)); unsigned iter = 0; while (1) { 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 62600c97dd..3217a466aa 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 @@ -703,7 +703,8 @@ RAFT_DEVICE_INLINE_FUNCTION void search_core( const std::uint32_t small_hash_bitlen, const std::uint32_t small_hash_reset_interval, const std::uint32_t query_id, - SAMPLE_FILTER_T sample_filter) + SAMPLE_FILTER_T sample_filter, + const typename DATASET_DESCRIPTOR_T::INDEX_T graph_size = 0) { using LOAD_T = device::LOAD_128BIT_T; @@ -792,7 +793,10 @@ RAFT_DEVICE_INLINE_FUNCTION void search_core( local_visited_hashmap_ptr, hash_bitlen, (INDEX_T*)nullptr, - 0); + 0, + 0, + 1, + graph_size); __syncthreads(); _CLK_REC(clk_compute_1st_distance); @@ -1125,7 +1129,8 @@ RAFT_KERNEL __launch_bounds__(1024, 1) search_kernel( const std::uint32_t hash_bitlen, const std::uint32_t small_hash_bitlen, const std::uint32_t small_hash_reset_interval, - SAMPLE_FILTER_T sample_filter) + SAMPLE_FILTER_T sample_filter, + const typename DATASET_DESCRIPTOR_T::INDEX_T graph_size = 0) { const auto query_id = blockIdx.y; search_core void { - 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_launch_kernel_with_smem_size( - kernel, smem_size, kernel_launcher); + 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, + static_cast(graph.extent(0))); RAFT_CUDA_TRY(cudaPeekAtLastError()); } } diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index 62fbf960ab..f5757da423 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -159,6 +159,7 @@ 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_graph_smaller_than_dataset.cu neighbors/ann_cagra/bug_iterative_cagra_build.cu neighbors/ann_cagra/bug_issue_93_reproducer.cu GPUS 1 diff --git a/cpp/tests/neighbors/ann_cagra/bug_graph_smaller_than_dataset.cu b/cpp/tests/neighbors/ann_cagra/bug_graph_smaller_than_dataset.cu new file mode 100644 index 0000000000..adeb774a8b --- /dev/null +++ b/cpp/tests/neighbors/ann_cagra/bug_graph_smaller_than_dataset.cu @@ -0,0 +1,160 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include + +#include + +#include +#include +#include +#include +#include + +#include + +namespace cuvs::neighbors::cagra { + +/** + * @brief Test verifying graph.extent(0) is used for random seed selection + * + * This test ensures that CAGRA search kernels correctly use graph.extent(0) + * (graph size) rather than dataset.size for random seed node selection. + * + * The bug: random seed selection previously used dataset_desc.size, which + * could cause OOB access if the graph size differed from dataset size + * (e.g., in CAGRA-Q iterative builds with compression). + * + * The fix: kernels now receive graph.extent(0) as graph_size parameter, + * ensuring seeds are always within valid graph node range [0, graph_size). + */ +class cagra_graph_smaller_than_dataset_test : public ::testing::Test { + public: + using data_type = float; + using index_type = uint32_t; + + protected: + void run() + { + // Create a dataset with 1000 points + constexpr int64_t n_dataset = 1000; + constexpr int64_t n_dim = 128; + constexpr int64_t n_queries = 100; + constexpr int64_t k = 10; + + // Build index normally + auto dataset = raft::make_device_matrix(res, n_dataset, n_dim); + raft::random::RngState r(1234ULL); + raft::random::uniform( + res, r, dataset.data_handle(), n_dataset * n_dim, data_type(-1), data_type(1)); + + cagra::index_params index_params; + index_params.graph_degree = 32; + index_params.intermediate_graph_degree = 64; + + auto index = cagra::build(res, index_params, raft::make_const_mdspan(dataset.view())); + raft::resource::sync_stream(res); + + // Get the graph from the index + auto original_graph = index.graph(); + ASSERT_EQ(original_graph.extent(0), n_dataset); + + // Recreate the bug scenario: LARGE dataset, SMALL graph + // (like iterative_build_graph does in intermediate iterations) + constexpr int64_t n_graph = n_dataset / 2; // Only 500 nodes in graph + + // Step 1: Build index on SMALL subset (500 points) + auto small_dataset_view = raft::make_device_matrix_view( + dataset.data_handle(), n_graph, n_dim); + + cagra::index_params small_index_params; + small_index_params.graph_degree = 32; + auto small_index = cagra::build(res, small_index_params, small_dataset_view); + raft::resource::sync_stream(res); + + // Step 2: Update to FULL dataset (1000 points) but keep small graph (500 nodes) + // This creates the exact bug scenario: dataset.size=1000, graph.extent(0)=500 + small_index.update_dataset(res, raft::make_const_mdspan(dataset.view())); + + // Verify the mismatch - THIS IS THE BUG SCENARIO! + ASSERT_EQ(small_index.graph().extent(0), n_graph); // Graph has 500 nodes + ASSERT_EQ(small_index.size(), n_dataset); // Dataset has 1000 points + ASSERT_NE(small_index.graph().extent(0), small_index.size()); // Mismatch! + + // Create queries + auto queries = raft::make_device_matrix(res, n_queries, n_dim); + raft::random::uniform( + res, r, queries.data_handle(), n_queries * n_dim, data_type(-1), data_type(1)); + + // Allocate output + auto neighbors = raft::make_device_matrix(res, n_queries, k); + auto distances = raft::make_device_matrix(res, n_queries, k); + + // Setup search params + cagra::search_params search_params; + search_params.itopk_size = 64; + search_params.search_width = 1; + search_params.max_iterations = 10; + search_params.algo = cagra::search_algo::SINGLE_CTA; + + // THIS SHOULD NOT CRASH OR CAUSE OOB ACCESS + // Before fix: random seeds use dataset.size (1000) -> tries to access graph[700] -> CRASH! + // After fix: random seeds use graph.extent(0) (500) -> only accesses graph[0-499] -> SAFE! + cagra::search(res, + search_params, + small_index, + raft::make_const_mdspan(queries.view()), + neighbors.view(), + distances.view()); + + raft::resource::sync_stream(res); + + // Verify results are valid (neighbors should be < graph size) + auto neighbors_host = raft::make_host_matrix(n_queries, k); + raft::copy(neighbors_host.data_handle(), + neighbors.data_handle(), + n_queries * k, + raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + + // All neighbor indices should be valid (< n_graph) + for (int64_t i = 0; i < n_queries * k; i++) { + ASSERT_LT(neighbors_host.data_handle()[i], n_graph) + << "Neighbor index " << neighbors_host.data_handle()[i] << " is >= graph size " << n_graph; + } + + // Test with MULTI_CTA algorithm as well (also had the same bug) + search_params.algo = cagra::search_algo::MULTI_CTA; + + cagra::search(res, + search_params, + small_index, + raft::make_const_mdspan(queries.view()), + neighbors.view(), + distances.view()); + + raft::resource::sync_stream(res); + + // Verify again + raft::copy(neighbors_host.data_handle(), + neighbors.data_handle(), + n_queries * k, + raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + + for (int64_t i = 0; i < n_queries * k; i++) { + ASSERT_LT(neighbors_host.data_handle()[i], n_graph) + << "Neighbor index " << neighbors_host.data_handle()[i] << " is >= graph size " << n_graph + << " (MULTI_CTA)"; + } + } + + private: + raft::resources res; +}; + +TEST_F(cagra_graph_smaller_than_dataset_test, search_with_smaller_graph) { this->run(); } + +} // namespace cuvs::neighbors::cagra