diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 7e4ee04c27..747a4af1d2 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -84,7 +84,12 @@ option(DISABLE_OPENMP "Disable OpenMP" OFF) option(CUVS_COMPILE_DYNAMIC_ONLY "Only build the shared library and skip the static library." OFF) option(CUVS_NVTX "Enable nvtx markers" OFF) option(CUVS_RAFT_CLONE_ON_PIN "Explicitly clone RAFT branch when pinned to non-feature branch" ON) +option(CUVS_ENABLE_FLOWANN_SEARCH "Enable experimental FlowANN tiered-graph search" OFF) +option(CUVS_FLOWANN_USE_GDRCOPY "Use GDRCopy for FlowANN GPU-to-CPU command queues" ON) +if(CUVS_ENABLE_FLOWANN_SEARCH AND BUILD_CPU_ONLY) + message(FATAL_ERROR "CUVS_ENABLE_FLOWANN_SEARCH requires a CUDA-enabled build") +endif() if(BUILD_CPU_ONLY) set(BUILD_SHARED_LIBS OFF) set(BUILD_TESTS OFF) @@ -113,7 +118,6 @@ elseif(NOT BUILD_SHARED_LIBS) set(cuvs_compile_mode "static_only") set(_cuvs_lib_targets cuvs_static) endif() - # Needed because GoogleBenchmark changes the state of FindThreads.cmake, causing subsequent runs to # have different values for the `Threads::Threads` target. Setting this flag ensures # `Threads::Threads` is the same value across all builds so that cache hits occur @@ -127,6 +131,8 @@ message(VERBOSE "cuVS: Build cuVS unit-tests: ${BUILD_TESTS}") message(VERBOSE "cuVS: Build CPU only components: ${BUILD_CPU_ONLY}") message(VERBOSE "cuVS: Build ANN benchmarks: ${BUILD_CUVS_BENCH}") message(VERBOSE "cuVS: Build only the shared library: ${CUVS_COMPILE_DYNAMIC_ONLY}") +message(VERBOSE "cuVS: Enable FlowANN tiered-graph search: ${CUVS_ENABLE_FLOWANN_SEARCH}") +message(VERBOSE "cuVS: Use GDRCopy for FlowANN queues: ${CUVS_FLOWANN_USE_GDRCOPY}") message(VERBOSE "cuVS: Enable detection of conda environment for dependencies: ${DETECT_CONDA_ENV}") message(VERBOSE "cuVS: Disable deprecation warnings " ${DISABLE_DEPRECATION_WARNINGS}) message(VERBOSE "cuVS: Disable OpenMP: ${DISABLE_OPENMP}") @@ -169,6 +175,22 @@ if(NOT BUILD_CPU_ONLY) BUILD_EXPORT_SET cuvs-exports INSTALL_EXPORT_SET cuvs-exports ) + + if(CUVS_ENABLE_FLOWANN_SEARCH AND CUVS_FLOWANN_USE_GDRCOPY) + rapids_find_generate_module( + GDRCopy + HEADER_NAMES gdrapi.h + LIBRARY_NAMES gdrapi + BUILD_EXPORT_SET cuvs-exports + INSTALL_EXPORT_SET cuvs-exports + ) + rapids_find_package( + GDRCopy REQUIRED + BUILD_EXPORT_SET cuvs-exports + INSTALL_EXPORT_SET cuvs-exports + ) + set(CUVS_FLOWANN_GDRCOPY_DEPENDENCIES GDRCopy::GDRCopy CUDA::cuda_driver) + endif() else() add_compile_definitions(BUILD_CPU_ONLY) endif() @@ -233,6 +255,9 @@ if(NOT BUILD_CPU_ONLY) "$" "$" ) + target_compile_definitions( + cuvs_cpp_headers INTERFACE $<$:CUVS_ENABLE_FLOWANN_SEARCH> + ) target_link_libraries(cuvs_cpp_headers INTERFACE raft::raft rmm::rmm) generate_inst_matrix( @@ -1486,6 +1511,8 @@ if(NOT BUILD_CPU_ONLY) src/neighbors/brute_force.cu src/neighbors/brute_force_serialize.cu src/neighbors/cagra.cpp + $<$:src/neighbors/flowann.cu> + $<$:src/neighbors/flowann_queue.cpp> ${cagra_build_inst_files} ${cagra_extend_inst_files} src/neighbors/cagra_optimize.cu @@ -1592,6 +1619,12 @@ if(NOT BUILD_CPU_ONLY) "$<$,$>:${CUVS_DEBUG_CUDA_FLAGS}>" ) + target_compile_definitions( + cuvs_objs + PRIVATE + $<$,$>:CUVS_FLOWANN_USE_GDRCOPY> + ) + target_compile_definitions( cuvs_objs PRIVATE CUVS_CUTILE_ENABLED=${CUVS_CUTILE_ENABLED} @@ -1608,6 +1641,7 @@ if(NOT BUILD_CPU_ONLY) nvidia::cutlass::cutlass rtcx::rtcx ${CUVS_CTK_MATH_DEPENDENCIES} + ${CUVS_FLOWANN_GDRCOPY_DEPENDENCIES} $ $ kvikio::kvikio @@ -1687,10 +1721,26 @@ if(NOT BUILD_CPU_ONLY) $> $> $<$:CUDA::nvtx3> - PRIVATE $ $ - $ CUDA::nvJitLink CUDA::nvrtc rtcx::rtcx kvikio::kvikio + PRIVATE $ + $ + $ + CUDA::nvJitLink + CUDA::nvrtc + rtcx::rtcx + kvikio::kvikio + ${CUVS_FLOWANN_GDRCOPY_DEPENDENCIES} INTERFACE $ ) + if(CUVS_ENABLE_FLOWANN_SEARCH) + target_link_libraries(cuvs INTERFACE $>) + endif() + if(CUVS_ENABLE_FLOWANN_SEARCH AND CUVS_FLOWANN_USE_GDRCOPY) + # Consumers must resolve this runtime dependency even when it is installed outside the + # linker's default search paths. Its headers remain private to the queue implementation. + target_link_libraries( + cuvs INTERFACE $ $ + ) + endif() endif() @@ -1735,6 +1785,7 @@ if(NOT BUILD_CPU_ONLY) CUDA::nvJitLink CUDA::nvrtc rtcx::rtcx + ${CUVS_FLOWANN_GDRCOPY_DEPENDENCIES} $<$:CUDA::nvtx3> $ $ diff --git a/cpp/src/neighbors/detail/cagra/jit_lto_kernels/search_multi_cta_jit.cuh b/cpp/src/neighbors/detail/cagra/jit_lto_kernels/search_multi_cta_jit.cuh index 0d16a32fc7..ec8f1abcdf 100644 --- a/cpp/src/neighbors/detail/cagra/jit_lto_kernels/search_multi_cta_jit.cuh +++ b/cpp/src/neighbors/detail/cagra/jit_lto_kernels/search_multi_cta_jit.cuh @@ -13,6 +13,7 @@ #include +#include #include #include #include @@ -31,6 +32,88 @@ using cuvs::neighbors::cagra::detail::device::compute_distance_to_child_nodes_ji using cuvs::neighbors::cagra::detail::device::compute_distance_to_random_nodes_jit; using cuvs::neighbors::detail::sample_filter; +template +struct device_graph_provider { + const IndexT* graph; + std::uint32_t degree; + + RAFT_DEVICE_INLINE_FUNCTION void initialize() const {} + + RAFT_DEVICE_INLINE_FUNCTION auto result_buffer_size(std::uint32_t base_size, std::uint32_t) const + -> std::uint32_t + { + return base_size; + } + + RAFT_DEVICE_INLINE_FUNCTION auto seed_pointer(const IndexT* seeds, + std::uint32_t num_seeds, + std::uint32_t query_id) const -> const IndexT* + { + return seeds == nullptr ? nullptr : seeds + static_cast(num_seeds) * query_id; + } + + RAFT_DEVICE_INLINE_FUNCTION auto seed_block_id(std::uint32_t cta_id, + std::uint32_t num_cta_per_query, + std::uint32_t row) const -> std::uint32_t + { + return cta_id + num_cta_per_query * row; + } + + RAFT_DEVICE_INLINE_FUNCTION auto seed_num_blocks(std::uint32_t num_cta_per_query, + std::uint32_t num_queries, + std::uint32_t num_partitions) const + -> std::uint32_t + { + return num_cta_per_query * num_queries * num_partitions; + } + + RAFT_DEVICE_INLINE_FUNCTION auto initial_candidate_count(std::uint32_t candidate_count, + std::uint32_t, + std::uint32_t, + std::uint32_t) const -> std::uint32_t + { + return candidate_count; + } + + RAFT_DEVICE_INLINE_FUNCTION auto initial_distillation_count(std::uint32_t count, + std::uint32_t) const -> std::uint32_t + { + return count; + } + + RAFT_DEVICE_INLINE_FUNCTION auto expand( + IndexT* result_indices, + DistanceT* result_distances, + const dataset_descriptor_base_t* dataset_desc, + IndexT* visited_hashmap, + std::uint32_t visited_hash_bitlen, + IndexT* traversed_hashmap, + std::uint32_t traversed_hash_bitlen, + const IndexT* parent_indices, + const IndexT* internal_topk, + int* result_position, + int max_result_position) const -> std::uint32_t + { + compute_distance_to_child_nodes_jit(result_indices, + result_distances, + dataset_desc, + graph, + degree, + visited_hashmap, + visited_hash_bitlen, + traversed_hashmap, + traversed_hash_bitlen, + parent_indices, + internal_topk, + 1, + result_position, + max_result_position); + return degree; + } + + RAFT_DEVICE_INLINE_FUNCTION void finalize() const {} +}; + // Shared MULTI_CTA search body for both single-partition and multi-partition kernels. // // The multi-partition launch is a strict generalization of the single-partition one: it adds a @@ -49,16 +132,17 @@ template > RAFT_DEVICE_INLINE_FUNCTION void search_core( IndexT* const result_indices_ptr, // [num_queries, num_cta_per_query, itopk_size] DistanceT* const result_distances_ptr, // [num_queries, num_cta_per_query, itopk_size] const dataset_descriptor_base_t* dataset_desc, const DataT* const queries_ptr, // [num_queries, dataset_dim] - const IndexT* const knn_graph, // [dataset_size, graph_degree] + GraphProviderT graph_provider, const uint32_t max_elements, const uint32_t graph_degree, - const uint32_t buffer_graph_degree, + const uint32_t buffer_candidate_count, const SourceIndexT* source_indices_ptr, // [num_queries, search_width] const unsigned num_distilation, const uint64_t rand_xor_mask, @@ -116,7 +200,9 @@ RAFT_DEVICE_INLINE_FUNCTION void search_core( // | | upto 32 | | // +----------------+---------+---------------------------+ // |<--- result_buffer_size_32 --->| - const auto result_buffer_size = itopk_size + buffer_graph_degree; + const auto base_result_buffer_size = itopk_size + buffer_candidate_count; + const auto result_buffer_size = + graph_provider.result_buffer_size(base_result_buffer_size, num_seeds); const auto result_buffer_size_32 = raft::round_up_safe(result_buffer_size, 32); assert(result_buffer_size_32 <= max_elements); @@ -147,20 +233,25 @@ RAFT_DEVICE_INLINE_FUNCTION void search_core( result_distances_buffer[i] = utils::get_max_value(); } hashmap::init(local_visited_hashmap_ptr, visited_hash_bitlen); + graph_provider.initialize(); __syncthreads(); _CLK_REC(clk_init); // compute distance to randomly selecting nodes using JIT version _CLK_START(); - const INDEX_T* const local_seed_ptr = seed_ptr ? seed_ptr + (num_seeds * query_id) : nullptr; - uint32_t block_id = cta_id + (num_cta_per_query * row); - uint32_t num_blocks = num_cta_per_query * num_queries * gridDim.z; + const INDEX_T* const local_seed_ptr = graph_provider.seed_pointer(seed_ptr, num_seeds, query_id); + uint32_t block_id = graph_provider.seed_block_id(cta_id, num_cta_per_query, row); + uint32_t num_blocks = graph_provider.seed_num_blocks(num_cta_per_query, num_queries, gridDim.z); + auto const initial_candidate_count = + graph_provider.initial_candidate_count(graph_degree, num_seeds, block_id, num_blocks); + auto const initial_distillation_count = + graph_provider.initial_distillation_count(num_distilation, num_seeds); compute_distance_to_random_nodes_jit(result_indices_buffer, result_distances_buffer, smem_desc, - graph_degree, - num_distilation, + initial_candidate_count, + initial_distillation_count, rand_xor_mask, local_seed_ptr, num_seeds, @@ -245,30 +336,26 @@ RAFT_DEVICE_INLINE_FUNCTION void search_core( hashmap::insert(local_visited_hashmap_ptr, visited_hash_bitlen, index); } } - // Initialize buffer for compute_distance_to_child_nodes. + // Initialize buffer for graph expansion. if (threadIdx.x == blockDim.x - 1) { result_position[0] = result_buffer_size_32; } __syncthreads(); - // Compute the norms between child nodes and query node using JIT version - compute_distance_to_child_nodes_jit(result_indices_buffer, - result_distances_buffer, - smem_desc, - knn_graph, - graph_degree, - local_visited_hashmap_ptr, - visited_hash_bitlen, - local_traversed_hashmap_ptr, - traversed_hash_bitlen, - parent_indices_buffer, - result_indices_buffer, - 1, - result_position, - result_buffer_size_32); + // Compute the distances between child nodes and the query node. + graph_provider.expand(result_indices_buffer, + result_distances_buffer, + smem_desc, + local_visited_hashmap_ptr, + visited_hash_bitlen, + local_traversed_hashmap_ptr, + traversed_hash_bitlen, + parent_indices_buffer, + result_indices_buffer, + result_position, + result_buffer_size_32); __syncthreads(); - // Check the state of the nodes in the result buffer which were not updated - // by the compute_distance_to_child_nodes above, and if it cannot be used as - // a parent node, it is deactivated. + // Check the state of result-buffer nodes not updated by graph expansion. Deactivate any node + // that cannot be used as a parent. for (uint32_t i = threadIdx.x; i < result_position[0]; i += blockDim.x) { INDEX_T index = result_indices_buffer[i]; if (index == invalid_index || index & index_msb_1_mask) { continue; } @@ -298,6 +385,8 @@ RAFT_DEVICE_INLINE_FUNCTION void search_core( iter++; } + graph_provider.finalize(); + // Filtering - use extern sample_filter function (linked via JIT LTO) for (uint32_t i = threadIdx.x; i < result_buffer_size_32; i += blockDim.x) { INDEX_T index = result_indices_buffer[i]; @@ -414,29 +503,30 @@ __device__ void search_kernel_jit( const uint32_t query_id_offset, // Offset to add to query_id when calling filter cagra_sample_filter filter_payload) { - search_core(result_indices_ptr, - result_distances_ptr, - dataset_desc, - queries_ptr, - knn_graph, - max_elements, - graph_degree, - /*buffer_graph_degree=*/graph_degree, - source_indices_ptr, - num_distilation, - rand_xor_mask, - seed_ptr, - num_seeds, - visited_hash_bitlen, - traversed_hashmap_ptr, - traversed_hash_bitlen, - itopk_size, - min_iteration, - max_iteration, - num_executed_iterations, - graph_size, - query_id_offset, - filter_payload); + search_core( + result_indices_ptr, + result_distances_ptr, + dataset_desc, + queries_ptr, + device_graph_provider{knn_graph, graph_degree}, + max_elements, + graph_degree, + /*buffer_candidate_count=*/graph_degree, + source_indices_ptr, + num_distilation, + rand_xor_mask, + seed_ptr, + num_seeds, + visited_hash_bitlen, + traversed_hashmap_ptr, + traversed_hash_bitlen, + itopk_size, + min_iteration, + max_iteration, + num_executed_iterations, + graph_size, + query_id_offset, + filter_payload); } // Multi-partition variant of search_kernel_jit. Grid is (num_cta_per_query, num_queries, @@ -474,29 +564,30 @@ __device__ void search_multi_cta_mp_jit( cagra_sample_filter filter_payload{}; filter_payload.filter_data = (part.bitset_ptr != nullptr) ? &part_bitset : nullptr; - search_core(result_indices_ptr, - result_distances_ptr, - part.dataset_desc, - queries_ptr, - part.graph, - max_elements, - /*graph_degree=*/part.graph_degree, - /*buffer_graph_degree=*/max_graph_degree, - /*source_indices_ptr=*/nullptr, - num_distilation, - rand_xor_mask, - /*seed_ptr=*/nullptr, - /*num_seeds=*/0, - visited_hash_bitlen, - traversed_hashmap_ptr, - traversed_hash_bitlen, - itopk_size, - min_iteration, - max_iteration, - /*num_executed_iterations=*/nullptr, - /*graph_size=*/0, - query_id_offset, - filter_payload); + search_core( + result_indices_ptr, + result_distances_ptr, + part.dataset_desc, + queries_ptr, + device_graph_provider{part.graph, part.graph_degree}, + max_elements, + /*graph_degree=*/part.graph_degree, + /*buffer_candidate_count=*/max_graph_degree, + /*source_indices_ptr=*/nullptr, + num_distilation, + rand_xor_mask, + /*seed_ptr=*/nullptr, + /*num_seeds=*/0, + visited_hash_bitlen, + traversed_hashmap_ptr, + traversed_hash_bitlen, + itopk_size, + min_iteration, + max_iteration, + /*num_executed_iterations=*/nullptr, + /*graph_size=*/0, + query_id_offset, + filter_payload); } } // namespace cuvs::neighbors::cagra::detail::multi_cta_search diff --git a/cpp/src/neighbors/detail/cagra/jit_lto_kernels/search_single_cta_jit.cuh b/cpp/src/neighbors/detail/cagra/jit_lto_kernels/search_single_cta_jit.cuh index 47f02fb6ef..e40e54e056 100644 --- a/cpp/src/neighbors/detail/cagra/jit_lto_kernels/search_single_cta_jit.cuh +++ b/cpp/src/neighbors/detail/cagra/jit_lto_kernels/search_single_cta_jit.cuh @@ -46,6 +46,74 @@ using cuvs::neighbors::detail::sample_filter; using cuvs::neighbors::cagra::detail::device::compute_distance_to_child_nodes_jit; using cuvs::neighbors::cagra::detail::device::compute_distance_to_random_nodes_jit; +template +struct device_graph_provider { + const IndexT* graph; + std::uint32_t degree; + + RAFT_DEVICE_INLINE_FUNCTION void initialize() const {} + + RAFT_DEVICE_INLINE_FUNCTION auto candidate_capacity(std::uint32_t search_width) const + -> std::uint32_t + { + return search_width * degree; + } + + RAFT_DEVICE_INLINE_FUNCTION auto result_buffer_size(std::uint32_t base_size, std::uint32_t) const + -> std::uint32_t + { + return base_size; + } + + RAFT_DEVICE_INLINE_FUNCTION auto seed_pointer(const IndexT* seeds, + std::uint32_t num_seeds, + std::uint32_t query_id) const -> const IndexT* + { + return seeds == nullptr ? nullptr : seeds + static_cast(num_seeds) * query_id; + } + + RAFT_DEVICE_INLINE_FUNCTION auto initial_candidate_count(std::uint32_t candidate_count, + std::uint32_t, + std::uint32_t, + std::uint32_t) const -> std::uint32_t + { + return candidate_count; + } + + RAFT_DEVICE_INLINE_FUNCTION auto initial_distillation_count(std::uint32_t count, + std::uint32_t) const -> std::uint32_t + { + return count; + } + + RAFT_DEVICE_INLINE_FUNCTION auto expand( + IndexT* result_indices, + DistanceT* result_distances, + const dataset_descriptor_base_t* dataset_desc, + IndexT* visited_hashmap, + std::uint32_t visited_hash_bitlen, + const IndexT* parent_indices, + const IndexT* internal_topk, + std::uint32_t search_width) const -> std::uint32_t + { + compute_distance_to_child_nodes_jit(result_indices, + result_distances, + dataset_desc, + graph, + degree, + visited_hashmap, + visited_hash_bitlen, + static_cast(nullptr), + 0u, + parent_indices, + internal_topk, + search_width); + return search_width * degree; + } + + RAFT_DEVICE_INLINE_FUNCTION void finalize() const {} +}; + // JIT search_core - setup_workspace/compute_distance via function pointers template > RAFT_DEVICE_INLINE_FUNCTION void search_core( uintptr_t result_indices_ptr, DistanceT* const result_distances_ptr, const std::uint32_t top_k, const DataT* const queries_ptr, - const IndexT* const knn_graph, + GraphProviderT graph_provider, const std::uint32_t graph_degree, const SourceIndexT* source_indices_ptr, const unsigned num_distilation, @@ -111,7 +180,10 @@ RAFT_DEVICE_INLINE_FUNCTION void search_core( extern __shared__ uint8_t smem[]; // Layout of result_buffer - const auto result_buffer_size = internal_topk + (search_width * graph_degree); + const auto candidate_capacity = graph_provider.candidate_capacity(search_width); + const auto base_result_buffer_size = internal_topk + candidate_capacity; + const auto result_buffer_size = + graph_provider.result_buffer_size(base_result_buffer_size, num_seeds); const auto result_buffer_size_32 = raft::round_up_safe(result_buffer_size, 32); const auto small_hash_size = hashmap::get_size(small_hash_bitlen); @@ -150,19 +222,24 @@ RAFT_DEVICE_INLINE_FUNCTION void search_core( local_visited_hashmap_ptr = visited_hashmap_ptr + (hashmap::get_size(hash_bitlen) * blockIdx.y); } hashmap::init(local_visited_hashmap_ptr, hash_bitlen, 0); + graph_provider.initialize(); __syncthreads(); _CLK_REC(clk_init); // compute distance to randomly selecting nodes using JIT version _CLK_START(); - const IndexT* const local_seed_ptr = seed_ptr ? seed_ptr + (num_seeds * query_id) : nullptr; + const IndexT* const local_seed_ptr = graph_provider.seed_pointer(seed_ptr, num_seeds, query_id); + auto const initial_candidate_count = + graph_provider.initial_candidate_count(result_buffer_size, num_seeds, 0, 1); + auto const initial_distillation_count = + graph_provider.initial_distillation_count(num_distilation, num_seeds); // Get dataset_size directly from base descriptor IndexT dataset_size = smem_desc->size; compute_distance_to_random_nodes_jit(result_indices_buffer, result_distances_buffer, smem_desc, - result_buffer_size, - num_distilation, + initial_candidate_count, + initial_distillation_count, rand_xor_mask, local_seed_ptr, num_seeds, @@ -173,10 +250,16 @@ RAFT_DEVICE_INLINE_FUNCTION void search_core( 0, 1, graph_size); + for (std::uint32_t i = initial_candidate_count + threadIdx.x; i < result_buffer_size; + i += blockDim.x) { + result_indices_buffer[i] = raft::upper_bound(); + result_distances_buffer[i] = raft::upper_bound(); + } __syncthreads(); _CLK_REC(clk_compute_1st_distance); - std::uint32_t iter = 0; + std::uint32_t iter = 0; + std::uint32_t computed_children_this_iteration = result_buffer_size - internal_topk; while (1) { // sort if constexpr (TOPK_BY_BITONIC_SORT) { @@ -225,7 +308,7 @@ RAFT_DEVICE_INLINE_FUNCTION void search_core( result_distances_buffer + internal_topk, result_indices_buffer + internal_topk, max_candidates, - search_width * graph_degree, + computed_children_this_iteration, topk_ws, (iter == 0)); __syncthreads(); @@ -280,19 +363,15 @@ RAFT_DEVICE_INLINE_FUNCTION void search_core( __syncthreads(); // compute the norms between child nodes and query node using JIT version _CLK_START(); - compute_distance_to_child_nodes_jit( - result_indices_buffer + internal_topk, - result_distances_buffer + internal_topk, - smem_desc, - knn_graph, - graph_degree, - local_visited_hashmap_ptr, - hash_bitlen, - (IndexT*)nullptr, - 0u, - parent_list_buffer, - result_indices_buffer, - search_width); + computed_children_this_iteration = + graph_provider.expand(result_indices_buffer + internal_topk, + result_distances_buffer + internal_topk, + smem_desc, + local_visited_hashmap_ptr, + hash_bitlen, + parent_list_buffer, + result_indices_buffer, + search_width); // Critical: __syncthreads() must be reached by ALL threads // If any thread is stuck in compute_distance_to_child_nodes_jit, this will hang __syncthreads(); @@ -322,11 +401,13 @@ RAFT_DEVICE_INLINE_FUNCTION void search_core( iter++; } + graph_provider.finalize(); + // Post process for filtering - use extern sample_filter function constexpr IndexT index_msb_1_mask = utils::gen_index_msb_1_mask::value; const IndexT invalid_index = utils::get_max_value(); - for (unsigned i = threadIdx.x; i < internal_topk + search_width * graph_degree; i += blockDim.x) { + for (unsigned i = threadIdx.x; i < internal_topk + candidate_capacity; i += blockDim.x) { const auto node_id = result_indices_buffer[i] & ~index_msb_1_mask; if (node_id != (invalid_index & ~index_msb_1_mask) && !sample_filter(query_id + query_id_offset, @@ -387,7 +468,7 @@ RAFT_DEVICE_INLINE_FUNCTION void search_core( result_distances_buffer + internal_topk, result_indices_buffer + internal_topk, max_candidates, - search_width * graph_degree, + computed_children_this_iteration, topk_ws, (iter == 0)); } @@ -418,9 +499,13 @@ RAFT_DEVICE_INLINE_FUNCTION void search_core( if constexpr (TOPK_BY_BITONIC_SORT) { ii = device::swizzling(i); } if (result_distances_ptr != nullptr) { result_distances_ptr[j] = result_distances_buffer[ii]; } - auto internal_index = - result_indices_buffer[ii] & ~index_msb_1_mask; // clear most significant bit - auto source_index = to_source_index(internal_index); + auto const result_index = result_indices_buffer[ii]; + if (result_index == invalid_index) { + write_indices(result_indices_ptr, j, static_cast(invalid_index)); + continue; + } + auto internal_index = result_index & ~index_msb_1_mask; // clear most significant bit + auto source_index = to_source_index(internal_index); write_indices(result_indices_ptr, j, source_index); } if (threadIdx.x == 0 && num_executed_iterations != nullptr) { @@ -495,33 +580,34 @@ __device__ void search_kernel_jit( DataT, IndexT, DistanceT, - SourceIndexT>(result_indices_ptr, - result_distances_ptr, - top_k, - queries_ptr, - knn_graph, - graph_degree, - source_indices_ptr, - num_distilation, - rand_xor_mask, - seed_ptr, - num_seeds, - visited_hashmap_ptr, - max_candidates, - max_itopk, - internal_topk, - search_width, - min_iteration, - max_iteration, - num_executed_iterations, - hash_bitlen, - small_hash_bitlen, - small_hash_reset_interval, - query_id, - query_id_offset, - dataset_desc, - filter_payload, - graph_size); + SourceIndexT>( + result_indices_ptr, + result_distances_ptr, + top_k, + queries_ptr, + device_graph_provider{knn_graph, graph_degree}, + graph_degree, + source_indices_ptr, + num_distilation, + rand_xor_mask, + seed_ptr, + num_seeds, + visited_hashmap_ptr, + max_candidates, + max_itopk, + internal_topk, + search_width, + min_iteration, + max_iteration, + num_executed_iterations, + hash_bitlen, + small_hash_bitlen, + small_hash_reset_interval, + query_id, + query_id_offset, + dataset_desc, + filter_payload, + graph_size); } // JIT persistent device implementation - called from extern "C" __global__ entry in generated .cu @@ -605,33 +691,34 @@ __device__ void search_single_cta_p_impl( DataT, IndexT, DistanceT, - SourceIndexT>(result_indices_ptr, - result_distances_ptr, - top_k, - queries_ptr, - knn_graph, - graph_degree, - source_indices_ptr, - num_distilation, - rand_xor_mask, - seed_ptr, - num_seeds, - visited_hashmap_ptr, - max_candidates, - max_itopk, - internal_topk, - search_width, - min_iteration, - max_iteration, - num_executed_iterations, - hash_bitlen, - small_hash_bitlen, - small_hash_reset_interval, - query_id, - query_id_offset, - dataset_desc, - filter_payload, - graph_size); + SourceIndexT>( + result_indices_ptr, + result_distances_ptr, + top_k, + queries_ptr, + device_graph_provider{knn_graph, graph_degree}, + graph_degree, + source_indices_ptr, + num_distilation, + rand_xor_mask, + seed_ptr, + num_seeds, + visited_hashmap_ptr, + max_candidates, + max_itopk, + internal_topk, + search_width, + min_iteration, + max_iteration, + num_executed_iterations, + hash_bitlen, + small_hash_bitlen, + small_hash_reset_interval, + query_id, + query_id_offset, + dataset_desc, + filter_payload, + graph_size); // make sure all writes are visible even for the host // (e.g. when result buffers are in pinned memory) @@ -709,33 +796,34 @@ __device__ void search_single_cta_mp_impl( DataT, IndexT, DistanceT, - SourceIndexT>(tagged_indices_ptr, - part_result_distances, - top_k, - queries_ptr, - part.graph, - part.graph_degree, - static_cast(nullptr), - num_distilation, - rand_xor_mask, - static_cast(nullptr), - num_seeds, - part_hashmap_ptr, - max_candidates, - max_itopk, - internal_topk, - search_width, - min_iteration, - max_iteration, - num_executed_iterations, - hash_bitlen, - small_hash_bitlen, - small_hash_reset_interval, - query_id, - query_id_offset, - part.dataset_desc, - filter_payload, - /*graph_size=*/IndexT{0}); + SourceIndexT>( + tagged_indices_ptr, + part_result_distances, + top_k, + queries_ptr, + device_graph_provider{part.graph, part.graph_degree}, + part.graph_degree, + static_cast(nullptr), + num_distilation, + rand_xor_mask, + static_cast(nullptr), + num_seeds, + part_hashmap_ptr, + max_candidates, + max_itopk, + internal_topk, + search_width, + min_iteration, + max_iteration, + num_executed_iterations, + hash_bitlen, + small_hash_bitlen, + small_hash_reset_interval, + query_id, + query_id_offset, + part.dataset_desc, + filter_payload, + /*graph_size=*/IndexT{0}); } } // namespace cuvs::neighbors::cagra::detail::single_cta_search diff --git a/cpp/src/neighbors/detail/flowann/README.md b/cpp/src/neighbors/detail/flowann/README.md new file mode 100644 index 0000000000..52394332f5 --- /dev/null +++ b/cpp/src/neighbors/detail/flowann/README.md @@ -0,0 +1,5 @@ +# FlowANN implementation + +FlowANN is described in "Disentangling Graph Dependencies for Efficient Billion-Scale GPU Vector Search" (OSDI 2026). + +These are internal implementation headers. The integration is introduced in separate PRs; only the interfaces present in the current revision are available. diff --git a/cpp/src/neighbors/detail/flowann/flowann.hpp b/cpp/src/neighbors/detail/flowann/flowann.hpp new file mode 100644 index 0000000000..18c22227f5 --- /dev/null +++ b/cpp/src/neighbors/detail/flowann/flowann.hpp @@ -0,0 +1,736 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * FlowANN is described in "Disentangling Graph Dependencies for Efficient + * Billion-Scale GPU Vector Search" (OSDI 2026). + */ + +#pragma once + +#ifdef CUVS_ENABLE_FLOWANN_SEARCH + +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace cuvs::neighbors::cagra::experimental::flowann { + +namespace detail { +struct search_context_access; +} + +class search_session; + +/** Parameters controlling the host-serviced cross-graph queues. */ +struct queue_params { + /** Number of independent GPU-to-CPU command queues. */ + std::uint32_t num_queues = 1; + /** Number of CPU pause instructions issued after an empty poll. */ + std::uint32_t empty_pause = 64; + /** Collect aggregate queue and polling statistics. */ + bool collect_statistics = false; +}; + +/** Aggregate counters collected by the host-serviced cross-graph queues. */ +struct queue_statistics { + std::uint64_t polls{}; + std::uint64_t empty_polls{}; + std::uint64_t processed_commands{}; +}; + +/** + * @brief Owns FlowANN command queues and their CPU polling threads. + * + * The cross-graph storage passed to the constructor must outlive this context. A context can be + * reused across search calls for the same index. Construction and search must use the same current + * CUDA device; CPU pollers select that device before servicing requests. + */ +class CUVS_EXPORT search_context { + public: + search_context(raft::host_matrix_view cross_graph, + queue_params params = {}); + ~search_context(); + + search_context(search_context const&) = delete; + auto operator=(search_context const&) -> search_context& = delete; + search_context(search_context&&) noexcept; + auto operator=(search_context&&) noexcept -> search_context&; + + /** + * Pause all CPU pollers after in-flight work has completed. + * + * @throws raft::logic_error if a search_session is active. + */ + void pause(); + /** Resume all CPU pollers before submitting search work. */ + void resume(); + [[nodiscard]] auto num_queues() const noexcept -> std::uint32_t; + /** Return aggregate counters. Counters are zero unless statistics were enabled. */ + [[nodiscard]] auto statistics() const noexcept -> queue_statistics; + + private: + struct impl; + std::unique_ptr impl_; + + friend struct detail::search_context_access; + friend class search_session; +}; + +/** + * Keep a search context's pollers running for the lifetime of this session. + * + * Use a session for repeated low-latency search calls. The search context must outlive the session + * and must not be moved while the session is active. + */ +class CUVS_EXPORT search_session { + public: + explicit search_session(search_context& context); + ~search_session(); + + search_session(search_session const&) = delete; + auto operator=(search_session const&) -> search_session& = delete; + search_session(search_session&& other) noexcept; + auto operator=(search_session&& other) noexcept -> search_session&; + + private: + void reset() noexcept; + + search_context* context_{}; +}; + +/** + * @brief FlowANN tiered-graph index. + * + * The index stores a compressed inner graph in device memory and the cross-graph rows in host + * memory. The dataset is a non-owning view and its backing storage must outlive the index. + */ +template > +class index : public cuvs::neighbors::index { + public: + using value_type = T; + using index_type = IdxT; + using graph_index_type = std::uint32_t; + using inner_index_type = std::uint8_t; + using dataset_view_type = DatasetViewT; + + static_assert(!raft::is_narrowing_v, + "IdxT must be able to represent all uint32_t graph indices"); + + explicit index(raft::resources const& res, + cuvs::distance::DistanceType metric = cuvs::distance::DistanceType::L2Expanded, + std::uint16_t node_per_cacheline = 4, + std::uint16_t n_bits = 16) + : metric_(metric), + node_per_cacheline_(node_per_cacheline), + n_bits_(n_bits), + inner_graph_(raft::make_device_matrix(res, 0, 0)), + cross_graph_(raft::make_host_matrix(0, 0)), + dataset_norms_(std::nullopt) + { + validate_encoding_(); + } + + index(raft::resources const& res, + cuvs::distance::DistanceType metric, + DatasetViewT dataset, + std::uint16_t node_per_cacheline = 4, + std::uint16_t n_bits = 16) + : metric_(metric), + node_per_cacheline_(node_per_cacheline), + n_bits_(n_bits), + inner_graph_(raft::make_device_matrix(res, 0, 0)), + cross_graph_(raft::make_host_matrix(0, 0)), + dataset_(dataset), + dataset_norms_(std::nullopt) + { + validate_encoding_(); + if constexpr (cuvs::neighbors::is_device_padded_dataset_view_v) { + if (metric_ == cuvs::distance::DistanceType::CosineExpanded && dataset_.n_rows() > 0) { + compute_dataset_norms_(res); + raft::resource::sync_stream(res); + } + } + } + + index(index const&) = delete; + auto operator=(index const&) -> index& = delete; + index(index&&) = default; + auto operator=(index&&) -> index& = default; + ~index() = default; + + [[nodiscard]] constexpr auto metric() const noexcept -> cuvs::distance::DistanceType + { + return metric_; + } + + [[nodiscard]] auto size() const noexcept -> IdxT + { + return static_cast(cross_graph_.extent(0)); + } + + [[nodiscard]] auto dim() const noexcept -> std::uint32_t { return dataset_.dim(); } + + [[nodiscard]] auto graph_degree() const noexcept -> std::uint32_t { return full_graph_degree_; } + + /** Uniform number of neighbors stored in the packed resident graph. Zero for legacy layouts. */ + [[nodiscard]] constexpr auto resident_degree() const noexcept -> std::uint32_t + { + return resident_degree_; + } + + /** Maximum number of host-resident neighbors stored after the count column. */ + [[nodiscard]] auto cross_degree() const noexcept -> std::uint32_t + { + return cross_graph_.extent(1) == 0 ? 0 : static_cast(cross_graph_.extent(1) - 1); + } + + [[nodiscard]] constexpr auto node_per_cacheline() const noexcept -> std::uint16_t + { + return node_per_cacheline_; + } + + [[nodiscard]] constexpr auto n_bits() const noexcept -> std::uint16_t { return n_bits_; } + + [[nodiscard]] auto dataset() const noexcept -> DatasetViewT const& { return dataset_; } + + /** Dataset norms used by cosine distance. */ + [[nodiscard]] auto dataset_norms() const noexcept + -> std::optional> + { + if (!dataset_norms_.has_value()) { return std::nullopt; } + return raft::make_const_mdspan(dataset_norms_->view()); + } + + void update_dataset(raft::resources const& res, DatasetViewT dataset) + { + if (cross_graph_.extent(0) > 0) { + RAFT_EXPECTS(dataset.n_rows() == cross_graph_.extent(0), + "Dataset and FlowANN graph must have the same number of rows"); + } + dataset_ = dataset; + dataset_norms_.reset(); + if constexpr (cuvs::neighbors::is_device_padded_dataset_view_v) { + if (metric_ == cuvs::distance::DistanceType::CosineExpanded && dataset_.n_rows() > 0) { + compute_dataset_norms_(res); + raft::resource::sync_stream(res); + } + } + } + + [[nodiscard]] auto inner_graph() const noexcept + -> raft::device_matrix_view + { + return raft::make_const_mdspan(inner_graph_.view()); + } + + /** + * Cross-graph rows have shape `[size, cross_degree + 1]`. Element zero stores the number of valid + * cross-graph neighbors in the rest of the row. `cross_degree` may be smaller than the complete + * graph degree when every node has a uniform resident degree. + */ + [[nodiscard]] auto cross_graph() const noexcept + -> raft::host_matrix_view + { + return raft::make_const_mdspan(cross_graph_.view()); + } + + [[nodiscard]] auto source_indices() const noexcept + -> std::optional> + { + if (!source_indices_.has_value()) { return std::nullopt; } + return raft::make_const_mdspan(source_indices_->view()); + } + + [[nodiscard]] auto subgraph_offsets() const noexcept + -> std::optional> + { + if (!subgraph_offsets_.has_value()) { return std::nullopt; } + return raft::make_const_mdspan(subgraph_offsets_->view()); + } + + /** Subgraph identifier for every internal node. */ + [[nodiscard]] auto subgraph_ids() const noexcept + -> std::optional> + { + if (!subgraph_ids_.has_value()) { return std::nullopt; } + return raft::make_const_mdspan(subgraph_ids_->view()); + } + + /** Optional medoid seed nodes shared by all queries. */ + [[nodiscard]] auto seeds() const noexcept + -> std::optional> + { + if (!seeds_.has_value()) { return std::nullopt; } + return raft::make_const_mdspan(seeds_->view()); + } + + /** Cached cross-graph rows for the fixed medoid seeds. */ + [[nodiscard]] auto seed_cross_graph() const noexcept + -> std::optional> + { + if (!seed_cross_graph_.has_value()) { return std::nullopt; } + return raft::make_const_mdspan(seed_cross_graph_->view()); + } + + [[nodiscard]] auto seed_lookup_keys() const noexcept + -> std::optional> + { + if (!seed_lookup_keys_.has_value()) { return std::nullopt; } + return raft::make_const_mdspan(seed_lookup_keys_->view()); + } + + [[nodiscard]] auto seed_lookup_values() const noexcept + -> std::optional> + { + if (!seed_lookup_values_.has_value()) { return std::nullopt; } + return raft::make_const_mdspan(seed_lookup_values_->view()); + } + + void update_graph( + raft::resources const& res, + raft::host_matrix_view inner_graph, + raft::host_matrix_view cross_graph) + { + auto const inferred_degree = + cross_graph.extent(1) == 0 ? 0u : static_cast(cross_graph.extent(1) - 1); + update_graph(res, inner_graph, cross_graph, inferred_degree, 0); + } + + /** Replace the graph and explicitly describe a compact uniform-residency layout. */ + void update_graph( + raft::resources const& res, + raft::host_matrix_view inner_graph, + raft::host_matrix_view cross_graph, + std::uint32_t full_graph_degree, + std::uint32_t resident_degree) + { + validate_graph_(inner_graph, cross_graph, full_graph_degree, resident_degree); + + auto new_inner = raft::make_device_matrix( + res, inner_graph.extent(0), inner_graph.extent(1)); + raft::copy(new_inner.data_handle(), + inner_graph.data_handle(), + inner_graph.size(), + raft::resource::get_cuda_stream(res)); + + auto new_cross = raft::make_host_matrix(cross_graph.extent(0), + cross_graph.extent(1)); + std::copy_n(cross_graph.data_handle(), cross_graph.size(), new_cross.data_handle()); + + inner_graph_ = std::move(new_inner); + cross_graph_ = std::move(new_cross); + full_graph_degree_ = full_graph_degree; + resident_degree_ = resident_degree; + rebuild_seed_cross_graph_(res); + } + + /** Replace the graph while taking ownership of the host cross-graph allocation. */ + void update_graph(raft::resources const& res, + raft::host_matrix&& inner_graph, + raft::host_matrix&& cross_graph) + { + auto const inferred_degree = + cross_graph.extent(1) == 0 ? 0u : static_cast(cross_graph.extent(1) - 1); + update_graph(res, std::move(inner_graph), std::move(cross_graph), inferred_degree, 0); + } + + /** Replace the graph while taking ownership of a compact uniform-residency layout. */ + void update_graph(raft::resources const& res, + raft::host_matrix&& inner_graph, + raft::host_matrix&& cross_graph, + std::uint32_t full_graph_degree, + std::uint32_t resident_degree) + { + validate_graph_(raft::make_const_mdspan(inner_graph.view()), + raft::make_const_mdspan(cross_graph.view()), + full_graph_degree, + resident_degree); + + auto new_inner = raft::make_device_matrix( + res, inner_graph.extent(0), inner_graph.extent(1)); + raft::copy(new_inner.data_handle(), + inner_graph.data_handle(), + inner_graph.size(), + raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + + inner_graph_ = std::move(new_inner); + cross_graph_ = std::move(cross_graph); + full_graph_degree_ = full_graph_degree; + resident_degree_ = resident_degree; + rebuild_seed_cross_graph_(res); + } + + /** Replace a variable-residency grouped graph while retaining the packed rows on the device. */ + void update_graph(raft::resources const& res, + raft::device_matrix&& inner_graph, + raft::host_matrix&& cross_graph, + bool validate_edges) + { + RAFT_EXPECTS(cross_graph.extent(0) > 0 && cross_graph.extent(1) > 1, + "FlowANN grouped cross graph must contain rows and neighbors"); + auto const full_degree = static_cast(cross_graph.extent(1) - 1); + auto const expected_rows = + raft::div_rounding_up_safe(cross_graph.extent(0), node_per_cacheline_); + auto const has_inner_graph = inner_graph.extent(0) != 0 || inner_graph.extent(1) != 0; + if (has_inner_graph) { + RAFT_EXPECTS(inner_graph.extent(0) == expected_rows, + "FlowANN grouped inner graph has an invalid row count"); + RAFT_EXPECTS(inner_graph.extent(1) >= node_per_cacheline_, + "FlowANN grouped inner graph rows cannot hold their count headers"); + } + if (dataset_.n_rows() > 0) { + RAFT_EXPECTS(dataset_.n_rows() == cross_graph.extent(0), + "Dataset and FlowANN grouped graph must have the same number of rows"); + } + for (int64_t node = 0; node < cross_graph.extent(0); ++node) { + auto const cross_count = cross_graph(node, 0); + RAFT_EXPECTS(cross_count <= full_degree, + "FlowANN grouped cross-neighbor count exceeds graph degree at row %ld", + static_cast(node)); + if (validate_edges) { + for (std::uint32_t edge = 0; edge < cross_count; ++edge) { + RAFT_EXPECTS(cross_graph(node, edge + 1) < cross_graph.extent(0), + "FlowANN grouped cross graph contains an out-of-range node at row %ld", + static_cast(node)); + } + } + } + + inner_graph_ = std::move(inner_graph); + cross_graph_ = std::move(cross_graph); + full_graph_degree_ = full_degree; + resident_degree_ = 0; + rebuild_seed_cross_graph_(res); + } + + /** Replace a compact graph while taking ownership of both device and host allocations. */ + void update_graph(raft::resources const& res, + raft::device_matrix&& inner_graph, + raft::host_matrix&& cross_graph, + std::uint32_t full_graph_degree, + std::uint32_t resident_degree) + { + RAFT_EXPECTS(cross_graph.extent(0) > 0 && cross_graph.extent(1) > 0, + "FlowANN cross graph must contain rows and a neighbor-count column"); + RAFT_EXPECTS(full_graph_degree > 0 && resident_degree <= full_graph_degree, + "FlowANN compact graph degree metadata is invalid"); + auto const cross_degree = static_cast(cross_graph.extent(1) - 1); + RAFT_EXPECTS(resident_degree + cross_degree == full_graph_degree, + "FlowANN compact graph degrees do not add up to the complete graph degree"); + if (resident_degree == 0) { + RAFT_EXPECTS(inner_graph.extent(0) == 0 && inner_graph.extent(1) == 0, + "FlowANN zero-residency graph must not allocate packed rows"); + } else { + auto const expected_rows = + raft::div_rounding_up_safe(cross_graph.extent(0), node_per_cacheline_); + auto const required_bits = + static_cast(node_per_cacheline_) * 8 + + static_cast(node_per_cacheline_) * resident_degree * n_bits_; + RAFT_EXPECTS(inner_graph.extent(0) == expected_rows, + "FlowANN compact inner graph has an invalid row count"); + RAFT_EXPECTS(static_cast(inner_graph.extent(1)) * 8 >= required_bits, + "FlowANN compact inner graph rows are too short"); + } + if (dataset_.n_rows() > 0) { + RAFT_EXPECTS(dataset_.n_rows() == cross_graph.extent(0), + "Dataset and FlowANN compact graph must have the same number of rows"); + } + for (int64_t node = 0; node < cross_graph.extent(0); ++node) { + RAFT_EXPECTS(cross_graph(node, 0) == cross_degree, + "FlowANN uniform compact graph has an invalid cross count at row %ld", + static_cast(node)); + for (std::uint32_t edge = 0; edge < cross_degree; ++edge) { + RAFT_EXPECTS(cross_graph(node, edge + 1) < cross_graph.extent(0), + "FlowANN compact cross graph contains an out-of-range node at row %ld", + static_cast(node)); + } + } + + inner_graph_ = std::move(inner_graph); + cross_graph_ = std::move(cross_graph); + full_graph_degree_ = full_graph_degree; + resident_degree_ = resident_degree; + rebuild_seed_cross_graph_(res); + } + + void update_source_indices( + raft::resources const& res, + raft::host_vector_view source_indices) + { + RAFT_EXPECTS(source_indices.extent(0) == static_cast(size()), + "FlowANN source-index mapping must have one entry per graph row"); + source_indices_.emplace(raft::make_device_vector(res, source_indices.extent(0))); + raft::copy(source_indices_->data_handle(), + source_indices.data_handle(), + source_indices.extent(0), + raft::resource::get_cuda_stream(res)); + } + + void update_seeds(raft::resources const& res, + raft::host_vector_view seeds) + { + RAFT_EXPECTS(seeds.extent(0) > 0, "FlowANN seeds must not be empty"); + RAFT_EXPECTS(seeds.extent(0) <= static_cast(size()), + "FlowANN seed count must not exceed the index size"); + for (int64_t i = 0; i < seeds.extent(0); ++i) { + RAFT_EXPECTS( + seeds(i) < size(), "FlowANN seed at position %ld is out of range", static_cast(i)); + } + seeds_.emplace(raft::make_device_vector(res, seeds.extent(0))); + raft::copy(seeds_->data_handle(), + seeds.data_handle(), + seeds.extent(0), + raft::resource::get_cuda_stream(res)); + seed_ids_host_.assign(seeds.data_handle(), seeds.data_handle() + seeds.extent(0)); + rebuild_seed_lookup_(res); + rebuild_seed_cross_graph_(res); + } + + void update_subgraph_layout( + raft::resources const& res, + raft::host_vector_view offsets, + raft::host_vector_view subgraph_ids) + { + RAFT_EXPECTS(offsets.extent(0) > 0, "FlowANN subgraph offsets must not be empty"); + RAFT_EXPECTS(offsets.extent(0) <= static_cast(size()), + "FlowANN subgraph count must not exceed the index size"); + RAFT_EXPECTS(subgraph_ids.extent(0) == static_cast(size()), + "FlowANN must store one subgraph identifier per graph row"); + RAFT_EXPECTS(offsets(0) == 0, "FlowANN first subgraph offset must be zero"); + for (int64_t i = 1; i < offsets.extent(0); ++i) { + RAFT_EXPECTS(offsets(i - 1) < offsets(i), + "FlowANN subgraph offsets must be strictly increasing"); + } + RAFT_EXPECTS(offsets(offsets.extent(0) - 1) < static_cast(size()), + "FlowANN subgraph offsets exceed the index size"); + for (int64_t i = 0; i < subgraph_ids.extent(0); ++i) { + RAFT_EXPECTS(subgraph_ids(i) < offsets.extent(0), + "FlowANN subgraph identifier is out of range at row %ld", + static_cast(i)); + } + + subgraph_offsets_.emplace( + raft::make_device_vector(res, offsets.extent(0))); + raft::copy(subgraph_offsets_->data_handle(), + offsets.data_handle(), + offsets.extent(0), + raft::resource::get_cuda_stream(res)); + subgraph_ids_.emplace( + raft::make_device_vector(res, subgraph_ids.extent(0))); + raft::copy(subgraph_ids_->data_handle(), + subgraph_ids.data_handle(), + subgraph_ids.extent(0), + raft::resource::get_cuda_stream(res)); + } + + private: + void validate_graph_( + raft::host_matrix_view inner_graph, + raft::host_matrix_view cross_graph, + std::uint32_t full_graph_degree, + std::uint32_t resident_degree) const + { + RAFT_EXPECTS(cross_graph.extent(1) > 0, + "FlowANN cross graph must include a neighbor-count column"); + RAFT_EXPECTS(static_cast(cross_graph.extent(0)) <= + static_cast(std::numeric_limits::max()), + "FlowANN graph row count exceeds the index type"); + auto const expected_inner_rows = + raft::div_rounding_up_safe(cross_graph.extent(0), node_per_cacheline_); + auto const has_inner_graph = inner_graph.extent(0) != 0 || inner_graph.extent(1) != 0; + if (has_inner_graph) { + RAFT_EXPECTS(inner_graph.extent(0) == expected_inner_rows, + "FlowANN inner graph must contain one packed row per node cacheline"); + RAFT_EXPECTS(inner_graph.extent(1) >= node_per_cacheline_, + "FlowANN inner graph rows are too short to contain neighbor counts"); + } else { + RAFT_EXPECTS(resident_degree == 0, + "FlowANN may omit the packed graph only when resident degree is zero"); + } + RAFT_EXPECTS(full_graph_degree > 0, "FlowANN complete graph degree must be positive"); + RAFT_EXPECTS(resident_degree <= full_graph_degree, + "FlowANN resident graph degree exceeds the complete graph degree"); + auto const cross_degree = static_cast(cross_graph.extent(1) - 1); + RAFT_EXPECTS(cross_degree <= full_graph_degree, + "FlowANN cross graph degree exceeds the complete graph degree"); + if (resident_degree > 0 || cross_degree < full_graph_degree) { + RAFT_EXPECTS(resident_degree + cross_degree == full_graph_degree, + "FlowANN compact graph degrees do not add up to the complete graph degree"); + } + if (dataset_.n_rows() > 0) { + RAFT_EXPECTS(dataset_.n_rows() == cross_graph.extent(0), + "Dataset and FlowANN cross graph must have the same number of rows"); + } + + for (int64_t node = 0; node < cross_graph.extent(0); ++node) { + auto const cross_count = cross_graph(node, 0); + RAFT_EXPECTS(cross_count <= cross_degree, + "FlowANN cross-neighbor count exceeds graph degree at row %ld", + static_cast(node)); + for (std::uint32_t edge = 0; edge < cross_count; ++edge) { + RAFT_EXPECTS(cross_graph(node, edge + 1) < cross_graph.extent(0), + "FlowANN cross graph contains an out-of-range node at row %ld", + static_cast(node)); + } + + auto const packed_row = node / node_per_cacheline_; + auto const packed_col = node % node_per_cacheline_; + auto const inner_count = + has_inner_graph ? static_cast(inner_graph(packed_row, packed_col)) : 0u; + RAFT_EXPECTS(inner_count + cross_count <= full_graph_degree, + "FlowANN inner and cross neighbors exceed graph degree at row %ld", + static_cast(node)); + } + for (int64_t packed_row = 0; packed_row < inner_graph.extent(0); ++packed_row) { + std::uint64_t payload_bits = static_cast(node_per_cacheline_) * 8; + for (std::uint16_t packed_col = 0; packed_col < node_per_cacheline_; ++packed_col) { + auto const node = packed_row * node_per_cacheline_ + packed_col; + if (node < cross_graph.extent(0)) { + payload_bits += static_cast(inner_graph(packed_row, packed_col)) * n_bits_; + } + } + RAFT_EXPECTS(payload_bits <= static_cast(inner_graph.extent(1)) * 8, + "FlowANN inner graph payload exceeds its packed row at row %ld", + static_cast(packed_row)); + } + } + + static auto seed_hash_(IdxT key, std::uint32_t capacity) -> std::uint32_t + { + std::uint64_t value = static_cast(key); + value ^= value >> 16; + value *= 0x85ebca77u; + value ^= value >> 13; + value *= 0xc2b2ae3du; + value ^= value >> 16; + return static_cast(value % capacity); + } + + void rebuild_seed_lookup_(raft::resources const& res) + { + if (seed_ids_host_.empty()) { + seed_lookup_keys_.reset(); + seed_lookup_values_.reset(); + return; + } + RAFT_EXPECTS(seed_ids_host_.size() <= + (static_cast(std::numeric_limits::max()) - 1) / 2, + "FlowANN seed lookup exceeds uint32 capacity"); + + auto const capacity = static_cast(seed_ids_host_.size() * 2 + 1); + auto const empty_key = std::numeric_limits::max(); + auto const empty_value = std::numeric_limits::max(); + std::vector host_keys(capacity, empty_key); + std::vector host_values(capacity, empty_value); + + for (std::uint32_t seed = 0; seed < seed_ids_host_.size(); ++seed) { + auto const key = seed_ids_host_[seed]; + auto slot = seed_hash_(key, capacity); + while (host_keys[slot] != empty_key) { + slot = (slot + 1) % capacity; + } + host_keys[slot] = key; + host_values[slot] = seed; + } + + seed_lookup_keys_.emplace(raft::make_device_vector(res, capacity)); + seed_lookup_values_.emplace(raft::make_device_vector(res, capacity)); + auto const stream = raft::resource::get_cuda_stream(res); + raft::copy(seed_lookup_keys_->data_handle(), host_keys.data(), capacity, stream); + raft::copy(seed_lookup_values_->data_handle(), host_values.data(), capacity, stream); + } + + void rebuild_seed_cross_graph_(raft::resources const& res) + { + if (seed_ids_host_.empty()) { + seed_cross_graph_.reset(); + return; + } + RAFT_EXPECTS(cross_graph_.extent(0) > 0 && cross_graph_.extent(1) > 0, + "FlowANN seed caching requires a non-empty cross graph"); + + auto host_cache = raft::make_host_matrix( + static_cast(seed_ids_host_.size()), cross_graph_.extent(1)); + for (std::size_t seed = 0; seed < seed_ids_host_.size(); ++seed) { + auto const node = seed_ids_host_[seed]; + RAFT_EXPECTS(node < size(), "FlowANN seed at position %zu is out of range", seed); + std::copy_n( + cross_graph_.data_handle() + static_cast(node) * cross_graph_.extent(1), + cross_graph_.extent(1), + host_cache.data_handle() + seed * cross_graph_.extent(1)); + } + + seed_cross_graph_.emplace(raft::make_device_matrix( + res, host_cache.extent(0), host_cache.extent(1))); + raft::copy(seed_cross_graph_->data_handle(), + host_cache.data_handle(), + host_cache.size(), + raft::resource::get_cuda_stream(res)); + } + + void validate_encoding_() const + { + RAFT_EXPECTS(node_per_cacheline_ > 0, "FlowANN node_per_cacheline must be positive"); + RAFT_EXPECTS(n_bits_ > 0 && n_bits_ <= 32, "FlowANN n_bits must be in [1, 32]"); + } + + CUVS_EXPORT void compute_dataset_norms_(raft::resources const& res); + + cuvs::distance::DistanceType metric_; + std::uint16_t node_per_cacheline_; + std::uint16_t n_bits_; + std::uint32_t full_graph_degree_ = 0; + std::uint32_t resident_degree_ = 0; + raft::device_matrix inner_graph_; + raft::host_matrix cross_graph_; + DatasetViewT dataset_{}; + std::optional> dataset_norms_; + std::optional> source_indices_; + std::optional> subgraph_offsets_; + std::optional> subgraph_ids_; + std::optional> seeds_; + std::vector seed_ids_host_; + std::optional> seed_lookup_keys_; + std::optional> seed_lookup_values_; + std::optional> seed_cross_graph_; +}; + +template +using device_padded_index = index>; + +/** Owns the dataset backing storage required by a FlowANN index. */ +template +struct index_bundle { + std::unique_ptr dataset; + IndexT index; +}; + +template +using device_padded_index_bundle = index_bundle, + flowann::device_padded_index>; + +} // namespace cuvs::neighbors::cagra::experimental::flowann + +#endif // CUVS_ENABLE_FLOWANN_SEARCH diff --git a/cpp/src/neighbors/detail/flowann/queue.cuh b/cpp/src/neighbors/detail/flowann/queue.cuh new file mode 100644 index 0000000000..ce7edc8873 --- /dev/null +++ b/cpp/src/neighbors/detail/flowann/queue.cuh @@ -0,0 +1,262 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * FlowANN is described in "Disentangling Graph Dependencies for Efficient + * Billion-Scale GPU Vector Search" (OSDI 2026). + */ + +#pragma once + +#include + +#include + +#include + +#include +#include + +namespace cuvs::neighbors::cagra::experimental::flowann::detail { + +inline constexpr std::uint32_t queue_capacity = 1u << 16; +inline constexpr std::uint32_t queue_mask = queue_capacity - 1; +inline constexpr unsigned long long empty_command = 0ull; + +static_assert((queue_capacity & queue_mask) == 0, "FlowANN queue capacity must be a power of two"); + +// Each word carries the row ID and its complement. Empty or inconsistent +// words fail validation; every uint32 row ID remains representable. +RAFT_INLINE_FUNCTION constexpr auto make_command(std::uint32_t row_id) -> unsigned long long +{ + return (static_cast(~row_id) << 32) | row_id; +} + +RAFT_INLINE_FUNCTION constexpr auto command_ready(unsigned long long word) -> bool +{ + return static_cast(word >> 32) == ~static_cast(word); +} + +struct alignas(128) queue_state { + alignas(128) unsigned long long reservation_head; + alignas(128) unsigned long long cpu_tail; + alignas(128) unsigned long long gpu_tail; + std::uint32_t capacity; + std::uint32_t row_length; +}; + +/** Device pointers for one host-serviced cross-graph queue. */ +struct queue_view { + queue_state* state; + std::uint32_t* rows; + unsigned long long* command_words; + unsigned long long* result_sequence; + unsigned long long* released_sequence; +}; + +struct queue_reservation { + unsigned long long sequence; + std::uint32_t slot; +}; + +#ifdef __CUDACC__ + +RAFT_DEVICE_INLINE_FUNCTION auto load_counter(unsigned long long const* value) -> unsigned long long +{ + cuda::atomic_ref counter( + *const_cast(value)); + return counter.load(cuda::memory_order_relaxed); +} + +[[nodiscard]] RAFT_DEVICE_INLINE_FUNCTION auto reservation_at(queue_reservation first, + std::uint32_t offset) + -> queue_reservation +{ + auto const sequence = first.sequence + offset; + return {sequence, static_cast(sequence) & queue_mask}; +} + +RAFT_DEVICE_INLINE_FUNCTION auto load_published_counter(unsigned long long const* value) + -> unsigned long long +{ + cuda::atomic_ref counter( + *const_cast(value)); + return counter.load(cuda::memory_order_acquire); +} + +/** + * Reclaim only the contiguous released prefix. A failed CAS means another reclaimer advanced + * the tail; restart from that tail so no slot is returned to producers twice. + * This scalar fallback is used when a deferred producer runs out of queue capacity. + */ +RAFT_DEVICE_INLINE_FUNCTION void reclaim_released(queue_view queue) +{ + while (true) { + auto const tail = load_counter(&queue.state->gpu_tail); + auto const head = load_counter(&queue.state->reservation_head); + auto end = tail; + while (end < head) { + auto const slot = static_cast(end) & queue_mask; + cuda::atomic_ref sequence( + queue.released_sequence[slot]); + if (sequence.load(cuda::memory_order_relaxed) != end + 1) { break; } + ++end; + } + if (end == tail) { return; } + if (atomicCAS(&queue.state->gpu_tail, tail, end) == tail) { return; } + } +} + +/** + * Called by all 32 lanes of a warp after the CTA has consumed its pending rows. The first + * unreleased sequence stops reclamation, including when the scan crosses the ring boundary. + * The tail CAS also permits concurrent reclaimers and users of the ordinary release path. + */ +RAFT_DEVICE_INLINE_FUNCTION void reclaim_released_warp(queue_view queue) +{ + auto const lane = threadIdx.x & 31u; + while (true) { + auto const tail = + __shfl_sync(0xffffffffu, lane == 0 ? load_counter(&queue.state->gpu_tail) : 0ull, 0); + auto const head = + __shfl_sync(0xffffffffu, lane == 0 ? load_counter(&queue.state->reservation_head) : 0ull, 0); + auto end = tail; + while (end < head) { + auto const sequence = end + lane; + auto const slot = static_cast(sequence) & queue_mask; + cuda::atomic_ref released( + queue.released_sequence[slot]); + auto const ready = + sequence < head && released.load(cuda::memory_order_relaxed) == sequence + 1; + auto const mask = __ballot_sync(0xffffffffu, ready); + auto const count = mask == 0xffffffffu ? 32u : static_cast(__ffs(~mask) - 1); + end += count; + if (count < 32) { break; } + } + if (end == tail) { return; } + int success = lane == 0 && atomicCAS(&queue.state->gpu_tail, tail, end) == tail; + if (__shfl_sync(0xffffffffu, success, 0)) { return; } + } +} + +/** Mark a consumed row; the caller must reclaim at capacity pressure and on completion. */ +RAFT_DEVICE_INLINE_FUNCTION void release_deferred(queue_view queue, queue_reservation reservation) +{ + atomicExch(&queue.released_sequence[reservation.slot], reservation.sequence + 1); +} + +/** Try to reserve a contiguous batch of queue slots atomically. */ +RAFT_DEVICE_INLINE_FUNCTION auto try_reserve_batch(queue_view queue, + std::uint32_t count, + queue_reservation* first) -> bool +{ + auto* state = queue.state; + if (count == 0 || count > state->capacity) { return false; } + + auto const head = load_counter(&state->reservation_head); + auto const tail = load_counter(&state->gpu_tail); + if (head - tail > state->capacity - count) { return false; } + if (atomicCAS(&state->reservation_head, head, head + count) != head) { return false; } + + first->sequence = head; + first->slot = static_cast(head) & queue_mask; + return true; +} + +RAFT_DEVICE_INLINE_FUNCTION auto try_reserve_deferred(queue_view queue, + std::uint32_t count, + queue_reservation* first) -> bool +{ + auto* state = queue.state; + if (count == 0 || count > state->capacity) { return false; } + + auto const head = load_counter(&state->reservation_head); + auto const tail = load_counter(&state->gpu_tail); + if (head - tail > state->capacity - count) { + reclaim_released(queue); + return false; + } + if (atomicCAS(&state->reservation_head, head, head + count) != head) { return false; } + + first->sequence = head; + first->slot = static_cast(head) & queue_mask; + return true; +} + +/** Try to reserve one queue slot without blocking the CTA indefinitely. */ +RAFT_DEVICE_INLINE_FUNCTION auto try_reserve(queue_view queue, queue_reservation* reservation) + -> bool +{ + return try_reserve_batch(queue, 1, reservation); +} + +/** Atomically publish each complete request in a reserved batch. */ +RAFT_DEVICE_INLINE_FUNCTION void submit_batch(queue_view queue, + queue_reservation first, + std::uint32_t const* row_ids, + std::uint32_t count) +{ + // The atomically published word contains the complete request. No separate + // payload needs a system fence. The host clears this word before publishing + // the result; result acquisition and deferred reclamation gate slot reuse. + for (std::uint32_t i = 0; i < count; ++i) { + auto const reservation = reservation_at(first, i); + cuda::atomic_ref command( + queue.command_words[reservation.slot]); + static_cast(command.exchange(make_command(row_ids[i]), cuda::memory_order_relaxed)); + } +} + +/** Publish one complete cross-graph row request atomically. */ +RAFT_DEVICE_INLINE_FUNCTION void submit(queue_view queue, + queue_reservation reservation, + std::uint32_t row_id) +{ + submit_batch(queue, reservation, &row_id, 1); +} + +RAFT_DEVICE_INLINE_FUNCTION auto result_ready(queue_view queue, queue_reservation reservation) + -> bool +{ + return load_published_counter(&queue.result_sequence[reservation.slot]) == + reservation.sequence + 1; +} + +RAFT_DEVICE_INLINE_FUNCTION auto result_row(queue_view queue, queue_reservation reservation) + -> std::uint32_t const* +{ + return queue.rows + static_cast(reservation.slot) * queue.state->row_length; +} + +/** + * Mark a consumed result as released and advance the global GPU tail over consecutive releases. + */ +RAFT_DEVICE_INLINE_FUNCTION void release(queue_view queue, queue_reservation reservation) +{ + atomicExch(&queue.released_sequence[reservation.slot], reservation.sequence + 1); + + while (true) { + auto const tail = load_counter(&queue.state->gpu_tail); + auto const slot = static_cast(tail) & queue_mask; + if (load_counter(&queue.released_sequence[slot]) != tail + 1) { break; } + if (atomicCAS(&queue.state->gpu_tail, tail, tail + 1) != tail) { continue; } + } +} + +#endif // __CUDACC__ + +#ifdef CUVS_ENABLE_FLOWANN_SEARCH +struct CUVS_EXPORT search_context_access { + static auto device_queues(search_context const& context) noexcept -> queue_view*; + static auto device_id(search_context const& context) noexcept -> int; + static auto num_queues(search_context const& context) noexcept -> std::uint32_t; + static auto row_count(search_context const& context) noexcept -> std::uint64_t; + static auto row_length(search_context const& context) noexcept -> std::uint32_t; + static void acquire_search(search_context& context); + static void release_search(search_context& context); + static auto is_paused(search_context const& context) noexcept -> bool; + static auto pending_pause_requests(search_context const& context) noexcept -> std::size_t; +}; +#endif + +} // namespace cuvs::neighbors::cagra::experimental::flowann::detail diff --git a/cpp/src/neighbors/flowann.cu b/cpp/src/neighbors/flowann.cu new file mode 100644 index 0000000000..6f7e5ee2c1 --- /dev/null +++ b/cpp/src/neighbors/flowann.cu @@ -0,0 +1,60 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * FlowANN is described in "Disentangling Graph Dependencies for Efficient + * Billion-Scale GPU Vector Search" (OSDI 2026). + */ + +#include + +#include "detail/ann_utils.cuh" + +#include +#include + +#include +#include +#include +#include + +namespace cuvs::neighbors::cagra::experimental::flowann { + +template +void index::compute_dataset_norms_(raft::resources const& res) +{ + if constexpr (cuvs::neighbors::is_device_padded_dataset_view_v) { + auto dataset = dataset_.view(); + if (!dataset_norms_.has_value() || dataset_norms_->extent(0) != dataset.extent(0)) { + dataset_norms_.reset(); + dataset_norms_.emplace(raft::make_device_vector(res, dataset.extent(0))); + } + + constexpr float scale = cuvs::spatial::knn::detail::utils::config::kDivisor / + cuvs::spatial::knn::detail::utils::config::kDivisor; + auto scaled_sq_op = + raft::compose_op(raft::sq_op{}, raft::div_const_op{scale}, raft::cast_op()); + raft::linalg::reduce(res, + dataset, + dataset_norms_->view(), + 0.0f, + false, + scaled_sq_op, + raft::add_op(), + raft::sqrt_op{}); + } +} + +#define CUVS_FLOWANN_DEFINE_STORAGE(T) \ + template CUVS_EXPORT void \ + index>:: \ + compute_dataset_norms_(raft::resources const& res) + +CUVS_FLOWANN_DEFINE_STORAGE(float); +CUVS_FLOWANN_DEFINE_STORAGE(half); +CUVS_FLOWANN_DEFINE_STORAGE(std::int8_t); +CUVS_FLOWANN_DEFINE_STORAGE(std::uint8_t); + +#undef CUVS_FLOWANN_DEFINE_STORAGE + +} // namespace cuvs::neighbors::cagra::experimental::flowann diff --git a/cpp/src/neighbors/flowann_queue.cpp b/cpp/src/neighbors/flowann_queue.cpp new file mode 100644 index 0000000000..87f212b4d3 --- /dev/null +++ b/cpp/src/neighbors/flowann_queue.cpp @@ -0,0 +1,918 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * FlowANN is described in "Disentangling Graph Dependencies for Efficient + * Billion-Scale GPU Vector Search" (OSDI 2026). + */ + +#include + +#include "detail/flowann/queue.cuh" + +#include +#include +#include + +#ifdef CUVS_FLOWANN_USE_GDRCOPY +#include +#include +#endif + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#if defined(__linux__) +#include +#include +#include +#endif + +namespace cuvs::neighbors::cagra::experimental::flowann { +namespace { + +constexpr auto align_up(std::size_t value, std::size_t alignment) -> std::size_t +{ + return (value + alignment - 1) & ~(alignment - 1); +} + +struct queue_layout { + std::size_t state; + std::size_t rows; + std::size_t command_words; + std::size_t result_sequence; + std::size_t released_sequence; + std::size_t total_size; + + explicit queue_layout(std::uint32_t row_length) + { + constexpr std::size_t alignment = 128; + state = 0; + rows = align_up(sizeof(detail::queue_state), alignment); + + RAFT_EXPECTS(row_length <= std::numeric_limits::max() / detail::queue_capacity / + sizeof(std::uint32_t), + "FlowANN queue row storage size overflows"); + auto const row_bytes = + static_cast(detail::queue_capacity) * row_length * sizeof(std::uint32_t); + command_words = align_up(rows + row_bytes, alignment); + result_sequence = + align_up(command_words + detail::queue_capacity * sizeof(unsigned long long), alignment); + released_sequence = + align_up(result_sequence + detail::queue_capacity * sizeof(unsigned long long), alignment); + total_size = + align_up(released_sequence + detail::queue_capacity * sizeof(unsigned long long), alignment); + } +}; + +template +auto add_bytes(void* base, std::size_t offset) -> T* +{ + return reinterpret_cast(static_cast(base) + offset); +} + +class cuda_allocation { + public: + cuda_allocation() = default; + cuda_allocation(cuda_allocation const&) = delete; + auto operator=(cuda_allocation const&) -> cuda_allocation& = delete; + + ~cuda_allocation() + { + if (allocation_ != nullptr) { RAFT_CUDA_TRY_NO_THROW(cudaFree(allocation_)); } + } + + void allocate(std::size_t size, std::size_t alignment = 1) + { + RAFT_EXPECTS(allocation_ == nullptr, "FlowANN queue allocation is already initialized"); + RAFT_EXPECTS(alignment > 0 && (alignment & (alignment - 1)) == 0, + "FlowANN queue allocation alignment must be a power of two"); + RAFT_EXPECTS(size <= std::numeric_limits::max() - (alignment - 1), + "FlowANN queue allocation size overflows"); + RAFT_CUDA_TRY(cudaMalloc(&allocation_, size + alignment - 1)); + auto const address = reinterpret_cast(allocation_); + ptr_ = reinterpret_cast(align_up(address, alignment)); + } + + [[nodiscard]] auto get() const noexcept -> void* { return ptr_; } + + private: + void* allocation_{}; + void* ptr_{}; +}; + +class cuda_stream { + public: + cuda_stream() = default; + cuda_stream(cuda_stream const&) = delete; + auto operator=(cuda_stream const&) -> cuda_stream& = delete; + + ~cuda_stream() + { + if (stream_ != nullptr) { RAFT_CUDA_TRY_NO_THROW(cudaStreamDestroy(stream_)); } + } + + void create() + { + RAFT_EXPECTS(stream_ == nullptr, "FlowANN queue stream is already initialized"); + RAFT_CUDA_TRY(cudaStreamCreateWithFlags(&stream_, cudaStreamNonBlocking)); + } + + [[nodiscard]] auto get() const noexcept -> cudaStream_t { return stream_; } + + private: + cudaStream_t stream_{}; +}; + +#ifndef CUVS_FLOWANN_USE_GDRCOPY +class cuda_host_buffer { + public: + cuda_host_buffer() = default; + cuda_host_buffer(cuda_host_buffer const&) = delete; + auto operator=(cuda_host_buffer const&) -> cuda_host_buffer& = delete; + + ~cuda_host_buffer() + { + if (ptr_ != nullptr) { RAFT_CUDA_TRY_NO_THROW(cudaFreeHost(ptr_)); } + } + + void allocate(std::size_t size) + { + RAFT_EXPECTS(ptr_ == nullptr, "FlowANN host copy buffer is already initialized"); + RAFT_CUDA_TRY(cudaMallocHost(&ptr_, size)); + size_ = size; + } + + [[nodiscard]] auto get(std::size_t size) const -> void* + { + RAFT_EXPECTS(size <= size_, "FlowANN transfer exceeds the host copy buffer"); + return ptr_; + } + + private: + void* ptr_{}; + std::size_t size_{}; +}; +#endif + +#ifdef CUVS_FLOWANN_USE_GDRCOPY +class gdrcopy_mapping { + public: + gdrcopy_mapping(void* device_ptr, std::size_t size) : size_(size) + { + try { + unsigned int sync_memops = 1; + RAFT_EXPECTS(cuPointerSetAttribute(&sync_memops, + CU_POINTER_ATTRIBUTE_SYNC_MEMOPS, + reinterpret_cast(device_ptr)) == CUDA_SUCCESS, + "Failed to enable synchronous memory operations for a FlowANN GDRCopy queue"); + + handle_ = gdr_open(); + RAFT_EXPECTS(handle_ != nullptr, + "Failed to open GDRCopy; ensure the gdrdrv kernel module is loaded"); + RAFT_EXPECTS( + gdr_pin_buffer( + handle_, reinterpret_cast(device_ptr), size_, 0, 0, &memory_handle_) == 0, + "Failed to pin a FlowANN queue with GDRCopy"); + pinned_ = true; + RAFT_EXPECTS(gdr_map(handle_, memory_handle_, &mapped_base_, size_) == 0, + "Failed to map a FlowANN queue with GDRCopy"); + mapped_ = true; + + gdr_info_t info{}; + RAFT_EXPECTS(gdr_get_info(handle_, memory_handle_, &info) == 0, + "Failed to query a FlowANN GDRCopy mapping"); + auto const device_address = reinterpret_cast(device_ptr); + RAFT_EXPECTS(device_address >= info.va, + "GDRCopy returned an invalid FlowANN mapping address"); + auto const offset = device_address - info.va; + RAFT_EXPECTS(offset <= size_, "GDRCopy returned an invalid FlowANN mapping offset"); + mapped_device_ptr_ = static_cast(mapped_base_) + offset; + } catch (...) { + reset(); + throw; + } + } + + gdrcopy_mapping(gdrcopy_mapping const&) = delete; + auto operator=(gdrcopy_mapping const&) -> gdrcopy_mapping& = delete; + + ~gdrcopy_mapping() { reset(); } + + void read(void* dst, std::size_t offset, std::size_t size) + { + RAFT_EXPECTS(offset <= size_ && size <= size_ - offset, + "FlowANN GDRCopy read exceeds its queue mapping"); + RAFT_EXPECTS(gdr_copy_from_mapping(memory_handle_, dst, mapped_device_ptr_ + offset, size) == 0, + "Failed to read a FlowANN GDRCopy queue"); + } + + void write(std::size_t offset, void const* src, std::size_t size) + { + RAFT_EXPECTS(offset <= size_ && size <= size_ - offset, + "FlowANN GDRCopy write exceeds its queue mapping"); + RAFT_EXPECTS(gdr_copy_to_mapping(memory_handle_, mapped_device_ptr_ + offset, src, size) == 0, + "Failed to write a FlowANN GDRCopy queue"); + } + + private: + void reset() noexcept + { + if (mapped_) { + static_cast(gdr_unmap(handle_, memory_handle_, mapped_base_, size_)); + mapped_ = false; + } + if (pinned_) { + static_cast(gdr_unpin_buffer(handle_, memory_handle_)); + pinned_ = false; + } + if (handle_ != nullptr) { + static_cast(gdr_close(handle_)); + handle_ = nullptr; + } + } + + gdr_t handle_{}; + gdr_mh_t memory_handle_{}; + void* mapped_base_{}; + std::byte* mapped_device_ptr_{}; + std::size_t size_{}; + bool pinned_{}; + bool mapped_{}; +}; +#endif + +#if defined(__linux__) +auto read_topology_value(int cpu, char const* name) -> int +{ + std::ifstream input("/sys/devices/system/cpu/cpu" + std::to_string(cpu) + "/topology/" + name); + int value = -1; + input >> value; + return input ? value : -1; +} + +auto physical_cpus() -> std::vector +{ + cpu_set_t allowed; + CPU_ZERO(&allowed); + if (sched_getaffinity(0, sizeof(allowed), &allowed) != 0) { return {}; } + + std::vector> cores; + std::vector cpus; + for (int cpu = 0; cpu < CPU_SETSIZE; ++cpu) { + if (!CPU_ISSET(cpu, &allowed)) { continue; } + auto package = read_topology_value(cpu, "physical_package_id"); + auto core = read_topology_value(cpu, "core_id"); + if (package < 0 || core < 0) { + package = 0; + core = cpu; + } + auto const identity = std::tuple{package, core}; + if (std::find(cores.begin(), cores.end(), identity) == cores.end()) { + cores.push_back(identity); + cpus.push_back(cpu); + } + } + return cpus; +} + +void bind_to_physical_cpu(std::uint32_t queue_id) +{ + static auto const cpus = physical_cpus(); + if (cpus.empty()) { return; } + auto const cpu = cpus[queue_id % cpus.size()]; + cpu_set_t selected; + CPU_ZERO(&selected); + CPU_SET(cpu, &selected); + auto const error = pthread_setaffinity_np(pthread_self(), sizeof(selected), &selected); + if (error != 0) { + RAFT_LOG_WARN( + "FlowANN queue %u could not bind its poller to CPU %d (error=%d)", queue_id, cpu, error); + } + if (queue_id >= cpus.size()) { + RAFT_LOG_WARN("FlowANN queue %u reuses physical CPU %d because only %zu are available", + queue_id, + cpu, + cpus.size()); + } +} +#else +void bind_to_physical_cpu(std::uint32_t) {} +#endif + +class host_queue { + public: + static constexpr std::uint32_t max_batch_size = 128; + + host_queue(std::uint32_t const* cross_graph, + std::uint64_t row_count, + std::uint32_t row_length, + bool collect_statistics) + : cross_graph_(cross_graph), + row_count_(row_count), + row_length_(row_length), + collect_statistics_(collect_statistics), + layout_(row_length), + empty_commands_(max_batch_size, detail::empty_command), + command_staging_(max_batch_size), + result_sequence_staging_(max_batch_size), + row_staging_(static_cast(max_batch_size) * row_length) + { +#ifdef CUVS_FLOWANN_USE_GDRCOPY + // gdrdrv rejects mappings whose device address is not aligned to its 64 KiB GPU page. + // cudaMalloc can return a suballocation after the graph build, so align within an owning + // allocation instead of relying on cudaMalloc's current address. + storage_.allocate(layout_.total_size, 64u * 1024u); +#else + storage_.allocate(layout_.total_size); +#endif + RAFT_CUDA_TRY(cudaMemset(storage_.get(), 0, layout_.total_size)); + + detail::queue_state state{}; + state.capacity = detail::queue_capacity; + state.row_length = row_length_; + RAFT_CUDA_TRY(cudaMemcpy(add_bytes(storage_.get(), layout_.state), + &state, + sizeof(state), + cudaMemcpyHostToDevice)); + + view_.state = add_bytes(storage_.get(), layout_.state); + view_.rows = add_bytes(storage_.get(), layout_.rows); + view_.command_words = add_bytes(storage_.get(), layout_.command_words); + view_.result_sequence = add_bytes(storage_.get(), layout_.result_sequence); + view_.released_sequence = + add_bytes(storage_.get(), layout_.released_sequence); + +#ifdef CUVS_FLOWANN_USE_GDRCOPY + gdrcopy_.emplace(storage_.get(), layout_.total_size); +#else + copy_stream_.create(); + // Pageable copies may synchronize inside the driver and contend with lazy + // kernel initialization. Allocate once, before starting any pollers, and + // use pinned memory for every transfer on the nonblocking copy stream. + copy_staging_.allocate(std::max(row_staging_.size() * sizeof(std::uint32_t), + command_staging_.size() * sizeof(unsigned long long))); +#endif + } + + host_queue(host_queue const&) = delete; + auto operator=(host_queue const&) -> host_queue& = delete; + + [[nodiscard]] auto view() const noexcept -> detail::queue_view { return view_; } + + [[nodiscard]] auto poll_batch() -> bool + { + if (collect_statistics_) { polls_.fetch_add(1, std::memory_order_relaxed); } + auto const first_slot = static_cast(cpu_tail_) & detail::queue_mask; + // Sparse traffic needs only a short prefix. Read the rest only when the + // whole prefix is ready, retaining full service batches under load. + constexpr std::uint32_t prefix_size = 8; + read_ring(command_staging_.data(), layout_.command_words, first_slot, prefix_size); + if (!detail::command_ready(command_staging_[0])) { + if (collect_statistics_) { empty_polls_.fetch_add(1, std::memory_order_relaxed); } + return false; + } + + std::uint32_t count = 1; + while (count < prefix_size && detail::command_ready(command_staging_[count])) { + ++count; + } + if (count == prefix_size) { + auto const next_slot = (first_slot + prefix_size) & detail::queue_mask; + read_ring(command_staging_.data() + prefix_size, + layout_.command_words, + next_slot, + max_batch_size - prefix_size); + while (count < max_batch_size && detail::command_ready(command_staging_[count])) { + ++count; + } + } + + for (std::uint32_t i = 0; i < count; ++i) { + auto const row_id = static_cast(command_staging_[i]); + RAFT_EXPECTS( + row_id < row_count_, "FlowANN queue received out-of-range cross-graph row %u", row_id); + std::copy_n(cross_graph_ + static_cast(row_id) * row_length_, + row_length_, + row_staging_.data() + static_cast(i) * row_length_); + result_sequence_staging_[i] = cpu_tail_ + i + 1; + } + write_ring(row_staging_.data(), layout_.rows, first_slot, count, row_length_); + // Clearing the request is ordered before result publication. A producer + // cannot reuse this slot until it acquires and releases that result. + write_ring(empty_commands_.data(), layout_.command_words, first_slot, count); + std::atomic_thread_fence(std::memory_order_release); + write_ring(result_sequence_staging_.data(), layout_.result_sequence, first_slot, count); + + cpu_tail_ += count; + if (collect_statistics_) { processed_commands_.fetch_add(count, std::memory_order_relaxed); } + return true; + } + + void publish_progress() + { + // cpu_tail is an observational snapshot, not a GPU synchronization input: + // consumers use result_sequence and producers reclaim capacity via gpu_tail. + // Publish before acknowledging a pause instead of writing on every batch. + write(layout_.state + offsetof(detail::queue_state, cpu_tail), &cpu_tail_, sizeof(cpu_tail_)); + } + + [[nodiscard]] auto statistics() const noexcept -> queue_statistics + { + return {polls_.load(std::memory_order_relaxed), + empty_polls_.load(std::memory_order_relaxed), + processed_commands_.load(std::memory_order_relaxed)}; + } + + private: + template + void read_ring(T* destination, + std::size_t base_offset, + std::uint32_t first_slot, + std::uint32_t count) + { + auto const first_count = std::min(count, detail::queue_capacity - first_slot); + read(destination, + base_offset + static_cast(first_slot) * sizeof(T), + static_cast(first_count) * sizeof(T)); + if (first_count < count) { + read(destination + first_count, + base_offset, + static_cast(count - first_count) * sizeof(T)); + } + } + + template + void write_ring(T const* source, + std::size_t base_offset, + std::uint32_t first_slot, + std::uint32_t count, + std::uint32_t elements_per_slot = 1) + { + auto const first_count = std::min(count, detail::queue_capacity - first_slot); + auto const slot_bytes = static_cast(elements_per_slot) * sizeof(T); + write(base_offset + static_cast(first_slot) * slot_bytes, + source, + static_cast(first_count) * slot_bytes); + if (first_count < count) { + write(base_offset, + source + static_cast(first_count) * elements_per_slot, + static_cast(count - first_count) * slot_bytes); + } + } + + void read(void* dst, std::size_t offset, std::size_t size) + { +#ifdef CUVS_FLOWANN_USE_GDRCOPY + gdrcopy_->read(dst, offset, size); +#else + auto* staging = copy_staging_.get(size); + RAFT_CUDA_TRY(cudaMemcpyAsync(staging, + add_bytes(storage_.get(), offset), + size, + cudaMemcpyDeviceToHost, + copy_stream_.get())); + RAFT_CUDA_TRY(cudaStreamSynchronize(copy_stream_.get())); + std::memcpy(dst, staging, size); +#endif + } + + void write(std::size_t offset, void const* src, std::size_t size) + { +#ifdef CUVS_FLOWANN_USE_GDRCOPY + gdrcopy_->write(offset, src, size); +#else + auto* staging = copy_staging_.get(size); + std::memcpy(staging, src, size); + RAFT_CUDA_TRY(cudaMemcpyAsync(add_bytes(storage_.get(), offset), + staging, + size, + cudaMemcpyHostToDevice, + copy_stream_.get())); + RAFT_CUDA_TRY(cudaStreamSynchronize(copy_stream_.get())); +#endif + } + + std::uint32_t const* cross_graph_; + std::uint64_t row_count_; + std::uint32_t row_length_; + bool collect_statistics_; + queue_layout layout_; + cuda_allocation storage_; + detail::queue_view view_{}; + unsigned long long cpu_tail_{}; + std::atomic polls_{}; + std::atomic empty_polls_{}; + std::atomic processed_commands_{}; + std::vector empty_commands_; + std::vector command_staging_; + std::vector result_sequence_staging_; + std::vector row_staging_; +#ifdef CUVS_FLOWANN_USE_GDRCOPY + std::optional gdrcopy_; +#else + cuda_stream copy_stream_; + cuda_host_buffer copy_staging_; +#endif +}; + +} // namespace + +struct search_context::impl { + explicit impl(raft::host_matrix_view cross_graph, + queue_params params) + : params_(params), + row_count_(static_cast(cross_graph.extent(0))), + row_length_(static_cast(cross_graph.extent(1))) + { + RAFT_CUDA_TRY(cudaGetDevice(&device_id_)); + RAFT_EXPECTS(cross_graph.extent(0) > 0, "FlowANN cross graph must not be empty"); + RAFT_EXPECTS(cross_graph.extent(1) > 0, + "FlowANN cross graph must contain a neighbor-count column"); + RAFT_EXPECTS(params.num_queues > 0, "FlowANN must use at least one command queue"); + RAFT_EXPECTS(cross_graph.extent(1) <= std::numeric_limits::max(), + "FlowANN cross-graph row length exceeds uint32"); + + queues_.reserve(params.num_queues); + std::vector views; + views.reserve(params.num_queues); + for (std::uint32_t i = 0; i < params.num_queues; ++i) { + queues_.push_back( + std::make_unique(cross_graph.data_handle(), + static_cast(cross_graph.extent(0)), + static_cast(cross_graph.extent(1)), + params.collect_statistics)); + views.push_back(queues_.back()->view()); + } + + RAFT_CUDA_TRY(cudaMalloc(&device_views_, views.size() * sizeof(detail::queue_view))); + try { + RAFT_CUDA_TRY(cudaMemcpy(device_views_, + views.data(), + views.size() * sizeof(detail::queue_view), + cudaMemcpyHostToDevice)); + + pollers_.reserve(queues_.size()); + for (std::uint32_t i = 0; i < queues_.size(); ++i) { + pollers_.emplace_back([this, queue = queues_[i].get(), i](std::stop_token stop_token) { + poll(stop_token, *queue, i); + }); + } + + std::unique_lock lock(pause_mutex_); + pause_cv_.wait(lock, [this] { return paused_pollers_ + failed_pollers_ == pollers_.size(); }); + auto error = poller_error_; + lock.unlock(); + if (error != nullptr) { std::rethrow_exception(error); } + } catch (...) { + stop(); + throw; + } + } + + ~impl() { stop(); } + + void throw_if_failed() + { + if (poller_error_ != nullptr) { std::rethrow_exception(poller_error_); } + } + + [[nodiscard]] auto should_run() const noexcept -> bool + { + return manually_resumed_ || active_sessions_ > 0 || active_searches_ > 0; + } + + void update_pollers(std::unique_lock& lock) + { + while (true) { + throw_if_failed(); + if (should_run()) { + if (paused_.load(std::memory_order_acquire)) { + paused_.store(false, std::memory_order_release); + pause_cv_.notify_all(); + } + pause_cv_.wait(lock, [this] { + return !should_run() || paused_pollers_ == 0 || poller_error_ != nullptr; + }); + throw_if_failed(); + if (should_run() && paused_pollers_ == 0) { return; } + } else { + if (!paused_.load(std::memory_order_acquire)) { + paused_.store(true, std::memory_order_release); + pause_cv_.notify_all(); + } + pause_cv_.wait(lock, [this] { + return should_run() || paused_pollers_ + failed_pollers_ == pollers_.size() || + poller_error_ != nullptr; + }); + throw_if_failed(); + if (!should_run() && paused_pollers_ + failed_pollers_ == pollers_.size()) { return; } + } + } + } + + void stop() noexcept + { + for (auto& poller : pollers_) { + poller.request_stop(); + } + { + std::lock_guard lock(pause_mutex_); + paused_.store(false, std::memory_order_release); + } + pause_cv_.notify_all(); + pollers_.clear(); + if (device_views_ != nullptr) { + RAFT_CUDA_TRY_NO_THROW(cudaFree(device_views_)); + device_views_ = nullptr; + } + } + + void pause() + { + std::unique_lock lock(pause_mutex_); + RAFT_EXPECTS(active_sessions_ == 0, "Cannot pause while a FlowANN search session is active"); + manually_resumed_ = false; + ++pending_pause_requests_; + try { + pause_cv_.wait(lock, [this] { return active_searches_ == 0 || poller_error_ != nullptr; }); + throw_if_failed(); + update_pollers(lock); + } catch (...) { + --pending_pause_requests_; + pause_cv_.notify_all(); + throw; + } + --pending_pause_requests_; + pause_cv_.notify_all(); + } + + void resume() + { + std::unique_lock lock(pause_mutex_); + pause_cv_.wait(lock, + [this] { return pending_pause_requests_ == 0 || poller_error_ != nullptr; }); + throw_if_failed(); + manually_resumed_ = true; + update_pollers(lock); + } + + void acquire_search() + { + std::unique_lock lock(pause_mutex_); + pause_cv_.wait(lock, + [this] { return pending_pause_requests_ == 0 || poller_error_ != nullptr; }); + throw_if_failed(); + manually_resumed_ = false; + ++active_searches_; + try { + update_pollers(lock); + } catch (...) { + --active_searches_; + pause_cv_.notify_all(); + throw; + } + } + + void release_search() + { + std::unique_lock lock(pause_mutex_); + RAFT_EXPECTS(active_searches_ > 0, "Cannot release an inactive FlowANN search lease"); + --active_searches_; + if (active_searches_ == 0) { + manually_resumed_ = false; + pause_cv_.notify_all(); + update_pollers(lock); + } + } + + void begin_session() + { + std::unique_lock lock(pause_mutex_); + pause_cv_.wait(lock, + [this] { return pending_pause_requests_ == 0 || poller_error_ != nullptr; }); + throw_if_failed(); + ++active_sessions_; + try { + update_pollers(lock); + } catch (...) { + --active_sessions_; + throw; + } + } + + void end_session() + { + std::unique_lock lock(pause_mutex_); + RAFT_EXPECTS(active_sessions_ > 0, "Cannot end an inactive FlowANN search session"); + --active_sessions_; + update_pollers(lock); + } + + void poll(std::stop_token stop_token, host_queue& queue, std::uint32_t queue_id) noexcept + { + try { + RAFT_CUDA_TRY(cudaSetDevice(device_id_)); + bind_to_physical_cpu(queue_id); + while (!stop_token.stop_requested()) { + if (paused_.load(std::memory_order_acquire)) { + queue.publish_progress(); + std::unique_lock lock(pause_mutex_); + ++paused_pollers_; + pause_cv_.notify_all(); + pause_cv_.wait(lock, [this, &stop_token] { + return stop_token.stop_requested() || !paused_.load(std::memory_order_acquire); + }); + --paused_pollers_; + pause_cv_.notify_all(); + if (stop_token.stop_requested()) { break; } + } + + auto const processed = queue.poll_batch(); + if (!processed) { +#if defined(__x86_64__) || defined(__i386__) + for (std::uint32_t i = 0; i < params_.empty_pause; ++i) { + __asm__ __volatile__("pause"); + } +#else + std::this_thread::yield(); +#endif + } + } + } catch (...) { + std::lock_guard lock(pause_mutex_); + if (poller_error_ == nullptr) { poller_error_ = std::current_exception(); } + ++failed_pollers_; + pause_cv_.notify_all(); + } + } + + [[nodiscard]] auto statistics() const noexcept -> queue_statistics + { + queue_statistics total{}; + for (auto const& queue : queues_) { + auto const value = queue->statistics(); + total.polls += value.polls; + total.empty_polls += value.empty_polls; + total.processed_commands += value.processed_commands; + } + return total; + } + + queue_params params_; + int device_id_{}; + std::uint64_t row_count_; + std::uint32_t row_length_; + std::vector> queues_; + std::vector pollers_; + detail::queue_view* device_views_{}; + std::atomic paused_{true}; + mutable std::mutex pause_mutex_; + std::condition_variable pause_cv_; + std::size_t paused_pollers_{}; + std::size_t failed_pollers_{}; + std::size_t active_searches_{}; + std::size_t active_sessions_{}; + bool manually_resumed_{}; + std::size_t pending_pause_requests_{}; + std::exception_ptr poller_error_{}; +}; + +search_context::search_context( + raft::host_matrix_view cross_graph, + queue_params params) + : impl_(std::make_unique(cross_graph, params)) +{ +} + +search_context::~search_context() = default; + +search_context::search_context(search_context&&) noexcept = default; + +auto search_context::operator=(search_context&&) noexcept -> search_context& = default; + +search_session::search_session(search_context& context) : context_(&context) +{ + RAFT_EXPECTS(context.impl_ != nullptr, "Cannot start a session on a moved-from search context"); + context.impl_->begin_session(); +} + +search_session::~search_session() { reset(); } + +search_session::search_session(search_session&& other) noexcept + : context_(std::exchange(other.context_, nullptr)) +{ +} + +auto search_session::operator=(search_session&& other) noexcept -> search_session& +{ + if (this != &other) { + reset(); + context_ = std::exchange(other.context_, nullptr); + } + return *this; +} + +void search_session::reset() noexcept +{ + if (context_ == nullptr) { return; } + auto* context = std::exchange(context_, nullptr); + try { + context->impl_->end_session(); + } catch (...) { + } +} + +void search_context::pause() +{ + RAFT_EXPECTS(impl_ != nullptr, "Cannot pause a moved-from FlowANN search context"); + impl_->pause(); +} + +void search_context::resume() +{ + RAFT_EXPECTS(impl_ != nullptr, "Cannot resume a moved-from FlowANN search context"); + impl_->resume(); +} + +auto search_context::num_queues() const noexcept -> std::uint32_t +{ + return impl_ == nullptr ? 0 : static_cast(impl_->queues_.size()); +} + +auto search_context::statistics() const noexcept -> queue_statistics +{ + return impl_ == nullptr ? queue_statistics{} : impl_->statistics(); +} + +auto detail::search_context_access::device_queues(search_context const& context) noexcept + -> queue_view* +{ + return context.impl_ == nullptr ? nullptr : context.impl_->device_views_; +} + +auto detail::search_context_access::device_id(search_context const& context) noexcept -> int +{ + return context.impl_ == nullptr ? -1 : context.impl_->device_id_; +} + +auto detail::search_context_access::num_queues(search_context const& context) noexcept + -> std::uint32_t +{ + return context.num_queues(); +} + +auto detail::search_context_access::row_count(search_context const& context) noexcept + -> std::uint64_t +{ + return context.impl_ == nullptr ? 0 : context.impl_->row_count_; +} + +auto detail::search_context_access::row_length(search_context const& context) noexcept + -> std::uint32_t +{ + return context.impl_ == nullptr ? 0 : context.impl_->row_length_; +} + +void detail::search_context_access::acquire_search(search_context& context) +{ + RAFT_EXPECTS(context.impl_ != nullptr, "Cannot search with a moved-from FlowANN search context"); + context.impl_->acquire_search(); +} + +void detail::search_context_access::release_search(search_context& context) +{ + RAFT_EXPECTS(context.impl_ != nullptr, "Cannot search with a moved-from FlowANN search context"); + context.impl_->release_search(); +} + +auto detail::search_context_access::is_paused(search_context const& context) noexcept -> bool +{ + return context.impl_ == nullptr || context.impl_->paused_.load(std::memory_order_acquire); +} + +auto detail::search_context_access::pending_pause_requests(search_context const& context) noexcept + -> std::size_t +{ + if (context.impl_ == nullptr) { return 0; } + std::lock_guard lock(context.impl_->pause_mutex_); + return context.impl_->pending_pause_requests_; +} + +} // namespace cuvs::neighbors::cagra::experimental::flowann diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index 81a499419c..4319c28288 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -266,6 +266,34 @@ ConfigureTest( PERCENT 100 ) +if(CUVS_ENABLE_FLOWANN_SEARCH) + set(CUVS_FLOWANN_TEST_SOURCES + neighbors/ann_cagra/flowann_queue.cu + neighbors/ann_cagra/flowann_storage.cu + ) + ConfigureTest( + NAME NEIGHBORS_FLOWANN_TEST + PATH ${CUVS_FLOWANN_TEST_SOURCES} + GPUS 1 + PERCENT 100 + ) + target_include_directories( + NEIGHBORS_FLOWANN_TEST PRIVATE "$" + ) + # Isolate this case so earlier tests cannot preload the kernel module. + rapids_test_add( + NAME NEIGHBORS_FLOWANN_COLD_START_TEST + COMMAND NEIGHBORS_FLOWANN_TEST --gtest_filter=FlowannQueue.ColdLaunchWithManyPollers + GPUS 1 + PERCENT 100 + INSTALL_COMPONENT_SET testing + ) + set_tests_properties( + NEIGHBORS_FLOWANN_COLD_START_TEST + PROPERTIES TIMEOUT 45 ENVIRONMENT "CUDA_MODULE_LOADING=LAZY;CUDA_MODULE_DATA_LOADING=LAZY" + ) +endif() + ConfigureTest( NAME NEIGHBORS_ANN_CAGRA_MERGE_TEST PATH neighbors/ann_cagra/test_merge_fastener.cu diff --git a/cpp/tests/neighbors/ann_cagra/flowann_queue.cu b/cpp/tests/neighbors/ann_cagra/flowann_queue.cu new file mode 100644 index 0000000000..4150516828 --- /dev/null +++ b/cpp/tests/neighbors/ann_cagra/flowann_queue.cu @@ -0,0 +1,716 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * FlowANN is described in "Disentangling Graph Dependencies for Efficient + * Billion-Scale GPU Vector Search" (OSDI 2026). + */ + +#include + +#include + +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include + +namespace flowann = cuvs::neighbors::cagra::experimental::flowann; + +namespace { + +RAFT_KERNEL queue_round_trip(flowann::detail::queue_view* queues, + std::uint32_t num_queues, + std::uint32_t row_count, + std::uint32_t row_length, + std::uint32_t* output) +{ + if (threadIdx.x != 0) { return; } + + auto const command = blockIdx.x % row_count; + auto queue = queues[blockIdx.x % num_queues]; + flowann::detail::queue_reservation reservation{}; + while (!flowann::detail::try_reserve(queue, &reservation)) {} + flowann::detail::submit(queue, reservation, command); + while (!flowann::detail::result_ready(queue, reservation)) {} + auto const* row = flowann::detail::result_row(queue, reservation); + for (std::uint32_t i = 0; i < row_length; ++i) { + output[static_cast(blockIdx.x) * row_length + i] = row[i]; + } + flowann::detail::release(queue, reservation); +} + +RAFT_KERNEL mixed_release_queue_round_trip(flowann::detail::queue_view* queues, + std::uint32_t num_queues, + std::uint32_t row_count, + std::uint32_t row_length, + std::uint32_t round, + std::uint32_t batch, + std::uint32_t* output) +{ + auto queue = queues[blockIdx.x % num_queues]; + if (threadIdx.x == 0) { + std::uint32_t ids[7]; + for (std::uint32_t i = 0; i < batch; ++i) { + ids[i] = (blockIdx.x * batch + i + round) % row_count; + } + flowann::detail::queue_reservation first{}; + while (!flowann::detail::try_reserve_deferred(queue, batch, &first)) {} + flowann::detail::submit_batch(queue, first, ids, batch); + for (std::uint32_t i = 0; i < batch; ++i) { + auto const reservation = flowann::detail::reservation_at(first, i); + while (!flowann::detail::result_ready(queue, reservation)) {} + auto const* row = flowann::detail::result_row(queue, reservation); + for (std::uint32_t col = 0; col < row_length; ++col) { + output[(static_cast(blockIdx.x) * batch + i) * row_length + col] = row[col]; + } + if (((blockIdx.x + i) & 1u) == 0) { + flowann::detail::release_deferred(queue, reservation); + } else { + flowann::detail::release(queue, reservation); + } + } + } + __syncthreads(); + flowann::detail::reclaim_released_warp(queue); +} + +TEST(FlowannQueue, CommandEncodingPreservesRowIdsAndRejectsPartialPublication) +{ + EXPECT_FALSE(flowann::detail::command_ready(flowann::detail::empty_command)); + for (std::uint32_t row : {0u, 1u, 65535u, 65536u, 0xabcdefu, 0x7fffffffu, 0xffffffffu}) { + auto const word = flowann::detail::make_command(row); + ASSERT_TRUE(flowann::detail::command_ready(word)); + EXPECT_EQ(static_cast(word), row); + // Reading either 32-bit half before the other is published must not + // produce a ready word carrying a different request. + for (auto partial : {word & 0xffffffff00000000ull, word & 0xffffffffull}) { + if (flowann::detail::command_ready(partial)) { EXPECT_EQ(partial, word); } + } + for (unsigned bit = 0; bit < 64; ++bit) { + EXPECT_FALSE(flowann::detail::command_ready(word ^ (1ull << bit))); + } + } +} + +void check_round_trip_pause_resume(std::uint32_t num_queues) +{ + raft::resources res; + constexpr std::uint32_t row_count = 17; + constexpr std::uint32_t row_length = 9; + constexpr std::uint32_t commands = 257; + + auto cross_graph = raft::make_host_matrix(row_count, row_length); + for (std::uint32_t row = 0; row < row_count; ++row) { + for (std::uint32_t col = 0; col < row_length; ++col) { + cross_graph(row, col) = row * 1000 + col; + } + } + + flowann::queue_params params{}; + params.num_queues = num_queues; + params.empty_pause = 8; + params.collect_statistics = true; + flowann::search_context context(raft::make_const_mdspan(cross_graph.view()), params); + auto output = raft::make_device_matrix(res, commands, row_length); + + context.resume(); + queue_round_trip<<>>( + flowann::detail::search_context_access::device_queues(context), + context.num_queues(), + row_count, + row_length, + output.data_handle()); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + raft::resource::sync_stream(res); + context.pause(); + + auto host_output = raft::make_host_matrix(commands, row_length); + raft::copy(host_output.data_handle(), + output.data_handle(), + output.size(), + raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + for (std::uint32_t command = 0; command < commands; ++command) { + auto const row = command % row_count; + for (std::uint32_t col = 0; col < row_length; ++col) { + EXPECT_EQ(host_output(command, col), cross_graph(row, col)); + } + } + EXPECT_EQ(context.statistics().processed_commands, commands); + + context.resume(); + context.pause(); +} + +TEST(FlowannQueue, RoundTripPauseResume) { check_round_trip_pause_resume(2); } + +// Also run this test alone in a fresh process to exercise lazy kernel loading +// while many idle CUDA-copy pollers are already issuing transfers. +TEST(FlowannQueue, ColdLaunchWithManyPollers) { check_round_trip_pause_resume(44); } + +TEST(FlowannQueue, PollersUseContextCudaDevice) +{ + if (raft::device_setter::get_device_count() < 2) { + GTEST_SKIP() << "Requires at least two CUDA devices"; + } + + raft::device_setter device_scope(1); + raft::resources res; + constexpr std::uint32_t row_count = 4; + constexpr std::uint32_t row_length = 3; + + auto cross_graph = raft::make_host_matrix(row_count, row_length); + for (std::uint32_t row = 0; row < row_count; ++row) { + for (std::uint32_t col = 0; col < row_length; ++col) { + cross_graph(row, col) = row * 10 + col; + } + } + + flowann::search_context context(raft::make_const_mdspan(cross_graph.view())); + auto output = raft::make_device_matrix(res, row_count, row_length); + context.resume(); + queue_round_trip<<>>( + flowann::detail::search_context_access::device_queues(context), + context.num_queues(), + row_count, + row_length, + output.data_handle()); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + raft::resource::sync_stream(res); + context.pause(); + + auto host_output = raft::make_host_matrix(row_count, row_length); + raft::copy(host_output.data_handle(), + output.data_handle(), + output.size(), + raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + for (std::uint32_t row = 0; row < row_count; ++row) { + for (std::uint32_t col = 0; col < row_length; ++col) { + EXPECT_EQ(host_output(row, col), cross_graph(row, col)); + } + } +} + +TEST(FlowannQueue, RealPollersPreservePayloadsAcrossWrapWithMixedReleases) +{ + raft::resources res; + auto const stream = raft::resource::get_cuda_stream(res); + constexpr std::uint32_t row_count = 17; + constexpr std::uint32_t row_length = 9; + constexpr std::uint32_t blocks = 4099; + constexpr std::uint32_t max_batch = 7; + constexpr std::uint32_t rounds = 65; + static_assert(static_cast(blocks) * rounds > + static_cast(flowann::detail::queue_capacity) * 4); + + auto cross_graph = raft::make_host_matrix(row_count, row_length); + for (std::uint32_t row = 0; row < row_count; ++row) { + for (std::uint32_t col = 0; col < row_length; ++col) { + cross_graph(row, col) = row * 1000 + col; + } + } + auto output = + raft::make_device_matrix(res, blocks * max_batch, row_length); + auto host_output = raft::make_host_matrix(blocks * max_batch, row_length); + + for (std::uint32_t num_queues : {1u, 4u}) { + for (std::uint32_t batch : {1u, max_batch}) { + auto const commands = blocks * batch; + flowann::queue_params params{}; + params.num_queues = num_queues; + params.empty_pause = 0; + params.collect_statistics = true; + flowann::search_context context(raft::make_const_mdspan(cross_graph.view()), params); + { + flowann::search_session session(context); + for (std::uint32_t round = 0; round < rounds; ++round) { + mixed_release_queue_round_trip<<>>( + flowann::detail::search_context_access::device_queues(context), + num_queues, + row_count, + row_length, + round, + batch, + output.data_handle()); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + // Finish the host-serviced kernel before copying into pageable host memory. + // A pageable D2H copy can block inside the driver while waiting for this + // kernel, preventing the CUDA-copy poller from submitting its transfers. + raft::resource::sync_stream(res); + raft::copy(host_output.data_handle(), + output.data_handle(), + static_cast(commands) * row_length, + stream); + raft::resource::sync_stream(res); + for (std::uint32_t command = 0; command < commands; ++command) { + auto const row = (command + round) % row_count; + for (std::uint32_t col = 0; col < row_length; ++col) { + if (host_output(command, col) != cross_graph(row, col)) { + FAIL() << "payload mismatch with " << num_queues << " queues at round " << round + << ", command " << command << ", column " << col; + } + } + } + } + } + + EXPECT_TRUE(flowann::detail::search_context_access::is_paused(context)); + EXPECT_EQ(context.statistics().processed_commands, + static_cast(commands) * rounds); + std::vector queues(num_queues); + RAFT_CUDA_TRY(cudaMemcpy(queues.data(), + flowann::detail::search_context_access::device_queues(context), + queues.size() * sizeof(flowann::detail::queue_view), + cudaMemcpyDeviceToHost)); + for (auto const& queue : queues) { + flowann::detail::queue_state state{}; + RAFT_CUDA_TRY(cudaMemcpy(&state, queue.state, sizeof(state), cudaMemcpyDeviceToHost)); + EXPECT_EQ(state.gpu_tail, state.reservation_head); + EXPECT_EQ(state.cpu_tail, state.reservation_head); + } + } + } +} + +TEST(FlowannQueue, PollerStopsAtUnpublishedCommandsAndResumesFullBursts) +{ + constexpr std::uint32_t row_count = 17; + constexpr std::uint32_t row_length = 33; + constexpr std::uint32_t commands = 140; + auto graph = raft::make_host_matrix(row_count, row_length); + for (std::uint32_t row = 0; row < row_count; ++row) { + for (std::uint32_t col = 0; col < row_length; ++col) { + graph(row, col) = row * 1000 + col; + } + } + + // Exercise holes within the short prefix, at its boundary, and in a full batch. + for (std::uint32_t hole : {5u, 8u, 127u}) { + SCOPED_TRACE(hole); + flowann::queue_params params{}; + params.num_queues = 1; + params.empty_pause = 0; + params.collect_statistics = true; + flowann::search_context context(raft::make_const_mdspan(graph.view()), params); + flowann::detail::queue_view queue{}; + RAFT_CUDA_TRY(cudaMemcpy(&queue, + flowann::detail::search_context_access::device_queues(context), + sizeof(queue), + cudaMemcpyDeviceToHost)); + unsigned long long const reserved = commands; + RAFT_CUDA_TRY(cudaMemcpy( + &queue.state->reservation_head, &reserved, sizeof(reserved), cudaMemcpyHostToDevice)); + std::vector ids(commands); + std::vector sequences(commands); + for (std::uint32_t i = 0; i < commands; ++i) { + ids[i] = i % row_count; + sequences[i] = + i == hole ? flowann::detail::empty_command : flowann::detail::make_command(ids[i]); + } + RAFT_CUDA_TRY(cudaMemcpy(queue.command_words, + sequences.data(), + sequences.size() * sizeof(sequences[0]), + cudaMemcpyHostToDevice)); + auto wait_for = [&](std::uint64_t expected) { + auto const deadline = std::chrono::steady_clock::now() + std::chrono::seconds(5); + while (context.statistics().processed_commands < expected && + std::chrono::steady_clock::now() < deadline) { + std::this_thread::yield(); + } + return context.statistics().processed_commands; + }; + + context.resume(); + auto const first_count = wait_for(hole); + context.pause(); + ASSERT_EQ(first_count, hole); + ASSERT_EQ(context.statistics().processed_commands, hole); + flowann::detail::queue_state state{}; + RAFT_CUDA_TRY(cudaMemcpy(&state, queue.state, sizeof(state), cudaMemcpyDeviceToHost)); + EXPECT_EQ(state.cpu_tail, hole); + std::vector published(commands); + RAFT_CUDA_TRY(cudaMemcpy(published.data(), + queue.result_sequence, + published.size() * sizeof(published[0]), + cudaMemcpyDeviceToHost)); + for (std::uint32_t i = 0; i < commands; ++i) { + EXPECT_EQ(published[i], i < hole ? i + 1ull : 0ull); + } + + auto const ready = flowann::detail::make_command(ids[hole]); + RAFT_CUDA_TRY( + cudaMemcpy(queue.command_words + hole, &ready, sizeof(ready), cudaMemcpyHostToDevice)); + context.resume(); + auto const final_count = wait_for(commands); + context.pause(); + ASSERT_EQ(final_count, commands); + RAFT_CUDA_TRY(cudaMemcpy(&state, queue.state, sizeof(state), cudaMemcpyDeviceToHost)); + EXPECT_EQ(state.cpu_tail, commands); + std::vector cleared(commands); + RAFT_CUDA_TRY(cudaMemcpy(cleared.data(), + queue.command_words, + cleared.size() * sizeof(cleared[0]), + cudaMemcpyDeviceToHost)); + for (auto word : cleared) { + EXPECT_EQ(word, flowann::detail::empty_command); + } + std::vector returned(commands * row_length); + RAFT_CUDA_TRY(cudaMemcpy( + returned.data(), queue.rows, returned.size() * sizeof(returned[0]), cudaMemcpyDeviceToHost)); + for (std::uint32_t i = 0; i < commands; ++i) { + for (std::uint32_t col = 0; col < row_length; ++col) { + EXPECT_EQ(returned[i * row_length + col], graph(ids[i], col)); + } + } + } +} + +RAFT_KERNEL deferred_capacity_reclaim(flowann::detail::queue_view q, unsigned* errors) +{ + if (threadIdx.x != 0 || blockIdx.x != 0) return; + for (unsigned i = 0; i < flowann::detail::queue_capacity * 3 + 17; i++) { + flowann::detail::queue_reservation r{}; + unsigned tries = 0; + while (!flowann::detail::try_reserve_deferred(q, 1, &r)) { + if (++tries > 4) { + *errors = 1; + return; + } + } + if (r.sequence != i) { + *errors = 2; + return; + } + flowann::detail::release_deferred(q, r); + } + flowann::detail::reclaim_released(q); + if (q.state->gpu_tail != q.state->reservation_head) *errors = 3; +} +RAFT_KERNEL deferred_release_hole(flowann::detail::queue_view q, unsigned* errors) +{ + if (threadIdx.x == 0) { + q.state->gpu_tail = 0; + q.state->reservation_head = 3; + q.released_sequence[0] = 0; + flowann::detail::release_deferred(q, {1, 1}); + flowann::detail::release_deferred(q, {2, 2}); + } + __syncthreads(); + flowann::detail::reclaim_released_warp(q); + if (threadIdx.x == 0 && q.state->gpu_tail != 0) *errors = 4; + if (threadIdx.x == 0) flowann::detail::release_deferred(q, {0, 0}); + __syncthreads(); + flowann::detail::reclaim_released_warp(q); + if (threadIdx.x == 0 && q.state->gpu_tail != 3) *errors = 5; +} + +RAFT_KERNEL reserve_deferred_concurrently(flowann::detail::queue_view queue, + std::uint64_t expected_base, + std::uint32_t reservations_per_block, + flowann::detail::queue_reservation* reservations, + std::uint32_t* valid, + std::uint32_t* sequence_owners, + std::uint32_t* errors) +{ + auto const reservation_index = blockIdx.x; + if (threadIdx.x == 0) { + flowann::detail::queue_reservation first{}; + bool reserved = false; + for (std::uint32_t attempt = 0; attempt < 4096 && !reserved; ++attempt) { + reserved = flowann::detail::try_reserve_deferred(queue, reservations_per_block, &first); + } + reservations[reservation_index] = first; + valid[reservation_index] = reserved; + if (!reserved) { atomicOr(errors, 1u); } + } + __syncthreads(); + + if (valid[reservation_index] == 0) { return; } + auto const first = reservations[reservation_index]; + for (std::uint32_t i = threadIdx.x; i < reservations_per_block; i += blockDim.x) { + auto const reservation = flowann::detail::reservation_at(first, i); + if (reservation.sequence < expected_base || + reservation.sequence >= + expected_base + static_cast(gridDim.x) * reservations_per_block) { + atomicOr(errors, 2u); + continue; + } + auto const offset = static_cast(reservation.sequence - expected_base); + if (atomicCAS(&sequence_owners[offset], ~std::uint32_t{0}, reservation_index) != + ~std::uint32_t{0}) { + atomicOr(errors, 4u); + } + } +} + +RAFT_KERNEL release_deferred_except_first(flowann::detail::queue_view queue, + std::uint64_t first_sequence, + std::uint32_t reservations_per_block, + flowann::detail::queue_reservation const* reservations, + std::uint32_t const* valid) +{ + if (valid[blockIdx.x] != 0) { + auto const first = reservations[blockIdx.x]; + for (std::uint32_t i = threadIdx.x; i < reservations_per_block; i += blockDim.x) { + auto const reservation = flowann::detail::reservation_at(first, i); + if (reservation.sequence != first_sequence) { + flowann::detail::release_deferred(queue, reservation); + } + } + } + __syncthreads(); + if (threadIdx.x < 32) { flowann::detail::reclaim_released_warp(queue); } +} + +RAFT_KERNEL release_first_and_reclaim(flowann::detail::queue_view queue, + std::uint64_t first_sequence, + std::uint32_t reservations_per_block, + flowann::detail::queue_reservation const* reservations, + std::uint32_t const* valid) +{ + if (valid[blockIdx.x] != 0) { + auto const first = reservations[blockIdx.x]; + if (first.sequence <= first_sequence && + first_sequence < first.sequence + reservations_per_block && threadIdx.x == 0) { + flowann::detail::release_deferred( + queue, + flowann::detail::reservation_at( + first, static_cast(first_sequence - first.sequence))); + } + } + __syncthreads(); + if (threadIdx.x < 32) { flowann::detail::reclaim_released_warp(queue); } +} + +RAFT_KERNEL check_deferred_queue_counters(flowann::detail::queue_view queue, + std::uint64_t expected_tail, + std::uint32_t error_bit, + std::uint32_t* errors) +{ + if (threadIdx.x == 0 && blockIdx.x == 0 && + (queue.state->gpu_tail != expected_tail || + (error_bit != 0 && queue.state->reservation_head != expected_tail))) { + atomicOr(errors, error_bit == 0 ? 8u : error_bit); + } +} + +TEST(FlowannQueue, DeferredReclamationHandlesCapacityWrapAndUnreleasedHoles) +{ + raft::resources res; + auto const stream = raft::resource::get_cuda_stream(res); + auto cross = raft::make_host_matrix(1, 2); + cross(0, 0) = 1; + cross(0, 1) = 0; + flowann::search_context context(raft::make_const_mdspan(cross.view())); + flowann::detail::queue_view queue; + RAFT_CUDA_TRY(cudaMemcpy(&queue, + flowann::detail::search_context_access::device_queues(context), + sizeof(queue), + cudaMemcpyDeviceToHost)); + auto errors = raft::make_device_matrix(res, 1, 1); + RAFT_CUDA_TRY(cudaMemsetAsync(errors.data_handle(), 0, sizeof(std::uint32_t), stream.get())); + deferred_capacity_reclaim<<<1, 1, 0, stream.get()>>>(queue, errors.data_handle()); + deferred_release_hole<<<1, 32, 0, stream.get()>>>(queue, errors.data_handle()); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + std::uint32_t host_error{}; + raft::copy(&host_error, errors.data_handle(), 1, stream); + raft::resource::sync_stream(res); + EXPECT_EQ(host_error, 0u); +} + +TEST(FlowannQueue, ConcurrentDeferredReclamationPreservesOwnershipAcrossWrap) +{ + raft::resources res; + auto const stream = raft::resource::get_cuda_stream(res); + auto cross = raft::make_host_matrix(1, 2); + cross(0, 0) = 1; + cross(0, 1) = 0; + flowann::search_context context(raft::make_const_mdspan(cross.view())); + flowann::detail::queue_view queue; + RAFT_CUDA_TRY(cudaMemcpy(&queue, + flowann::detail::search_context_access::device_queues(context), + sizeof(queue), + cudaMemcpyDeviceToHost)); + + constexpr std::uint32_t blocks = 255; + constexpr std::uint32_t reservations_per_block = 257; + constexpr std::uint32_t reservations_per_wave = blocks * reservations_per_block; + static_assert(reservations_per_wave == flowann::detail::queue_capacity - 1); + + auto reservations = + raft::make_device_matrix(res, blocks, 1); + auto valid = raft::make_device_matrix(res, blocks, 1); + auto owners = raft::make_device_matrix(res, reservations_per_wave, 1); + auto errors = raft::make_device_matrix(res, 1, 1); + RAFT_CUDA_TRY(cudaMemsetAsync(errors.data_handle(), 0, sizeof(std::uint32_t), stream.get())); + + std::uint64_t expected_base = 0; + for (std::uint32_t wave = 0; wave < 3; ++wave) { + RAFT_CUDA_TRY( + cudaMemsetAsync(valid.data_handle(), 0, valid.size() * sizeof(std::uint32_t), stream.get())); + RAFT_CUDA_TRY(cudaMemsetAsync( + owners.data_handle(), 0xff, owners.size() * sizeof(std::uint32_t), stream.get())); + reserve_deferred_concurrently<<>>(queue, + expected_base, + reservations_per_block, + reservations.data_handle(), + valid.data_handle(), + owners.data_handle(), + errors.data_handle()); + release_deferred_except_first<<>>(queue, + expected_base, + reservations_per_block, + reservations.data_handle(), + valid.data_handle()); + check_deferred_queue_counters<<<1, 1, 0, stream.get()>>>( + queue, expected_base, 0, errors.data_handle()); + release_first_and_reclaim<<>>(queue, + expected_base, + reservations_per_block, + reservations.data_handle(), + valid.data_handle()); + expected_base += reservations_per_wave; + check_deferred_queue_counters<<<1, 1, 0, stream.get()>>>( + queue, expected_base, 16u, errors.data_handle()); + } + RAFT_CUDA_TRY(cudaPeekAtLastError()); + + std::uint32_t host_error{}; + raft::copy(&host_error, errors.data_handle(), 1, stream); + raft::resource::sync_stream(res); + EXPECT_EQ(host_error, 0u); +} + +TEST(FlowannQueue, SearchLeaseConsumesRawResume) +{ + constexpr std::uint32_t row_count = 4; + constexpr std::uint32_t row_length = 3; + auto cross_graph = raft::make_host_matrix(row_count, row_length); + std::fill_n(cross_graph.data_handle(), cross_graph.size(), 0u); + + flowann::search_context context(raft::make_const_mdspan(cross_graph.view())); + EXPECT_TRUE(flowann::detail::search_context_access::is_paused(context)); + + context.resume(); + EXPECT_FALSE(flowann::detail::search_context_access::is_paused(context)); + flowann::detail::search_context_access::acquire_search(context); + flowann::detail::search_context_access::release_search(context); + EXPECT_TRUE(flowann::detail::search_context_access::is_paused(context)); + + context.resume(); + context.pause(); + EXPECT_TRUE(flowann::detail::search_context_access::is_paused(context)); + + flowann::detail::search_context_access::acquire_search(context); + context.resume(); + flowann::detail::search_context_access::release_search(context); + EXPECT_TRUE(flowann::detail::search_context_access::is_paused(context)); +} + +TEST(FlowannQueue, SearchSessionKeepsPollersRunningAcrossSearchLease) +{ + constexpr std::uint32_t row_count = 4; + constexpr std::uint32_t row_length = 3; + auto cross_graph = raft::make_host_matrix(row_count, row_length); + std::fill_n(cross_graph.data_handle(), cross_graph.size(), 0u); + + flowann::search_context context(raft::make_const_mdspan(cross_graph.view())); + { + flowann::search_session session(context); + EXPECT_FALSE(flowann::detail::search_context_access::is_paused(context)); + flowann::detail::search_context_access::acquire_search(context); + flowann::detail::search_context_access::release_search(context); + EXPECT_FALSE(flowann::detail::search_context_access::is_paused(context)); + } + EXPECT_TRUE(flowann::detail::search_context_access::is_paused(context)); +} + +TEST(FlowannQueue, PauseBlocksNewSearchesUntilPollersReachQuiescence) +{ + using namespace std::chrono_literals; + constexpr std::uint32_t row_count = 4; + constexpr std::uint32_t row_length = 3; + auto cross_graph = raft::make_host_matrix(row_count, row_length); + std::fill_n(cross_graph.data_handle(), cross_graph.size(), 0u); + + flowann::search_context context(raft::make_const_mdspan(cross_graph.view())); + flowann::detail::search_context_access::acquire_search(context); + std::atomic pauses_returned{}; + std::jthread first_pauser([&] { + context.pause(); + pauses_returned.fetch_add(1, std::memory_order_release); + }); + + auto const deadline = std::chrono::steady_clock::now() + 2s; + while (flowann::detail::search_context_access::pending_pause_requests(context) < 1 && + std::chrono::steady_clock::now() < deadline) { + std::this_thread::yield(); + } + std::jthread second_pauser([&] { + context.pause(); + pauses_returned.fetch_add(1, std::memory_order_release); + }); + auto const second_deadline = std::chrono::steady_clock::now() + 2s; + while (flowann::detail::search_context_access::pending_pause_requests(context) < 2 && + std::chrono::steady_clock::now() < second_deadline) { + std::this_thread::yield(); + } + auto const pending_pauses = + flowann::detail::search_context_access::pending_pause_requests(context); + + std::atomic later_search_started{}; + std::atomic later_search_acquired{}; + std::jthread later_search([&] { + later_search_started.store(true, std::memory_order_release); + flowann::detail::search_context_access::acquire_search(context); + later_search_acquired.store(true, std::memory_order_release); + flowann::detail::search_context_access::release_search(context); + }); + while (!later_search_started.load(std::memory_order_acquire)) { + std::this_thread::yield(); + } + std::this_thread::sleep_for(20ms); + if (pending_pauses == 2) { EXPECT_FALSE(later_search_acquired.load(std::memory_order_acquire)); } + + flowann::detail::search_context_access::release_search(context); + first_pauser.join(); + second_pauser.join(); + later_search.join(); + EXPECT_EQ(pending_pauses, 2); + EXPECT_EQ(pauses_returned.load(std::memory_order_acquire), 2); + EXPECT_TRUE(later_search_acquired.load(std::memory_order_acquire)); + EXPECT_TRUE(flowann::detail::search_context_access::is_paused(context)); +} + +TEST(FlowannQueue, SearchLeasePausesAfterLastSearch) +{ + constexpr std::uint32_t row_count = 4; + constexpr std::uint32_t row_length = 3; + auto cross_graph = raft::make_host_matrix(row_count, row_length); + std::fill_n(cross_graph.data_handle(), cross_graph.size(), 0u); + + flowann::search_context context(raft::make_const_mdspan(cross_graph.view())); + flowann::detail::search_context_access::acquire_search(context); + flowann::detail::search_context_access::acquire_search(context); + EXPECT_FALSE(flowann::detail::search_context_access::is_paused(context)); + + flowann::detail::search_context_access::release_search(context); + EXPECT_FALSE(flowann::detail::search_context_access::is_paused(context)); + flowann::detail::search_context_access::release_search(context); + EXPECT_TRUE(flowann::detail::search_context_access::is_paused(context)); +} + +} // namespace diff --git a/cpp/tests/neighbors/ann_cagra/flowann_storage.cu b/cpp/tests/neighbors/ann_cagra/flowann_storage.cu new file mode 100644 index 0000000000..f2f85b7af9 --- /dev/null +++ b/cpp/tests/neighbors/ann_cagra/flowann_storage.cu @@ -0,0 +1,76 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * FlowANN is described in "Disentangling Graph Dependencies for Efficient + * Billion-Scale GPU Vector Search" (OSDI 2026). + */ + +#include + +#include +#include + +#include + +namespace flowann = cuvs::neighbors::cagra::experimental::flowann; + +TEST(FlowannStorage, OwnsCrossGraphAndPreservesItAcrossMove) +{ + raft::resources res; + flowann::device_padded_index index(res); + auto inner = raft::make_host_matrix(0, 0); + auto cross = raft::make_host_matrix(2, 2); + cross(0, 0) = 1; + cross(0, 1) = 1; + cross(1, 0) = 1; + cross(1, 1) = 0; + index.update_graph( + res, raft::make_const_mdspan(inner.view()), raft::make_const_mdspan(cross.view()), 1, 0); + cross(0, 1) = 0; + auto moved = std::move(index); + EXPECT_EQ(moved.size(), 2); + EXPECT_EQ(moved.graph_degree(), 1); + EXPECT_EQ(moved.resident_degree(), 0); + EXPECT_EQ(moved.cross_graph()(0, 1), 1); + raft::resource::sync_stream(res); +} + +TEST(FlowannStorage, RejectsInvalidCrossEdgesBeforeReplacingGraph) +{ + raft::resources res; + flowann::device_padded_index index(res); + auto inner = raft::make_host_matrix(0, 0); + auto cross = raft::make_host_matrix(2, 2); + cross(0, 0) = 1; + cross(0, 1) = 1; + cross(1, 0) = 1; + cross(1, 1) = 0; + index.update_graph( + res, raft::make_const_mdspan(inner.view()), raft::make_const_mdspan(cross.view()), 1, 0); + cross(0, 1) = 2; + EXPECT_ANY_THROW(index.update_graph( + res, raft::make_const_mdspan(inner.view()), raft::make_const_mdspan(cross.view()), 1, 0)); + EXPECT_EQ(index.cross_graph()(0, 1), 1); + EXPECT_EQ(index.size(), 2); + raft::resource::sync_stream(res); +} + +TEST(FlowannStorage, ComputesCosineNormsWithoutSearchImplementation) +{ + raft::resources res; + auto host = raft::make_host_matrix(2, 2); + host(0, 0) = 3.0f; + host(0, 1) = 4.0f; + host(1, 0) = 0.0f; + host(1, 1) = 2.0f; + auto dataset = cuvs::neighbors::make_device_padded_dataset(res, host.view()); + flowann::device_padded_index index( + res, cuvs::distance::DistanceType::CosineExpanded, dataset->as_dataset_view()); + ASSERT_TRUE(index.dataset_norms().has_value()); + float norms[2]{}; + raft::copy(norms, index.dataset_norms()->data_handle(), 2, raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + EXPECT_FLOAT_EQ(norms[0], 5.0f); + EXPECT_FLOAT_EQ(norms[1], 2.0f); +}