diff --git a/c/CMakeLists.txt b/c/CMakeLists.txt index 59932c32ee..56417602f3 100644 --- a/c/CMakeLists.txt +++ b/c/CMakeLists.txt @@ -32,6 +32,8 @@ option(BUILD_CAGRA_HNSWLIB "Build CAGRA+hnswlib interface" ON) option(BUILD_MG_ALGOS "Build with multi-GPU support" ON) option(CUVSC_STATIC_CUVS_LIBRARY "Link against statical version of libcuvs" OFF) +set(CUVS_FLOWANN_GDRCOPY_INSTALL_RPATH "$ORIGIN" CACHE STRING + "Install RPATH for external GDRCopy runtime libraries") # Check if cuVS is already available. If so, it is the user's responsibility to ensure that the # CMake package is also available at build time of the Python cuvs package. @@ -123,6 +125,13 @@ set_target_properties( VISIBILITY_INLINES_HIDDEN ON ) +if(CUVSC_STATIC_CUVS_LIBRARY) + # Imported cuvs_static can carry GDRCopy without exporting its build-time feature flags. + # Keep the runtime deployment setting available to standalone C builds too. + set_property(TARGET cuvs_c APPEND PROPERTY INSTALL_RPATH + "${CUVS_FLOWANN_GDRCOPY_INSTALL_RPATH}") +endif() + # Setup the VERSION and SOVERSION of our library # # Long term goal is that every release that breaks the C ABI will update `cuvs_abi_version` allowing diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 78bde6f657..bc1bcc363a 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -84,7 +84,14 @@ 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) +set(CUVS_FLOWANN_GDRCOPY_INSTALL_RPATH "\$ORIGIN" CACHE STRING + "Runtime search paths for externally deployed GDRCopy; never inferred from the build host") +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 +120,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 +133,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 +177,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 +257,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( @@ -959,6 +986,27 @@ if(NOT BUILD_CPU_ONLY) OUTPUT_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/cagra/search_single_cta" KERNEL_LINK_LIBRARIES jit_lto_kernel_usage_requirements ) + if(CUVS_ENABLE_FLOWANN_SEARCH) + set(flowann_ns "cuvs::neighbors::cagra::experimental::flowann::detail") + generate_jit_lto_kernels( + jit_lto_files + NAME_FORMAT + "flowann_search_single_cta_@topk_by_bitonic_sort_str@_@bitonic_sort_and_merge_multi_warps_str@_data_@data_abbrev@" + MATRIX_JSON_FILE + "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/detail/flowann/search_single_cta_matrix.json" + KERNEL_INPUT_FILE + "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/detail/flowann/search_single_cta_kernel.cu.in" + FRAGMENT_TAG_FORMAT + "${flowann_ns}::fragment_tag_search_single_cta<${neighbors_ns}::tag_@data_abbrev@, ${neighbors_ns}::tag_index_@source_index_abbrev@, ${neighbors_ns}::tag_index_@index_abbrev@, ${cagra_ns}::tag_dist_@distance_abbrev@, @topk_by_bitonic_sort@, @bitonic_sort_and_merge_multi_warps@>" + FRAGMENT_TAG_HEADER_FILES + "" + "" + "" + OUTPUT_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/flowann/search_single_cta" + KERNEL_LINK_LIBRARIES jit_lto_kernel_usage_requirements + KERNEL_EXTRA_COMPILE_OPTIONS --maxrregcount=64 + ) + endif() generate_jit_lto_kernels( jit_lto_files NAME_FORMAT @@ -1486,6 +1534,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 +1642,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 +1664,7 @@ if(NOT BUILD_CPU_ONLY) nvidia::cutlass::cutlass rtcx::rtcx ${CUVS_CTK_MATH_DEPENDENCIES} + ${CUVS_FLOWANN_GDRCOPY_DEPENDENCIES} $ $ kvikio::kvikio @@ -1687,10 +1744,30 @@ 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) + # GDRCopy is an external runtime dependency. Default to an install-relative path; + # packagers may supply an explicit deployment path, independent of the build prefix. + set_property(TARGET cuvs APPEND PROPERTY INSTALL_RPATH + "${CUVS_FLOWANN_GDRCOPY_INSTALL_RPATH}") + # 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 +1812,7 @@ if(NOT BUILD_CPU_ONLY) CUDA::nvJitLink CUDA::nvrtc rtcx::rtcx + ${CUVS_FLOWANN_GDRCOPY_DEPENDENCIES} $<$:CUDA::nvtx3> $ $ diff --git a/cpp/include/cuvs/detail/jit_lto/flowann/flowann_fragments.hpp b/cpp/include/cuvs/detail/jit_lto/flowann/flowann_fragments.hpp new file mode 100644 index 0000000000..8b07cd6401 --- /dev/null +++ b/cpp/include/cuvs/detail/jit_lto/flowann/flowann_fragments.hpp @@ -0,0 +1,21 @@ +/* + * 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 + +namespace cuvs::neighbors::cagra::experimental::flowann::detail { + +template +struct fragment_tag_search_single_cta {}; + +} // namespace cuvs::neighbors::cagra::experimental::flowann::detail 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..5d11f4f9f5 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, @@ -81,7 +150,9 @@ RAFT_DEVICE_INLINE_FUNCTION void search_core( const std::uint32_t query_id_offset, // Offset to add to query_id when calling filter const dataset_descriptor_base_t* dataset_desc, cagra_sample_filter filter_payload, - const IndexT graph_size = 0) // Original number of bits + const IndexT graph_size = 0, + const std::uint32_t query_stride = 0, // Zero retains contiguous query addressing. + const bool full_width_invalid = false) { using LOAD_T = device::LOAD_128BIT_T; @@ -111,7 +182,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); @@ -119,8 +193,12 @@ RAFT_DEVICE_INLINE_FUNCTION void search_core( uint32_t dim = dataset_desc->args.dim; uint32_t smem_ws_size_in_bytes = dataset_desc->smem_ws_size_in_bytes(); - auto smem_desc = - setup_workspace(dataset_desc, smem, queries_ptr, query_id); + auto smem_desc = setup_workspace( + dataset_desc, + smem, + query_stride == 0 ? queries_ptr + : queries_ptr + static_cast(query_id) * query_stride, + query_stride == 0 ? query_id : 0); auto* __restrict__ result_indices_buffer = reinterpret_cast(smem + smem_ws_size_in_bytes); @@ -150,19 +228,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 +256,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 +314,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,24 +369,30 @@ 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(); _CLK_REC(clk_compute_distance); + // Radix selection reads the entire result buffer. A provider may produce fewer + // children than its capacity, so old candidates must not survive in the tail. + for (std::uint32_t i = internal_topk + computed_children_this_iteration + threadIdx.x; + i < result_buffer_size; + i += blockDim.x) { + result_indices_buffer[i] = raft::upper_bound(); + result_distances_buffer[i] = raft::upper_bound(); + } + + // The filter-flag barrier below also orders these tail writes. // Filtering - use extern sample_filter function if (threadIdx.x == 0) { *filter_flag = 0; } __syncthreads(); @@ -322,11 +417,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 < result_buffer_size; 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 +484,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 +515,22 @@ 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) { + // Regular CAGRA retains its historical sentinel. FlowANN requests an output-width + // sentinel even when no source map is installed. + if (source_indices_ptr == nullptr && !full_width_invalid) { + write_indices( + result_indices_ptr, j, static_cast(invalid_index & ~index_msb_1_mask)); + } else if (index_element_tag == 3) { + reinterpret_cast(result_indices_ptr)[j] = ~uint64_t{0}; + } else { + write_indices(result_indices_ptr, j, raft::upper_bound()); + } + 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 +605,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 +716,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 +821,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..74caa6535e --- /dev/null +++ b/cpp/src/neighbors/detail/flowann/README.md @@ -0,0 +1,38 @@ +# 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. + + +## Host service thread lifetime + +The experimental queue runtime uses `std::jthread` for persistent CPU pollers. They must +continue servicing device requests while the calling thread waits for CUDA completion, +and an explicit search session can keep them alive between calls. An OpenMP parallel +region scoped to one search call cannot provide that lifetime. Threads select the +context's CUDA device and are stopped and joined before queue storage is released. +This threading choice needs an explicit exception to the repository's OpenMP-only host +threading guideline during maintainer review; this rationale does not claim that approval. + +Pollers inherit the caller's CPU affinity by default. Internal callers may opt into +physical-core pinning with `queue_params::pin_pollers`; pinned contexts share a process-wide +allocation counter rather than all starting at CPU zero. + + +## Deploying a GDRCopy-enabled library + +GDRCopy remains an external runtime dependency; installing cuVS does not bundle `libgdrapi`. +The runtime environment must provide a compatible `libgdrapi` and its SONAME links. Place them +beside the installed cuVS libraries, install them in the system loader paths, or set +`LD_LIBRARY_PATH` to the runtime GDRCopy directory before loading cuVS. The kernel driver is +also required for GDRCopy transfers. + +The installed RPATH defaults to `$ORIGIN`, so moving an installation does not retain the +build host's GDRCopy prefix. Packagers with a fixed runtime layout may set +`CUVS_FLOWANN_GDRCOPY_INSTALL_RPATH` to their deployment search paths. The same setting applies +to `libcuvs_c` when `CUVSC_STATIC_CUVS_LIBRARY=ON`, since that library then carries the GDRCopy +dependency directly. A standalone C build accepts the same cache setting independently of the +C++ build's feature flags. Set it in each separate build when using a custom deployment path. +Compile-time package discovery still uses `GDRCopy_ROOT` or the usual +CMake search paths and does not establish the runtime layout. diff --git a/cpp/src/neighbors/detail/flowann/flowann.hpp b/cpp/src/neighbors/detail/flowann/flowann.hpp new file mode 100644 index 0000000000..cc7ea971c2 --- /dev/null +++ b/cpp/src/neighbors/detail/flowann/flowann.hpp @@ -0,0 +1,827 @@ +/* + * 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; + /** Opt in to pinning pollers; otherwise inherit caller CPU affinity. */ + bool pin_pollers = 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{}; +}; + +/** FlowANN search parameters. */ +struct search_params : cuvs::neighbors::cagra::search_params { + /** Maximum number of preloaded medoid seed nodes used to initialize each query. */ + std::uint32_t num_seeds = 0; + /** Scale used by the adaptive deferred-cross-edge synchronization window. */ + float sync_window_scale = 10.0f; + /** Parent-position drop that forces synchronization for a submitted cross-edge request. */ + std::uint32_t sync_drop_threshold = 51; +}; + +/** + * @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::shared_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. Moving a context preserves active + * sessions. The host cross graph must outlive both the context and its sessions. + */ +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; + + std::shared_ptr 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; + using search_params_type = flowann::search_params; + + 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)); + // A non-owning host view may refer to pinned memory released by the caller on return. + raft::resource::sync_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()); + + auto new_seed_cross_graph = + prepare_graph_update_(res, raft::make_const_mdspan(new_cross.view())); + inner_graph_ = std::move(new_inner); + cross_graph_ = std::move(new_cross); + full_graph_degree_ = full_graph_degree; + resident_degree_ = resident_degree; + seed_cross_graph_ = std::move(new_seed_cross_graph); + } + + /** 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); + + auto new_seed_cross_graph = + prepare_graph_update_(res, raft::make_const_mdspan(cross_graph.view())); + inner_graph_ = std::move(new_inner); + cross_graph_ = std::move(cross_graph); + full_graph_degree_ = full_graph_degree; + resident_degree_ = resident_degree; + seed_cross_graph_ = std::move(new_seed_cross_graph); + } + + /** 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)); + } + } + } + + auto new_seed_cross_graph = + prepare_graph_update_(res, raft::make_const_mdspan(cross_graph.view())); + inner_graph_ = std::move(inner_graph); + cross_graph_ = std::move(cross_graph); + full_graph_degree_ = full_degree; + resident_degree_ = 0; + seed_cross_graph_ = std::move(new_seed_cross_graph); + } + + /** 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)); + } + } + + auto new_seed_cross_graph = + prepare_graph_update_(res, raft::make_const_mdspan(cross_graph.view())); + inner_graph_ = std::move(inner_graph); + cross_graph_ = std::move(cross_graph); + full_graph_degree_ = full_graph_degree; + resident_degree_ = resident_degree; + seed_cross_graph_ = std::move(new_seed_cross_graph); + } + + 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); + } + + // Prepare all potentially failing work before committing a graph replacement. + auto prepare_graph_update_( + raft::resources const& res, + raft::host_matrix_view incoming) + -> std::optional> + { + for (std::size_t seed = 0; seed < seed_ids_host_.size(); ++seed) { + RAFT_EXPECTS(seed_ids_host_[seed] < incoming.extent(0), + "FlowANN seed at position %zu is out of range for the new graph", + seed); + } + if (incoming.extent(0) != cross_graph_.extent(0)) { + source_indices_.reset(); + subgraph_offsets_.reset(); + subgraph_ids_.reset(); + seeds_.reset(); + seed_ids_host_.clear(); + seed_lookup_keys_.reset(); + seed_lookup_values_.reset(); + return std::nullopt; + } + if (seed_ids_host_.empty()) { return std::nullopt; } + auto host_cache = raft::make_host_matrix( + static_cast(seed_ids_host_.size()), incoming.extent(1)); + for (std::size_t seed = 0; seed < seed_ids_host_.size(); ++seed) { + std::copy_n(incoming.data_handle() + + static_cast(seed_ids_host_[seed]) * incoming.extent(1), + incoming.extent(1), + host_cache.data_handle() + seed * incoming.extent(1)); + } + auto device_cache = raft::make_device_matrix( + res, host_cache.extent(0), host_cache.extent(1)); + raft::copy(device_cache.data_handle(), + host_cache.data_handle(), + host_cache.size(), + raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + return device_cache; + } + + 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_{cuvs::neighbors::detail::make_empty_dataset_view(0)}; + 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>; + +/** + * Search a FlowANN index. The overload set matches the data and output-index types supported by + * CAGRA search. The graph and source-index mapping remain uint32_t internally. + */ +#define CUVS_FLOWANN_DECLARE_SEARCH(T, DatasetAlias, OutputIdxT) \ + CUVS_EXPORT void search( \ + raft::resources const& res, \ + search_params const& params, \ + DatasetAlias const& index, \ + search_context& context, \ + raft::device_matrix_view queries, \ + raft::device_matrix_view neighbors, \ + raft::device_matrix_view distances, \ + cuvs::neighbors::filtering::base_filter const& sample_filter = \ + cuvs::neighbors::filtering::none_sample_filter{}) + +#define CUVS_FLOWANN_DECLARE_SEARCH_FOR_TYPE(T) \ + CUVS_FLOWANN_DECLARE_SEARCH(T, device_padded_index, std::uint32_t); \ + CUVS_FLOWANN_DECLARE_SEARCH(T, device_padded_index, std::int64_t); + +CUVS_FLOWANN_DECLARE_SEARCH_FOR_TYPE(float); +CUVS_FLOWANN_DECLARE_SEARCH_FOR_TYPE(half); +CUVS_FLOWANN_DECLARE_SEARCH_FOR_TYPE(std::int8_t); +CUVS_FLOWANN_DECLARE_SEARCH_FOR_TYPE(std::uint8_t); + +#undef CUVS_FLOWANN_DECLARE_SEARCH_FOR_TYPE +#undef CUVS_FLOWANN_DECLARE_SEARCH + +} // namespace cuvs::neighbors::cagra::experimental::flowann + +#endif // CUVS_ENABLE_FLOWANN_SEARCH diff --git a/cpp/src/neighbors/detail/flowann/graph_state.cuh b/cpp/src/neighbors/detail/flowann/graph_state.cuh new file mode 100644 index 0000000000..d1ffa4e3e7 --- /dev/null +++ b/cpp/src/neighbors/detail/flowann/graph_state.cuh @@ -0,0 +1,138 @@ +/* + * 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 "kernel.cuh" + +#include + +namespace cuvs::neighbors::cagra::experimental::flowann::detail { + +inline constexpr std::uint32_t pending_capacity = 512; +inline constexpr std::uint32_t pending_mask = pending_capacity - 1; + +static_assert((pending_capacity & pending_mask) == 0, + "FlowANN pending capacity must be a power of two"); + +struct pending_entry { + std::uint32_t parent; + queue_reservation reservation; + std::uint16_t sync_window; +}; + +struct pending_queue { + std::uint16_t head; + std::uint16_t tail; + std::uint32_t consumed; + pending_entry entries[pending_capacity]; + + RAFT_DEVICE_INLINE_FUNCTION void initialize() + { + head = 0; + tail = 0; + consumed = 0; + } + + [[nodiscard]] RAFT_DEVICE_INLINE_FUNCTION auto size() const -> std::uint16_t + { + return static_cast((head - tail) & pending_mask); + } + + [[nodiscard]] RAFT_DEVICE_INLINE_FUNCTION auto available() const -> std::uint16_t + { + return static_cast(pending_mask - size()); + } + + [[nodiscard]] RAFT_DEVICE_INLINE_FUNCTION auto empty() const -> bool { return head == tail; } + + RAFT_DEVICE_INLINE_FUNCTION void push(pending_entry entry) + { + entries[head] = entry; + head = static_cast((head + 1) & pending_mask); + } + + [[nodiscard]] RAFT_DEVICE_INLINE_FUNCTION auto front() -> pending_entry& { return entries[tail]; } + + RAFT_DEVICE_INLINE_FUNCTION void pop() + { + tail = static_cast((tail + 1) & pending_mask); + } +}; + +template +struct graph_state { + pending_queue pending; + IndexT parent_ids[max_search_width]; + bool parent_valid[max_search_width]; + std::uint32_t parent_seed_positions[max_search_width]; + std::uint32_t parent_inner_counts[max_search_width]; + std::uint64_t parent_bit_offsets[max_search_width]; + std::uint32_t parent_subgraph_offsets[max_search_width]; + const std::uint32_t* ready_cross_data[max_search_width]; + std::uint32_t ready_cross_edge_counts[max_search_width]; + queue_reservation ready_reservations[max_search_width]; + std::uint32_t ready_cross_count; + std::uint32_t total_inner_children; + std::uint32_t total_children; + std::uint32_t previous_parent_position; + std::uint32_t iteration; + bool use_seed_cache; +}; + +#ifdef __CUDACC__ +/** Submit a complete parent batch or restore its traversal state for retry. */ +template +RAFT_DEVICE_INLINE_FUNCTION void submit_parent_requests(queue_view const& queue, + graph_state* state, + const IndexT* parent_indices, + const IndexT* internal_topk, + std::uint32_t search_width, + std::uint16_t sync_window, + IndexT* traversed_hashmap, + std::uint32_t traversed_hash_bitlen) +{ + constexpr IndexT index_msb = + cuvs::neighbors::cagra::detail::utils::gen_index_msb_1_mask::value; + std::uint32_t commands[max_search_width]; + std::uint32_t command_count = 0; + for (std::uint32_t i = 0; i < search_width; ++i) { + if (!state->parent_valid[i]) { continue; } + commands[command_count++] = static_cast(state->parent_ids[i]); + } + + queue_reservation first{}; + bool reserved = command_count == 0; + // One slot is reserved to distinguish a full pending ring from an empty one. + // Do not publish commands unless every reservation can be tracked until release. + if (command_count <= state->pending.available()) { + for (std::uint32_t attempt = 0; attempt < 200 && !reserved; ++attempt) { + reserved = try_reserve_deferred(queue, command_count, &first); + } + } + if (command_count > 0 && reserved) { + submit_batch(queue, first, commands, command_count); + for (std::uint32_t i = 0; i < command_count; ++i) { + state->pending.push({commands[i], reservation_at(first, i), sync_window}); + } + } else if (!reserved) { + // Parent selection marked these nodes as traversed. Retry their cross rows + // after making progress on older requests rather than silently losing edges. + for (std::uint32_t i = 0; i < search_width; ++i) { + if (!state->parent_valid[i]) { continue; } + const_cast(internal_topk)[parent_indices[i]] &= ~index_msb; + if (traversed_hashmap != nullptr) { + cuvs::neighbors::cagra::detail::hashmap::remove( + traversed_hashmap, traversed_hash_bitlen, state->parent_ids[i]); + } + } + } +} +#endif + +} // namespace cuvs::neighbors::cagra::experimental::flowann::detail diff --git a/cpp/src/neighbors/detail/flowann/kernel.cuh b/cpp/src/neighbors/detail/flowann/kernel.cuh new file mode 100644 index 0000000000..0b827f7073 --- /dev/null +++ b/cpp/src/neighbors/detail/flowann/kernel.cuh @@ -0,0 +1,106 @@ +/* + * 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 "queue.cuh" + +#include +#include + +#include + +namespace cuvs::neighbors::cagra::experimental::flowann::detail { + +inline constexpr std::uint32_t max_search_width = 8; +inline constexpr std::uint32_t additional_buffer_width = 2; + +template +using search_single_cta_kernel_func_t = + void(uintptr_t, + DistanceT*, + std::uint32_t, + const DataT*, + std::uint32_t, + const std::uint8_t*, + std::uint32_t, + std::uint32_t, + std::uint32_t, + const std::uint32_t*, + const std::uint32_t*, + queue_view*, + std::uint32_t, + std::uint32_t, + std::uint32_t, + const SourceIndexT*, + unsigned, + std::uint64_t, + const IndexT*, + std::uint32_t, + const std::uint32_t*, + const IndexT*, + const std::uint32_t*, + std::uint32_t, + IndexT*, + std::uint32_t, + std::uint32_t, + std::uint32_t, + std::uint32_t, + std::uint32_t, + std::uint32_t, + std::uint32_t*, + std::uint32_t, + std::uint32_t, + std::uint32_t, + std::uint32_t, + const cuvs::neighbors::cagra::detail::dataset_descriptor_base_t*, + IndexT, + float, + std::uint32_t, + cuvs::neighbors::cagra::detail::cagra_sample_filter); + +template +using search_multi_cta_kernel_func_t = + void(IndexT*, + DistanceT*, + const cuvs::neighbors::cagra::detail::dataset_descriptor_base_t*, + const DataT*, + const std::uint8_t*, + std::uint32_t, + std::uint32_t, + std::uint32_t, + const std::uint32_t*, + const std::uint32_t*, + queue_view*, + std::uint32_t, + std::uint32_t, + std::uint32_t, + std::uint32_t, + const SourceIndexT*, + unsigned, + std::uint64_t, + const IndexT*, + std::uint32_t, + const std::uint32_t*, + const IndexT*, + const std::uint32_t*, + std::uint32_t, + std::uint32_t, + IndexT*, + std::uint32_t, + std::uint32_t, + std::uint32_t, + std::uint32_t, + std::uint32_t*, + IndexT, + std::uint32_t, + float, + std::uint32_t, + cuvs::neighbors::cagra::detail::cagra_sample_filter); + +} // namespace cuvs::neighbors::cagra::experimental::flowann::detail diff --git a/cpp/src/neighbors/detail/flowann/launcher.hpp b/cpp/src/neighbors/detail/flowann/launcher.hpp new file mode 100644 index 0000000000..49cd6021bf --- /dev/null +++ b/cpp/src/neighbors/detail/flowann/launcher.hpp @@ -0,0 +1,115 @@ +/* + * 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 "kernel.cuh" +#include "planner.hpp" + +#include +#include + +#include +#include + +namespace cuvs::neighbors::cagra::experimental::flowann::detail { + +template +auto build_single_cta_launcher( + const cuvs::neighbors::cagra::detail::dataset_descriptor_host& + dataset_desc, + bool topk_by_bitonic_sort, + bool bitonic_sort_and_merge_multi_warps, + std::unique_ptr sample_filter_udf_fragment) + -> std::shared_ptr +{ + single_cta_planner + planner; + + planner.add_setup_workspace_device_function(dataset_desc.team_size, + dataset_desc.dataset_block_dim); + planner.add_compute_distance_device_function( + dataset_desc.metric, dataset_desc.team_size, dataset_desc.dataset_block_dim); + + planner.add_search_kernel(topk_by_bitonic_sort, bitonic_sort_and_merge_multi_warps); + planner.add_sample_filter_device_function(std::move(sample_filter_udf_fragment)); + return planner.get_launcher(); +} + +template +auto make_single_cta_launcher( + const cuvs::neighbors::cagra::detail::dataset_descriptor_host& + dataset_desc, + bool topk_by_bitonic_sort, + bool bitonic_sort_and_merge_multi_warps, + std::unique_ptr sample_filter_udf_fragment = nullptr) + -> std::shared_ptr +{ + using namespace cuvs::neighbors::cagra::detail; + using data_tag = decltype(get_data_type_tag()); + using index_tag = decltype(get_index_type_tag()); + using distance_tag = decltype(get_distance_type_tag()); + using source_index_tag = decltype(get_source_index_type_tag()); + + using codebook_tag = codebook_tag_standard_t; + if (dataset_desc.metric == cuvs::distance::DistanceType::BitwiseHamming) { + using query_tag = + query_type_tag_standard_t; + return build_single_cta_launcher(dataset_desc, + topk_by_bitonic_sort, + bitonic_sort_and_merge_multi_warps, + std::move(sample_filter_udf_fragment)); + } + using query_tag = query_type_tag_standard_t; + return build_single_cta_launcher(dataset_desc, + topk_by_bitonic_sort, + bitonic_sort_and_merge_multi_warps, + std::move(sample_filter_udf_fragment)); +} + +} // namespace cuvs::neighbors::cagra::experimental::flowann::detail diff --git a/cpp/src/neighbors/detail/flowann/planner.hpp b/cpp/src/neighbors/detail/flowann/planner.hpp new file mode 100644 index 0000000000..10589b6575 --- /dev/null +++ b/cpp/src/neighbors/detail/flowann/planner.hpp @@ -0,0 +1,67 @@ +/* + * 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 + +namespace cuvs::neighbors::cagra::experimental::flowann::detail { + +template +struct single_cta_planner + : cuvs::neighbors::cagra::detail:: + CagraPlannerBase { + using base_type = cuvs::neighbors::cagra::detail:: + CagraPlannerBase; + static inline rtcx::launcher_jit_cache launcher_cache{}; + + single_cta_planner() : base_type("flowann_search_single_cta", launcher_cache) {} + + void add_search_kernel(bool topk_by_bitonic_sort, bool bitonic_sort_and_merge_multi_warps) + { + if (topk_by_bitonic_sort && bitonic_sort_and_merge_multi_warps) { + this->template add_static_fragment>(); + } else if (topk_by_bitonic_sort) { + this->template add_static_fragment>(); + } else if (bitonic_sort_and_merge_multi_warps) { + this->template add_static_fragment>(); + } else { + this->template add_static_fragment>(); + } + } +}; + +} // namespace cuvs::neighbors::cagra::experimental::flowann::detail diff --git a/cpp/src/neighbors/detail/flowann/queue.cuh b/cpp/src/neighbors/detail/flowann/queue.cuh new file mode 100644 index 0000000000..6e32f49a88 --- /dev/null +++ b/cpp/src/neighbors/detail/flowann/queue.cuh @@ -0,0 +1,286 @@ +/* + * 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 +#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; + std::uint32_t const* abort_flag; + // Private to each device-side queue view. Avoid a PCIe read on every readiness poll. + mutable unsigned long long next_abort_check{}; + mutable bool aborted{}; +}; + +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 const& queue, + queue_reservation reservation) -> bool +{ + if (load_published_counter(&queue.result_sequence[reservation.slot]) == + reservation.sequence + 1) { + return true; + } + if (queue.aborted) { return true; } + if (queue.abort_flag == nullptr) { return false; } + auto const now = clock64(); + if (queue.next_abort_check == 0) { + queue.next_abort_check = now + (1ull << 24); + return false; + } + if (now < queue.next_abort_check) { return false; } + // Read the host-written emergency signal at a bounded interval. Ordinary completion + // stays in device memory; continuously loading mapped host memory stalls queue traffic. + queue.next_abort_check = now + (1ull << 24); + cuda::atomic_ref aborted( + *const_cast(queue.abort_flag)); + queue.aborted = aborted.load(cuda::memory_order_acquire) != 0; + return queue.aborted; +} + +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); + // Internal failure injection uses the same path as an unexpected poller exception. + static void report_failure(search_context& context, std::exception_ptr error); + 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/detail/flowann/search_single_cta.cuh b/cpp/src/neighbors/detail/flowann/search_single_cta.cuh new file mode 100644 index 0000000000..9cc3ae6986 --- /dev/null +++ b/cpp/src/neighbors/detail/flowann/search_single_cta.cuh @@ -0,0 +1,542 @@ +/* + * 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 "graph_state.cuh" +#include "launcher.hpp" +#include "queue.cuh" + +#include +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace cuvs::neighbors::cagra::experimental::flowann::detail { + +inline auto automatic_max_iterations(std::uint32_t itopk_size, + std::uint32_t search_width, + std::uint64_t dataset_size, + std::uint32_t graph_degree, + std::uint32_t min_iterations) -> std::uint32_t +{ + auto iterations = itopk_size / search_width; + std::uint64_t reachable = 1; + while (reachable < dataset_size) { + auto const expansion = std::max(2, graph_degree / 2); + if (reachable > std::numeric_limits::max() / expansion) { break; } + reachable *= expansion; + ++iterations; + } + return std::max(iterations, min_iterations); +} + +template +struct single_cta_plan + : cuvs::neighbors::cagra::detail:: + search_plan_impl { + using base_type = cuvs::neighbors::cagra::detail:: + search_plan_impl; + + std::uint32_t num_itopk_candidates{}; + + static auto planning_params(flowann::search_params const& input) + -> cuvs::neighbors::cagra::search_params + { + cuvs::neighbors::cagra::search_params result = input; + result.search_width += additional_buffer_width; + result.algo = cuvs::neighbors::cagra::search_algo::SINGLE_CTA; + result.persistent = false; + return result; + } + + single_cta_plan( + raft::resources const& res, + flowann::search_params const& params, + const cuvs::neighbors::cagra::detail::dataset_descriptor_host& + dataset_desc, + int64_t dim, + int64_t dataset_size, + int64_t graph_degree, + std::uint32_t topk) + : base_type(res, planning_params(params), dataset_desc, dim, dataset_size, graph_degree, topk) + { + // The inherited CAGRA plan sizes hashmaps and candidate storage for two additional graph rows, + // while the actual traversal width remains the user-requested value. + this->search_width = params.search_width; + if (params.max_iterations == 0) { + this->max_iterations = automatic_max_iterations(static_cast(this->itopk_size), + params.search_width, + dataset_size, + graph_degree, + params.min_iterations); + } else { + this->max_iterations = std::max(params.max_iterations, params.min_iterations); + } + set_launch_params(res, params.num_seeds); + } + + void set_launch_params(raft::resources const& res, std::uint32_t num_seeds) + { + // Seed initialization inserts every requested seed before selecting the internal top-k. + // The inherited CAGRA plan only budgets for the retained top-k and graph expansion. + if (num_seeds > this->itopk_size) { + namespace hashmap = cuvs::neighbors::cagra::detail::hashmap; + if (this->small_hash_bitlen != 0) { + auto bits = this->small_hash_bitlen; + while (bits <= 13 && num_seeds > hashmap::get_size(bits) * this->hashmap_max_fill_rate) { + ++bits; + } + if (bits <= 13) { + this->hash_bitlen = this->small_hash_bitlen = bits; + } else { + RAFT_EXPECTS(this->hashmap_mode != cuvs::neighbors::cagra::hash_mode::SMALL, + "FlowANN seeds exceed the small visited hashmap capacity"); + this->small_hash_bitlen = 0; + this->small_hash_reset_interval = 1024 * 1024; + this->hash_bitlen = 11; + } + } + if (this->small_hash_bitlen == 0) { + auto const max_visited = + static_cast(num_seeds) + + (static_cast(this->search_width) + additional_buffer_width) * + this->graph_degree * this->max_iterations; + while (this->hash_bitlen <= 20 && + max_visited > hashmap::get_size(this->hash_bitlen) * this->hashmap_max_fill_rate) { + ++this->hash_bitlen; + } + RAFT_EXPECTS(this->hash_bitlen <= 20, + "FlowANN seeds and traversal exceed the visited hashmap capacity"); + } + } + auto const expansion_candidates = + (static_cast(this->search_width) + additional_buffer_width) * + static_cast(this->graph_degree); + auto const base_result_buffer_size = + static_cast(this->itopk_size) + expansion_candidates; + this->result_buffer_size = std::max(base_result_buffer_size, num_seeds); + num_itopk_candidates = this->result_buffer_size - static_cast(this->itopk_size); + auto const result_buffer_size_32 = + raft::round_up_safe(this->result_buffer_size, 32); + constexpr std::uint32_t max_itopk = 512; + RAFT_EXPECTS( + this->itopk_size <= max_itopk, "FlowANN itopk_size cannot be larger than %u", max_itopk); + + auto const base_smem_size = + this->dataset_desc.smem_ws_size_in_bytes + + (sizeof(IndexT) + sizeof(DistanceT)) * result_buffer_size_32 + + sizeof(IndexT) * cuvs::neighbors::cagra::detail::hashmap::get_size(this->small_hash_bitlen) + + sizeof(IndexT) * this->search_width + 4 * sizeof(std::uint32_t); + auto additional_smem_size = std::uint32_t{0}; + if (num_itopk_candidates > 256) { + auto const radix_itopk = this->itopk_size <= 256 ? 256u : 512u; + additional_smem_size = + cuvs::neighbors::cagra::detail::single_cta_search::topk_by_radix_sort::smem_size( + radix_itopk) * + sizeof(std::uint32_t); + } + if constexpr (!std::is_same_v) { + using scan_op_t = cub::WarpScan; + additional_smem_size = + std::max(additional_smem_size, sizeof(typename scan_op_t::TempStorage)); + } + this->smem_size = base_smem_size + additional_smem_size; + // graph_state is declared with __shared__ in the JIT kernel, in addition to + // the dynamic workspace passed to the launch API. Include alignment padding. + auto const total_smem_size = + raft::round_up_safe(sizeof(graph_state), std::size_t{128}) + + this->smem_size; + auto const properties = raft::resource::get_device_properties(res); + RAFT_EXPECTS(total_smem_size <= properties.sharedMemPerBlockOptin, + "FlowANN search needs %zu bytes of shared memory, device limit is %zu", + total_smem_size, + properties.sharedMemPerBlockOptin); + + constexpr std::uint32_t min_block_size = 64; + constexpr std::uint32_t min_radix_block_size = 256; + constexpr std::uint32_t max_block_size = 1024; + auto block_size = static_cast(this->thread_block_size); + if (block_size == 0) { + block_size = num_itopk_candidates > 256 ? min_radix_block_size : min_block_size; + while (num_itopk_candidates > 256 && block_size < max_block_size && + max_itopk / block_size > 4) { + block_size *= 2; + } + constexpr std::uint32_t smem_per_warp_limit = 4096; + while (block_size < max_block_size && + total_smem_size > smem_per_warp_limit / 32 * block_size) { + block_size *= 2; + } + auto const properties = raft::resource::get_device_properties(res); + while (block_size < max_block_size && + this->graph_degree * this->search_width * this->team_size >= block_size * 2 && + this->max_queries <= (1024 / (block_size * 2)) * properties.multiProcessorCount) { + block_size *= 2; + } + } + RAFT_EXPECTS(block_size >= min_block_size && block_size <= max_block_size, + "FlowANN thread_block_size must be in [64, 1024]"); + this->thread_block_size = block_size; + + this->hashmap_size = 0; + if (this->small_hash_bitlen == 0) { + this->hashmap_size = + this->max_queries * cuvs::neighbors::cagra::detail::hashmap::get_size(this->hash_bitlen); + this->hashmap.resize(this->hashmap_size, raft::resource::get_cuda_stream(res)); + } + } +}; + +template +auto make_dataset_descriptor(raft::resources const& res, + flowann::search_params& params, + flowann::index const& index) +{ + if constexpr (cuvs::neighbors::is_device_vpq_f16_dataset_view_v) { + if (params.smem_dtype == cuvs::neighbors::cagra::internal_dtype::E5M2 && + raft::getComputeCapability().first < 9) { + RAFT_LOG_WARN( + "FlowANN VPQ E5M2 smem_dtype requires native FP8 support on SM90+. Falling back to F16."); + params.smem_dtype = cuvs::neighbors::cagra::internal_dtype::F16; + } + return cuvs::neighbors::cagra::detail:: + dataset_descriptor_init_with_cache( + res, params, index.dataset(), index.metric(), nullptr); + } else { + static_assert(cuvs::neighbors::is_device_padded_dataset_view_v, + "FlowANN search supports device-padded or float16 VPQ datasets"); + if (params.smem_dtype != cuvs::neighbors::cagra::internal_dtype::F16) { + RAFT_LOG_WARN("FlowANN dense search supports only F16 smem_dtype. Falling back to F16."); + params.smem_dtype = cuvs::neighbors::cagra::internal_dtype::F16; + } + auto const dataset_norms = index.dataset_norms(); + RAFT_EXPECTS( + index.metric() != cuvs::distance::DistanceType::CosineExpanded || dataset_norms.has_value(), + "FlowANN dataset norms must be available for CosineExpanded metric"); + auto const* dataset_norms_ptr = + dataset_norms.has_value() ? dataset_norms->data_handle() : nullptr; + return cuvs::neighbors::cagra::detail:: + dataset_descriptor_init_with_cache( + res, params, index.dataset(), index.metric(), dataset_norms_ptr); + } +} + +template +void postprocess_distances(raft::resources const& res, + cuvs::distance::DistanceType metric, + raft::device_matrix_view queries, + raft::device_matrix_view distances) +{ + static_assert(std::is_same_v, + "FlowANN currently supports only float distances"); + constexpr float scale = cuvs::spatial::knn::detail::utils::config::kDivisor / + cuvs::spatial::knn::detail::utils::config::kDivisor; + + if (metric == cuvs::distance::DistanceType::CosineExpanded) { + auto query_norms = raft::make_device_vector(res, queries.extent(0)); + auto scaled_sq_op = raft::compose_op( + raft::sq_op{}, raft::div_const_op{DistanceT(scale)}, raft::cast_op()); + raft::linalg::reduce(res, + queries, + query_norms.view(), + DistanceT{0}, + false, + scaled_sq_op, + raft::add_op(), + raft::sqrt_op{}); + raft::linalg::matrix_vector_op( + res, + raft::make_const_mdspan(distances), + raft::make_const_mdspan(query_norms.view()), + distances, + raft::compose_op(raft::add_const_op{DistanceT{1}}, raft::div_checkzero_op{})); + } else { + cuvs::neighbors::ivf::detail::postprocess_distances(res, + distances.data_handle(), + distances.data_handle(), + metric, + distances.extent(0), + distances.extent(1), + scale, + true); + } +} + +class context_run_guard { + public: + explicit context_run_guard(search_context& context, cudaStream_t stream) + : context_(context), stream_(stream) + { + search_context_access::acquire_search(context_); + } + context_run_guard(context_run_guard const&) = delete; + auto operator=(context_run_guard const&) -> context_run_guard& = delete; + + ~context_run_guard() + { + if (active_) { + // A failed launch or an exception during a later batch must not stop pollers + // while an earlier kernel still waits for its host-serviced requests. + RAFT_CUDA_TRY_NO_THROW(cudaStreamSynchronize(stream_)); + try { + active_ = false; + search_context_access::release_search(context_); + } catch (...) { + } + } + } + + void pause() + { + active_ = false; + search_context_access::release_search(context_); + } + + private: + search_context& context_; + cudaStream_t stream_; + bool active_{true}; +}; + +template +void search_single_cta(raft::resources const& res, + flowann::search_params params, + flowann::index const& index, + search_context& context, + raft::device_matrix_view queries, + raft::device_matrix_view neighbors, + raft::device_matrix_view distances, + SampleFilterT sample_filter) +{ + using data_t = DataT; + using index_t = std::uint32_t; + using distance_t = float; + using source_index_t = std::uint32_t; + using filter_t = + typename cuvs::neighbors::cagra::detail::CagraSampleFilterT_Selector::type; + + static_assert( + std::is_same_v || std::is_same_v, + "FlowANN output indices must be uint32_t or int64_t"); + RAFT_EXPECTS(index.size() > 0 && index.graph_degree() > 0, + "FlowANN search requires a non-empty tiered graph"); + RAFT_EXPECTS(index.dataset().n_rows() == static_cast(index.size()), + "FlowANN search requires a dataset row for every graph node"); + RAFT_EXPECTS(index.subgraph_offsets().has_value() && index.subgraph_ids().has_value(), + "FlowANN search requires subgraph offsets and per-node subgraph IDs"); + RAFT_EXPECTS(queries.extent(1) == index.dim(), + "FlowANN query dimension must match the index dimension"); + RAFT_EXPECTS(neighbors.extent(0) == queries.extent(0) && distances.extent(0) == queries.extent(0), + "FlowANN output row count must match the number of queries"); + RAFT_EXPECTS(neighbors.extent(1) == distances.extent(1) && neighbors.extent(1) > 0, + "FlowANN neighbor and distance outputs must have the same positive k"); + RAFT_EXPECTS(queries.extent(0) <= std::numeric_limits::max(), + "FlowANN single-CTA query count exceeds uint32"); + RAFT_EXPECTS(params.search_width > 0 && params.search_width <= max_search_width, + "FlowANN search_width must be in [1, %u]", + max_search_width); + RAFT_EXPECTS(params.sync_window_scale >= 0.0f, "FlowANN sync_window_scale must not be negative"); + RAFT_EXPECTS(params.algo == cuvs::neighbors::cagra::search_algo::AUTO || + params.algo == cuvs::neighbors::cagra::search_algo::SINGLE_CTA, + "This FlowANN path supports SINGLE_CTA or AUTO"); + RAFT_EXPECTS(search_context_access::row_count(context) == index.size() && + search_context_access::row_length(context) == index.cross_degree() + 1, + "FlowANN search context does not match the index cross graph"); + + auto const requested_seeds = params.num_seeds; + auto const index_seeds = index.seeds(); + auto const seed_cross_graph = index.seed_cross_graph(); + auto const seed_lookup_keys = index.seed_lookup_keys(); + auto const seed_lookup_values = index.seed_lookup_values(); + RAFT_EXPECTS(requested_seeds == 0 || index_seeds.has_value(), + "FlowANN num_seeds requires preloaded index seeds"); + RAFT_EXPECTS(!index_seeds.has_value() || requested_seeds <= index_seeds->extent(0), + "FlowANN num_seeds exceeds the number of preloaded seeds"); + RAFT_EXPECTS(requested_seeds == 0 || seed_cross_graph.has_value(), + "FlowANN num_seeds requires cached seed cross-graph rows"); + RAFT_EXPECTS( + requested_seeds == 0 || (seed_lookup_keys.has_value() && seed_lookup_values.has_value()), + "FlowANN num_seeds requires a prebuilt seed lookup"); + auto const* seeds = requested_seeds == 0 ? nullptr : index_seeds->data_handle(); + auto const* seed_cross_ptr = requested_seeds == 0 ? nullptr : seed_cross_graph->data_handle(); + auto const* seed_lookup_keys_ptr = + requested_seeds == 0 ? nullptr : seed_lookup_keys->data_handle(); + auto const* seed_lookup_values_ptr = + requested_seeds == 0 ? nullptr : seed_lookup_values->data_handle(); + auto const seed_lookup_capacity = + requested_seeds == 0 ? 0u : static_cast(seed_lookup_keys->extent(0)); + + if (queries.extent(0) == 0) { return; } + + auto device_properties = raft::resource::get_device_properties(res); + if (params.max_queries == 0) { + params.max_queries = std::min(queries.extent(0), device_properties.maxGridSize[1]); + } + params.max_queries = std::min(params.max_queries, device_properties.maxGridSize[1]); + RAFT_EXPECTS(params.max_queries > 0, "FlowANN max_queries must be positive"); + + auto descriptor = make_dataset_descriptor(res, params, index); + auto const topk = static_cast(neighbors.extent(1)); + single_cta_plan plan( + res, params, descriptor, index.dim(), index.size(), index.graph_degree(), topk); + plan.check(topk); + plan.num_seeds = requested_seeds; + + auto launch_config = cuvs::neighbors::cagra::detail::single_cta_search::compute_launch_config( + plan.num_itopk_candidates, + static_cast(plan.itopk_size), + static_cast(plan.thread_block_size)); + auto launcher_filter = cuvs::neighbors::cagra::detail::set_offset(sample_filter, 0); + auto launcher = make_single_cta_launcher< + data_t, + index_t, + distance_t, + source_index_t, + cuvs::neighbors::cagra::detail::sample_filter_jit_tag_t>( + descriptor, + launch_config.topk_by_bitonic_sort, + launch_config.bitonic_sort_and_merge_multi_warps, + cuvs::neighbors::cagra::detail::make_cagra_sample_filter_udf_fragment( + launcher_filter)); + RAFT_EXPECTS(launcher != nullptr, "Failed to create the FlowANN single-CTA JIT launcher"); + + std::unique_ptr> padded_queries; + data_t const* query_data{}; + std::uint32_t query_stride{}; + if (cuvs::neighbors::matrix_row_width_matches_cagra_required(queries)) { + auto view = cuvs::neighbors::make_device_padded_dataset_view(res, queries); + query_data = view.data_view().data_handle(); + query_stride = view.data_view().stride(0); + } else { + padded_queries = cuvs::neighbors::make_device_padded_dataset(res, queries); + auto view = padded_queries->view(); + query_data = view.data_view().data_handle(); + query_stride = view.data_view().stride(0); + } + auto const max_batch_queries = plan.max_queries; + auto const stream = raft::resource::get_cuda_stream(res).get(); + auto const* descriptor_ptr = descriptor.dev_ptr(stream); + auto const inner_graph = index.inner_graph(); + auto const subgraph_offsets = index.subgraph_offsets().value(); + auto const subgraph_ids = index.subgraph_ids().value(); + auto const source_indices = index.source_indices(); + auto const* source_ptr = source_indices.has_value() ? source_indices->data_handle() : nullptr; + constexpr uintptr_t output_tag = raft::Pow2::Log2; + static_assert(output_tag <= 3, "FlowANN output index type cannot exceed 8 bytes"); + auto const num_queries = static_cast(queries.extent(0)); + using kernel_type = search_single_cta_kernel_func_t; + + context_run_guard running(context, stream); + for (std::uint32_t offset = 0; offset < num_queries; + offset += static_cast(max_batch_queries)) { + auto const count = + static_cast(std::min(max_batch_queries, num_queries - offset)); + auto const result_ptr = reinterpret_cast(neighbors.data_handle() + + static_cast(offset) * topk); + if constexpr (output_tag <= 1) { + RAFT_EXPECTS((result_ptr & 0x3) == 0, + "FlowANN result indices must be at least 4-byte aligned"); + } + auto const* query_ptr = query_data + static_cast(offset) * query_stride; + auto batch_filter = cuvs::neighbors::cagra::detail::set_offset(sample_filter, offset); + auto filter_payload = + cuvs::neighbors::cagra::detail::extract_cagra_sample_filter(batch_filter, + stream); + dim3 grid(1, count, 1); + dim3 block(static_cast(plan.thread_block_size), 1, 1); + + auto kernel_launcher = [&]() { + launcher->template dispatch( + stream, + grid, + block, + static_cast(plan.smem_size), + result_ptr | output_tag, + distances.data_handle() + static_cast(offset) * topk, + topk, + query_ptr, + query_stride, + inner_graph.data_handle(), + static_cast(inner_graph.extent(1)), + static_cast(index.node_per_cacheline()), + static_cast(index.n_bits()), + subgraph_offsets.data_handle(), + subgraph_ids.data_handle(), + search_context_access::device_queues(context), + search_context_access::num_queues(context), + index.graph_degree(), + index.cross_degree(), + source_ptr, + static_cast(plan.num_random_samplings), + plan.rand_xor_mask, + seeds, + requested_seeds, + seed_cross_ptr, + seed_lookup_keys_ptr, + seed_lookup_values_ptr, + seed_lookup_capacity, + plan.hashmap.data(), + launch_config.max_candidates, + launch_config.max_itopk, + static_cast(plan.itopk_size), + static_cast(plan.search_width), + static_cast(plan.min_iterations), + static_cast(plan.max_iterations), + static_cast(nullptr), + static_cast(plan.hash_bitlen), + static_cast(plan.small_hash_bitlen), + static_cast(plan.small_hash_reset_interval), + filter_payload.query_id_offset, + descriptor_ptr, + index.size(), + params.sync_window_scale, + params.sync_drop_threshold, + filter_payload); + }; + cuvs::neighbors::detail::safely_launch_kernel_with_smem_size( + plan.smem_size, kernel_launcher, launcher->get_kernel()); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + } + raft::resource::sync_stream(res); + running.pause(); + + postprocess_distances(res, index.metric(), queries, distances); +} + +} // namespace cuvs::neighbors::cagra::experimental::flowann::detail diff --git a/cpp/src/neighbors/detail/flowann/search_single_cta_jit.cuh b/cpp/src/neighbors/detail/flowann/search_single_cta_jit.cuh new file mode 100644 index 0000000000..8dd67283ce --- /dev/null +++ b/cpp/src/neighbors/detail/flowann/search_single_cta_jit.cuh @@ -0,0 +1,637 @@ +/* + * 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 "graph_state.cuh" + +#include +#include + +#include +#include +#include + +#include +#include +#include + +namespace cuvs::neighbors::cagra::experimental::flowann::detail { + +RAFT_DEVICE_INLINE_FUNCTION auto load_packed_neighbor(const std::uint8_t* inner_graph, + std::uint64_t bit_position, + std::uint32_t neighbor, + std::uint32_t n_bits) -> std::uint32_t +{ + bit_position += static_cast(neighbor) * n_bits; + auto const byte_position = bit_position / 8; + auto const bit_offset = static_cast(bit_position % 8); + auto const byte_count = (bit_offset + n_bits + 7) / 8; + + // The legacy 24-bit layout is byte aligned. Issue the three independent byte + // reads without a runtime loop; never read beyond the encoded neighbor. + if (n_bits == 24 && bit_offset == 0) { + return static_cast(inner_graph[byte_position]) | + (static_cast(inner_graph[byte_position + 1]) << 8) | + (static_cast(inner_graph[byte_position + 2]) << 16); + } + + std::uint64_t packed = 0; + for (std::uint32_t i = 0; i < byte_count; ++i) { + packed |= static_cast(inner_graph[byte_position + i]) << (8 * i); + } + auto const mask = n_bits == 32 ? std::uint64_t{0xffffffffu} : ((std::uint64_t{1} << n_bits) - 1); + return static_cast((packed >> bit_offset) & mask); +} + +template +struct tiered_graph_provider { + const std::uint8_t* inner_graph; + std::uint32_t inner_row_length; + std::uint32_t node_per_cacheline; + std::uint32_t n_bits; + const std::uint32_t* subgraph_offsets; + const std::uint32_t* subgraph_ids; + const IndexT* seeds; + const std::uint32_t* seed_cross_graph; + const IndexT* seed_lookup_keys; + const std::uint32_t* seed_lookup_values; + std::uint32_t seed_lookup_capacity; + std::uint32_t num_seeds; + queue_view queue; + std::uint32_t graph_degree; + std::uint32_t cross_graph_degree; + IndexT graph_size; + float sync_window_scale; + std::uint32_t sync_drop_threshold; + graph_state* state; + + RAFT_DEVICE_INLINE_FUNCTION void initialize() const + { + if (threadIdx.x == 0) { + state->pending.initialize(); + state->previous_parent_position = 0; + state->iteration = 0; + state->use_seed_cache = false; + } + } + + [[nodiscard]] RAFT_DEVICE_INLINE_FUNCTION auto candidate_capacity( + std::uint32_t search_width) const -> std::uint32_t + { + return (search_width + additional_buffer_width) * graph_degree; + } + + [[nodiscard]] RAFT_DEVICE_INLINE_FUNCTION auto result_buffer_size( + std::uint32_t base_size, std::uint32_t requested_seeds) const -> std::uint32_t + { + return seeds == nullptr ? base_size : std::max(base_size, requested_seeds); + } + + [[nodiscard]] RAFT_DEVICE_INLINE_FUNCTION auto seed_pointer(const IndexT* seeds, + std::uint32_t, + std::uint32_t) const -> const IndexT* + { + return seeds; + } + + 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 seeds != nullptr && num_seeds > 0 ? 0 : 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 seeds != nullptr && num_seeds > 0 ? 1 : num_cta_per_query * num_queries * num_partitions; + } + + [[nodiscard]] RAFT_DEVICE_INLINE_FUNCTION auto seed_position(IndexT key) const -> std::uint32_t + { + if (seed_lookup_keys == nullptr || seed_lookup_values == nullptr || seed_lookup_capacity == 0) { + return num_seeds; + } + + std::uint64_t hash = static_cast(key); + hash ^= hash >> 16; + hash *= 0x85ebca77u; + hash ^= hash >> 13; + hash *= 0xc2b2ae3du; + hash ^= hash >> 16; + auto slot = static_cast(hash % seed_lookup_capacity); + + constexpr auto empty_key = ~IndexT{0}; + for (std::uint32_t probe = 0; probe < seed_lookup_capacity; ++probe) { + auto const stored = seed_lookup_keys[slot]; + if (stored == empty_key) { return num_seeds; } + if (stored == key) { return seed_lookup_values[slot]; } + slot = (slot + 1) % seed_lookup_capacity; + } + return num_seeds; + } + + RAFT_DEVICE_INLINE_FUNCTION auto initial_candidate_count(std::uint32_t candidate_count, + std::uint32_t requested_seeds, + std::uint32_t block_id, + std::uint32_t num_blocks) const + -> std::uint32_t + { + if (seeds == nullptr || requested_seeds == 0) { return candidate_count; } + if (block_id >= requested_seeds) { return 0; } + auto const seeds_for_block = (requested_seeds - block_id + num_blocks - 1) / num_blocks; + // The result buffer includes space for the explicit seeds. The graph degree only limits + // random initialization; applying it here silently discards later medoids. + return seeds_for_block; + } + + RAFT_DEVICE_INLINE_FUNCTION auto initial_distillation_count(std::uint32_t count, + std::uint32_t requested_seeds) const + -> std::uint32_t + { + return seeds != nullptr && requested_seeds > 0 ? 1u : count; + } + + RAFT_DEVICE_INLINE_FUNCTION auto expand( + IndexT* result_indices, + DistanceT* result_distances, + const cuvs::neighbors::cagra::detail::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, + IndexT* traversed_hashmap = nullptr, + std::uint32_t traversed_hash_bitlen = 0, + int* result_position = nullptr, + int max_result_position = 0) const -> std::uint32_t + { + using cuvs::neighbors::cagra::detail::device::team_sum; + using cuvs::neighbors::cagra::detail::utils::gen_index_msb_1_mask; + constexpr IndexT invalid_index = ~static_cast(0); + constexpr IndexT index_msb = gen_index_msb_1_mask::value; + + auto store_child = [&](IndexT child, std::uint32_t i) { + if (child != invalid_index && cuvs::neighbors::cagra::detail::hashmap::insert( + visited_hashmap, visited_hash_bitlen, child) == 0) { + child = invalid_index; + } else if (child != invalid_index && traversed_hashmap != nullptr && + cuvs::neighbors::cagra::detail::hashmap::search( + traversed_hashmap, traversed_hash_bitlen, child)) { + child = invalid_index; + } + if (result_position == nullptr) { + result_indices[i] = child; + } else if (child != invalid_index) { + auto const position = atomicSub(result_position, 1) - 1; + result_indices[position] = child; + } + }; + + // A one-parent compacted expansion fits in one warp. Decode and deduplicate + // its inner neighbors while warp 0 publishes requests. Compaction writes + // only the candidate tail, beyond the retained parent pool. + bool const overlap_inner = + result_position != nullptr && search_width == 1 && graph_degree <= 32; + + // Read packed-row metadata independently while warp 0 publishes cross-graph requests. + // These lanes derive parents directly; they do not depend on warp 0's shared writes. + if (threadIdx.x >= 32 && threadIdx.x - 32 < search_width) { + auto const p = threadIdx.x - 32; + auto const position = parent_indices[p]; + auto const valid = + position != invalid_index && inner_graph != nullptr && inner_row_length != 0; + std::uint32_t count = 0; + std::uint64_t bits = 0; + std::uint32_t base = 0; + if (valid) { + auto const parent = static_cast(internal_topk[position] & ~index_msb); + auto const row = static_cast(parent) / node_per_cacheline; + auto const col = static_cast(parent % node_per_cacheline); + auto const start = row * inner_row_length; + count = inner_graph[start + col]; + bits = (start + node_per_cacheline) * 8; + for (std::uint32_t j = 0; j < col; ++j) { + bits += static_cast(inner_graph[start + j]) * n_bits; + } + if (count > 0) { base = subgraph_offsets[subgraph_ids[parent]]; } + } + state->parent_inner_counts[p] = count; + state->parent_bit_offsets[p] = bits; + state->parent_subgraph_offsets[p] = base; + } + + if (overlap_inner && threadIdx.x >= 32 && threadIdx.x < 64) { + __syncwarp(); + auto const lane = threadIdx.x - 32; + if (lane < state->parent_inner_counts[0]) { + auto const local = + load_packed_neighbor(inner_graph, state->parent_bit_offsets[0], lane, n_bits); + auto const decoded = static_cast(state->parent_subgraph_offsets[0]) + local; + if (decoded < graph_size) { store_child(static_cast(decoded), lane); } + } + } + + if (threadIdx.x == 0) { + auto const current_parent_position = + parent_indices[0] == invalid_index ? 0u : static_cast(parent_indices[0]); + auto const sharp_drop = + state->previous_parent_position > current_parent_position && + state->previous_parent_position - current_parent_position >= sync_drop_threshold; + auto const window_value = + sharp_drop + ? 0u + : std::max(search_width, + static_cast(sync_window_scale * current_parent_position)); + auto const sync_window = static_cast(std::min(window_value, 0xffffu)); + + for (std::uint32_t i = 0; i < search_width; ++i) { + auto const parent_position = parent_indices[i]; + auto const valid = parent_position != invalid_index; + auto const parent = + valid ? static_cast(internal_topk[parent_position] & ~index_msb) : IndexT{}; + state->parent_ids[i] = parent; + state->parent_valid[i] = valid; + } + + state->use_seed_cache = state->iteration == 0 && seeds != nullptr && + seed_cross_graph != nullptr && seed_lookup_keys != nullptr && + seed_lookup_values != nullptr && seed_lookup_capacity > 0 && + num_seeds > 0; + if (state->use_seed_cache) { + for (std::uint32_t i = 0; i < search_width; ++i) { + if (!state->parent_valid[i]) { continue; } + auto const seed_position = this->seed_position(state->parent_ids[i]); + state->parent_seed_positions[i] = seed_position; + if (seed_position == num_seeds) { + state->use_seed_cache = false; + break; + } + } + } + + if (!state->use_seed_cache) { + submit_parent_requests(queue, + state, + parent_indices, + internal_topk, + search_width, + sync_window, + traversed_hashmap, + traversed_hash_bitlen); + } + } + // The one-parent path can consume ready rows while warp 1 prepares inner + // neighbors. Warp 0 reads the raw inner count independently below; the + // following join still protects all candidate writes and row descriptors. + if (!overlap_inner) { __syncthreads(); } + + auto decode_inner = [&](std::uint32_t i) { + IndexT child = invalid_index; + auto local = i; + std::uint32_t parent_index = 0; + for (; parent_index < search_width; ++parent_index) { + auto const count = state->parent_inner_counts[parent_index]; + if (local < count) { break; } + local -= count; + } + if (parent_index < search_width && state->parent_valid[parent_index]) { + auto const local_neighbor = + load_packed_neighbor(inner_graph, state->parent_bit_offsets[parent_index], local, n_bits); + auto const decoded = + static_cast(state->parent_subgraph_offsets[parent_index]) + local_neighbor; + if (decoded < graph_size) { child = static_cast(decoded); } + } + return child; + }; + + // The caller has acquired this row's ready sequence. A partial row remains + // pending until the rest of its neighbors fit in a later expansion. + auto consume_cross_row = [&](std::uint32_t& remaining) { + auto const& entry = state->pending.front(); + auto const* row = result_row(queue, entry.reservation); + auto edge_count = std::min(row[0], cross_graph_degree); + auto const consumed = state->pending.consumed; + edge_count = edge_count > consumed ? edge_count - consumed : 0; + auto const take = std::min(edge_count, remaining); + auto const ready = state->ready_cross_count; + state->ready_cross_data[ready] = row + 1 + consumed; + state->ready_cross_edge_counts[ready] = take; + ++state->ready_cross_count; + remaining -= take; + + if (take < edge_count) { + state->pending.consumed += take; + return false; + } + state->pending.consumed = 0; + state->ready_reservations[ready] = entry.reservation; + state->pending.pop(); + return true; + }; + + if (threadIdx.x == 0) { + std::uint32_t inner_count = 0; + if (overlap_inner) { + if (state->parent_valid[0] && inner_graph != nullptr && inner_row_length != 0) { + auto const parent = state->parent_ids[0]; + auto const row = static_cast(parent) / node_per_cacheline; + auto const col = static_cast(parent % node_per_cacheline); + inner_count = inner_graph[row * inner_row_length + col]; + } + } else { + for (std::uint32_t i = 0; i < search_width; ++i) { + inner_count += state->parent_inner_counts[i]; + } + } + state->total_inner_children = inner_count; + + auto remaining = candidate_capacity(search_width) - inner_count; + state->ready_cross_count = 0; + + if (state->use_seed_cache) { + auto const row_length = cross_graph_degree + 1; + for (std::uint32_t i = 0; + i < search_width && state->ready_cross_count < max_search_width && remaining > 0; + ++i) { + if (!state->parent_valid[i]) { continue; } + auto const* row = seed_cross_graph + + static_cast(state->parent_seed_positions[i]) * row_length; + auto const take = std::min(std::min(row[0], cross_graph_degree), remaining); + if (take == 0) { continue; } + auto const ready = state->ready_cross_count; + state->ready_cross_data[ready] = row + 1; + state->ready_cross_edge_counts[ready] = take; + ++state->ready_cross_count; + remaining -= take; + } + } else { + auto available = state->pending.available(); + auto deferred = state->pending.size(); + while (!state->pending.empty() && state->ready_cross_count < max_search_width && + remaining > 0) { + auto& entry = state->pending.front(); + auto need_wait = available < search_width || deferred > entry.sync_window; + if (need_wait) { + while (!result_ready(queue, entry.reservation)) {} + ++available; + --deferred; + } else { + if (!result_ready(queue, entry.reservation)) { break; } + } + + if (!consume_cross_row(remaining)) { break; } + } + } + state->total_children = candidate_capacity(search_width) - remaining; + state->previous_parent_position = + parent_indices[0] == invalid_index ? 0u : static_cast(parent_indices[0]); + } else if (!overlap_inner && threadIdx.x >= 32) { + std::uint32_t inner_count = 0; + for (std::uint32_t p = 0; p < search_width; ++p) { + inner_count += state->parent_inner_counts[p]; + } + for (std::uint32_t i = threadIdx.x - 32; i < inner_count; i += blockDim.x - 32) { + store_child(decode_inner(i), i); + } + } + __syncthreads(); + + // In the compacted multi-CTA path, include rows that became ready during + // inner-neighbor work. This pass never waits or expands the candidate budget. + // The single-CTA path retains its single readiness pass. + if (result_position != nullptr) { + if (threadIdx.x == 0 && !state->use_seed_cache) { + auto remaining = candidate_capacity(search_width) - state->total_children; + while (!state->pending.empty() && state->ready_cross_count < max_search_width && + remaining > 0) { + if (!result_ready(queue, state->pending.front().reservation)) { break; } + if (!consume_cross_row(remaining)) { break; } + } + state->total_children = candidate_capacity(search_width) - remaining; + } + __syncthreads(); + } + + for (std::uint32_t i = threadIdx.x + state->total_inner_children; i < state->total_children; + i += blockDim.x) { + IndexT child = invalid_index; + auto local = i - state->total_inner_children; + for (std::uint32_t ready = 0; ready < state->ready_cross_count; ++ready) { + auto const count = state->ready_cross_edge_counts[ready]; + if (local < count) { + auto const decoded = state->ready_cross_data[ready][local]; + if (decoded < graph_size) { child = static_cast(decoded); } + break; + } + local -= count; + } + store_child(child, i); + } + __syncthreads(); + + if (threadIdx.x == 0) { + if (!state->use_seed_cache) { + auto release_count = state->ready_cross_count; + if (state->pending.consumed != 0 && release_count > 0) { --release_count; } + for (std::uint32_t i = 0; i < release_count; ++i) { + release_deferred(queue, state->ready_reservations[i]); + } + } + ++state->iteration; + } + // PQ uses only the copied child IDs. The caller joins the CTA after expansion, + // before any thread can reuse iteration state or consume another queue row. + + auto const team_size_bits = dataset_desc->team_size_bitshift_from_smem(); + auto const offset = + result_position == nullptr ? 0u : static_cast(result_position[0]); + auto const max_i = raft::round_up_safe( + state->total_children, cuvs::neighbors::cagra::detail::device::warp_size >> team_size_bits); + auto const args = dataset_desc->args.load(); + auto const lead_lane = (threadIdx.x & ((1u << team_size_bits) - 1u)) == 0; + for (std::uint32_t i = threadIdx.x >> team_size_bits; i < max_i; + i += blockDim.x >> team_size_bits) { + auto const position = i + offset; + auto const valid = i < state->total_children && + (result_position == nullptr || position < max_result_position) && + result_indices[position] != invalid_index; + auto const distance = team_sum( + valid + ? cuvs::neighbors::cagra::detail::compute_distance_per_thread( + args, result_indices[position]) + : (lead_lane ? raft::upper_bound() : DistanceT{}), + team_size_bits); + __syncwarp(); + if (i < state->total_children && + (result_position == nullptr || position < max_result_position) && lead_lane) { + result_distances[position] = distance; + } + } + return state->total_children; + } + + RAFT_DEVICE_INLINE_FUNCTION auto expand( + IndexT* result_indices, + DistanceT* result_distances, + const cuvs::neighbors::cagra::detail::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 + { + return expand(result_indices, + result_distances, + dataset_desc, + visited_hashmap, + visited_hash_bitlen, + parent_indices, + internal_topk, + 1, + traversed_hashmap, + traversed_hash_bitlen, + result_position, + max_result_position); + } + + RAFT_DEVICE_INLINE_FUNCTION void finalize() const + { + if (threadIdx.x == 0) { + while (!state->pending.empty()) { + auto const reservation = state->pending.front().reservation; + while (!result_ready(queue, reservation)) {} + release_deferred(queue, reservation); + state->pending.pop(); + } + state->pending.consumed = 0; + } + __syncthreads(); + if (threadIdx.x < 32) { reclaim_released_warp(queue); } + __syncthreads(); + } +}; + +template +__device__ void search_kernel_jit( + uintptr_t result_indices_ptr, + DistanceT* result_distances_ptr, + std::uint32_t top_k, + const DataT* queries_ptr, + std::uint32_t query_stride, + const std::uint8_t* inner_graph, + std::uint32_t inner_row_length, + std::uint32_t node_per_cacheline, + std::uint32_t n_bits, + const std::uint32_t* subgraph_offsets, + const std::uint32_t* subgraph_ids, + queue_view* queues, + std::uint32_t num_queues, + std::uint32_t graph_degree, + std::uint32_t cross_graph_degree, + const SourceIndexT* source_indices, + unsigned num_distillation, + std::uint64_t rand_xor_mask, + const IndexT* seeds, + std::uint32_t num_seeds, + const std::uint32_t* seed_cross_graph, + const IndexT* seed_lookup_keys, + const std::uint32_t* seed_lookup_values, + std::uint32_t seed_lookup_capacity, + IndexT* visited_hashmap, + std::uint32_t max_candidates, + std::uint32_t max_itopk, + std::uint32_t internal_topk, + std::uint32_t search_width, + std::uint32_t min_iterations, + std::uint32_t max_iterations, + std::uint32_t* num_executed_iterations, + std::uint32_t hash_bitlen, + std::uint32_t small_hash_bitlen, + std::uint32_t small_hash_reset_interval, + std::uint32_t query_id_offset, + const cuvs::neighbors::cagra::detail::dataset_descriptor_base_t* + dataset_desc, + IndexT graph_size, + float sync_window_scale, + std::uint32_t sync_drop_threshold, + cuvs::neighbors::cagra::detail::cagra_sample_filter filter_payload) +{ + __shared__ graph_state state; + auto const query_id = blockIdx.y; + auto const queue_id = query_id % num_queues; + tiered_graph_provider provider{inner_graph, + inner_row_length, + node_per_cacheline, + n_bits, + subgraph_offsets, + subgraph_ids, + seeds, + seed_cross_graph, + seed_lookup_keys, + seed_lookup_values, + seed_lookup_capacity, + num_seeds, + queues[queue_id], + graph_degree, + cross_graph_degree, + graph_size, + sync_window_scale, + sync_drop_threshold, + &state}; + cuvs::neighbors::cagra::detail::single_cta_search::search_core( + result_indices_ptr, + result_distances_ptr, + top_k, + queries_ptr, + provider, + graph_degree, + source_indices, + num_distillation, + rand_xor_mask, + seeds, + num_seeds, + visited_hashmap, + max_candidates, + max_itopk, + internal_topk, + search_width, + min_iterations, + max_iterations, + num_executed_iterations, + hash_bitlen, + small_hash_bitlen, + small_hash_reset_interval, + query_id, + query_id_offset, + dataset_desc, + filter_payload, + graph_size, + query_stride, + true); +} + +} // namespace cuvs::neighbors::cagra::experimental::flowann::detail diff --git a/cpp/src/neighbors/detail/flowann/search_single_cta_kernel.cu.in b/cpp/src/neighbors/detail/flowann/search_single_cta_kernel.cu.in new file mode 100644 index 0000000000..7e0fea9789 --- /dev/null +++ b/cpp/src/neighbors/detail/flowann/search_single_cta_kernel.cu.in @@ -0,0 +1,127 @@ +/* + * 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 + +namespace { + +constexpr bool k_topk_by_bitonic_sort = @topk_by_bitonic_sort@; +constexpr bool k_bitonic_sort_and_merge_multi_warps = @bitonic_sort_and_merge_multi_warps@; + +using data_t = @data_type@; +using index_t = @index_type@; +using distance_t = @distance_type@; +using source_index_t = @source_index_type@; +using dataset_desc_base = + cuvs::neighbors::cagra::detail::dataset_descriptor_base_t; +using filter_payload_t = cuvs::neighbors::cagra::detail::cagra_sample_filter; +using queue_view = cuvs::neighbors::cagra::experimental::flowann::detail::queue_view; + +} // namespace + +namespace cuvs::neighbors::cagra::experimental::flowann::detail { + +extern "C" __global__ __launch_bounds__(1024, 1) void flowann_search_single_cta( + uintptr_t result_indices, + distance_t* distances, + std::uint32_t top_k, + const data_t* queries, + std::uint32_t query_stride, + const std::uint8_t* inner_graph, + std::uint32_t inner_row_length, + std::uint32_t node_per_cacheline, + std::uint32_t n_bits, + const std::uint32_t* subgraph_offsets, + const std::uint32_t* subgraph_ids, + queue_view* queues, + std::uint32_t num_queues, + std::uint32_t graph_degree, + std::uint32_t cross_graph_degree, + const source_index_t* source_indices, + unsigned num_distillation, + std::uint64_t rand_xor_mask, + const index_t* seeds, + std::uint32_t num_seeds, + const std::uint32_t* seed_cross_graph, + const index_t* seed_lookup_keys, + const std::uint32_t* seed_lookup_values, + std::uint32_t seed_lookup_capacity, + index_t* visited_hashmap, + std::uint32_t max_candidates, + std::uint32_t max_itopk, + std::uint32_t internal_topk, + std::uint32_t search_width, + std::uint32_t min_iterations, + std::uint32_t max_iterations, + std::uint32_t* num_executed_iterations, + std::uint32_t hash_bitlen, + std::uint32_t small_hash_bitlen, + std::uint32_t small_hash_reset_interval, + std::uint32_t query_id_offset, + const dataset_desc_base* dataset_desc, + index_t graph_size, + float sync_window_scale, + std::uint32_t sync_drop_threshold, + filter_payload_t filter_payload) +{ + search_kernel_jit(result_indices, + distances, + top_k, + queries, + query_stride, + inner_graph, + inner_row_length, + node_per_cacheline, + n_bits, + subgraph_offsets, + subgraph_ids, + queues, + num_queues, + graph_degree, + cross_graph_degree, + source_indices, + num_distillation, + rand_xor_mask, + seeds, + num_seeds, + seed_cross_graph, + seed_lookup_keys, + seed_lookup_values, + seed_lookup_capacity, + visited_hashmap, + max_candidates, + max_itopk, + internal_topk, + search_width, + min_iterations, + max_iterations, + num_executed_iterations, + hash_bitlen, + small_hash_bitlen, + small_hash_reset_interval, + query_id_offset, + dataset_desc, + graph_size, + sync_window_scale, + sync_drop_threshold, + filter_payload); +} + +static_assert( + std::is_same_v>); + +} // namespace cuvs::neighbors::cagra::experimental::flowann::detail diff --git a/cpp/src/neighbors/detail/flowann/search_single_cta_matrix.json b/cpp/src/neighbors/detail/flowann/search_single_cta_matrix.json new file mode 100644 index 0000000000..ba1334f11d --- /dev/null +++ b/cpp/src/neighbors/detail/flowann/search_single_cta_matrix.json @@ -0,0 +1,21 @@ +[ + { + "_data": [ + {"data_type": "float", "data_abbrev": "f"}, + {"data_type": "__half", "data_abbrev": "h"}, + {"data_type": "uint8_t", "data_abbrev": "u8"}, + {"data_type": "int8_t", "data_abbrev": "i8"} + ], + "_source_index": [{"source_index_type": "uint32_t", "source_index_abbrev": "u32"}], + "_index": [{"index_type": "uint32_t", "index_abbrev": "u32"}], + "_distance": [{"distance_type": "float", "distance_abbrev": "f"}], + "_topk_by_bitonic": [ + {"topk_by_bitonic_sort": "true", "topk_by_bitonic_sort_str": "topk_by_bitonic_sort"}, + {"topk_by_bitonic_sort": "false", "topk_by_bitonic_sort_str": "no_topk_by_bitonic_sort"} + ], + "_bitonic_sort_and_merge_multi_warps": [ + {"bitonic_sort_and_merge_multi_warps": "true", "bitonic_sort_and_merge_multi_warps_str": "bitonic_sort_and_merge_multi_warps"}, + {"bitonic_sort_and_merge_multi_warps": "false", "bitonic_sort_and_merge_multi_warps_str": "no_bitonic_sort_and_merge_multi_warps"} + ] + } +] diff --git a/cpp/src/neighbors/flowann.cu b/cpp/src/neighbors/flowann.cu new file mode 100644 index 0000000000..8abd3b20ab --- /dev/null +++ b/cpp/src/neighbors/flowann.cu @@ -0,0 +1,190 @@ +/* + * 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 "detail/flowann/search_single_cta.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 = cuvs::neighbors::dense_row_major_view(dataset_); + 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{}); + } +} + +namespace detail { + +template +void search_dispatch(raft::resources const& res, + flowann::search_params params, + flowann::index const& index, + search_context& context, + raft::device_matrix_view queries, + raft::device_matrix_view neighbors, + raft::device_matrix_view distances, + SampleFilterT sample_filter) +{ + int current_device{}; + RAFT_CUDA_TRY(cudaGetDevice(¤t_device)); + RAFT_EXPECTS(current_device == search_context_access::device_id(context), + "FlowANN search context must be used on the CUDA device where it was created"); + RAFT_EXPECTS(!params.persistent, + "FlowANN persistent search is not available for the host-serviced graph provider"); + RAFT_EXPECTS(params.algo == cuvs::neighbors::cagra::search_algo::SINGLE_CTA || + params.algo == cuvs::neighbors::cagra::search_algo::AUTO, + "This revision supports SINGLE_CTA or AUTO; multi-CTA follows separately"); + params.algo = cuvs::neighbors::cagra::search_algo::SINGLE_CTA; + search_single_cta(res, params, index, context, queries, neighbors, distances, sample_filter); +} + +template +void search_with_filtering(raft::resources const& res, + search_params const& params, + flowann::index const& index, + search_context& context, + raft::device_matrix_view queries, + raft::device_matrix_view neighbors, + raft::device_matrix_view distances, + SampleFilterT sample_filter) +{ + search_dispatch(res, params, index, context, queries, neighbors, distances, sample_filter); +} + +template +void search_with_base_filter( + raft::resources const& res, + search_params const& params, + flowann::index const& index, + search_context& context, + raft::device_matrix_view queries, + raft::device_matrix_view neighbors, + raft::device_matrix_view distances, + cuvs::neighbors::filtering::base_filter const& sample_filter_ref) +{ + try { + using filter_type = cuvs::neighbors::filtering::none_sample_filter; + auto const& filter = dynamic_cast(sample_filter_ref); + auto adjusted_params = params; + if (adjusted_params.filtering_rate < 0.0f) { adjusted_params.filtering_rate = 0.0f; } + return search_with_filtering( + res, adjusted_params, index, context, queries, neighbors, distances, filter); + } catch (std::bad_cast const&) { + } + + try { + using filter_type = cuvs::neighbors::filtering::bitset_filter; + auto const& filter = dynamic_cast(sample_filter_ref); + auto adjusted_params = params; + if (adjusted_params.filtering_rate < 0.0f) { + auto const num_set_bits = filter.bitset_view_.count(res); + auto const filtering_rate = + static_cast(index.dataset().n_rows() - num_set_bits) / index.dataset().n_rows(); + adjusted_params.filtering_rate = std::clamp(filtering_rate, 0.0f, 0.999f); + } + return search_with_filtering( + res, adjusted_params, index, context, queries, neighbors, distances, filter); + } catch (std::bad_cast const&) { + } + + try { + using filter_type = cuvs::neighbors::filtering::bloom_filter; + auto const& filter = dynamic_cast(sample_filter_ref); + auto adjusted_params = params; + if (adjusted_params.filtering_rate < 0.0f) { + auto const* bloom = static_cast(filter.filter_data); + RAFT_EXPECTS(bloom != nullptr, + "bloom_filter must carry a valid cuvs::core::bloom_filter handle"); + adjusted_params.filtering_rate = bloom->estimate_filtering_rate(); + } + return search_with_filtering( + res, adjusted_params, index, context, queries, neighbors, distances, filter); + } catch (std::bad_cast const&) { + } + + try { + using filter_type = cuvs::neighbors::filtering::udf_filter; + auto const& filter = dynamic_cast(sample_filter_ref); + auto adjusted_params = params; + if (adjusted_params.filtering_rate < 0.0f) { + adjusted_params.filtering_rate = + filter.filtering_rate < 0.0f ? 0.0f : std::clamp(filter.filtering_rate, 0.0f, 0.999f); + } + return search_with_filtering( + res, adjusted_params, index, context, queries, neighbors, distances, filter); + } catch (std::bad_cast const&) { + RAFT_FAIL("Unsupported FlowANN sample filter type"); + } +} + +} // namespace detail + +#define CUVS_FLOWANN_DEFINE_SEARCH(T, DatasetAlias, OutputIdxT) \ + void search(raft::resources const& res, \ + search_params const& params, \ + DatasetAlias const& index, \ + search_context& context, \ + raft::device_matrix_view queries, \ + raft::device_matrix_view neighbors, \ + raft::device_matrix_view distances, \ + cuvs::neighbors::filtering::base_filter const& sample_filter) \ + { \ + detail::search_with_base_filter( \ + res, params, index, context, queries, neighbors, distances, sample_filter); \ + } + +#define CUVS_FLOWANN_DEFINE_SEARCH_FOR_TYPE(T) \ + CUVS_FLOWANN_DEFINE_SEARCH(T, device_padded_index, std::uint32_t) \ + CUVS_FLOWANN_DEFINE_SEARCH(T, device_padded_index, std::int64_t) \ + template CUVS_EXPORT void \ + index>:: \ + compute_dataset_norms_(raft::resources const& res) + +CUVS_FLOWANN_DEFINE_SEARCH_FOR_TYPE(float); +CUVS_FLOWANN_DEFINE_SEARCH_FOR_TYPE(half); +CUVS_FLOWANN_DEFINE_SEARCH_FOR_TYPE(std::int8_t); +CUVS_FLOWANN_DEFINE_SEARCH_FOR_TYPE(std::uint8_t); + +#undef CUVS_FLOWANN_DEFINE_SEARCH_FOR_TYPE +#undef CUVS_FLOWANN_DEFINE_SEARCH + +} // 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..f9fc171601 --- /dev/null +++ b/cpp/src/neighbors/flowann_queue.cpp @@ -0,0 +1,999 @@ +/* + * 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_{}; +}; + +// A host-written flag remains publishable even if CUDA copies fail in a poller. +class mapped_abort_flag { + public: + mapped_abort_flag() + { + RAFT_CUDA_TRY(cudaHostAlloc(&host_, sizeof(std::uint32_t), cudaHostAllocMapped)); + try { + *host_ = 0; + RAFT_CUDA_TRY(cudaHostGetDevicePointer(&device_, host_, 0)); + } catch (...) { + RAFT_CUDA_TRY_NO_THROW(cudaFreeHost(host_)); + throw; + } + } + mapped_abort_flag(mapped_abort_flag const&) = delete; + auto operator=(mapped_abort_flag const&) -> mapped_abort_flag& = delete; + ~mapped_abort_flag() { RAFT_CUDA_TRY_NO_THROW(cudaFreeHost(host_)); } + void abort() noexcept + { + std::atomic_ref(*host_).store(1, std::memory_order_release); + } + auto device_ptr() const noexcept -> std::uint32_t const* { return device_; } + + private: + std::uint32_t* host_{}; + std::uint32_t* device_{}; +}; + +#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. + constexpr std::size_t gpu_page_size = 64u * 1024u; + RAFT_EXPECTS( + layout_.total_size <= std::numeric_limits::max() - (gpu_page_size - 1), + "FlowANN GDRCopy mapping size overflows"); + mapped_size_ = align_up(layout_.total_size, gpu_page_size); + storage_.allocate(mapped_size_, gpu_page_size); +#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); + view_.abort_flag = abort_flag_.device_ptr(); + +#ifdef CUVS_FLOWANN_USE_GDRCOPY + gdrcopy_.emplace(storage_.get(), mapped_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_; } + void abort() noexcept { abort_flag_.abort(); } + auto consume_invalid_requests() noexcept -> std::uint64_t + { + return invalid_requests_.exchange(0, std::memory_order_relaxed); + } + + [[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]); + auto* destination = row_staging_.data() + static_cast(i) * row_length_; + if (row_id < row_count_) { + std::copy_n( + cross_graph_ + static_cast(row_id) * row_length_, row_length_, destination); + } else { + // Complete malformed requests with an empty row so GPU waiters can finish. + std::fill_n(destination, row_length_, 0u); + invalid_requests_.fetch_add(1, std::memory_order_relaxed); + } + 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_; + mapped_abort_flag abort_flag_; + cuda_allocation storage_; + detail::queue_view view_{}; + unsigned long long cpu_tail_{}; + std::atomic polls_{}; + std::atomic empty_polls_{}; + std::atomic processed_commands_{}; + std::atomic invalid_requests_{}; + std::vector empty_commands_; + std::vector command_staging_; + std::vector result_sequence_staging_; + std::vector row_staging_; +#ifdef CUVS_FLOWANN_USE_GDRCOPY + std::size_t mapped_size_{}; + 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 report_failure(std::exception_ptr error) noexcept + { + std::lock_guard lock(pause_mutex_); + for (auto const& queue : queues_) { + queue->abort(); + } + if (poller_error_ == nullptr) { poller_error_ = error; } + pause_cv_.notify_all(); + } + + void throw_if_failed() + { + if (poller_error_ != nullptr) { std::rethrow_exception(poller_error_); } + } + + void throw_if_invalid_requests() const + { + bool invalid = false; + for (auto const& queue : queues_) { + invalid |= queue->consume_invalid_requests() != 0; + } + RAFT_EXPECTS(!invalid, "FlowANN queue received an out-of-range cross-graph row"); + } + + [[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(); + throw_if_invalid_requests(); + } + + void resume() + { + std::unique_lock lock(pause_mutex_); + pause_cv_.wait(lock, + [this] { return pending_pause_requests_ == 0 || poller_error_ != nullptr; }); + throw_if_failed(); + throw_if_invalid_requests(); + 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(); + throw_if_invalid_requests(); + 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); + } + throw_if_failed(); + throw_if_invalid_requests(); + } + + 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(); + throw_if_invalid_requests(); + ++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_)); + if (params_.pin_pollers) { + static std::atomic next_poller{0}; + bind_to_physical_cpu(next_poller.fetch_add(1, std::memory_order_relaxed)); + } + 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 (...) { + report_failure(std::current_exception()); + std::lock_guard lock(pause_mutex_); + ++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_shared(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.impl_) +{ + RAFT_EXPECTS(context.impl_ != nullptr, "Cannot start a session on a moved-from search context"); + context_->begin_session(); +} + +search_session::~search_session() { reset(); } + +search_session::search_session(search_session&& other) noexcept + : context_(std::move(other.context_)) +{ +} + +auto search_session::operator=(search_session&& other) noexcept -> search_session& +{ + if (this != &other) { + reset(); + context_ = std::move(other.context_); + } + return *this; +} + +void search_session::reset() noexcept +{ + if (context_ == nullptr) { return; } + auto context = std::move(context_); + try { + context->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(); +} + +void detail::search_context_access::report_failure(search_context& context, + std::exception_ptr error) +{ + RAFT_EXPECTS(context.impl_ != nullptr, "Cannot fail a moved-from FlowANN search context"); + context.impl_->report_failure(std::move(error)); +} + +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 f132bb21a3..49425c64bd 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -267,6 +267,35 @@ 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 + neighbors/ann_cagra/flowann_search.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..e8991c79c0 --- /dev/null +++ b/cpp/tests/neighbors/ann_cagra/flowann_queue.cu @@ -0,0 +1,839 @@ +/* + * 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 + +TEST(FlowannQueue, SessionSurvivesContextMove) +{ + auto graph = raft::make_host_matrix(2, 2); + graph(0, 0) = graph(1, 0) = 1; + graph(0, 1) = 1; + graph(1, 1) = 0; + flowann::search_context context(raft::make_const_mdspan(graph.view())); + auto session = std::make_unique(context); + flowann::search_context moved(std::move(context)); + session.reset(); + EXPECT_NO_THROW(moved.pause()); + EXPECT_TRUE(flowann::detail::search_context_access::is_paused(moved)); + + flowann::search_context destination(raft::make_const_mdspan(graph.view())); + session = std::make_unique(moved); + destination = std::move(moved); + session.reset(); + EXPECT_NO_THROW(destination.pause()); + EXPECT_TRUE(flowann::detail::search_context_access::is_paused(destination)); +} + +RAFT_KERNEL invalid_queue_request(flowann::detail::queue_view* queues, std::uint32_t* output) +{ + if (threadIdx.x != 0) { return; } + auto queue = queues[0]; + flowann::detail::queue_reservation reservation{}; + if (!flowann::detail::try_reserve(queue, &reservation)) { + output[0] = 0; + return; + } + flowann::detail::submit(queue, reservation, 0xffffffffu); + auto const start = clock64(); + while (!flowann::detail::result_ready(queue, reservation) && clock64() - start < 100000000ull) {} + auto const ready = flowann::detail::result_ready(queue, reservation); + output[0] = ready; + output[1] = ready ? flowann::detail::result_row(queue, reservation)[0] : 0xffffffffu; + if (ready) { flowann::detail::release(queue, reservation); } +} + +TEST(FlowannQueue, InvalidRequestsCompleteAndReportAnError) +{ + raft::resources res; + auto graph = raft::make_host_matrix(2, 2); + graph(0, 0) = graph(1, 0) = 1; + graph(0, 1) = 1; + graph(1, 1) = 0; + flowann::search_context context(raft::make_const_mdspan(graph.view())); + auto output = raft::make_device_vector(res, 2); + context.resume(); + invalid_queue_request<<<1, 1, 0, raft::resource::get_cuda_stream(res).get()>>>( + flowann::detail::search_context_access::device_queues(context), output.data_handle()); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + raft::resource::sync_stream(res); + std::uint32_t observed[2]{}; + raft::copy(observed, output.data_handle(), 2, raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + EXPECT_EQ(observed[0], 1); + EXPECT_EQ(observed[1], 0); + EXPECT_ANY_THROW(context.pause()); + // Reporting consumes the malformed-request count. Reuse this same queue/context. + EXPECT_NO_THROW(context.resume()); + EXPECT_NO_THROW(context.pause()); + EXPECT_NO_THROW(flowann::detail::search_context_access::acquire_search(context)); + EXPECT_NO_THROW(flowann::detail::search_context_access::release_search(context)); +} + +TEST(FlowannQueue, EveryLeaseReportsPollerFailure) +{ + auto graph = raft::make_host_matrix(4, 1); + std::fill_n(graph.data_handle(), graph.size(), 0u); + flowann::search_context context(raft::make_const_mdspan(graph.view())); + flowann::detail::search_context_access::acquire_search(context); + flowann::detail::search_context_access::acquire_search(context); + flowann::detail::search_context_access::report_failure( + context, std::make_exception_ptr(std::runtime_error("injected poller failure"))); + EXPECT_THROW(flowann::detail::search_context_access::release_search(context), std::runtime_error); + EXPECT_THROW(flowann::detail::search_context_access::release_search(context), std::runtime_error); +} + +namespace { +RAFT_KERNEL wait_for_host_abort(flowann::detail::queue_view queue, std::uint32_t* completed) +{ + if (threadIdx.x != 0) { return; } + flowann::detail::queue_reservation reservation{0, 0}; + auto const start = clock64(); + bool ready = false; + while (!(ready = flowann::detail::result_ready(queue, reservation)) && + clock64() - start < 100000000ull) {} + // One queue view must drain every outstanding reservation after the first abort read. + for (std::uint32_t i = 1; i < 8; ++i) { + ready = flowann::detail::result_ready(queue, {i, i}) && ready; + } + *completed = ready; +} +} // namespace + +TEST(FlowannQueue, HostAbortCompletesAWaitWithoutCudaCopies) +{ + raft::resources res; + auto graph = raft::make_host_matrix(4, 1); + std::fill_n(graph.data_handle(), graph.size(), 0u); + flowann::search_context context(raft::make_const_mdspan(graph.view())); + flowann::detail::queue_view queue{}; + RAFT_CUDA_TRY(cudaMemcpy(&queue, + flowann::detail::search_context_access::device_queues(context), + sizeof(queue), + cudaMemcpyDeviceToHost)); + cudaPointerAttributes attributes{}; + RAFT_CUDA_TRY(cudaPointerGetAttributes(&attributes, queue.abort_flag)); + auto* host_flag = static_cast(attributes.hostPointer); + ASSERT_NE(host_flag, nullptr); + auto completed = raft::make_device_vector(res, 1); + std::atomic_ref(*host_flag).store(1, std::memory_order_release); + wait_for_host_abort<<<1, 1, 0, raft::resource::get_cuda_stream(res).get()>>>( + queue, completed.data_handle()); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + raft::resource::sync_stream(res); + std::uint32_t actual{}; + raft::copy(&actual, completed.data_handle(), 1, raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + EXPECT_EQ(actual, 1u); +} diff --git a/cpp/tests/neighbors/ann_cagra/flowann_search.cu b/cpp/tests/neighbors/ann_cagra/flowann_search.cu new file mode 100644 index 0000000000..80edd9557e --- /dev/null +++ b/cpp/tests/neighbors/ann_cagra/flowann_search.cu @@ -0,0 +1,652 @@ +/* + * 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 +#include + +namespace flowann = cuvs::neighbors::cagra::experimental::flowann; + +namespace { + +void pack_two_nodes(raft::host_matrix& inner, + int64_t row, + std::uint8_t first_neighbor, + std::uint8_t second_neighbor, + std::uint32_t packed_bits) +{ + inner(row, 0) = 1; + inner(row, 1) = 1; + for (int64_t col = 2; col < inner.extent(1); ++col) { + inner(row, col) = 0; + } + for (std::uint32_t node = 0; node < 2; ++node) { + auto const value = node == 0 ? first_neighbor : second_neighbor; + for (std::uint32_t bit = 0; bit < packed_bits; ++bit) { + auto const position = node * packed_bits + bit; + inner(row, 2 + position / 8) |= ((std::uint32_t{value} >> bit) & 1u) << (position % 8); + } + } +} + +struct search_result { + raft::host_matrix neighbors; + raft::host_matrix distances; + flowann::queue_statistics queue_statistics; + bool pollers_paused; +}; + +auto run_dense_search(cuvs::neighbors::cagra::search_algo algo, + std::vector const& query_values, + std::vector const& seeds = {}, + std::uint32_t max_iterations = 16, + bool keep_pollers_running = false, + std::uint32_t max_batch_queries = 0, + std::uint32_t itopk_size = 32, + std::uint32_t search_width = 1, + std::uint32_t num_queues = 2, + bool alternate_modes = false, + std::uint32_t packed_bits = 2, + std::uint32_t row_padding = 0, + std::uint32_t query_dim = 4, + bool use_tail_feature = false) -> search_result +{ + raft::resources res; + constexpr int64_t rows = 8; + auto const dim = static_cast(query_dim); + + auto host_dataset = raft::make_host_matrix(rows, dim); + for (int64_t row = 0; row < rows; ++row) { + auto const group = row < 4 ? 0.0f : 10.0f; + host_dataset(row, 0) = group + static_cast(row % 4); + for (int64_t col = 1; col < dim; ++col) { + host_dataset(row, col) = + use_tail_feature && col == dim - 1 ? static_cast(row % 4) : 0.0f; + } + } + auto device_dataset = raft::make_device_matrix(res, rows, dim); + raft::copy(device_dataset.data_handle(), + host_dataset.data_handle(), + host_dataset.size(), + raft::resource::get_cuda_stream(res)); + std::unique_ptr> dataset_owner; + auto dataset = [&]() -> cuvs::neighbors::device_padded_dataset_view { + if (dim % 4 == 0) { + return cuvs::neighbors::make_device_padded_dataset_view(res, device_dataset.view()); + } + dataset_owner = cuvs::neighbors::make_device_padded_dataset(res, device_dataset.view()); + return dataset_owner->view(); + }(); + + flowann::device_padded_index index( + res, cuvs::distance::DistanceType::L2Expanded, dataset, 2, packed_bits); + auto inner = + raft::make_host_matrix(4, 2 + (2 * packed_bits + 7) / 8 + row_padding); + pack_two_nodes(inner, 0, 1, 2, packed_bits); + pack_two_nodes(inner, 1, 3, 0, packed_bits); + pack_two_nodes(inner, 2, 1, 2, packed_bits); + pack_two_nodes(inner, 3, 3, 0, packed_bits); + + auto cross = raft::make_host_matrix(rows, 2); + for (std::uint32_t row = 0; row < rows; ++row) { + cross(row, 0) = 1; + cross(row, 1) = row < 4 ? row + 4 : row - 4; + } + index.update_graph( + res, raft::make_const_mdspan(inner.view()), raft::make_const_mdspan(cross.view()), 2, 1); + + auto offsets = raft::make_host_vector(2); + offsets(0) = 0; + offsets(1) = 4; + auto subgraph_ids = raft::make_host_vector(rows); + auto source_ids = raft::make_host_vector(rows); + for (std::uint32_t row = 0; row < rows; ++row) { + subgraph_ids(row) = row / 4; + source_ids(row) = rows - 1 - row; + } + index.update_subgraph_layout( + res, + raft::make_host_vector_view(offsets.data_handle(), 2), + raft::make_host_vector_view(subgraph_ids.data_handle(), rows)); + index.update_source_indices( + res, raft::make_host_vector_view(source_ids.data_handle(), rows)); + + if (!seeds.empty()) { + index.update_seeds(res, + raft::make_host_vector_view( + seeds.data(), static_cast(seeds.size()))); + } + + auto const query_count = static_cast(query_values.size()); + auto host_queries = raft::make_host_matrix(query_count, dim); + for (int64_t row = 0; row < query_count; ++row) { + host_queries(row, 0) = query_values[row]; + for (int64_t col = 1; col < dim; ++col) { + host_queries(row, col) = use_tail_feature && col == dim - 1 ? 0.5f : 0.0f; + } + } + auto queries = raft::make_device_matrix(res, query_count, dim); + raft::copy(queries.data_handle(), + host_queries.data_handle(), + host_queries.size(), + raft::resource::get_cuda_stream(res)); + auto neighbors = raft::make_device_matrix(res, query_count, 2); + auto distances = raft::make_device_matrix(res, query_count, 2); + + flowann::queue_params queue_params{}; + queue_params.num_queues = num_queues; + queue_params.collect_statistics = true; + flowann::search_context context(index.cross_graph(), queue_params); + flowann::search_params params{}; + params.algo = algo; + params.itopk_size = itopk_size; + params.search_width = search_width; + params.max_iterations = max_iterations; + params.max_queries = max_batch_queries == 0 ? query_count : max_batch_queries; + params.num_seeds = static_cast(seeds.size()); + std::optional session; + if (keep_pollers_running) { session.emplace(context); } + if (alternate_modes) { + for (auto alternate_algo : {cuvs::neighbors::cagra::search_algo::SINGLE_CTA}) { + params.algo = alternate_algo; + for (auto alternate_width : {1u, 2u}) { + params.search_width = alternate_width; + for (int64_t count : {int64_t{1}, query_count}) { + flowann::search( + res, + params, + index, + context, + raft::make_device_matrix_view(queries.data_handle(), count, dim), + raft::make_device_matrix_view( + neighbors.data_handle(), count, 2), + raft::make_device_matrix_view(distances.data_handle(), count, 2)); + auto observed_neighbors = raft::make_host_matrix(count, 2); + auto observed_distances = raft::make_host_matrix(count, 2); + raft::copy(observed_neighbors.data_handle(), + neighbors.data_handle(), + observed_neighbors.size(), + raft::resource::get_cuda_stream(res)); + raft::copy(observed_distances.data_handle(), + distances.data_handle(), + observed_distances.size(), + raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + for (int64_t query = 0; query < count; ++query) { + int64_t nearest = 0; + for (int64_t row = 1; row < rows; ++row) { + if (std::abs(query_values[query] - host_dataset(row, 0)) < + std::abs(query_values[query] - host_dataset(nearest, 0))) { + nearest = row; + } + } + auto const delta = query_values[query] - host_dataset(nearest, 0); + EXPECT_EQ(observed_neighbors(query, 0), source_ids(nearest)); + EXPECT_NEAR(observed_distances(query, 0), delta * delta, 1e-4f); + } + } + } + } + params.algo = algo; + params.search_width = search_width; + } + flowann::search(res, + params, + index, + context, + raft::make_const_mdspan(queries.view()), + neighbors.view(), + distances.view()); + raft::resource::sync_stream(res); + auto const pollers_paused = flowann::detail::search_context_access::is_paused(context); + session.reset(); + + auto host_neighbors = raft::make_host_matrix(query_count, 2); + auto host_distances = raft::make_host_matrix(query_count, 2); + raft::copy(host_neighbors.data_handle(), + neighbors.data_handle(), + neighbors.size(), + raft::resource::get_cuda_stream(res)); + raft::copy(host_distances.data_handle(), + distances.data_handle(), + distances.size(), + raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + + return { + std::move(host_neighbors), std::move(host_distances), context.statistics(), pollers_paused}; +} + +class FlowannSearchModesTest + : public ::testing::TestWithParam {}; + +TEST_P(FlowannSearchModesTest, TraversesInnerAndCrossEdgesAndMapsSourceIds) +{ + auto result = run_dense_search(GetParam(), {0.1f, 12.2f}); + + EXPECT_EQ(result.neighbors(0, 0), 7); + EXPECT_EQ(result.neighbors(1, 0), 1); + EXPECT_NEAR(result.distances(0, 0), 0.01f, 1e-4f); + EXPECT_NEAR(result.distances(1, 0), 0.04f, 1e-4f); + EXPECT_GT(result.queue_statistics.processed_commands, 0); +} + +TEST_P(FlowannSearchModesTest, DecodesPackedRowsAcrossByteBoundaries) +{ + for (std::uint32_t bits : {2u, 3u, 7u, 8u, 15u, 24u, 31u, 32u}) { + SCOPED_TRACE(bits); + auto result = + run_dense_search(GetParam(), {0.1f, 12.2f}, {}, 16, true, 2, 256, 8, 1, false, bits); + EXPECT_EQ(result.neighbors(0, 0), 7u); + EXPECT_EQ(result.neighbors(1, 0), 1u); + EXPECT_NEAR(result.distances(0, 0), 0.01f, 1e-4f); + EXPECT_NEAR(result.distances(1, 0), 0.04f, 1e-4f); + EXPECT_GT(result.queue_statistics.processed_commands, 0u); + } +} + +TEST_P(FlowannSearchModesTest, Decodes24BitRowsWithEveryRowAlignment) +{ + for (std::uint32_t padding : {0u, 1u, 2u, 3u}) { + for (std::uint32_t width : {1u, 8u}) { + SCOPED_TRACE(padding); + SCOPED_TRACE(width); + auto result = run_dense_search( + GetParam(), {0.1f, 12.2f}, {}, 16, true, 2, 256, width, 1, false, 24, padding); + EXPECT_EQ(result.neighbors(0, 0), 7u); + EXPECT_EQ(result.neighbors(1, 0), 1u); + EXPECT_NEAR(result.distances(0, 0), 0.01f, 1e-4f); + EXPECT_NEAR(result.distances(1, 0), 0.04f, 1e-4f); + EXPECT_GT(result.queue_statistics.processed_commands, 0u); + } + } +} + +TEST_P(FlowannSearchModesTest, UsesFixedSeedCrossEdgeCacheOnFirstExpansion) +{ + auto result = run_dense_search(GetParam(), {10.1f}, {0}, 2); + + EXPECT_EQ(result.neighbors(0, 0), 3); + EXPECT_NEAR(result.distances(0, 0), 0.01f, 1e-4f); + EXPECT_EQ(result.queue_statistics.processed_commands, 0); +} + +TEST_P(FlowannSearchModesTest, PreservesSearchSessionPollers) +{ + auto result = run_dense_search(GetParam(), {0.1f, 12.2f}, {}, 16, true); + + EXPECT_FALSE(result.pollers_paused); +} + +TEST(FlowannSearch, RejectsContextFromDifferentCudaDevice) +{ + if (raft::device_setter::get_device_count() < 2) { + GTEST_SKIP() << "Requires at least two CUDA devices"; + } + + raft::device_setter device_zero(0); + raft::resources res; + flowann::device_padded_index index(res); + auto queries = raft::make_device_matrix(res, 1, 1); + auto neighbors = raft::make_device_matrix(res, 1, 1); + auto distances = raft::make_device_matrix(res, 1, 1); + auto cross = raft::make_host_matrix(1, 1); + cross(0, 0) = 0; + + raft::device_setter device_one(1); + flowann::search_context context(raft::make_const_mdspan(cross.view())); + raft::device_setter restore_device_zero(0); + try { + flowann::search(res, + flowann::search_params{}, + index, + context, + raft::make_const_mdspan(queries.view()), + neighbors.view(), + distances.view()); + FAIL() << "Expected a mismatched-device FlowANN context to be rejected"; + } catch (raft::exception const& error) { + EXPECT_NE(std::string{error.what()}.find( + "FlowANN search context must be used on the CUDA device where it was created"), + std::string::npos); + } +} + +TEST_P(FlowannSearchModesTest, UniversalProviderPreservesSharedQueueBatchesAndWidths) +{ + std::vector const all_queries = {0.1f, 12.2f, 1.1f, 11.2f, 2.1f, 10.2f, 3.1f, 13.2f}; + std::vector const expected_sources = {7, 1, 6, 2, 5, 3, 4, 0}; + std::vector const expected_distances = { + 0.01f, 0.04f, 0.01f, 0.04f, 0.01f, 0.04f, 0.01f, 0.04f}; + for (std::uint32_t batch : {1u, 2u, 8u}) { + auto queries = std::vector(all_queries.begin(), all_queries.begin() + batch); + for (std::uint32_t width : {1u, 2u, 4u}) { + auto result = run_dense_search(GetParam(), queries, {}, 16, true, batch, 256, width, 1); + for (std::uint32_t query = 0; query < batch; ++query) { + EXPECT_EQ(result.neighbors(query, 0), expected_sources[query]); + EXPECT_NEAR(result.distances(query, 0), expected_distances[query], 1e-4f); + } + EXPECT_GT(result.queue_statistics.processed_commands, 0); + EXPECT_FALSE(result.pollers_paused); + } + } +} + +INSTANTIATE_TEST_SUITE_P(SingleAndMultiCta, + FlowannSearchModesTest, + ::testing::Values(cuvs::neighbors::cagra::search_algo::SINGLE_CTA)); + +} // namespace + +namespace { +RAFT_KERNEL check_pending_backpressure(flowann::detail::queue_view queue, + std::uint32_t mode, + std::uint32_t* output) +{ + if (threadIdx.x != 0) { return; } + __shared__ flowann::detail::graph_state state; + __shared__ std::uint32_t traversed[32]; + std::uint32_t topk[8]; + std::uint32_t positions[8]; + for (std::uint32_t i = 0; i < 32; ++i) { + traversed[i] = ~std::uint32_t{0}; + } + state.pending.initialize(); + // Mode 0 has one free pending slot; mode 1 has a full device queue; + // mode 2 has capacity in both and must publish the complete batch. + state.pending.head = mode == 0 ? 510 : 0; + queue.state->reservation_head = mode == 1 ? flowann::detail::queue_capacity : 0; + queue.state->gpu_tail = 0; + queue.state->capacity = flowann::detail::queue_capacity; + queue.released_sequence[0] = 0; + for (std::uint32_t i = 0; i < 8; ++i) { + queue.command_words[i] = flowann::detail::empty_command; + topk[i] = (i + 1) | 0x80000000u; + positions[i] = i; + state.parent_ids[i] = i + 1; + state.parent_valid[i] = true; + cuvs::neighbors::cagra::detail::hashmap::insert(traversed, 5, i + 1); + } + flowann::detail::submit_parent_requests(queue, &state, positions, topk, 8, 1, traversed, 5); + output[0] = state.pending.size(); + output[1] = queue.state->reservation_head; + output[2] = 0; + output[3] = 0; + for (std::uint32_t i = 0; i < 8; ++i) { + output[2] += (topk[i] & 0x80000000u) != 0; + output[3] += + cuvs::neighbors::cagra::detail::hashmap::search(traversed, 5, i + 1); + } + output[4] = 0; + for (std::uint32_t i = 0; i < 8; ++i) { + output[4] += queue.command_words[i] == flowann::detail::make_command(i + 1); + } +} +} // namespace + +TEST(FlowannSearch, RetriesParentsWithoutOverwritingPendingRequests) +{ + raft::resources res; + auto queue_state = raft::make_device_vector(res, 1); + auto commands = raft::make_device_vector(res, 8); + auto released = raft::make_device_vector(res, 1); + auto output = raft::make_device_vector(res, 5); + flowann::detail::queue_view queue{}; + queue.state = queue_state.data_handle(); + queue.command_words = commands.data_handle(); + queue.released_sequence = released.data_handle(); + for (std::uint32_t mode = 0; mode < 3; ++mode) { + check_pending_backpressure<<<1, 1, 0, raft::resource::get_cuda_stream(res).get()>>>( + queue, mode, output.data_handle()); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + std::uint32_t actual[5]; + raft::copy(actual, output.data_handle(), 5, raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + EXPECT_EQ(actual[0], mode == 0 ? 510 : mode == 1 ? 0 : 8); + EXPECT_EQ(actual[1], mode == 1 ? flowann::detail::queue_capacity : mode == 2 ? 8 : 0); + EXPECT_EQ(actual[2], mode == 2 ? 8 : 0); + EXPECT_EQ(actual[3], mode == 2 ? 8 : 0); + EXPECT_EQ(actual[4], mode == 2 ? 8 : 0); + } +} + +TEST(FlowannSearch, EmptyQueriesAreANoOp) +{ + EXPECT_NO_THROW(run_dense_search(cuvs::neighbors::cagra::search_algo::SINGLE_CTA, {})); +} + +TEST(FlowannSearch, PaddedSingleCtaQueriesMatchAcrossBatches) +{ + for (std::uint32_t dim : {5u, 33u}) { + auto one = run_dense_search(cuvs::neighbors::cagra::search_algo::SINGLE_CTA, + {0.1f, 12.2f, 1.1f, 10.3f, 2.2f}, + {}, + 16, + false, + 1, + 32, + 1, + 2, + false, + 2, + 0, + dim, + true); + auto batched = run_dense_search(cuvs::neighbors::cagra::search_algo::SINGLE_CTA, + {0.1f, 12.2f, 1.1f, 10.3f, 2.2f}, + {}, + 16, + false, + 3, + 32, + 1, + 2, + false, + 2, + 0, + dim, + true); + std::uint32_t const expected_ids[] = {7, 2, 6, 3, 6}; + float const expected_distances[] = {0.26f, 1.69f, 0.26f, 0.34f, 1.69f}; + for (int64_t row = 0; row < 5; ++row) { + EXPECT_EQ(batched.neighbors(row, 0), expected_ids[row]); + EXPECT_NEAR(batched.distances(row, 0), expected_distances[row], 1e-4f); + for (int64_t rank = 0; rank < 2; ++rank) { + EXPECT_EQ(batched.neighbors(row, rank), one.neighbors(row, rank)); + EXPECT_NEAR(batched.distances(row, rank), one.distances(row, rank), 1e-4f); + } + } + } +} + +TEST_P(FlowannSearchModesTest, InvalidIdsMatchTheOutputWidthWithAndWithoutSourceMaps) +{ + for (bool mapped : {false, true}) { + raft::resources res; + auto host = raft::make_host_matrix(4, 4); + std::fill_n(host.data_handle(), host.size(), 0.0f); + auto dataset = cuvs::neighbors::make_device_padded_dataset(res, host.view()); + flowann::device_padded_index index( + res, cuvs::distance::DistanceType::L2Expanded, dataset->view()); + auto inner = raft::make_host_matrix(0, 0); + auto cross = raft::make_host_matrix(4, 2); + auto ids = raft::make_host_vector(4); + for (std::uint32_t row = 0; row < 4; ++row) { + cross(row, 0) = 1; + cross(row, 1) = (row + 1) % 4; + ids(row) = 3 - row; + } + index.update_graph( + res, raft::make_const_mdspan(inner.view()), raft::make_const_mdspan(cross.view()), 1, 0); + if (mapped) { index.update_source_indices(res, raft::make_const_mdspan(ids.view())); } + auto offsets = raft::make_host_vector(1); + auto groups = raft::make_host_vector(4); + offsets(0) = 0; + std::fill_n(groups.data_handle(), groups.size(), 0u); + index.update_subgraph_layout( + res, raft::make_const_mdspan(offsets.view()), raft::make_const_mdspan(groups.view())); + auto queries = raft::make_device_matrix(res, 2, 4); + auto neighbors = raft::make_device_matrix(res, 2, 2); + auto distances = raft::make_device_matrix(res, 2, 2); + auto bits = raft::make_device_vector(res, 1); + RAFT_CUDA_TRY(cudaMemsetAsync(queries.data_handle(), + 0, + queries.size() * sizeof(float), + raft::resource::get_cuda_stream(res).get())); + RAFT_CUDA_TRY(cudaMemsetAsync( + bits.data_handle(), 0, sizeof(std::uint32_t), raft::resource::get_cuda_stream(res).get())); + cuvs::neighbors::filtering::bitset_filter filter( + cuvs::core::bitset_view(bits.data_handle(), 4)); + flowann::search_context context(index.cross_graph()); + flowann::search_params params; + params.algo = GetParam(); + params.itopk_size = 32; + params.max_iterations = 16; + flowann::search(res, + params, + index, + context, + raft::make_const_mdspan(queries.view()), + neighbors.view(), + distances.view(), + filter); + int64_t actual[4]; + raft::copy(actual, neighbors.data_handle(), 4, raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + for (auto id : actual) { + EXPECT_EQ(id, int64_t{-1}); + } + } +} + +TEST_P(FlowannSearchModesTest, ExplicitBatchSizeBeyondGridLimitIsSplit) +{ + raft::resources res; + auto const count = + static_cast(raft::resource::get_device_properties(res).maxGridSize[1]) + 1; + auto result = run_dense_search( + GetParam(), std::vector(count, 0.1f), {0, 1, 2, 3, 4, 5, 6, 7}, 1, false, count); + for (auto row : {0u, count - 2, count - 1}) { + EXPECT_EQ(result.neighbors(row, 0), 7u); + EXPECT_NEAR(result.distances(row, 0), 0.01f, 1e-4f); + } +} + +TEST(FlowannSearch, AllRequestedSeedsParticipateInInitialSelection) +{ + // 300 seeds exceed the inherited 256-slot small hashmap for itopk=64, degree=16. + // Every query exactly matches a seed, including seeds outside the first 256 slots. + using hash_mode = cuvs::neighbors::cagra::hash_mode; + // Independently exercise shared growth, global growth, AUTO fallback, and the explicit limit. + for (auto const [rows, mode] : {std::pair{int64_t{300}, hash_mode::AUTO}, + std::pair{int64_t{300}, hash_mode::SMALL}, + std::pair{int64_t{2200}, hash_mode::HASH}, + std::pair{int64_t{4200}, hash_mode::AUTO}, + std::pair{int64_t{4200}, hash_mode::SMALL}}) { + raft::resources res; + auto host = raft::make_host_matrix(rows, 4); + std::fill_n(host.data_handle(), host.size(), 0.0f); + for (int64_t row = 0; row < rows; ++row) { + host(row, 0) = static_cast(row); + } + auto dataset = cuvs::neighbors::make_device_padded_dataset(res, host.view()); + flowann::device_padded_index index( + res, cuvs::distance::DistanceType::L2Expanded, dataset->view()); + auto inner = raft::make_host_matrix(0, 0); + auto cross = raft::make_host_matrix(rows, 17); + auto ids = raft::make_host_vector(rows); + auto groups = raft::make_host_vector(rows); + auto offsets = raft::make_host_vector(1); + offsets(0) = 0; + for (std::uint32_t row = 0; row < rows; ++row) { + cross(row, 0) = 16; + for (int64_t col = 1; col < 17; ++col) { + cross(row, col) = (row + col) % rows; + } + ids(row) = row; + groups(row) = 0; + } + index.update_graph( + res, raft::make_const_mdspan(inner.view()), raft::make_const_mdspan(cross.view()), 16, 0); + index.update_subgraph_layout( + res, raft::make_const_mdspan(offsets.view()), raft::make_const_mdspan(groups.view())); + index.update_seeds(res, raft::make_const_mdspan(ids.view())); + auto const query_count = rows == 300 ? rows : int64_t{3}; + auto queries = raft::make_device_matrix(res, query_count, 4); + auto host_queries = raft::make_host_matrix(query_count, 4); + std::fill_n(host_queries.data_handle(), host_queries.size(), 0.0f); + for (int64_t query = 0; query < query_count; ++query) { + host_queries(query, 0) = static_cast(rows == 300 ? query : query * (rows - 1) / 2); + } + raft::copy(queries.data_handle(), + host_queries.data_handle(), + host_queries.size(), + raft::resource::get_cuda_stream(res)); + auto neighbors = raft::make_device_matrix(res, query_count, 1); + auto distances = raft::make_device_matrix(res, query_count, 1); + flowann::search_context context(index.cross_graph()); + flowann::search_params params; + params.algo = cuvs::neighbors::cagra::search_algo::SINGLE_CTA; + params.itopk_size = 64; + params.max_iterations = 1; + params.num_seeds = rows; + params.hashmap_mode = mode; + if (mode == hash_mode::SMALL && rows > 4096) { + EXPECT_ANY_THROW(flowann::search(res, + params, + index, + context, + raft::make_const_mdspan(queries.view()), + neighbors.view(), + distances.view())); + EXPECT_TRUE(flowann::detail::search_context_access::is_paused(context)); + raft::resource::sync_stream(res); + continue; + } + flowann::search(res, + params, + index, + context, + raft::make_const_mdspan(queries.view()), + neighbors.view(), + distances.view()); + auto actual = raft::make_host_matrix(query_count, 1); + auto actual_distances = raft::make_host_matrix(query_count, 1); + raft::copy(actual.data_handle(), + neighbors.data_handle(), + query_count, + raft::resource::get_cuda_stream(res)); + raft::copy(actual_distances.data_handle(), + distances.data_handle(), + query_count, + raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + for (std::uint32_t row = 0; row < query_count; ++row) { + EXPECT_EQ(actual(row, 0), static_cast(host_queries(row, 0))) + << "hash mode " << static_cast(mode) << ", seeds " << rows; + EXPECT_FLOAT_EQ(actual_distances(row, 0), 0.0f); + } + } +} 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..03f418c29a --- /dev/null +++ b/cpp/tests/neighbors/ann_cagra/flowann_storage.cu @@ -0,0 +1,161 @@ +/* + * 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->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); +} + +TEST(FlowannStorage, RejectsStoredSeedsBeforeReplacingGraph) +{ + for (int overload = 0; overload < 4; ++overload) { + raft::resources res; + flowann::device_padded_index index(res); + auto inner = raft::make_host_matrix(0, 0); + auto original = raft::make_host_matrix(4, 2); + for (std::uint32_t row = 0; row < 4; ++row) { + original(row, 0) = 1; + original(row, 1) = (row + 1) % 4; + } + index.update_graph( + res, raft::make_const_mdspan(inner.view()), raft::make_const_mdspan(original.view()), 1, 0); + std::uint32_t seed = 3; + index.update_seeds(res, raft::make_host_vector_view(&seed, 1)); + auto replacement = raft::make_host_matrix(2, 2); + replacement(0, 0) = replacement(1, 0) = 1; + replacement(0, 1) = 1; + replacement(1, 1) = 0; + if (overload == 0) { + EXPECT_ANY_THROW(index.update_graph(res, + raft::make_const_mdspan(inner.view()), + raft::make_const_mdspan(replacement.view()), + 1, + 0)); + } else if (overload == 1) { + EXPECT_ANY_THROW(index.update_graph(res, std::move(inner), std::move(replacement), 1, 0)); + } else { + auto device_inner = raft::make_device_matrix(res, 0, 0); + if (overload == 2) { + EXPECT_ANY_THROW( + index.update_graph(res, std::move(device_inner), std::move(replacement), true)); + } else { + EXPECT_ANY_THROW( + index.update_graph(res, std::move(device_inner), std::move(replacement), 1, 0)); + } + } + ASSERT_EQ(index.size(), 4) << "overload " << overload; + EXPECT_EQ(index.cross_graph()(3, 1), 0); + ASSERT_TRUE(index.seeds().has_value()); + raft::resource::sync_stream(res); + } +} + +TEST(FlowannStorage, RowCountChangeClearsDependentMappings) +{ + raft::resources res; + flowann::device_padded_index index(res); + auto inner = raft::make_host_matrix(0, 0); + auto graph = raft::make_host_matrix(4, 2); + for (std::uint32_t row = 0; row < 4; ++row) { + graph(row, 0) = 1; + graph(row, 1) = (row + 1) % 4; + } + index.update_graph( + res, raft::make_const_mdspan(inner.view()), raft::make_const_mdspan(graph.view()), 1, 0); + std::uint32_t ids[4] = {3, 2, 1, 0}; + std::uint32_t groups[4] = {0, 0, 0, 0}; + std::uint32_t offset = 0; + std::uint32_t seed = 1; + index.update_source_indices(res, + raft::make_host_vector_view(ids, 4)); + index.update_subgraph_layout( + res, + raft::make_host_vector_view(&offset, 1), + raft::make_host_vector_view(groups, 4)); + index.update_seeds(res, raft::make_host_vector_view(&seed, 1)); + auto larger = raft::make_host_matrix(6, 2); + for (std::uint32_t row = 0; row < 6; ++row) { + larger(row, 0) = 1; + larger(row, 1) = (row + 1) % 6; + } + index.update_graph( + res, raft::make_const_mdspan(inner.view()), raft::make_const_mdspan(larger.view()), 1, 0); + EXPECT_EQ(index.size(), 6); + EXPECT_FALSE(index.source_indices().has_value()); + EXPECT_FALSE(index.subgraph_offsets().has_value()); + EXPECT_FALSE(index.subgraph_ids().has_value()); + EXPECT_FALSE(index.seeds().has_value()); + EXPECT_FALSE(index.seed_cross_graph().has_value()); + EXPECT_FALSE(index.seed_lookup_keys().has_value()); + EXPECT_FALSE(index.seed_lookup_values().has_value()); + raft::resource::sync_stream(res); +}