diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 7e4ee04c27..586bc72edb 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -84,7 +84,16 @@ 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_ENABLE_FLOWANN_BUILD "Enable the experimental integrated FlowANN graph builder" OFF) +option(CUVS_FLOWANN_USE_GDRCOPY "Use GDRCopy for FlowANN GPU-to-CPU command queues" ON) +if(CUVS_ENABLE_FLOWANN_SEARCH AND BUILD_CPU_ONLY) + message(FATAL_ERROR "CUVS_ENABLE_FLOWANN_SEARCH requires a CUDA-enabled build") +endif() +if(CUVS_ENABLE_FLOWANN_BUILD AND NOT CUVS_ENABLE_FLOWANN_SEARCH) + message(FATAL_ERROR "CUVS_ENABLE_FLOWANN_BUILD requires CUVS_ENABLE_FLOWANN_SEARCH") +endif() if(BUILD_CPU_ONLY) set(BUILD_SHARED_LIBS OFF) set(BUILD_TESTS OFF) @@ -113,7 +122,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 +135,9 @@ 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: Enable FlowANN integrated build: ${CUVS_ENABLE_FLOWANN_BUILD}") +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 +180,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 +260,10 @@ if(NOT BUILD_CPU_ONLY) "$" "$" ) + target_compile_definitions( + cuvs_cpp_headers INTERFACE $<$:CUVS_ENABLE_FLOWANN_SEARCH> + $<$:CUVS_ENABLE_FLOWANN_BUILD> + ) target_link_libraries(cuvs_cpp_headers INTERFACE raft::raft rmm::rmm) generate_inst_matrix( @@ -959,6 +990,44 @@ 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 + ) + generate_jit_lto_kernels( + jit_lto_files + NAME_FORMAT "flowann_search_multi_cta_data_@data_abbrev@" + MATRIX_JSON_FILE + "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/detail/flowann/search_multi_cta_matrix.json" + KERNEL_INPUT_FILE + "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/detail/flowann/search_multi_cta_kernel.cu.in" + FRAGMENT_TAG_FORMAT + "${flowann_ns}::fragment_tag_search_multi_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@>" + FRAGMENT_TAG_HEADER_FILES + "" + "" + "" + OUTPUT_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/flowann/search_multi_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 +1555,9 @@ 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_build.cu> + $<$:src/neighbors/flowann_queue.cpp> ${cagra_build_inst_files} ${cagra_extend_inst_files} src/neighbors/cagra_optimize.cu @@ -1592,6 +1664,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 +1686,7 @@ if(NOT BUILD_CPU_ONLY) nvidia::cutlass::cutlass rtcx::rtcx ${CUVS_CTK_MATH_DEPENDENCIES} + ${CUVS_FLOWANN_GDRCOPY_DEPENDENCIES} $ $ kvikio::kvikio @@ -1687,10 +1766,26 @@ if(NOT BUILD_CPU_ONLY) $> $> $<$:CUDA::nvtx3> - PRIVATE $ $ - $ CUDA::nvJitLink CUDA::nvrtc rtcx::rtcx kvikio::kvikio + PRIVATE $ + $ + $ + CUDA::nvJitLink + CUDA::nvrtc + rtcx::rtcx + kvikio::kvikio + ${CUVS_FLOWANN_GDRCOPY_DEPENDENCIES} INTERFACE $ ) + if(CUVS_ENABLE_FLOWANN_SEARCH) + target_link_libraries(cuvs INTERFACE $>) + endif() + if(CUVS_ENABLE_FLOWANN_SEARCH AND CUVS_FLOWANN_USE_GDRCOPY) + # Consumers must resolve this runtime dependency even when it is installed outside the + # linker's default search paths. Its headers remain private to the queue implementation. + target_link_libraries( + cuvs INTERFACE $ $ + ) + endif() endif() @@ -1735,6 +1830,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..3dacb8f6d7 --- /dev/null +++ b/cpp/include/cuvs/detail/jit_lto/flowann/flowann_fragments.hpp @@ -0,0 +1,24 @@ +/* + * 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 {}; + +template +struct fragment_tag_search_multi_cta {}; + +} // namespace cuvs::neighbors::cagra::experimental::flowann::detail diff --git a/cpp/src/neighbors/detail/cagra/cagra_build.cuh b/cpp/src/neighbors/detail/cagra/cagra_build.cuh index d4b874dd80..172abdc5ed 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_build.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_build.cuh @@ -1057,6 +1057,56 @@ auto build_cagra_host_graph_from_knn_params(raft::resources const& res, return cagra_graph; } +/** Build only the final optimized host graph from a dense dataset view. */ +template + requires cuvs::neighbors::is_dense_row_major_dataset_view_v +auto build_graph(raft::resources const& res, + index_params const& params, + DatasetViewT const& dataset) -> raft::host_matrix +{ + auto const n_rows = static_cast(dataset.n_rows()); + auto const dim = static_cast(dataset.dim()); + auto intermediate_degree = params.intermediate_graph_degree; + auto graph_degree = params.graph_degree; + common::nvtx::range function_scope( + "cagra::detail::build_graph(%zu, %zu)", intermediate_degree, graph_degree); + check_graph_degree(intermediate_degree, graph_degree, n_rows); + + auto const dataset_extents = + raft::matrix_extent(static_cast(n_rows), static_cast(dim)); + auto knn_build_params = resolve_cagra_default_knn_graph_build_params( + res, params, dataset_extents, intermediate_degree); + validate_cagra_knn_graph_build_constraints(params, knn_build_params); + + if (std::holds_alternative( + knn_build_params)) { + auto device_graph = [&]() { + if constexpr (cuvs::neighbors::is_host_dataset_view_v) { + auto padded = cuvs::neighbors::make_device_padded_dataset(res, dataset.view()); + auto graph = iterative_build_graph(res, params, padded->as_dataset_view()); + raft::resource::sync_stream(res); + return graph; + } else { + return iterative_build_graph(res, params, dataset); + } + }(); + auto host_graph = raft::make_host_matrix(n_rows, graph_degree); + raft::copy(host_graph.data_handle(), + device_graph.data_handle(), + host_graph.size(), + raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + return host_graph; + } + return build_cagra_host_graph_from_knn_params(res, + params, + knn_build_params, + static_cast(n_rows), + intermediate_degree, + graph_degree, + dataset.view()); +} + /** * Build from a host row-major matrix without uploading the full dataset early when IVF-PQ graph * construction can consume host batches directly. Iterative CAGRA search needs the rows on device diff --git a/cpp/src/neighbors/detail/cagra/jit_lto_kernels/search_multi_cta_jit.cuh b/cpp/src/neighbors/detail/cagra/jit_lto_kernels/search_multi_cta_jit.cuh index 0d16a32fc7..ec8f1abcdf 100644 --- a/cpp/src/neighbors/detail/cagra/jit_lto_kernels/search_multi_cta_jit.cuh +++ b/cpp/src/neighbors/detail/cagra/jit_lto_kernels/search_multi_cta_jit.cuh @@ -13,6 +13,7 @@ #include +#include #include #include #include @@ -31,6 +32,88 @@ using cuvs::neighbors::cagra::detail::device::compute_distance_to_child_nodes_ji using cuvs::neighbors::cagra::detail::device::compute_distance_to_random_nodes_jit; using cuvs::neighbors::detail::sample_filter; +template +struct device_graph_provider { + const IndexT* graph; + std::uint32_t degree; + + RAFT_DEVICE_INLINE_FUNCTION void initialize() const {} + + RAFT_DEVICE_INLINE_FUNCTION auto result_buffer_size(std::uint32_t base_size, std::uint32_t) const + -> std::uint32_t + { + return base_size; + } + + RAFT_DEVICE_INLINE_FUNCTION auto seed_pointer(const IndexT* seeds, + std::uint32_t num_seeds, + std::uint32_t query_id) const -> const IndexT* + { + return seeds == nullptr ? nullptr : seeds + static_cast(num_seeds) * query_id; + } + + RAFT_DEVICE_INLINE_FUNCTION auto seed_block_id(std::uint32_t cta_id, + std::uint32_t num_cta_per_query, + std::uint32_t row) const -> std::uint32_t + { + return cta_id + num_cta_per_query * row; + } + + RAFT_DEVICE_INLINE_FUNCTION auto seed_num_blocks(std::uint32_t num_cta_per_query, + std::uint32_t num_queries, + std::uint32_t num_partitions) const + -> std::uint32_t + { + return num_cta_per_query * num_queries * num_partitions; + } + + RAFT_DEVICE_INLINE_FUNCTION auto initial_candidate_count(std::uint32_t candidate_count, + std::uint32_t, + std::uint32_t, + std::uint32_t) const -> std::uint32_t + { + return candidate_count; + } + + RAFT_DEVICE_INLINE_FUNCTION auto initial_distillation_count(std::uint32_t count, + std::uint32_t) const -> std::uint32_t + { + return count; + } + + RAFT_DEVICE_INLINE_FUNCTION auto expand( + IndexT* result_indices, + DistanceT* result_distances, + const dataset_descriptor_base_t* dataset_desc, + IndexT* visited_hashmap, + std::uint32_t visited_hash_bitlen, + IndexT* traversed_hashmap, + std::uint32_t traversed_hash_bitlen, + const IndexT* parent_indices, + const IndexT* internal_topk, + int* result_position, + int max_result_position) const -> std::uint32_t + { + compute_distance_to_child_nodes_jit(result_indices, + result_distances, + dataset_desc, + graph, + degree, + visited_hashmap, + visited_hash_bitlen, + traversed_hashmap, + traversed_hash_bitlen, + parent_indices, + internal_topk, + 1, + result_position, + max_result_position); + return degree; + } + + RAFT_DEVICE_INLINE_FUNCTION void finalize() const {} +}; + // Shared MULTI_CTA search body for both single-partition and multi-partition kernels. // // The multi-partition launch is a strict generalization of the single-partition one: it adds a @@ -49,16 +132,17 @@ template > RAFT_DEVICE_INLINE_FUNCTION void search_core( IndexT* const result_indices_ptr, // [num_queries, num_cta_per_query, itopk_size] DistanceT* const result_distances_ptr, // [num_queries, num_cta_per_query, itopk_size] const dataset_descriptor_base_t* dataset_desc, const DataT* const queries_ptr, // [num_queries, dataset_dim] - const IndexT* const knn_graph, // [dataset_size, graph_degree] + GraphProviderT graph_provider, const uint32_t max_elements, const uint32_t graph_degree, - const uint32_t buffer_graph_degree, + const uint32_t buffer_candidate_count, const SourceIndexT* source_indices_ptr, // [num_queries, search_width] const unsigned num_distilation, const uint64_t rand_xor_mask, @@ -116,7 +200,9 @@ RAFT_DEVICE_INLINE_FUNCTION void search_core( // | | upto 32 | | // +----------------+---------+---------------------------+ // |<--- result_buffer_size_32 --->| - const auto result_buffer_size = itopk_size + buffer_graph_degree; + const auto base_result_buffer_size = itopk_size + buffer_candidate_count; + const auto result_buffer_size = + graph_provider.result_buffer_size(base_result_buffer_size, num_seeds); const auto result_buffer_size_32 = raft::round_up_safe(result_buffer_size, 32); assert(result_buffer_size_32 <= max_elements); @@ -147,20 +233,25 @@ RAFT_DEVICE_INLINE_FUNCTION void search_core( result_distances_buffer[i] = utils::get_max_value(); } hashmap::init(local_visited_hashmap_ptr, visited_hash_bitlen); + graph_provider.initialize(); __syncthreads(); _CLK_REC(clk_init); // compute distance to randomly selecting nodes using JIT version _CLK_START(); - const INDEX_T* const local_seed_ptr = seed_ptr ? seed_ptr + (num_seeds * query_id) : nullptr; - uint32_t block_id = cta_id + (num_cta_per_query * row); - uint32_t num_blocks = num_cta_per_query * num_queries * gridDim.z; + const INDEX_T* const local_seed_ptr = graph_provider.seed_pointer(seed_ptr, num_seeds, query_id); + uint32_t block_id = graph_provider.seed_block_id(cta_id, num_cta_per_query, row); + uint32_t num_blocks = graph_provider.seed_num_blocks(num_cta_per_query, num_queries, gridDim.z); + auto const initial_candidate_count = + graph_provider.initial_candidate_count(graph_degree, num_seeds, block_id, num_blocks); + auto const initial_distillation_count = + graph_provider.initial_distillation_count(num_distilation, num_seeds); compute_distance_to_random_nodes_jit(result_indices_buffer, result_distances_buffer, smem_desc, - graph_degree, - num_distilation, + initial_candidate_count, + initial_distillation_count, rand_xor_mask, local_seed_ptr, num_seeds, @@ -245,30 +336,26 @@ RAFT_DEVICE_INLINE_FUNCTION void search_core( hashmap::insert(local_visited_hashmap_ptr, visited_hash_bitlen, index); } } - // Initialize buffer for compute_distance_to_child_nodes. + // Initialize buffer for graph expansion. if (threadIdx.x == blockDim.x - 1) { result_position[0] = result_buffer_size_32; } __syncthreads(); - // Compute the norms between child nodes and query node using JIT version - compute_distance_to_child_nodes_jit(result_indices_buffer, - result_distances_buffer, - smem_desc, - knn_graph, - graph_degree, - local_visited_hashmap_ptr, - visited_hash_bitlen, - local_traversed_hashmap_ptr, - traversed_hash_bitlen, - parent_indices_buffer, - result_indices_buffer, - 1, - result_position, - result_buffer_size_32); + // Compute the distances between child nodes and the query node. + graph_provider.expand(result_indices_buffer, + result_distances_buffer, + smem_desc, + local_visited_hashmap_ptr, + visited_hash_bitlen, + local_traversed_hashmap_ptr, + traversed_hash_bitlen, + parent_indices_buffer, + result_indices_buffer, + result_position, + result_buffer_size_32); __syncthreads(); - // Check the state of the nodes in the result buffer which were not updated - // by the compute_distance_to_child_nodes above, and if it cannot be used as - // a parent node, it is deactivated. + // Check the state of result-buffer nodes not updated by graph expansion. Deactivate any node + // that cannot be used as a parent. for (uint32_t i = threadIdx.x; i < result_position[0]; i += blockDim.x) { INDEX_T index = result_indices_buffer[i]; if (index == invalid_index || index & index_msb_1_mask) { continue; } @@ -298,6 +385,8 @@ RAFT_DEVICE_INLINE_FUNCTION void search_core( iter++; } + graph_provider.finalize(); + // Filtering - use extern sample_filter function (linked via JIT LTO) for (uint32_t i = threadIdx.x; i < result_buffer_size_32; i += blockDim.x) { INDEX_T index = result_indices_buffer[i]; @@ -414,29 +503,30 @@ __device__ void search_kernel_jit( const uint32_t query_id_offset, // Offset to add to query_id when calling filter cagra_sample_filter filter_payload) { - search_core(result_indices_ptr, - result_distances_ptr, - dataset_desc, - queries_ptr, - knn_graph, - max_elements, - graph_degree, - /*buffer_graph_degree=*/graph_degree, - source_indices_ptr, - num_distilation, - rand_xor_mask, - seed_ptr, - num_seeds, - visited_hash_bitlen, - traversed_hashmap_ptr, - traversed_hash_bitlen, - itopk_size, - min_iteration, - max_iteration, - num_executed_iterations, - graph_size, - query_id_offset, - filter_payload); + search_core( + result_indices_ptr, + result_distances_ptr, + dataset_desc, + queries_ptr, + device_graph_provider{knn_graph, graph_degree}, + max_elements, + graph_degree, + /*buffer_candidate_count=*/graph_degree, + source_indices_ptr, + num_distilation, + rand_xor_mask, + seed_ptr, + num_seeds, + visited_hash_bitlen, + traversed_hashmap_ptr, + traversed_hash_bitlen, + itopk_size, + min_iteration, + max_iteration, + num_executed_iterations, + graph_size, + query_id_offset, + filter_payload); } // Multi-partition variant of search_kernel_jit. Grid is (num_cta_per_query, num_queries, @@ -474,29 +564,30 @@ __device__ void search_multi_cta_mp_jit( cagra_sample_filter filter_payload{}; filter_payload.filter_data = (part.bitset_ptr != nullptr) ? &part_bitset : nullptr; - search_core(result_indices_ptr, - result_distances_ptr, - part.dataset_desc, - queries_ptr, - part.graph, - max_elements, - /*graph_degree=*/part.graph_degree, - /*buffer_graph_degree=*/max_graph_degree, - /*source_indices_ptr=*/nullptr, - num_distilation, - rand_xor_mask, - /*seed_ptr=*/nullptr, - /*num_seeds=*/0, - visited_hash_bitlen, - traversed_hashmap_ptr, - traversed_hash_bitlen, - itopk_size, - min_iteration, - max_iteration, - /*num_executed_iterations=*/nullptr, - /*graph_size=*/0, - query_id_offset, - filter_payload); + search_core( + result_indices_ptr, + result_distances_ptr, + part.dataset_desc, + queries_ptr, + device_graph_provider{part.graph, part.graph_degree}, + max_elements, + /*graph_degree=*/part.graph_degree, + /*buffer_candidate_count=*/max_graph_degree, + /*source_indices_ptr=*/nullptr, + num_distilation, + rand_xor_mask, + /*seed_ptr=*/nullptr, + /*num_seeds=*/0, + visited_hash_bitlen, + traversed_hashmap_ptr, + traversed_hash_bitlen, + itopk_size, + min_iteration, + max_iteration, + /*num_executed_iterations=*/nullptr, + /*graph_size=*/0, + query_id_offset, + filter_payload); } } // namespace cuvs::neighbors::cagra::detail::multi_cta_search diff --git a/cpp/src/neighbors/detail/cagra/jit_lto_kernels/search_single_cta_jit.cuh b/cpp/src/neighbors/detail/cagra/jit_lto_kernels/search_single_cta_jit.cuh index 47f02fb6ef..e40e54e056 100644 --- a/cpp/src/neighbors/detail/cagra/jit_lto_kernels/search_single_cta_jit.cuh +++ b/cpp/src/neighbors/detail/cagra/jit_lto_kernels/search_single_cta_jit.cuh @@ -46,6 +46,74 @@ using cuvs::neighbors::detail::sample_filter; using cuvs::neighbors::cagra::detail::device::compute_distance_to_child_nodes_jit; using cuvs::neighbors::cagra::detail::device::compute_distance_to_random_nodes_jit; +template +struct device_graph_provider { + const IndexT* graph; + std::uint32_t degree; + + RAFT_DEVICE_INLINE_FUNCTION void initialize() const {} + + RAFT_DEVICE_INLINE_FUNCTION auto candidate_capacity(std::uint32_t search_width) const + -> std::uint32_t + { + return search_width * degree; + } + + RAFT_DEVICE_INLINE_FUNCTION auto result_buffer_size(std::uint32_t base_size, std::uint32_t) const + -> std::uint32_t + { + return base_size; + } + + RAFT_DEVICE_INLINE_FUNCTION auto seed_pointer(const IndexT* seeds, + std::uint32_t num_seeds, + std::uint32_t query_id) const -> const IndexT* + { + return seeds == nullptr ? nullptr : seeds + static_cast(num_seeds) * query_id; + } + + RAFT_DEVICE_INLINE_FUNCTION auto initial_candidate_count(std::uint32_t candidate_count, + std::uint32_t, + std::uint32_t, + std::uint32_t) const -> std::uint32_t + { + return candidate_count; + } + + RAFT_DEVICE_INLINE_FUNCTION auto initial_distillation_count(std::uint32_t count, + std::uint32_t) const -> std::uint32_t + { + return count; + } + + RAFT_DEVICE_INLINE_FUNCTION auto expand( + IndexT* result_indices, + DistanceT* result_distances, + const dataset_descriptor_base_t* dataset_desc, + IndexT* visited_hashmap, + std::uint32_t visited_hash_bitlen, + const IndexT* parent_indices, + const IndexT* internal_topk, + std::uint32_t search_width) const -> std::uint32_t + { + compute_distance_to_child_nodes_jit(result_indices, + result_distances, + dataset_desc, + graph, + degree, + visited_hashmap, + visited_hash_bitlen, + static_cast(nullptr), + 0u, + parent_indices, + internal_topk, + search_width); + return search_width * degree; + } + + RAFT_DEVICE_INLINE_FUNCTION void finalize() const {} +}; + // JIT search_core - setup_workspace/compute_distance via function pointers template > RAFT_DEVICE_INLINE_FUNCTION void search_core( uintptr_t result_indices_ptr, DistanceT* const result_distances_ptr, const std::uint32_t top_k, const DataT* const queries_ptr, - const IndexT* const knn_graph, + GraphProviderT graph_provider, const std::uint32_t graph_degree, const SourceIndexT* source_indices_ptr, const unsigned num_distilation, @@ -111,7 +180,10 @@ RAFT_DEVICE_INLINE_FUNCTION void search_core( extern __shared__ uint8_t smem[]; // Layout of result_buffer - const auto result_buffer_size = internal_topk + (search_width * graph_degree); + const auto candidate_capacity = graph_provider.candidate_capacity(search_width); + const auto base_result_buffer_size = internal_topk + candidate_capacity; + const auto result_buffer_size = + graph_provider.result_buffer_size(base_result_buffer_size, num_seeds); const auto result_buffer_size_32 = raft::round_up_safe(result_buffer_size, 32); const auto small_hash_size = hashmap::get_size(small_hash_bitlen); @@ -150,19 +222,24 @@ RAFT_DEVICE_INLINE_FUNCTION void search_core( local_visited_hashmap_ptr = visited_hashmap_ptr + (hashmap::get_size(hash_bitlen) * blockIdx.y); } hashmap::init(local_visited_hashmap_ptr, hash_bitlen, 0); + graph_provider.initialize(); __syncthreads(); _CLK_REC(clk_init); // compute distance to randomly selecting nodes using JIT version _CLK_START(); - const IndexT* const local_seed_ptr = seed_ptr ? seed_ptr + (num_seeds * query_id) : nullptr; + const IndexT* const local_seed_ptr = graph_provider.seed_pointer(seed_ptr, num_seeds, query_id); + auto const initial_candidate_count = + graph_provider.initial_candidate_count(result_buffer_size, num_seeds, 0, 1); + auto const initial_distillation_count = + graph_provider.initial_distillation_count(num_distilation, num_seeds); // Get dataset_size directly from base descriptor IndexT dataset_size = smem_desc->size; compute_distance_to_random_nodes_jit(result_indices_buffer, result_distances_buffer, smem_desc, - result_buffer_size, - num_distilation, + initial_candidate_count, + initial_distillation_count, rand_xor_mask, local_seed_ptr, num_seeds, @@ -173,10 +250,16 @@ RAFT_DEVICE_INLINE_FUNCTION void search_core( 0, 1, graph_size); + for (std::uint32_t i = initial_candidate_count + threadIdx.x; i < result_buffer_size; + i += blockDim.x) { + result_indices_buffer[i] = raft::upper_bound(); + result_distances_buffer[i] = raft::upper_bound(); + } __syncthreads(); _CLK_REC(clk_compute_1st_distance); - std::uint32_t iter = 0; + std::uint32_t iter = 0; + std::uint32_t computed_children_this_iteration = result_buffer_size - internal_topk; while (1) { // sort if constexpr (TOPK_BY_BITONIC_SORT) { @@ -225,7 +308,7 @@ RAFT_DEVICE_INLINE_FUNCTION void search_core( result_distances_buffer + internal_topk, result_indices_buffer + internal_topk, max_candidates, - search_width * graph_degree, + computed_children_this_iteration, topk_ws, (iter == 0)); __syncthreads(); @@ -280,19 +363,15 @@ RAFT_DEVICE_INLINE_FUNCTION void search_core( __syncthreads(); // compute the norms between child nodes and query node using JIT version _CLK_START(); - compute_distance_to_child_nodes_jit( - result_indices_buffer + internal_topk, - result_distances_buffer + internal_topk, - smem_desc, - knn_graph, - graph_degree, - local_visited_hashmap_ptr, - hash_bitlen, - (IndexT*)nullptr, - 0u, - parent_list_buffer, - result_indices_buffer, - search_width); + computed_children_this_iteration = + graph_provider.expand(result_indices_buffer + internal_topk, + result_distances_buffer + internal_topk, + smem_desc, + local_visited_hashmap_ptr, + hash_bitlen, + parent_list_buffer, + result_indices_buffer, + search_width); // Critical: __syncthreads() must be reached by ALL threads // If any thread is stuck in compute_distance_to_child_nodes_jit, this will hang __syncthreads(); @@ -322,11 +401,13 @@ RAFT_DEVICE_INLINE_FUNCTION void search_core( iter++; } + graph_provider.finalize(); + // Post process for filtering - use extern sample_filter function constexpr IndexT index_msb_1_mask = utils::gen_index_msb_1_mask::value; const IndexT invalid_index = utils::get_max_value(); - for (unsigned i = threadIdx.x; i < internal_topk + search_width * graph_degree; i += blockDim.x) { + for (unsigned i = threadIdx.x; i < internal_topk + candidate_capacity; i += blockDim.x) { const auto node_id = result_indices_buffer[i] & ~index_msb_1_mask; if (node_id != (invalid_index & ~index_msb_1_mask) && !sample_filter(query_id + query_id_offset, @@ -387,7 +468,7 @@ RAFT_DEVICE_INLINE_FUNCTION void search_core( result_distances_buffer + internal_topk, result_indices_buffer + internal_topk, max_candidates, - search_width * graph_degree, + computed_children_this_iteration, topk_ws, (iter == 0)); } @@ -418,9 +499,13 @@ RAFT_DEVICE_INLINE_FUNCTION void search_core( if constexpr (TOPK_BY_BITONIC_SORT) { ii = device::swizzling(i); } if (result_distances_ptr != nullptr) { result_distances_ptr[j] = result_distances_buffer[ii]; } - auto internal_index = - result_indices_buffer[ii] & ~index_msb_1_mask; // clear most significant bit - auto source_index = to_source_index(internal_index); + auto const result_index = result_indices_buffer[ii]; + if (result_index == invalid_index) { + write_indices(result_indices_ptr, j, static_cast(invalid_index)); + continue; + } + auto internal_index = result_index & ~index_msb_1_mask; // clear most significant bit + auto source_index = to_source_index(internal_index); write_indices(result_indices_ptr, j, source_index); } if (threadIdx.x == 0 && num_executed_iterations != nullptr) { @@ -495,33 +580,34 @@ __device__ void search_kernel_jit( DataT, IndexT, DistanceT, - SourceIndexT>(result_indices_ptr, - result_distances_ptr, - top_k, - queries_ptr, - knn_graph, - graph_degree, - source_indices_ptr, - num_distilation, - rand_xor_mask, - seed_ptr, - num_seeds, - visited_hashmap_ptr, - max_candidates, - max_itopk, - internal_topk, - search_width, - min_iteration, - max_iteration, - num_executed_iterations, - hash_bitlen, - small_hash_bitlen, - small_hash_reset_interval, - query_id, - query_id_offset, - dataset_desc, - filter_payload, - graph_size); + SourceIndexT>( + result_indices_ptr, + result_distances_ptr, + top_k, + queries_ptr, + device_graph_provider{knn_graph, graph_degree}, + graph_degree, + source_indices_ptr, + num_distilation, + rand_xor_mask, + seed_ptr, + num_seeds, + visited_hashmap_ptr, + max_candidates, + max_itopk, + internal_topk, + search_width, + min_iteration, + max_iteration, + num_executed_iterations, + hash_bitlen, + small_hash_bitlen, + small_hash_reset_interval, + query_id, + query_id_offset, + dataset_desc, + filter_payload, + graph_size); } // JIT persistent device implementation - called from extern "C" __global__ entry in generated .cu @@ -605,33 +691,34 @@ __device__ void search_single_cta_p_impl( DataT, IndexT, DistanceT, - SourceIndexT>(result_indices_ptr, - result_distances_ptr, - top_k, - queries_ptr, - knn_graph, - graph_degree, - source_indices_ptr, - num_distilation, - rand_xor_mask, - seed_ptr, - num_seeds, - visited_hashmap_ptr, - max_candidates, - max_itopk, - internal_topk, - search_width, - min_iteration, - max_iteration, - num_executed_iterations, - hash_bitlen, - small_hash_bitlen, - small_hash_reset_interval, - query_id, - query_id_offset, - dataset_desc, - filter_payload, - graph_size); + SourceIndexT>( + result_indices_ptr, + result_distances_ptr, + top_k, + queries_ptr, + device_graph_provider{knn_graph, graph_degree}, + graph_degree, + source_indices_ptr, + num_distilation, + rand_xor_mask, + seed_ptr, + num_seeds, + visited_hashmap_ptr, + max_candidates, + max_itopk, + internal_topk, + search_width, + min_iteration, + max_iteration, + num_executed_iterations, + hash_bitlen, + small_hash_bitlen, + small_hash_reset_interval, + query_id, + query_id_offset, + dataset_desc, + filter_payload, + graph_size); // make sure all writes are visible even for the host // (e.g. when result buffers are in pinned memory) @@ -709,33 +796,34 @@ __device__ void search_single_cta_mp_impl( DataT, IndexT, DistanceT, - SourceIndexT>(tagged_indices_ptr, - part_result_distances, - top_k, - queries_ptr, - part.graph, - part.graph_degree, - static_cast(nullptr), - num_distilation, - rand_xor_mask, - static_cast(nullptr), - num_seeds, - part_hashmap_ptr, - max_candidates, - max_itopk, - internal_topk, - search_width, - min_iteration, - max_iteration, - num_executed_iterations, - hash_bitlen, - small_hash_bitlen, - small_hash_reset_interval, - query_id, - query_id_offset, - part.dataset_desc, - filter_payload, - /*graph_size=*/IndexT{0}); + SourceIndexT>( + tagged_indices_ptr, + part_result_distances, + top_k, + queries_ptr, + device_graph_provider{part.graph, part.graph_degree}, + part.graph_degree, + static_cast(nullptr), + num_distilation, + rand_xor_mask, + static_cast(nullptr), + num_seeds, + part_hashmap_ptr, + max_candidates, + max_itopk, + internal_topk, + search_width, + min_iteration, + max_iteration, + num_executed_iterations, + hash_bitlen, + small_hash_bitlen, + small_hash_reset_interval, + query_id, + query_id_offset, + part.dataset_desc, + filter_payload, + /*graph_size=*/IndexT{0}); } } // namespace cuvs::neighbors::cagra::detail::single_cta_search diff --git a/cpp/src/neighbors/detail/flowann/README.md b/cpp/src/neighbors/detail/flowann/README.md new file mode 100644 index 0000000000..52394332f5 --- /dev/null +++ b/cpp/src/neighbors/detail/flowann/README.md @@ -0,0 +1,5 @@ +# FlowANN implementation + +FlowANN is described in "Disentangling Graph Dependencies for Efficient Billion-Scale GPU Vector Search" (OSDI 2026). + +These are internal implementation headers. The integration is introduced in separate PRs; only the interfaces present in the current revision are available. diff --git a/cpp/src/neighbors/detail/flowann/build.hpp b/cpp/src/neighbors/detail/flowann/build.hpp new file mode 100644 index 0000000000..8190bdd1ce --- /dev/null +++ b/cpp/src/neighbors/detail/flowann/build.hpp @@ -0,0 +1,101 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * FlowANN is described in "Disentangling Graph Dependencies for Efficient + * Billion-Scale GPU Vector Search" (OSDI 2026). + */ + +#pragma once + +#include + +#include +#include +#include +#include + +namespace cuvs::neighbors::cagra::experimental::flowann::detail { + +struct uniform_layout_plan { + std::uint16_t n_bits; + std::uint32_t resident_degree; + std::uint32_t cross_degree; + std::uint32_t inner_row_bytes; + std::uint64_t inner_graph_bytes; +}; + +inline auto required_id_bits(std::uint64_t n_rows) -> std::uint16_t +{ + RAFT_EXPECTS(n_rows > 0, "FlowANN cannot plan a graph layout for zero rows"); + auto value = n_rows - 1; + auto bits = std::uint16_t{0}; + while (value != 0) { + ++bits; + value >>= 1; + } + return std::max(1, bits); +} + +inline auto packed_row_bytes(std::uint16_t node_per_cacheline, + std::uint16_t n_bits, + std::uint32_t resident_degree) -> std::uint32_t +{ + RAFT_EXPECTS(node_per_cacheline > 0, "FlowANN node_per_cacheline must be positive"); + RAFT_EXPECTS(n_bits > 0 && n_bits <= 32, "FlowANN packed ID width must be in [1, 32]"); + if (resident_degree == 0) { return 0; } + auto const header_bits = static_cast(node_per_cacheline) * 8; + auto const edge_bits = static_cast(node_per_cacheline) * resident_degree * + static_cast(n_bits); + auto const bytes = (header_bits + edge_bits + 7) / 8; + RAFT_EXPECTS(bytes <= std::numeric_limits::max(), + "FlowANN packed row length exceeds uint32"); + return static_cast(bytes); +} + +inline auto plan_uniform_layout(std::uint64_t n_rows, + std::uint32_t graph_degree, + std::uint16_t node_per_cacheline, + std::uint64_t device_graph_budget_bytes) -> uniform_layout_plan +{ + RAFT_EXPECTS(graph_degree > 0 && graph_degree <= std::numeric_limits::max(), + "FlowANN complete graph degree must be in [1, 255]"); + RAFT_EXPECTS(node_per_cacheline > 0, "FlowANN node_per_cacheline must be positive"); + RAFT_EXPECTS(n_rows <= std::numeric_limits::max(), + "FlowANN graph row count must fit uint32 IDs"); + + auto const n_bits = required_id_bits(n_rows); + auto const packed_rows = n_rows / node_per_cacheline + (n_rows % node_per_cacheline != 0); + + for (auto resident = graph_degree; resident > 0; --resident) { + auto const row_bytes = packed_row_bytes(node_per_cacheline, n_bits, resident); + RAFT_EXPECTS(packed_rows <= std::numeric_limits::max() / row_bytes, + "FlowANN packed graph size overflows uint64"); + auto const total_bytes = packed_rows * row_bytes; + if (total_bytes <= device_graph_budget_bytes) { + return {n_bits, resident, graph_degree - resident, row_bytes, total_bytes}; + } + } + + return {n_bits, 0, graph_degree, 0, 0}; +} + +/** Budgeted grouped rows share capacity between nodes; do not cap them at average local degree. */ +inline auto budgeted_grouped_row_bytes(std::uint64_t n_rows, + std::uint32_t graph_degree, + std::uint16_t nodes_per_row, + std::uint16_t n_bits, + std::uint64_t budget_bytes) -> std::uint32_t +{ + RAFT_EXPECTS(n_rows > 0 && nodes_per_row > 0, "FlowANN grouped layout must have rows and nodes"); + RAFT_EXPECTS(graph_degree > 0 && graph_degree <= std::numeric_limits::max(), + "FlowANN grouped degree must be in [1, 255]"); + auto const rows = n_rows / nodes_per_row + (n_rows % nodes_per_row != 0); + auto const bytes = std::min(budget_bytes / rows, + packed_row_bytes(nodes_per_row, n_bits, graph_degree)); + // A header without space for even one edge is equivalent to an empty resident graph. + if (bytes < nodes_per_row + (n_bits + 7u) / 8u) { return 0; } + return static_cast(bytes); +} + +} // namespace cuvs::neighbors::cagra::experimental::flowann::detail diff --git a/cpp/src/neighbors/detail/flowann/flowann.hpp b/cpp/src/neighbors/detail/flowann/flowann.hpp new file mode 100644 index 0000000000..bb72c2d5c1 --- /dev/null +++ b/cpp/src/neighbors/detail/flowann/flowann.hpp @@ -0,0 +1,775 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * FlowANN is described in "Disentangling Graph Dependencies for Efficient + * Billion-Scale GPU Vector Search" (OSDI 2026). + */ + +#pragma once + +#ifdef CUVS_ENABLE_FLOWANN_SEARCH + +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace cuvs::neighbors::cagra::experimental::flowann { + +namespace detail { +struct search_context_access; +} + +class search_session; + +/** Parameters controlling the host-serviced cross-graph queues. */ +struct queue_params { + /** Number of independent GPU-to-CPU command queues. */ + std::uint32_t num_queues = 1; + /** Number of CPU pause instructions issued after an empty poll. */ + std::uint32_t empty_pause = 64; + /** Collect aggregate queue and polling statistics. */ + bool collect_statistics = false; +}; + +/** Aggregate counters collected by the host-serviced cross-graph queues. */ +struct queue_statistics { + std::uint64_t polls{}; + std::uint64_t empty_polls{}; + std::uint64_t processed_commands{}; +}; + +/** 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::unique_ptr impl_; + + friend struct detail::search_context_access; + friend class search_session; +}; + +/** + * Keep a search context's pollers running for the lifetime of this session. + * + * Use a session for repeated low-latency search calls. The search context must outlive the session + * and must not be moved while the session is active. + */ +class CUVS_EXPORT search_session { + public: + explicit search_session(search_context& context); + ~search_session(); + + search_session(search_session const&) = delete; + auto operator=(search_session const&) -> search_session& = delete; + search_session(search_session&& other) noexcept; + auto operator=(search_session&& other) noexcept -> search_session&; + + private: + void reset() noexcept; + + search_context* context_{}; +}; + +/** + * @brief FlowANN tiered-graph index. + * + * The index stores a compressed inner graph in device memory and the cross-graph rows in host + * memory. The dataset is a non-owning view and its backing storage must outlive the index. + */ +template > +class index : public cuvs::neighbors::index { + public: + using value_type = T; + using index_type = IdxT; + using graph_index_type = std::uint32_t; + using inner_index_type = std::uint8_t; + using dataset_view_type = DatasetViewT; + 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)); + + auto new_cross = raft::make_host_matrix(cross_graph.extent(0), + cross_graph.extent(1)); + std::copy_n(cross_graph.data_handle(), cross_graph.size(), new_cross.data_handle()); + + inner_graph_ = std::move(new_inner); + cross_graph_ = std::move(new_cross); + full_graph_degree_ = full_graph_degree; + resident_degree_ = resident_degree; + rebuild_seed_cross_graph_(res); + } + + /** Replace the graph while taking ownership of the host cross-graph allocation. */ + void update_graph(raft::resources const& res, + raft::host_matrix&& inner_graph, + raft::host_matrix&& cross_graph) + { + auto const inferred_degree = + cross_graph.extent(1) == 0 ? 0u : static_cast(cross_graph.extent(1) - 1); + update_graph(res, std::move(inner_graph), std::move(cross_graph), inferred_degree, 0); + } + + /** Replace the graph while taking ownership of a compact uniform-residency layout. */ + void update_graph(raft::resources const& res, + raft::host_matrix&& inner_graph, + raft::host_matrix&& cross_graph, + std::uint32_t full_graph_degree, + std::uint32_t resident_degree) + { + validate_graph_(raft::make_const_mdspan(inner_graph.view()), + raft::make_const_mdspan(cross_graph.view()), + full_graph_degree, + resident_degree); + + auto new_inner = raft::make_device_matrix( + res, inner_graph.extent(0), inner_graph.extent(1)); + raft::copy(new_inner.data_handle(), + inner_graph.data_handle(), + inner_graph.size(), + raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + + inner_graph_ = std::move(new_inner); + cross_graph_ = std::move(cross_graph); + full_graph_degree_ = full_graph_degree; + resident_degree_ = resident_degree; + rebuild_seed_cross_graph_(res); + } + + /** Replace a variable-residency grouped graph while retaining the packed rows on the device. */ + void update_graph(raft::resources const& res, + raft::device_matrix&& inner_graph, + raft::host_matrix&& cross_graph, + bool validate_edges) + { + RAFT_EXPECTS(cross_graph.extent(0) > 0 && cross_graph.extent(1) > 1, + "FlowANN grouped cross graph must contain rows and neighbors"); + auto const full_degree = static_cast(cross_graph.extent(1) - 1); + auto const expected_rows = + raft::div_rounding_up_safe(cross_graph.extent(0), node_per_cacheline_); + auto const has_inner_graph = inner_graph.extent(0) != 0 || inner_graph.extent(1) != 0; + if (has_inner_graph) { + RAFT_EXPECTS(inner_graph.extent(0) == expected_rows, + "FlowANN grouped inner graph has an invalid row count"); + RAFT_EXPECTS(inner_graph.extent(1) >= node_per_cacheline_, + "FlowANN grouped inner graph rows cannot hold their count headers"); + } + if (dataset_.n_rows() > 0) { + RAFT_EXPECTS(dataset_.n_rows() == cross_graph.extent(0), + "Dataset and FlowANN grouped graph must have the same number of rows"); + } + for (int64_t node = 0; node < cross_graph.extent(0); ++node) { + auto const cross_count = cross_graph(node, 0); + RAFT_EXPECTS(cross_count <= full_degree, + "FlowANN grouped cross-neighbor count exceeds graph degree at row %ld", + static_cast(node)); + if (validate_edges) { + for (std::uint32_t edge = 0; edge < cross_count; ++edge) { + RAFT_EXPECTS(cross_graph(node, edge + 1) < cross_graph.extent(0), + "FlowANN grouped cross graph contains an out-of-range node at row %ld", + static_cast(node)); + } + } + } + + inner_graph_ = std::move(inner_graph); + cross_graph_ = std::move(cross_graph); + full_graph_degree_ = full_degree; + resident_degree_ = 0; + rebuild_seed_cross_graph_(res); + } + + /** Replace a compact graph while taking ownership of both device and host allocations. */ + void update_graph(raft::resources const& res, + raft::device_matrix&& inner_graph, + raft::host_matrix&& cross_graph, + std::uint32_t full_graph_degree, + std::uint32_t resident_degree) + { + RAFT_EXPECTS(cross_graph.extent(0) > 0 && cross_graph.extent(1) > 0, + "FlowANN cross graph must contain rows and a neighbor-count column"); + RAFT_EXPECTS(full_graph_degree > 0 && resident_degree <= full_graph_degree, + "FlowANN compact graph degree metadata is invalid"); + auto const cross_degree = static_cast(cross_graph.extent(1) - 1); + RAFT_EXPECTS(resident_degree + cross_degree == full_graph_degree, + "FlowANN compact graph degrees do not add up to the complete graph degree"); + if (resident_degree == 0) { + RAFT_EXPECTS(inner_graph.extent(0) == 0 && inner_graph.extent(1) == 0, + "FlowANN zero-residency graph must not allocate packed rows"); + } else { + auto const expected_rows = + raft::div_rounding_up_safe(cross_graph.extent(0), node_per_cacheline_); + auto const required_bits = + static_cast(node_per_cacheline_) * 8 + + static_cast(node_per_cacheline_) * resident_degree * n_bits_; + RAFT_EXPECTS(inner_graph.extent(0) == expected_rows, + "FlowANN compact inner graph has an invalid row count"); + RAFT_EXPECTS(static_cast(inner_graph.extent(1)) * 8 >= required_bits, + "FlowANN compact inner graph rows are too short"); + } + if (dataset_.n_rows() > 0) { + RAFT_EXPECTS(dataset_.n_rows() == cross_graph.extent(0), + "Dataset and FlowANN compact graph must have the same number of rows"); + } + for (int64_t node = 0; node < cross_graph.extent(0); ++node) { + RAFT_EXPECTS(cross_graph(node, 0) == cross_degree, + "FlowANN uniform compact graph has an invalid cross count at row %ld", + static_cast(node)); + for (std::uint32_t edge = 0; edge < cross_degree; ++edge) { + RAFT_EXPECTS(cross_graph(node, edge + 1) < cross_graph.extent(0), + "FlowANN compact cross graph contains an out-of-range node at row %ld", + static_cast(node)); + } + } + + inner_graph_ = std::move(inner_graph); + cross_graph_ = std::move(cross_graph); + full_graph_degree_ = full_graph_degree; + resident_degree_ = resident_degree; + rebuild_seed_cross_graph_(res); + } + + void update_source_indices( + raft::resources const& res, + raft::host_vector_view source_indices) + { + RAFT_EXPECTS(source_indices.extent(0) == static_cast(size()), + "FlowANN source-index mapping must have one entry per graph row"); + source_indices_.emplace(raft::make_device_vector(res, source_indices.extent(0))); + raft::copy(source_indices_->data_handle(), + source_indices.data_handle(), + source_indices.extent(0), + raft::resource::get_cuda_stream(res)); + } + + void update_seeds(raft::resources const& res, + raft::host_vector_view seeds) + { + RAFT_EXPECTS(seeds.extent(0) > 0, "FlowANN seeds must not be empty"); + RAFT_EXPECTS(seeds.extent(0) <= static_cast(size()), + "FlowANN seed count must not exceed the index size"); + for (int64_t i = 0; i < seeds.extent(0); ++i) { + RAFT_EXPECTS( + seeds(i) < size(), "FlowANN seed at position %ld is out of range", static_cast(i)); + } + seeds_.emplace(raft::make_device_vector(res, seeds.extent(0))); + raft::copy(seeds_->data_handle(), + seeds.data_handle(), + seeds.extent(0), + raft::resource::get_cuda_stream(res)); + seed_ids_host_.assign(seeds.data_handle(), seeds.data_handle() + seeds.extent(0)); + rebuild_seed_lookup_(res); + rebuild_seed_cross_graph_(res); + } + + void update_subgraph_layout( + raft::resources const& res, + raft::host_vector_view offsets, + raft::host_vector_view subgraph_ids) + { + RAFT_EXPECTS(offsets.extent(0) > 0, "FlowANN subgraph offsets must not be empty"); + RAFT_EXPECTS(offsets.extent(0) <= static_cast(size()), + "FlowANN subgraph count must not exceed the index size"); + RAFT_EXPECTS(subgraph_ids.extent(0) == static_cast(size()), + "FlowANN must store one subgraph identifier per graph row"); + RAFT_EXPECTS(offsets(0) == 0, "FlowANN first subgraph offset must be zero"); + for (int64_t i = 1; i < offsets.extent(0); ++i) { + RAFT_EXPECTS(offsets(i - 1) < offsets(i), + "FlowANN subgraph offsets must be strictly increasing"); + } + RAFT_EXPECTS(offsets(offsets.extent(0) - 1) < static_cast(size()), + "FlowANN subgraph offsets exceed the index size"); + for (int64_t i = 0; i < subgraph_ids.extent(0); ++i) { + RAFT_EXPECTS(subgraph_ids(i) < offsets.extent(0), + "FlowANN subgraph identifier is out of range at row %ld", + static_cast(i)); + } + + subgraph_offsets_.emplace( + raft::make_device_vector(res, offsets.extent(0))); + raft::copy(subgraph_offsets_->data_handle(), + offsets.data_handle(), + offsets.extent(0), + raft::resource::get_cuda_stream(res)); + subgraph_ids_.emplace( + raft::make_device_vector(res, subgraph_ids.extent(0))); + raft::copy(subgraph_ids_->data_handle(), + subgraph_ids.data_handle(), + subgraph_ids.extent(0), + raft::resource::get_cuda_stream(res)); + } + + private: + void validate_graph_( + raft::host_matrix_view inner_graph, + raft::host_matrix_view cross_graph, + std::uint32_t full_graph_degree, + std::uint32_t resident_degree) const + { + RAFT_EXPECTS(cross_graph.extent(1) > 0, + "FlowANN cross graph must include a neighbor-count column"); + RAFT_EXPECTS(static_cast(cross_graph.extent(0)) <= + static_cast(std::numeric_limits::max()), + "FlowANN graph row count exceeds the index type"); + auto const expected_inner_rows = + raft::div_rounding_up_safe(cross_graph.extent(0), node_per_cacheline_); + auto const has_inner_graph = inner_graph.extent(0) != 0 || inner_graph.extent(1) != 0; + if (has_inner_graph) { + RAFT_EXPECTS(inner_graph.extent(0) == expected_inner_rows, + "FlowANN inner graph must contain one packed row per node cacheline"); + RAFT_EXPECTS(inner_graph.extent(1) >= node_per_cacheline_, + "FlowANN inner graph rows are too short to contain neighbor counts"); + } else { + RAFT_EXPECTS(resident_degree == 0, + "FlowANN may omit the packed graph only when resident degree is zero"); + } + RAFT_EXPECTS(full_graph_degree > 0, "FlowANN complete graph degree must be positive"); + RAFT_EXPECTS(resident_degree <= full_graph_degree, + "FlowANN resident graph degree exceeds the complete graph degree"); + auto const cross_degree = static_cast(cross_graph.extent(1) - 1); + RAFT_EXPECTS(cross_degree <= full_graph_degree, + "FlowANN cross graph degree exceeds the complete graph degree"); + if (resident_degree > 0 || cross_degree < full_graph_degree) { + RAFT_EXPECTS(resident_degree + cross_degree == full_graph_degree, + "FlowANN compact graph degrees do not add up to the complete graph degree"); + } + if (dataset_.n_rows() > 0) { + RAFT_EXPECTS(dataset_.n_rows() == cross_graph.extent(0), + "Dataset and FlowANN cross graph must have the same number of rows"); + } + + for (int64_t node = 0; node < cross_graph.extent(0); ++node) { + auto const cross_count = cross_graph(node, 0); + RAFT_EXPECTS(cross_count <= cross_degree, + "FlowANN cross-neighbor count exceeds graph degree at row %ld", + static_cast(node)); + for (std::uint32_t edge = 0; edge < cross_count; ++edge) { + RAFT_EXPECTS(cross_graph(node, edge + 1) < cross_graph.extent(0), + "FlowANN cross graph contains an out-of-range node at row %ld", + static_cast(node)); + } + + auto const packed_row = node / node_per_cacheline_; + auto const packed_col = node % node_per_cacheline_; + auto const inner_count = + has_inner_graph ? static_cast(inner_graph(packed_row, packed_col)) : 0u; + RAFT_EXPECTS(inner_count + cross_count <= full_graph_degree, + "FlowANN inner and cross neighbors exceed graph degree at row %ld", + static_cast(node)); + } + for (int64_t packed_row = 0; packed_row < inner_graph.extent(0); ++packed_row) { + std::uint64_t payload_bits = static_cast(node_per_cacheline_) * 8; + for (std::uint16_t packed_col = 0; packed_col < node_per_cacheline_; ++packed_col) { + auto const node = packed_row * node_per_cacheline_ + packed_col; + if (node < cross_graph.extent(0)) { + payload_bits += static_cast(inner_graph(packed_row, packed_col)) * n_bits_; + } + } + RAFT_EXPECTS(payload_bits <= static_cast(inner_graph.extent(1)) * 8, + "FlowANN inner graph payload exceeds its packed row at row %ld", + static_cast(packed_row)); + } + } + + static auto seed_hash_(IdxT key, std::uint32_t capacity) -> std::uint32_t + { + std::uint64_t value = static_cast(key); + value ^= value >> 16; + value *= 0x85ebca77u; + value ^= value >> 13; + value *= 0xc2b2ae3du; + value ^= value >> 16; + return static_cast(value % capacity); + } + + void rebuild_seed_lookup_(raft::resources const& res) + { + if (seed_ids_host_.empty()) { + seed_lookup_keys_.reset(); + seed_lookup_values_.reset(); + return; + } + RAFT_EXPECTS(seed_ids_host_.size() <= + (static_cast(std::numeric_limits::max()) - 1) / 2, + "FlowANN seed lookup exceeds uint32 capacity"); + + auto const capacity = static_cast(seed_ids_host_.size() * 2 + 1); + auto const empty_key = std::numeric_limits::max(); + auto const empty_value = std::numeric_limits::max(); + std::vector host_keys(capacity, empty_key); + std::vector host_values(capacity, empty_value); + + for (std::uint32_t seed = 0; seed < seed_ids_host_.size(); ++seed) { + auto const key = seed_ids_host_[seed]; + auto slot = seed_hash_(key, capacity); + while (host_keys[slot] != empty_key) { + slot = (slot + 1) % capacity; + } + host_keys[slot] = key; + host_values[slot] = seed; + } + + seed_lookup_keys_.emplace(raft::make_device_vector(res, capacity)); + seed_lookup_values_.emplace(raft::make_device_vector(res, capacity)); + auto const stream = raft::resource::get_cuda_stream(res); + raft::copy(seed_lookup_keys_->data_handle(), host_keys.data(), capacity, stream); + raft::copy(seed_lookup_values_->data_handle(), host_values.data(), capacity, stream); + } + + void rebuild_seed_cross_graph_(raft::resources const& res) + { + if (seed_ids_host_.empty()) { + seed_cross_graph_.reset(); + return; + } + RAFT_EXPECTS(cross_graph_.extent(0) > 0 && cross_graph_.extent(1) > 0, + "FlowANN seed caching requires a non-empty cross graph"); + + auto host_cache = raft::make_host_matrix( + static_cast(seed_ids_host_.size()), cross_graph_.extent(1)); + for (std::size_t seed = 0; seed < seed_ids_host_.size(); ++seed) { + auto const node = seed_ids_host_[seed]; + RAFT_EXPECTS(node < size(), "FlowANN seed at position %zu is out of range", seed); + std::copy_n( + cross_graph_.data_handle() + static_cast(node) * cross_graph_.extent(1), + cross_graph_.extent(1), + host_cache.data_handle() + seed * cross_graph_.extent(1)); + } + + seed_cross_graph_.emplace(raft::make_device_matrix( + res, host_cache.extent(0), host_cache.extent(1))); + raft::copy(seed_cross_graph_->data_handle(), + host_cache.data_handle(), + host_cache.size(), + raft::resource::get_cuda_stream(res)); + } + + void validate_encoding_() const + { + RAFT_EXPECTS(node_per_cacheline_ > 0, "FlowANN node_per_cacheline must be positive"); + RAFT_EXPECTS(n_bits_ > 0 && n_bits_ <= 32, "FlowANN n_bits must be in [1, 32]"); + } + + CUVS_EXPORT void compute_dataset_norms_(raft::resources const& res); + + cuvs::distance::DistanceType metric_; + std::uint16_t node_per_cacheline_; + std::uint16_t n_bits_; + std::uint32_t full_graph_degree_ = 0; + std::uint32_t resident_degree_ = 0; + raft::device_matrix inner_graph_; + raft::host_matrix cross_graph_; + DatasetViewT dataset_{}; + std::optional> dataset_norms_; + std::optional> source_indices_; + std::optional> subgraph_offsets_; + std::optional> subgraph_ids_; + std::optional> seeds_; + std::vector seed_ids_host_; + std::optional> seed_lookup_keys_; + std::optional> seed_lookup_values_; + std::optional> seed_cross_graph_; +}; + +template +using device_padded_index = index>; + +/** Owns the dataset backing storage required by a FlowANN index. */ +template +struct index_bundle { + std::unique_ptr dataset; + IndexT index; +}; + +template +using device_padded_index_bundle = index_bundle, + flowann::device_padded_index>; + +/** + * 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/flowann_build.hpp b/cpp/src/neighbors/detail/flowann/flowann_build.hpp new file mode 100644 index 0000000000..880f7c5a63 --- /dev/null +++ b/cpp/src/neighbors/detail/flowann/flowann_build.hpp @@ -0,0 +1,181 @@ +/* + * 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 + +#if defined(CUVS_ENABLE_FLOWANN_SEARCH) && defined(CUVS_ENABLE_FLOWANN_BUILD) + +#include +#include + +#include +#include +#include + +#include +#include +#include +#include + +namespace cuvs::neighbors::cagra::experimental::flowann { + +/** Parameters for converting a CAGRA graph into FlowANN's tiered graph representation. */ +struct build_params { + /** Number of consecutive nodes sharing one packed inner-graph row. */ + std::uint16_t node_per_cacheline = 2; + /** Number of bits used for a node offset inside its partition. */ + std::uint16_t n_bits = 24; + /** Packed row size in bytes. Zero selects the original average-local-degree heuristic. */ + std::uint32_t inner_graph_row_bytes = 0; + /** Verify every encoded inner edge and every cross edge before returning. */ + bool validate = true; +}; + +/** Parameters controlling the locality groups used by the integrated FlowANN build. */ +struct grouping_params { + /** Build and store a grouped graph layout. Disable this to retain the legacy rank-only layout. */ + bool enabled = true; + /** Final group count. Zero uses one group if IDs fit, otherwise ceil(N * (1 + tolerance) / 2^b). + * Each k-means node has at most 16 children; this does not limit the final group count. + */ + std::uint32_t n_groups = 0; + /** Local-ID width. Zero uses min(24, max(4, align_up_4(ceil(log2(N))))). + * Explicit widths in [1, 32] remain available for fixed-layout builds. + */ + std::uint16_t n_bits = 0; + /** Allowed relative deviation from the average group size, in the open interval `(0, 1)`. */ + double balance_tolerance = 0.10; + /** Maximum sampled rows per k-means node. Zero shares a 160,000-row budget by subtree + * leaf count, capped at 10,000 rows per child and floored at 256 rows per child. + * Sampling is always capped at the actual node size. + */ + std::size_t training_rows = 0; + /** Number of rows assigned to groups per GPU batch. Zero targets a 256 MiB vector buffer; + * distance and label buffers are additional. Does not limit host capacity-repair storage. + */ + std::size_t assignment_batch_rows = 0; + /** Number of balanced k-means training iterations. */ + std::uint32_t kmeans_n_iters = 20; + /** Verify every encoded edge after rearrangement. Intended for tests and small builds. */ + bool validate = false; +}; + +/** Parameters for the integrated rank/budget FlowANN build. */ +struct index_params { + /** Parameters used to build the complete CAGRA graph before tiering. */ + cuvs::neighbors::cagra::index_params cagra_params{}; + /** Maximum bytes occupied by the final packed resident graph. Grouped builds derive uniform + * row capacity from this budget, capped at the full graph degree. Zero stores all edges on host. + */ + std::size_t device_graph_budget_bytes = 0; + /** Number of consecutive nodes sharing one packed inner-graph row. */ + std::uint16_t node_per_cacheline = 2; + /** Locality grouping and group-local ID compression parameters. */ + grouping_params grouping{}; + /** Number of medoid seeds generated during build. Zero disables seed generation. */ + std::uint32_t num_seeds = 0; + /** Number of sampled rows used to train seed centroids. Zero selects an automatic limit. */ + std::size_t seed_training_rows = 0; + /** Deterministic seed used to rotate the uniform k-means training sample. */ + std::uint64_t seed = 0x9e3779b97f4a7c15ULL; +}; + +/** Grouping output, reusable when the complete CAGRA graph is already available. */ +struct grouping_result { + /** One group ID in [0, n_groups) per original dataset row; owns its host storage. */ + raft::host_vector labels; + /** Resolved uniform local-ID width, including automatic width selection. */ + std::uint16_t n_bits; + /** Resolved final group count. */ + std::uint32_t n_groups; +}; + +/** + * @brief Assign locality groups through a predetermined hierarchy with capacity repair at each + * level. + * + * Grouping uses vector distances, not graph adjacency. Only the graph shape is inspected; this + * function does not build or change edges. K-means nodes have at most 16 children. Assignment is + * batched on the GPU, but capacity repair retains count * children float costs in host memory for + * the current node. Large inputs therefore still require substantial host memory. + * This operation synchronizes the resource stream before returning completed host labels. + * + * @param res Allocation, stream ordering, and GPU execution resources. + * @param params Grouping settings, sampling seed, and cagra_params.metric. Supported metrics are + * L2Expanded, InnerProduct, and CosineExpanded. Other index settings, including the resident graph + * budget and grouping.enabled, do not affect this standalone operation. + * @param graph Host graph with original dataset row IDs, N rows, and degree in [1, 255]. + * @param dataset Full-precision host vectors with N rows and positive dimension. N must be positive + * and fit uint32 graph IDs. Training iterations must be positive and computed distances finite. + * @return Owning host labels plus the resolved uniform ID width and group count. At most 65,536 + * groups are supported by the graph metadata. A single group bypasses k-means training. + */ +#define CUVS_FLOWANN_DECLARE_GROUP_GRAPH(T) \ + CUVS_EXPORT auto group_graph( \ + raft::resources const& res, \ + index_params const& params, \ + raft::host_matrix_view graph, \ + cuvs::neighbors::host_standard_dataset_view const& dataset) -> grouping_result + +CUVS_FLOWANN_DECLARE_GROUP_GRAPH(float); +CUVS_FLOWANN_DECLARE_GROUP_GRAPH(half); +CUVS_FLOWANN_DECLARE_GROUP_GRAPH(std::int8_t); +CUVS_FLOWANN_DECLARE_GROUP_GRAPH(std::uint8_t); +#undef CUVS_FLOWANN_DECLARE_GROUP_GRAPH + +/** + * @brief Build a FlowANN index by constructing and immediately tiering a final CAGRA graph. + * + * The device-padded overload copies the dense dataset into its owning output bundle. The dense + * overload retains every final CAGRA edge. Group-local edges selected by the packed-row budget are + * stored on the device, and all remaining edges are stored in the host cross graph. ACE graph build + * parameters are not supported. + * + * `res` provides allocation, stream ordering, and GPU execution resources. `params` controls CAGRA + * construction, balanced locality grouping, resident-graph budgeting, packing, and optional seeds. + * Grouped builds rank graph edges with the full-precision input before choosing resident edges, + * reorder graph and dataset rows together, and expose original row IDs through `source_indices()`. + * `dataset` is the full-precision input and its row count must fit uint32 graph IDs. The overload + * returns an owning dataset/index bundle whose index view remains valid for the bundle lifetime. + */ +#define CUVS_FLOWANN_DECLARE_INTEGRATED_BUILD(T) \ + CUVS_EXPORT auto build(raft::resources const& res, \ + index_params const& params, \ + cuvs::neighbors::device_padded_dataset_view const& dataset) \ + -> device_padded_index_bundle + +CUVS_FLOWANN_DECLARE_INTEGRATED_BUILD(float); +CUVS_FLOWANN_DECLARE_INTEGRATED_BUILD(half); +CUVS_FLOWANN_DECLARE_INTEGRATED_BUILD(std::int8_t); +CUVS_FLOWANN_DECLARE_INTEGRATED_BUILD(std::uint8_t); + +#undef CUVS_FLOWANN_DECLARE_INTEGRATED_BUILD + +/** + * Convert a CAGRA index into an owning FlowANN index. `partition_labels` contains one partition + * identifier per CAGRA node. + */ +#define CUVS_FLOWANN_DECLARE_BUILD(T) \ + CUVS_EXPORT auto build_from_cagra( \ + raft::resources const& res, \ + build_params const& params, \ + cuvs::neighbors::cagra::device_padded_index const& cagra_index, \ + raft::host_vector_view partition_labels, \ + std::optional> seeds = std::nullopt) \ + -> device_padded_index_bundle + +CUVS_FLOWANN_DECLARE_BUILD(float); +CUVS_FLOWANN_DECLARE_BUILD(half); +CUVS_FLOWANN_DECLARE_BUILD(std::int8_t); +CUVS_FLOWANN_DECLARE_BUILD(std::uint8_t); + +#undef CUVS_FLOWANN_DECLARE_BUILD + +} // namespace cuvs::neighbors::cagra::experimental::flowann + +#endif // CUVS_ENABLE_FLOWANN_SEARCH && CUVS_ENABLE_FLOWANN_BUILD diff --git a/cpp/src/neighbors/detail/flowann/grouping.hpp b/cpp/src/neighbors/detail/flowann/grouping.hpp new file mode 100644 index 0000000000..d6970c8cda --- /dev/null +++ b/cpp/src/neighbors/detail/flowann/grouping.hpp @@ -0,0 +1,275 @@ +/* + * 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 +#include +#include + +namespace cuvs::neighbors::cagra::experimental::flowann::detail { + +struct grouping_plan { + std::uint16_t n_bits; + std::uint32_t n_groups; + std::uint64_t capacity; + std::uint64_t lower; + std::uint64_t upper; +}; + +/** Plan once, independently of the graph memory budget. This arithmetic also accepts dataset + * sizes beyond the current uint32 graph implementation, for capacity planning only. + */ +inline auto plan_grouping(std::uint64_t n_rows, + std::uint16_t bits = 0, + std::uint32_t groups = 0, + double tolerance = 0.10) -> grouping_plan +{ + RAFT_EXPECTS(n_rows > 0, "FlowANN cannot group an empty dataset"); + RAFT_EXPECTS(tolerance > 0 && tolerance < 1, "Grouping tolerance must be in (0, 1)"); + if (bits == 0) { + std::uint16_t required = 0; + for (auto remaining = n_rows - 1; remaining != 0; remaining >>= 1) { + ++required; + } + bits = std::min(24, std::max(4, (required + 3) / 4 * 4)); + } + RAFT_EXPECTS(bits <= 32, "FlowANN grouping n_bits must be in [1, 32]"); + auto const capacity = std::uint64_t{1} << bits; + // Round away sub-ULP arithmetic noise before ceil at exact capacity boundaries. + auto ceil_count = [](double value) { + auto nearest = std::round(value); + if (std::abs(value - nearest) <= 4 * std::numeric_limits::epsilon() * value) { + value = nearest; + } + return std::ceil(value); + }; + if (groups == 0) { + auto const required = + n_rows <= capacity ? 1.0 + : ceil_count(static_cast(n_rows) * (1.0 + tolerance) / capacity); + RAFT_EXPECTS(required <= std::numeric_limits::max(), "Too many groups"); + groups = static_cast(required); + } + RAFT_EXPECTS(groups <= n_rows && groups * capacity >= n_rows, + "Group count cannot satisfy local-ID capacity"); + auto const average = static_cast(n_rows) / groups; + auto const lower = std::max(1, std::floor(average * (1.0 - tolerance))); + auto const upper = std::min(capacity, std::ceil(average * (1.0 + tolerance))); + RAFT_EXPECTS(lower * groups <= n_rows && upper * groups >= n_rows, + "Infeasible grouping size bounds"); + return {bits, groups, capacity, lower, upper}; +} + +/** Each entry is the number of final leaves beneath a child. Intermediate capacities are + * additive: child_leaves * leaf_capacity (and likewise for the balance bounds). + */ +inline auto grouping_children(std::uint32_t leaves) -> std::vector +{ + RAFT_EXPECTS(leaves > 1, "A leaf has no children"); + auto const branches = leaves <= 16 ? leaves : std::min(16, (leaves + 15) / 16); + auto result = std::vector(branches, leaves / branches); + for (std::uint32_t i = 0; i < leaves % branches; ++i) { + ++result[i]; + } + return result; +} + +/** Reassign nearest-centroid labels using distance-increase priority. This is a bounded greedy + * repair, not optimal transport. Costs are row-major, with smaller values always preferable. + * Every move preserves donor lower bounds and receiver upper bounds. Overflow is removed first, + * then underfull groups are filled; selection is repeated when a destination/source fills up. + */ +inline auto rebalance_group_distances(float const* costs, + std::vector& labels, + std::vector const& lower, + std::vector const& upper) + -> std::vector +{ + auto const k = lower.size(); + RAFT_EXPECTS(k > 0 && k <= 16 && upper.size() == k, "Invalid grouping branches"); + std::vector sizes(k); + for (auto label : labels) { + RAFT_EXPECTS(label < k, "Invalid group label"); + ++sizes[label]; + } + RAFT_EXPECTS(std::accumulate(lower.begin(), lower.end(), std::uint64_t{0}) <= labels.size() && + std::accumulate(upper.begin(), upper.end(), std::uint64_t{0}) >= labels.size(), + "Infeasible subtree capacities"); + for (std::size_t g = 0; g < k; ++g) { + RAFT_EXPECTS(lower[g] <= upper[g], "Invalid subtree bounds"); + } + struct move { + // collect() writes every slot. Avoid zero-filling large temporary arrays before that write. + move() noexcept {} + move(float cost, std::uint32_t id, std::uint8_t destination) + : penalty(cost), row(id), target(destination) + { + } + float penalty; + std::uint32_t row; + std::uint8_t target; + }; + auto less = [](move const& a, move const& b) { + return a.penalty < b.penalty || (a.penalty == b.penalty && a.row < b.row); + }; + // Compact eligible rows in parallel, retaining original row order for deterministic ties. + auto collect = [](std::size_t rows, auto eligible, auto make_move) { + constexpr std::size_t block_rows = 65536; + auto const blocks = (rows + block_rows - 1) / block_rows; + std::vector offsets(blocks + 1); +#ifdef _OPENMP +#pragma omp parallel for schedule(static) if (rows > block_rows) +#endif + for (std::size_t block = 0; block < blocks; ++block) { + std::size_t count = 0; + for (auto row = block * block_rows; row < std::min(rows, (block + 1) * block_rows); ++row) { + count += eligible(row); + } + offsets[block + 1] = count; + } + std::partial_sum(offsets.begin(), offsets.end(), offsets.begin()); + std::vector result(offsets.back()); +#ifdef _OPENMP +#pragma omp parallel for schedule(static) if (rows > block_rows) +#endif + for (std::size_t block = 0; block < blocks; ++block) { + auto position = offsets[block]; + for (auto row = block * block_rows; row < std::min(rows, (block + 1) * block_rows); ++row) { + if (eligible(row)) { result[position++] = make_move(row); } + } + } + return result; + }; + std::vector candidates; + auto apply = [&](std::uint64_t requested, bool overflow) { + auto count = std::min(requested, candidates.size()); + RAFT_EXPECTS(count > 0, "No feasible capacity repair"); + std::vector selected; +#ifdef _OPENMP + if (count < candidates.size() && candidates.size() >= (1u << 20)) { + // Exact radix selection parallelizes the large partition operation. The row ID + // completes a unique key, so equal distances preserve the same priority as `less`. + auto key = [](move const& value) { + auto penalty = value.penalty == 0 ? 0.0f : value.penalty; + std::uint32_t bits; + std::memcpy(&bits, &penalty, sizeof(bits)); + bits ^= (bits & 0x80000000u) ? 0xffffffffu : 0x80000000u; + return (std::uint64_t{bits} << 32) | value.row; + }; + std::uint64_t prefix = 0, mask = 0; + auto rank = count - 1; + for (int shift = 56; shift >= 0; shift -= 8) { + std::uint64_t histogram[256]{}; +#pragma omp parallel for schedule(static) reduction(+ : histogram[ : 256]) + for (std::size_t i = 0; i < candidates.size(); ++i) { + auto const value = key(candidates[i]); + if ((value & mask) == prefix) { ++histogram[(value >> shift) & 255]; } + } + std::size_t bucket = 0; + while (bucket < 255 && rank >= histogram[bucket]) { + rank -= histogram[bucket++]; + } + prefix |= std::uint64_t{bucket} << shift; + mask |= std::uint64_t{255} << shift; + if (histogram[bucket] == 1) { break; } + } + auto threshold = prefix | ~mask; + selected = collect( + candidates.size(), + [&](auto i) { return key(candidates[i]) <= threshold; }, + [&](auto i) { return candidates[i]; }); + RAFT_EXPECTS(selected.size() == count, "Invalid capacity-repair selection"); + } else +#endif + { + if (count < candidates.size()) { + std::nth_element(candidates.begin(), candidates.begin() + count, candidates.end(), less); + } + selected.assign(candidates.begin(), candidates.begin() + count); + } + // All moves share a donor (overflow) or receiver (underflow). Only the opposite + // endpoint can saturate. Select that endpoint's cheapest feasible subset, without + // sorting millions of candidates that can all be moved independently. + std::vector> buckets(k); + for (std::size_t i = 0; i < count; ++i) { + auto const& candidate = selected[i]; + auto const group = overflow ? candidate.target : labels[candidate.row]; + buckets[group].push_back(candidate); + } + std::uint64_t moved = 0; + for (std::size_t group = 0; group < k; ++group) { + auto& bucket = buckets[group]; + auto const available = overflow ? upper[group] - std::min(sizes[group], upper[group]) + : sizes[group] - std::min(sizes[group], lower[group]); + auto const take = std::min(available, bucket.size()); + if (take == 0) { continue; } + if (take < bucket.size()) { + std::nth_element(bucket.begin(), bucket.begin() + take, bucket.end(), less); + } + for (std::size_t i = 0; i < take; ++i) { + auto const& candidate = bucket[i]; + auto& source = labels[candidate.row]; + --sizes[source]; + source = candidate.target; + ++sizes[source]; + } + moved += take; + } + RAFT_EXPECTS(moved > 0, "Capacity repair made no progress"); + }; + for (std::size_t donor = 0; donor < k; ++donor) { + while (sizes[donor] > upper[donor]) { + candidates = collect( + labels.size(), + [&](auto row) { return labels[row] == donor; }, + [&](auto row) { + auto target = k; + for (std::size_t g = 0; g < k; ++g) { + if (g != donor && sizes[g] < upper[g] && + (target == k || costs[row * k + g] < costs[row * k + target])) { + target = g; + } + } + // Feasible aggregate bounds guarantee at least one receiver. + return move{costs[row * k + target] - costs[row * k + donor], + static_cast(row), + static_cast(target)}; + }); + apply(sizes[donor] - upper[donor], true); + } + } + for (std::size_t receiver = 0; receiver < k; ++receiver) { + while (sizes[receiver] < lower[receiver]) { + candidates = collect( + labels.size(), + [&](auto row) { + auto const donor = labels[row]; + return donor != receiver && sizes[donor] > lower[donor]; + }, + [&](auto row) { + return move{costs[row * k + receiver] - costs[row * k + labels[row]], + static_cast(row), + static_cast(receiver)}; + }); + // Restrict each donor to its available surplus. + apply(lower[receiver] - sizes[receiver], false); + } + } + for (std::size_t g = 0; g < k; ++g) { + RAFT_EXPECTS(sizes[g] >= lower[g] && sizes[g] <= upper[g], "Subtree capacity violated"); + } + return sizes; +} + +} // 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..8ed411852a --- /dev/null +++ b/cpp/src/neighbors/detail/flowann/kernel.cuh @@ -0,0 +1,105 @@ +/* + * 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*, + 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..cbb30b68d5 --- /dev/null +++ b/cpp/src/neighbors/detail/flowann/launcher.hpp @@ -0,0 +1,200 @@ +/* + * 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)); +} + +template +auto build_multi_cta_launcher( + const cuvs::neighbors::cagra::detail::dataset_descriptor_host& + dataset_desc, + std::unique_ptr sample_filter_udf_fragment) + -> std::shared_ptr +{ + multi_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(); + planner.add_sample_filter_device_function(std::move(sample_filter_udf_fragment)); + return planner.get_launcher(); +} + +template +auto make_multi_cta_launcher( + const cuvs::neighbors::cagra::detail::dataset_descriptor_host& + dataset_desc, + 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_multi_cta_launcher(dataset_desc, + std::move(sample_filter_udf_fragment)); + } + using query_tag = query_type_tag_standard_t; + return build_multi_cta_launcher(dataset_desc, + 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..a16d3131d8 --- /dev/null +++ b/cpp/src/neighbors/detail/flowann/planner.hpp @@ -0,0 +1,90 @@ +/* + * 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>(); + } + } +}; + +template +struct multi_cta_planner + : cuvs::neighbors::cagra::detail:: + CagraPlannerBase { + using base_type = cuvs::neighbors::cagra::detail:: + CagraPlannerBase; + static inline rtcx::launcher_jit_cache launcher_cache{}; + + multi_cta_planner() : base_type("flowann_search_multi_cta", launcher_cache) {} + + void add_search_kernel() + { + this->template add_static_fragment< + fragment_tag_search_multi_cta>(); + } +}; + +} // 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..ce7edc8873 --- /dev/null +++ b/cpp/src/neighbors/detail/flowann/queue.cuh @@ -0,0 +1,262 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * FlowANN is described in "Disentangling Graph Dependencies for Efficient + * Billion-Scale GPU Vector Search" (OSDI 2026). + */ + +#pragma once + +#include + +#include + +#include + +#include +#include + +namespace cuvs::neighbors::cagra::experimental::flowann::detail { + +inline constexpr std::uint32_t queue_capacity = 1u << 16; +inline constexpr std::uint32_t queue_mask = queue_capacity - 1; +inline constexpr unsigned long long empty_command = 0ull; + +static_assert((queue_capacity & queue_mask) == 0, "FlowANN queue capacity must be a power of two"); + +// Each word carries the row ID and its complement. Empty or inconsistent +// words fail validation; every uint32 row ID remains representable. +RAFT_INLINE_FUNCTION constexpr auto make_command(std::uint32_t row_id) -> unsigned long long +{ + return (static_cast(~row_id) << 32) | row_id; +} + +RAFT_INLINE_FUNCTION constexpr auto command_ready(unsigned long long word) -> bool +{ + return static_cast(word >> 32) == ~static_cast(word); +} + +struct alignas(128) queue_state { + alignas(128) unsigned long long reservation_head; + alignas(128) unsigned long long cpu_tail; + alignas(128) unsigned long long gpu_tail; + std::uint32_t capacity; + std::uint32_t row_length; +}; + +/** Device pointers for one host-serviced cross-graph queue. */ +struct queue_view { + queue_state* state; + std::uint32_t* rows; + unsigned long long* command_words; + unsigned long long* result_sequence; + unsigned long long* released_sequence; +}; + +struct queue_reservation { + unsigned long long sequence; + std::uint32_t slot; +}; + +#ifdef __CUDACC__ + +RAFT_DEVICE_INLINE_FUNCTION auto load_counter(unsigned long long const* value) -> unsigned long long +{ + cuda::atomic_ref counter( + *const_cast(value)); + return counter.load(cuda::memory_order_relaxed); +} + +[[nodiscard]] RAFT_DEVICE_INLINE_FUNCTION auto reservation_at(queue_reservation first, + std::uint32_t offset) + -> queue_reservation +{ + auto const sequence = first.sequence + offset; + return {sequence, static_cast(sequence) & queue_mask}; +} + +RAFT_DEVICE_INLINE_FUNCTION auto load_published_counter(unsigned long long const* value) + -> unsigned long long +{ + cuda::atomic_ref counter( + *const_cast(value)); + return counter.load(cuda::memory_order_acquire); +} + +/** + * Reclaim only the contiguous released prefix. A failed CAS means another reclaimer advanced + * the tail; restart from that tail so no slot is returned to producers twice. + * This scalar fallback is used when a deferred producer runs out of queue capacity. + */ +RAFT_DEVICE_INLINE_FUNCTION void reclaim_released(queue_view queue) +{ + while (true) { + auto const tail = load_counter(&queue.state->gpu_tail); + auto const head = load_counter(&queue.state->reservation_head); + auto end = tail; + while (end < head) { + auto const slot = static_cast(end) & queue_mask; + cuda::atomic_ref sequence( + queue.released_sequence[slot]); + if (sequence.load(cuda::memory_order_relaxed) != end + 1) { break; } + ++end; + } + if (end == tail) { return; } + if (atomicCAS(&queue.state->gpu_tail, tail, end) == tail) { return; } + } +} + +/** + * Called by all 32 lanes of a warp after the CTA has consumed its pending rows. The first + * unreleased sequence stops reclamation, including when the scan crosses the ring boundary. + * The tail CAS also permits concurrent reclaimers and users of the ordinary release path. + */ +RAFT_DEVICE_INLINE_FUNCTION void reclaim_released_warp(queue_view queue) +{ + auto const lane = threadIdx.x & 31u; + while (true) { + auto const tail = + __shfl_sync(0xffffffffu, lane == 0 ? load_counter(&queue.state->gpu_tail) : 0ull, 0); + auto const head = + __shfl_sync(0xffffffffu, lane == 0 ? load_counter(&queue.state->reservation_head) : 0ull, 0); + auto end = tail; + while (end < head) { + auto const sequence = end + lane; + auto const slot = static_cast(sequence) & queue_mask; + cuda::atomic_ref released( + queue.released_sequence[slot]); + auto const ready = + sequence < head && released.load(cuda::memory_order_relaxed) == sequence + 1; + auto const mask = __ballot_sync(0xffffffffu, ready); + auto const count = mask == 0xffffffffu ? 32u : static_cast(__ffs(~mask) - 1); + end += count; + if (count < 32) { break; } + } + if (end == tail) { return; } + int success = lane == 0 && atomicCAS(&queue.state->gpu_tail, tail, end) == tail; + if (__shfl_sync(0xffffffffu, success, 0)) { return; } + } +} + +/** Mark a consumed row; the caller must reclaim at capacity pressure and on completion. */ +RAFT_DEVICE_INLINE_FUNCTION void release_deferred(queue_view queue, queue_reservation reservation) +{ + atomicExch(&queue.released_sequence[reservation.slot], reservation.sequence + 1); +} + +/** Try to reserve a contiguous batch of queue slots atomically. */ +RAFT_DEVICE_INLINE_FUNCTION auto try_reserve_batch(queue_view queue, + std::uint32_t count, + queue_reservation* first) -> bool +{ + auto* state = queue.state; + if (count == 0 || count > state->capacity) { return false; } + + auto const head = load_counter(&state->reservation_head); + auto const tail = load_counter(&state->gpu_tail); + if (head - tail > state->capacity - count) { return false; } + if (atomicCAS(&state->reservation_head, head, head + count) != head) { return false; } + + first->sequence = head; + first->slot = static_cast(head) & queue_mask; + return true; +} + +RAFT_DEVICE_INLINE_FUNCTION auto try_reserve_deferred(queue_view queue, + std::uint32_t count, + queue_reservation* first) -> bool +{ + auto* state = queue.state; + if (count == 0 || count > state->capacity) { return false; } + + auto const head = load_counter(&state->reservation_head); + auto const tail = load_counter(&state->gpu_tail); + if (head - tail > state->capacity - count) { + reclaim_released(queue); + return false; + } + if (atomicCAS(&state->reservation_head, head, head + count) != head) { return false; } + + first->sequence = head; + first->slot = static_cast(head) & queue_mask; + return true; +} + +/** Try to reserve one queue slot without blocking the CTA indefinitely. */ +RAFT_DEVICE_INLINE_FUNCTION auto try_reserve(queue_view queue, queue_reservation* reservation) + -> bool +{ + return try_reserve_batch(queue, 1, reservation); +} + +/** Atomically publish each complete request in a reserved batch. */ +RAFT_DEVICE_INLINE_FUNCTION void submit_batch(queue_view queue, + queue_reservation first, + std::uint32_t const* row_ids, + std::uint32_t count) +{ + // The atomically published word contains the complete request. No separate + // payload needs a system fence. The host clears this word before publishing + // the result; result acquisition and deferred reclamation gate slot reuse. + for (std::uint32_t i = 0; i < count; ++i) { + auto const reservation = reservation_at(first, i); + cuda::atomic_ref command( + queue.command_words[reservation.slot]); + static_cast(command.exchange(make_command(row_ids[i]), cuda::memory_order_relaxed)); + } +} + +/** Publish one complete cross-graph row request atomically. */ +RAFT_DEVICE_INLINE_FUNCTION void submit(queue_view queue, + queue_reservation reservation, + std::uint32_t row_id) +{ + submit_batch(queue, reservation, &row_id, 1); +} + +RAFT_DEVICE_INLINE_FUNCTION auto result_ready(queue_view queue, queue_reservation reservation) + -> bool +{ + return load_published_counter(&queue.result_sequence[reservation.slot]) == + reservation.sequence + 1; +} + +RAFT_DEVICE_INLINE_FUNCTION auto result_row(queue_view queue, queue_reservation reservation) + -> std::uint32_t const* +{ + return queue.rows + static_cast(reservation.slot) * queue.state->row_length; +} + +/** + * Mark a consumed result as released and advance the global GPU tail over consecutive releases. + */ +RAFT_DEVICE_INLINE_FUNCTION void release(queue_view queue, queue_reservation reservation) +{ + atomicExch(&queue.released_sequence[reservation.slot], reservation.sequence + 1); + + while (true) { + auto const tail = load_counter(&queue.state->gpu_tail); + auto const slot = static_cast(tail) & queue_mask; + if (load_counter(&queue.released_sequence[slot]) != tail + 1) { break; } + if (atomicCAS(&queue.state->gpu_tail, tail, tail + 1) != tail) { continue; } + } +} + +#endif // __CUDACC__ + +#ifdef CUVS_ENABLE_FLOWANN_SEARCH +struct CUVS_EXPORT search_context_access { + static auto device_queues(search_context const& context) noexcept -> queue_view*; + static auto device_id(search_context const& context) noexcept -> int; + static auto num_queues(search_context const& context) noexcept -> std::uint32_t; + static auto row_count(search_context const& context) noexcept -> std::uint64_t; + static auto row_length(search_context const& context) noexcept -> std::uint32_t; + static void acquire_search(search_context& context); + static void release_search(search_context& context); + static auto is_paused(search_context const& context) noexcept -> bool; + static auto pending_pause_requests(search_context const& context) noexcept -> std::size_t; +}; +#endif + +} // namespace cuvs::neighbors::cagra::experimental::flowann::detail diff --git a/cpp/src/neighbors/detail/flowann/search_multi_cta.cuh b/cpp/src/neighbors/detail/flowann/search_multi_cta.cuh new file mode 100644 index 0000000000..001badf98a --- /dev/null +++ b/cpp/src/neighbors/detail/flowann/search_multi_cta.cuh @@ -0,0 +1,401 @@ +/* + * 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 "search_single_cta.cuh" + +#include +#include + +#include + +#include + +namespace cuvs::neighbors::cagra::experimental::flowann::detail { + +template +struct multi_cta_plan + : cuvs::neighbors::cagra::detail:: + search_plan_impl { + using base_type = cuvs::neighbors::cagra::detail:: + search_plan_impl; + + static constexpr bool needs_index_copy = sizeof(IndexT) != sizeof(OutputIndexT); + + std::uint32_t graph_degree; + std::uint32_t num_cta_per_query; + rmm::device_uvector intermediate_indices; + rmm::device_uvector intermediate_distances; + rmm::device_uvector topk_workspace; + std::size_t topk_workspace_size{}; + + static auto planning_params(flowann::search_params const& input) + -> cuvs::neighbors::cagra::search_params + { + cuvs::neighbors::cagra::search_params result = input; + result.algo = cuvs::neighbors::cagra::search_algo::MULTI_CTA; + result.persistent = false; + return result; + } + + multi_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, + std::uint32_t input_graph_degree, + std::uint32_t topk) + : base_type(res, + planning_params(params), + dataset_desc, + dim, + dataset_size, + (1 + additional_buffer_width) * input_graph_degree, + topk), + graph_degree(input_graph_degree), + num_cta_per_query(0), + intermediate_indices(0, raft::resource::get_cuda_stream(res)), + intermediate_distances(0, raft::resource::get_cuda_stream(res)), + topk_workspace(0, raft::resource::get_cuda_stream(res)) + { + auto const global_itopk = this->itopk_size; + constexpr std::uint32_t local_itopk = 32; + this->itopk_size = local_itopk; + this->search_width = 1; + num_cta_per_query = static_cast( + std::max(params.search_width, raft::ceildiv(global_itopk, std::size_t{local_itopk}))); + auto const base_result_buffer_size = + local_itopk + (1 + additional_buffer_width) * input_graph_degree; + this->result_buffer_size = std::max(base_result_buffer_size, params.num_seeds); + auto const result_buffer_size_32 = + raft::round_up_safe(this->result_buffer_size, 32); + RAFT_EXPECTS(result_buffer_size_32 <= 256, + "FlowANN multi-CTA result buffer cannot exceed 256 elements"); + this->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) + sizeof(int); + + constexpr std::uint32_t min_block_size = 64; + constexpr std::uint32_t max_block_size = 1024; + auto block_size = static_cast(this->thread_block_size); + if (block_size == 0) { + block_size = min_block_size; + constexpr std::uint32_t smem_per_warp_limit = 4096; + while (block_sizesmem_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 && + input_graph_degree * this->team_size >= block_size * 2 && + num_cta_per_query * 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 multi-CTA thread_block_size must be in [64, 1024]"); + this->thread_block_size = block_size; + + auto const intermediate_count = num_cta_per_query * local_itopk; + intermediate_indices.resize( + intermediate_count * this->max_queries + (needs_index_copy ? topk * this->max_queries : 0), + stream(res)); + intermediate_distances.resize(intermediate_count * this->max_queries, stream(res)); + this->hashmap.resize(this->hashmap_size, stream(res)); + topk_workspace_size = cuvs::neighbors::cagra::detail::_cuann_find_topk_bufferSize( + topk, this->max_queries, intermediate_count, CUDA_R_32F); + topk_workspace.resize(topk_workspace_size, stream(res)); + } + + private: + static auto stream(raft::resources const& res) -> rmm::cuda_stream_view + { + return raft::resource::get_cuda_stream(res); + } +}; + +template +void search_multi_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 metadata"); + RAFT_EXPECTS(queries.extent(1) == index.dim(), + "FlowANN query dimension must match the index dimension"); + RAFT_EXPECTS(neighbors.extents() == distances.extents() && + neighbors.extent(0) == queries.extent(0) && neighbors.extent(1) > 0, + "FlowANN outputs must have shape [n_queries, k]"); + RAFT_EXPECTS(queries.extent(0) <= std::numeric_limits::max(), + "FlowANN multi-CTA query count exceeds uint32"); + RAFT_EXPECTS(params.search_width > 0 && params.sync_window_scale >= 0.0f, + "FlowANN multi-CTA parameters are invalid"); + 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 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(params.num_seeds == 0 || index_seeds.has_value(), + "FlowANN num_seeds requires preloaded index seeds"); + RAFT_EXPECTS(!index_seeds.has_value() || params.num_seeds <= index_seeds->extent(0), + "FlowANN num_seeds exceeds the preloaded seed count"); + RAFT_EXPECTS(params.num_seeds == 0 || seed_cross_graph.has_value(), + "FlowANN num_seeds requires cached seed cross-graph rows"); + RAFT_EXPECTS( + params.num_seeds == 0 || (seed_lookup_keys.has_value() && seed_lookup_values.has_value()), + "FlowANN num_seeds requires a prebuilt seed lookup"); + auto const* seeds = params.num_seeds == 0 ? nullptr : index_seeds->data_handle(); + auto const* seed_cross_ptr = params.num_seeds == 0 ? nullptr : seed_cross_graph->data_handle(); + auto const* seed_lookup_keys_ptr = + params.num_seeds == 0 ? nullptr : seed_lookup_keys->data_handle(); + auto const* seed_lookup_values_ptr = + params.num_seeds == 0 ? nullptr : seed_lookup_values->data_handle(); + auto const seed_lookup_capacity = + params.num_seeds == 0 ? 0u : static_cast(seed_lookup_keys->extent(0)); + auto const properties = raft::resource::get_device_properties(res); + if (params.max_queries == 0) { + params.max_queries = std::min(queries.extent(0), properties.maxGridSize[1]); + } + + auto descriptor = make_dataset_descriptor(res, params, index); + auto const topk = static_cast(neighbors.extent(1)); + multi_cta_plan plan( + res, params, descriptor, index.dim(), index.size(), index.graph_degree(), topk); + RAFT_EXPECTS(plan.num_cta_per_query * plan.itopk_size >= topk, + "FlowANN multi-CTA candidate count is smaller than k"); + auto launcher_filter = cuvs::neighbors::cagra::detail::set_offset(sample_filter, 0); + auto launcher = make_multi_cta_launcher< + data_t, + index_t, + distance_t, + source_index_t, + cuvs::neighbors::cagra::detail::sample_filter_jit_tag_t>( + descriptor, + cuvs::neighbors::cagra::detail::make_cagra_sample_filter_udf_fragment( + launcher_filter)); + RAFT_EXPECTS(launcher != nullptr, "Failed to create the FlowANN multi-CTA JIT launcher"); + + auto const query_dim = static_cast(queries.extent(1)); + 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.view().data_handle(); + query_stride = view.stride(); + } else { + padded_queries = cuvs::neighbors::make_device_padded_dataset(res, queries); + auto view = padded_queries->as_dataset_view(); + query_data = view.view().data_handle(); + query_stride = view.stride(); + } + auto const max_batch_queries = query_stride == query_dim ? plan.max_queries : std::size_t{1}; + 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; + using kernel_type = search_multi_cta_kernel_func_t; + auto const max_elements = raft::round_up_safe(plan.result_buffer_size, 32); + auto const intermediate_count = plan.num_cta_per_query * plan.itopk_size; + auto const num_queries = static_cast(queries.extent(0)); + + context_run_guard running(context); + 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* 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); + auto const traversed_hash_size = + cuvs::neighbors::cagra::detail::hashmap::get_size(plan.hash_bitlen); + cuvs::neighbors::cagra::detail::set_value_batch( + plan.hashmap.data(), traversed_hash_size, ~index_t{0}, traversed_hash_size, count, stream); + dim3 grid(plan.num_cta_per_query, 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), + plan.intermediate_indices.data(), + plan.intermediate_distances.data(), + descriptor_ptr, + query_ptr, + 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), + max_elements, + index.graph_degree(), + index.cross_degree(), + source_ptr, + static_cast(plan.num_random_samplings), + plan.rand_xor_mask, + seeds, + params.num_seeds, + seed_cross_ptr, + seed_lookup_keys_ptr, + seed_lookup_values_ptr, + seed_lookup_capacity, + static_cast(plan.small_hash_bitlen), + plan.hashmap.data(), + static_cast(plan.hash_bitlen), + static_cast(plan.itopk_size), + static_cast(plan.min_iterations), + static_cast(plan.max_iterations), + static_cast(nullptr), + index.size(), + filter_payload.query_id_offset, + 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()); + +#ifndef CUVS_FLOWANN_USE_GDRCOPY + // Lazy loading the top-k kernel can wait for active kernels while holding + // a driver lock needed by the CUDA-copy pollers. Finish the host-serviced + // search before launching kernels from another module. + raft::resource::sync_stream(res); +#endif + + auto* output_indices = + plan.needs_index_copy + ? plan.intermediate_indices.data() + intermediate_count * plan.max_queries + : reinterpret_cast(neighbors.data_handle() + + static_cast(offset) * topk); + cuvs::neighbors::cagra::detail::_cuann_find_topk( + topk, + count, + intermediate_count, + plan.intermediate_distances.data(), + intermediate_count, + plan.intermediate_indices.data(), + intermediate_count, + distances.data_handle() + static_cast(offset) * topk, + topk, + output_indices, + topk, + plan.topk_workspace.data(), + true, + nullptr, + stream); + if (source_ptr != nullptr) { + auto output = raft::make_device_matrix_view( + neighbors.data_handle() + static_cast(offset) * topk, count, topk); + raft::linalg::map( + res, + output, + [source_ptr] __device__(index_t value) -> OutputIndexT { + return value == ~index_t{0} ? ~OutputIndexT{0} + : static_cast(source_ptr[value]); + }, + raft::make_device_matrix_view(output_indices, count, topk)); + } else if constexpr (sizeof(index_t) != sizeof(OutputIndexT)) { + raft::linalg::map( + res, + raft::make_device_matrix_view( + neighbors.data_handle() + static_cast(offset) * topk, count, topk), + [] __device__(index_t value) -> OutputIndexT { + return value == ~index_t{0} ? ~OutputIndexT{0} : static_cast(value); + }, + raft::make_device_matrix_view(output_indices, count, topk)); + } + } + raft::resource::sync_stream(res); + running.pause(); + postprocess_distances(res, index.metric(), queries, distances); +} + +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"); + auto algo = params.algo; + if (algo == cuvs::neighbors::cagra::search_algo::AUTO) { + auto const properties = raft::resource::get_device_properties(res); + auto const effective_max_queries = + params.max_queries == 0 ? static_cast(queries.extent(0)) : params.max_queries; + algo = params.itopk_size <= 512 && + effective_max_queries >= static_cast(properties.multiProcessorCount) * 2 + ? cuvs::neighbors::cagra::search_algo::SINGLE_CTA + : cuvs::neighbors::cagra::search_algo::MULTI_CTA; + } + if (algo == cuvs::neighbors::cagra::search_algo::SINGLE_CTA) { + params.algo = algo; + search_single_cta(res, params, index, context, queries, neighbors, distances, sample_filter); + return; + } + RAFT_EXPECTS(algo == cuvs::neighbors::cagra::search_algo::MULTI_CTA, + "FlowANN search supports SINGLE_CTA, MULTI_CTA, or AUTO"); + params.algo = algo; + search_multi_cta(res, params, index, context, queries, neighbors, distances, sample_filter); +} + +} // namespace cuvs::neighbors::cagra::experimental::flowann::detail diff --git a/cpp/src/neighbors/detail/flowann/search_multi_cta_jit.cuh b/cpp/src/neighbors/detail/flowann/search_multi_cta_jit.cuh new file mode 100644 index 0000000000..c7edf76b42 --- /dev/null +++ b/cpp/src/neighbors/detail/flowann/search_multi_cta_jit.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 "search_single_cta_jit.cuh" + +#include + +namespace cuvs::neighbors::cagra::experimental::flowann::detail { + +template +__device__ void search_multi_cta_kernel_jit( + IndexT* result_indices, + DistanceT* result_distances, + const cuvs::neighbors::cagra::detail::dataset_descriptor_base_t* + dataset_desc, + const DataT* queries, + 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 max_elements, + 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, + std::uint32_t visited_hash_bitlen, + IndexT* traversed_hashmap, + std::uint32_t traversed_hash_bitlen, + std::uint32_t itopk_size, + std::uint32_t min_iterations, + std::uint32_t max_iterations, + std::uint32_t* num_executed_iterations, + IndexT graph_size, + std::uint32_t query_id_offset, + 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::multi_cta_search:: + search_core( + result_indices, + result_distances, + dataset_desc, + queries, + provider, + max_elements, + graph_degree, + (1 + additional_buffer_width) * graph_degree, + source_indices, + num_distillation, + rand_xor_mask, + seeds, + num_seeds, + visited_hash_bitlen, + traversed_hashmap, + traversed_hash_bitlen, + itopk_size, + min_iterations, + max_iterations, + num_executed_iterations, + graph_size, + query_id_offset, + filter_payload); +} + +} // namespace cuvs::neighbors::cagra::experimental::flowann::detail diff --git a/cpp/src/neighbors/detail/flowann/search_multi_cta_kernel.cu.in b/cpp/src/neighbors/detail/flowann/search_multi_cta_kernel.cu.in new file mode 100644 index 0000000000..0f469be744 --- /dev/null +++ b/cpp/src/neighbors/detail/flowann/search_multi_cta_kernel.cu.in @@ -0,0 +1,107 @@ +/* + * 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 { +using data_t = @data_type@; +using index_t = @index_type@; +using distance_t = @distance_type@; +using source_index_t = @source_index_type@; +using descriptor_t = + 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_multi_cta( + index_t* result_indices, + distance_t* result_distances, + const descriptor_t* dataset_desc, + const data_t* queries, + 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 max_elements, + 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, + std::uint32_t visited_hash_bitlen, + index_t* traversed_hashmap, + std::uint32_t traversed_hash_bitlen, + std::uint32_t itopk_size, + std::uint32_t min_iterations, + std::uint32_t max_iterations, + std::uint32_t* num_executed_iterations, + index_t graph_size, + std::uint32_t query_id_offset, + float sync_window_scale, + std::uint32_t sync_drop_threshold, + filter_payload_t filter_payload) +{ + search_multi_cta_kernel_jit(result_indices, + result_distances, + dataset_desc, + queries, + inner_graph, + inner_row_length, + node_per_cacheline, + n_bits, + subgraph_offsets, + subgraph_ids, + queues, + num_queues, + max_elements, + 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_hash_bitlen, + traversed_hashmap, + traversed_hash_bitlen, + itopk_size, + min_iterations, + max_iterations, + num_executed_iterations, + graph_size, + query_id_offset, + 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_multi_cta_matrix.json b/cpp/src/neighbors/detail/flowann/search_multi_cta_matrix.json new file mode 100644 index 0000000000..adfdf1e78b --- /dev/null +++ b/cpp/src/neighbors/detail/flowann/search_multi_cta_matrix.json @@ -0,0 +1,13 @@ +[ + { + "_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"}] + } +] 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..7b4e2e43c3 --- /dev/null +++ b/cpp/src/neighbors/detail/flowann/search_single_cta.cuh @@ -0,0 +1,490 @@ +/* + * 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 "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) + { + 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; + + 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_sizesmem_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().dset(), 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) : context_(context) + { + 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_) { + try { + active_ = false; + search_context_access::release_search(context_); + } catch (...) { + } + } + } + + void pause() + { + active_ = false; + search_context_access::release_search(context_); + } + + private: + search_context& context_; + 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)); + + 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]); + } + 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"); + + auto const query_dim = static_cast(queries.extent(1)); + 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.view().data_handle(); + query_stride = view.stride(); + } else { + padded_queries = cuvs::neighbors::make_device_padded_dataset(res, queries); + auto view = padded_queries->as_dataset_view(); + query_data = view.view().data_handle(); + query_stride = view.stride(); + } + auto const max_batch_queries = query_stride == query_dim ? plan.max_queries : std::size_t{1}; + 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); + 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, + 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..527f747d39 --- /dev/null +++ b/cpp/src/neighbors/detail/flowann/search_single_cta_jit.cuh @@ -0,0 +1,715 @@ +/* + * 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 +#include +#include + +#include +#include +#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; +}; + +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) { + 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; + 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}); + } + } + } + } + // 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, + 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); +} + +} // 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..41ff3b2c89 --- /dev/null +++ b/cpp/src/neighbors/detail/flowann/search_single_cta_kernel.cu.in @@ -0,0 +1,125 @@ +/* + * 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, + 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, + 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..715fbcb980 --- /dev/null +++ b/cpp/src/neighbors/flowann.cu @@ -0,0 +1,164 @@ +/* + * 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_multi_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 = dataset_.view(); + if (!dataset_norms_.has_value() || dataset_norms_->extent(0) != dataset.extent(0)) { + dataset_norms_.reset(); + dataset_norms_.emplace(raft::make_device_vector(res, dataset.extent(0))); + } + + constexpr float scale = cuvs::spatial::knn::detail::utils::config::kDivisor / + cuvs::spatial::knn::detail::utils::config::kDivisor; + auto scaled_sq_op = + raft::compose_op(raft::sq_op{}, raft::div_const_op{scale}, raft::cast_op()); + raft::linalg::reduce(res, + dataset, + dataset_norms_->view(), + 0.0f, + false, + scaled_sq_op, + raft::add_op(), + raft::sqrt_op{}); + } +} + +namespace detail { + +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_build.cu b/cpp/src/neighbors/flowann_build.cu new file mode 100644 index 0000000000..69f8497f04 --- /dev/null +++ b/cpp/src/neighbors/flowann_build.cu @@ -0,0 +1,2392 @@ +/* + * 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 "detail/cagra/cagra_build.cuh" + +#include +#include +#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 { + +namespace { + +using graph_index_type = std::uint32_t; +using host_graph_view = raft::host_matrix_view; +using host_edge_order_view = raft::host_matrix_view; + +constexpr std::uint16_t max_nodes_per_packed_row = 64; +constexpr std::uint32_t max_partitions = 1u << 16; + +template +auto cuda_stream_handle(StreamT const& stream) -> cudaStream_t +{ + if constexpr (requires { stream.get(); }) { + return stream.get(); + } else if constexpr (requires { stream.value(); }) { + return stream.value(); + } else { + return static_cast(stream); + } +} + +struct conversion_result { + raft::host_matrix inner_graph; + raft::host_matrix cross_graph; + std::vector old_to_new; + std::vector new_to_old; + std::vector subgraph_offsets; + std::vector subgraph_ids; +}; + +struct device_conversion_result { + raft::device_matrix inner_graph; + raft::host_matrix cross_graph; + std::vector old_to_new; + std::vector new_to_old; + std::vector subgraph_offsets; + std::vector subgraph_ids; + bool validated; +}; + +struct uniform_conversion_result { + raft::device_matrix inner_graph; + raft::host_matrix cross_graph; + uniform_layout_plan layout; +}; + +template +RAFT_DEVICE_INLINE_FUNCTION auto rank_value(T value) -> float +{ + return static_cast(value); +} + +RAFT_DEVICE_INLINE_FUNCTION void set_packed_bits(std::uint8_t* row, + std::uint64_t bit_offset, + std::uint32_t value, + std::uint16_t n_bits) +{ + for (std::uint16_t bit = 0; bit < n_bits; ++bit) { + auto const position = bit_offset + bit; + auto const mask = static_cast(1u << (position & 7u)); + if ((value >> bit) & 1u) { row[position >> 3] |= mask; } + } +} + +template +__global__ void rank_and_split_kernel(T const* dataset, + std::uint32_t dataset_stride, + T const* staged_rows, + graph_index_type const* graph, + std::uint64_t batch_offset, + std::uint32_t batch_rows, + std::uint32_t dim, + std::uint32_t graph_degree, + std::uint32_t resident_degree, + std::uint32_t cross_degree, + std::uint16_t node_per_cacheline, + std::uint16_t n_bits, + std::uint32_t inner_row_bytes, + std::uint8_t* inner_graph, + graph_index_type* cross_graph, + cuvs::distance::DistanceType metric) +{ + extern __shared__ std::uint8_t shared_bytes[]; + auto* scores = reinterpret_cast(shared_bytes); + auto* selected = + shared_bytes + static_cast(node_per_cacheline) * graph_degree * sizeof(float); + auto resident_offset = static_cast(node_per_cacheline) * graph_degree * + (sizeof(float) + sizeof(std::uint8_t)); + resident_offset = + (resident_offset + alignof(graph_index_type) - 1) & ~(alignof(graph_index_type) - 1); + auto* resident = reinterpret_cast(shared_bytes + resident_offset); + + auto const lane = threadIdx.x % raft::WarpSize; + auto const packed_col = threadIdx.x / raft::WarpSize; + auto const local_node = static_cast(blockIdx.x) * node_per_cacheline + packed_col; + auto const valid_node = packed_col < node_per_cacheline && local_node < batch_rows; + + if (valid_node && lane == 0) { + auto* node_selected = selected + static_cast(packed_col) * graph_degree; + for (std::uint32_t edge = 0; edge < graph_degree; ++edge) { + node_selected[edge] = 0; + } + } + __syncthreads(); + + if (valid_node) { + auto const global_node = batch_offset + local_node; + auto const* source = + StagedRows ? staged_rows + static_cast(local_node) * (graph_degree + 1) * dim + : dataset + global_node * dataset_stride; + for (std::uint32_t edge = 0; edge < graph_degree; ++edge) { + auto const neighbor = graph[local_node * graph_degree + edge]; + auto const* target = + StagedRows ? staged_rows + + (static_cast(local_node) * (graph_degree + 1) + edge + 1) * dim + : dataset + static_cast(neighbor) * dataset_stride; + float primary = 0.0f; + float source_l2 = 0.0f; + float target_l2 = 0.0f; + for (std::uint32_t column = lane; column < dim; column += raft::WarpSize) { + auto const lhs = rank_value(source[column]); + auto const rhs = rank_value(target[column]); + switch (metric) { + case cuvs::distance::DistanceType::L2Expanded: { + auto const diff = lhs - rhs; + primary += diff * diff; + break; + } + case cuvs::distance::DistanceType::InnerProduct: primary -= lhs * rhs; break; + case cuvs::distance::DistanceType::CosineExpanded: + primary -= lhs * rhs; + source_l2 += lhs * lhs; + target_l2 += rhs * rhs; + break; + case cuvs::distance::DistanceType::L1: primary += fabsf(lhs - rhs); break; + case cuvs::distance::DistanceType::BitwiseHamming: + if constexpr (std::is_integral_v) { + primary += __popc((static_cast(source[column]) ^ + static_cast(target[column])) & + 0xffu); + } + break; + default: primary = std::numeric_limits::infinity(); break; + } + } + for (auto offset = raft::WarpSize / 2; offset > 0; offset /= 2) { + primary += __shfl_down_sync(0xffffffffu, primary, offset); + source_l2 += __shfl_down_sync(0xffffffffu, source_l2, offset); + target_l2 += __shfl_down_sync(0xffffffffu, target_l2, offset); + } + if (lane == 0) { + if (metric == cuvs::distance::DistanceType::CosineExpanded) { + auto const denominator = sqrtf(source_l2 * target_l2); + primary = denominator > 0.0f ? 1.0f + primary / denominator + : std::numeric_limits::infinity(); + } + if (isnan(primary)) { primary = std::numeric_limits::infinity(); } + scores[static_cast(packed_col) * graph_degree + edge] = primary; + } + } + } + __syncthreads(); + + if (valid_node && lane == 0) { + auto const* node_graph = graph + local_node * graph_degree; + auto* node_scores = scores + static_cast(packed_col) * graph_degree; + auto* node_selected = selected + static_cast(packed_col) * graph_degree; + auto* node_resident = resident + static_cast(packed_col) * resident_degree; + + for (std::uint32_t rank = 0; rank < resident_degree; ++rank) { + auto best_edge = graph_degree; + auto best_id = std::numeric_limits::max(); + auto best_score = std::numeric_limits::infinity(); + for (std::uint32_t edge = 0; edge < graph_degree; ++edge) { + if (node_selected[edge] != 0) { continue; } + auto const score = node_scores[edge]; + auto const neighbor = node_graph[edge]; + if (score < best_score || (score == best_score && neighbor < best_id)) { + best_edge = edge; + best_id = neighbor; + best_score = score; + } + } + if (best_edge < graph_degree) { + node_selected[best_edge] = 1; + node_resident[rank] = best_id; + } + } + + auto* node_cross = cross_graph + local_node * (cross_degree + 1); + node_cross[0] = cross_degree; + auto output = std::uint32_t{0}; + for (std::uint32_t edge = 0; edge < graph_degree; ++edge) { + if (node_selected[edge] == 0) { node_cross[++output] = node_graph[edge]; } + } + } + __syncthreads(); + + if (threadIdx.x == 0 && resident_degree > 0) { + auto const global_first_node = + batch_offset + static_cast(blockIdx.x) * node_per_cacheline; + auto* inner_row = inner_graph + (global_first_node / node_per_cacheline) * inner_row_bytes; + for (std::uint32_t byte = 0; byte < inner_row_bytes; ++byte) { + inner_row[byte] = 0; + } + auto bit_offset = static_cast(node_per_cacheline) * 8; + for (std::uint16_t column = 0; column < node_per_cacheline; ++column) { + auto const node = static_cast(blockIdx.x) * node_per_cacheline + column; + auto const count = node < batch_rows ? resident_degree : 0; + inner_row[column] = static_cast(count); + for (std::uint32_t rank = 0; rank < count; ++rank) { + set_packed_bits(inner_row, + bit_offset + static_cast(rank) * n_bits, + resident[static_cast(column) * resident_degree + rank], + n_bits); + } + bit_offset += static_cast(count) * n_bits; + } + } +} + +template +__global__ void rank_graph_rows_kernel(T const* dataset, + std::uint32_t dataset_stride, + T const* staged_rows, + graph_index_type const* input_graph, + std::uint32_t const* labels, + std::uint8_t* output_order, + std::uint64_t batch_offset, + std::uint32_t batch_rows, + std::uint32_t dim, + std::uint32_t graph_degree, + cuvs::distance::DistanceType metric) +{ + extern __shared__ std::uint8_t shared_bytes[]; + auto const warp = threadIdx.x / raft::WarpSize; + auto const lane = threadIdx.x % raft::WarpSize; + auto const warps = blockDim.x / raft::WarpSize; + auto const local_node = static_cast(blockIdx.x) * warps + warp; + auto* scores = reinterpret_cast(shared_bytes) + warp * graph_degree; + auto* selected = shared_bytes + static_cast(warps) * graph_degree * sizeof(float) + + warp * graph_degree; + if (local_node >= batch_rows) { return; } + + if (lane == 0) { + for (std::uint32_t edge = 0; edge < graph_degree; ++edge) { + selected[edge] = 0; + output_order[local_node * graph_degree + edge] = std::uint8_t{255}; + } + } + auto const global_node = batch_offset + local_node; + auto const* source = StagedRows ? staged_rows + local_node * (graph_degree + 1) * dim + : dataset + global_node * dataset_stride; + for (std::uint32_t edge = 0; edge < graph_degree; ++edge) { + auto const neighbor = input_graph[local_node * graph_degree + edge]; + auto const* target = StagedRows + ? staged_rows + (local_node * (graph_degree + 1) + edge + 1) * dim + : dataset + static_cast(neighbor) * dataset_stride; + auto score = 0.0f; + auto source_l2 = 0.0f; + auto target_l2 = 0.0f; + for (std::uint32_t column = lane; column < dim; column += raft::WarpSize) { + auto const lhs = rank_value(source[column]); + auto const rhs = rank_value(target[column]); + switch (metric) { + case cuvs::distance::DistanceType::L2Expanded: { + auto const diff = lhs - rhs; + score += diff * diff; + break; + } + case cuvs::distance::DistanceType::InnerProduct: score -= lhs * rhs; break; + case cuvs::distance::DistanceType::CosineExpanded: + score -= lhs * rhs; + source_l2 += lhs * lhs; + target_l2 += rhs * rhs; + break; + default: score = std::numeric_limits::infinity(); break; + } + } + for (auto offset = raft::WarpSize / 2; offset > 0; offset /= 2) { + score += __shfl_down_sync(0xffffffffu, score, offset); + source_l2 += __shfl_down_sync(0xffffffffu, source_l2, offset); + target_l2 += __shfl_down_sync(0xffffffffu, target_l2, offset); + } + if (lane == 0) { + if (metric == cuvs::distance::DistanceType::CosineExpanded) { + auto const denominator = sqrtf(source_l2 * target_l2); + score = + denominator > 0.0f ? 1.0f + score / denominator : std::numeric_limits::infinity(); + } + scores[edge] = isnan(score) ? std::numeric_limits::infinity() : score; + } + } + __syncwarp(); + + if (lane == 0) { + auto* output = output_order + local_node * graph_degree; + for (std::uint32_t rank = 0; rank < graph_degree; ++rank) { + auto best_edge = graph_degree; + auto best_id = std::numeric_limits::max(); + auto best_score = std::numeric_limits::infinity(); + for (std::uint32_t edge = 0; edge < graph_degree; ++edge) { + if (selected[edge] != 0) { continue; } + auto const neighbor = input_graph[local_node * graph_degree + edge]; + if (labels[neighbor] != labels[global_node]) { continue; } + if (scores[edge] < best_score || (scores[edge] == best_score && neighbor < best_id)) { + best_edge = edge; + best_id = neighbor; + best_score = scores[edge]; + } + } + if (best_edge == graph_degree) { break; } + selected[best_edge] = 1; + output[best_edge] = static_cast(rank); + } + } +} + +template +__global__ void gather_float_rows_kernel(float* output, + T const* input, + std::uint64_t const* row_ids, + std::uint64_t row_offset, + std::uint64_t n_rows, + std::uint32_t dim, + std::uint32_t stride) +{ + auto index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + auto length = n_rows * dim; + auto step = static_cast(blockDim.x) * gridDim.x; + for (; index < length; index += step) { + auto const output_row = index / dim; + auto const column = index % dim; + auto const input_row = row_ids == nullptr ? row_offset + output_row : row_ids[output_row]; + output[index] = rank_value(input[input_row * stride + column]); + } +} + +template +auto rank_graph_edge_order(raft::resources const& res, + cuvs::distance::DistanceType metric, + host_graph_view graph, + raft::host_vector_view labels, + DatasetViewT const& dataset) + -> raft::host_matrix +{ + using value_type = + std::remove_const_t().view())::value_type>; + constexpr bool host_dataset = cuvs::neighbors::is_host_dataset_view_v; + + auto const n_rows = static_cast(graph.extent(0)); + auto const degree = static_cast(graph.extent(1)); + auto const dim = dataset.dim(); + RAFT_EXPECTS(n_rows > 0 && degree > 0 && dim > 0, + "FlowANN graph ranking requires a non-empty graph and dataset"); + RAFT_EXPECTS(degree <= std::numeric_limits::max(), + "FlowANN graph ranking requires degree <= 255"); + RAFT_EXPECTS(dataset.n_rows() == graph.extent(0), + "FlowANN graph and full-precision dataset row counts differ"); + RAFT_EXPECTS(labels.extent(0) == graph.extent(0), + "FlowANN graph ranking requires one group label per row"); + RAFT_EXPECTS(metric == cuvs::distance::DistanceType::L2Expanded || + metric == cuvs::distance::DistanceType::InnerProduct || + metric == cuvs::distance::DistanceType::CosineExpanded, + "FlowANN grouped graph ranking supports L2, inner product, and cosine metrics"); + + auto first_invalid = std::atomic{std::numeric_limits::max()}; +#ifdef _OPENMP +#pragma omp parallel for schedule(static) +#endif + for (int64_t node = 0; node < graph.extent(0); ++node) { + for (int64_t edge = 0; edge < graph.extent(1); ++edge) { + if (graph(node, edge) >= n_rows) { + auto const position = static_cast(node) * degree + edge; + auto expected = first_invalid.load(std::memory_order_relaxed); + while (position < expected && + !first_invalid.compare_exchange_weak( + expected, position, std::memory_order_relaxed, std::memory_order_relaxed)) {} + } + } + } + auto const invalid = first_invalid.load(std::memory_order_relaxed); + RAFT_EXPECTS(invalid == std::numeric_limits::max(), + "FlowANN graph contains an out-of-range neighbor at row %llu, column %llu", + static_cast(invalid / degree), + static_cast(invalid % degree)); + + constexpr auto target_staging_bytes = std::uint64_t{256} << 20; + auto const bytes_per_staged_row = + static_cast(degree + 1) * dim * sizeof(value_type); + auto batch_rows = host_dataset + ? std::max(1, target_staging_bytes / bytes_per_staged_row) + : std::uint64_t{1} << 20; + batch_rows = std::min(batch_rows, n_rows); + RAFT_EXPECTS(batch_rows > 0 && batch_rows <= std::numeric_limits::max(), + "FlowANN could not derive a valid graph-ranking batch size"); + + auto edge_order = raft::make_host_matrix(graph.extent(0), graph.extent(1)); + auto device_input = raft::make_device_matrix( + res, static_cast(batch_rows), graph.extent(1)); + auto device_output = raft::make_device_matrix( + res, static_cast(batch_rows), graph.extent(1)); + auto device_labels = raft::make_device_vector(res, graph.extent(0)); + + std::optional> device_staging; + std::optional> host_staging; + if constexpr (host_dataset) { + RAFT_EXPECTS(batch_rows <= std::numeric_limits::max() / + (static_cast(degree) + 1) / dim, + "FlowANN graph-ranking staging size overflows size_t"); + auto const elements = static_cast(batch_rows) * (degree + 1) * dim; + host_staging.emplace(raft::make_pinned_vector(res, elements)); + device_staging.emplace(raft::make_device_matrix( + res, static_cast(batch_rows) * (degree + 1), dim)); + } + + constexpr auto warps_per_block = std::uint32_t{4}; + constexpr auto block_size = warps_per_block * raft::WarpSize; + auto const shared_bytes = + static_cast(warps_per_block) * degree * (sizeof(float) + sizeof(std::uint8_t)); + auto device_id = 0; + RAFT_CUDA_TRY(cudaGetDevice(&device_id)); + auto shared_memory_limit = 0; + RAFT_CUDA_TRY( + cudaDeviceGetAttribute(&shared_memory_limit, cudaDevAttrMaxSharedMemoryPerBlock, device_id)); + RAFT_EXPECTS(shared_bytes <= static_cast(shared_memory_limit), + "FlowANN graph-ranking kernel requires %zu bytes of shared memory, but the device " + "limit is %d bytes", + shared_bytes, + shared_memory_limit); + + auto const stream = raft::resource::get_cuda_stream(res); + auto const raw_stream = cuda_stream_handle(stream); + auto const dataset_rows = dataset.view(); + auto const dataset_stride = dataset.stride(); + raft::copy(device_labels.data_handle(), labels.data_handle(), labels.extent(0), stream); + for (std::uint64_t offset = 0; offset < n_rows; offset += batch_rows) { + auto const count = + static_cast(std::min(batch_rows, n_rows - offset)); + raft::copy(device_input.data_handle(), + graph.data_handle() + offset * degree, + static_cast(count) * degree, + stream); + + auto const blocks = raft::div_rounding_up_safe(count, warps_per_block); + if constexpr (host_dataset) { + auto* staging = host_staging->data_handle(); +#ifdef _OPENMP +#pragma omp parallel for schedule(static) +#endif + for (std::uint32_t local = 0; local < count; ++local) { + auto const global = offset + local; + auto* output = staging + static_cast(local) * (degree + 1) * dim; + std::copy_n(dataset_rows.data_handle() + global * dataset_stride, dim, output); + auto const* node_graph = graph.data_handle() + global * degree; + for (std::uint32_t edge = 0; edge < degree; ++edge) { + std::copy_n(dataset_rows.data_handle() + + static_cast(node_graph[edge]) * dataset_stride, + dim, + output + static_cast(edge + 1) * dim); + } + } + raft::copy(device_staging->data_handle(), + host_staging->data_handle(), + static_cast(count) * (degree + 1) * dim, + stream); + rank_graph_rows_kernel + <<>>(nullptr, + 0, + device_staging->data_handle(), + device_input.data_handle(), + device_labels.data_handle(), + device_output.data_handle(), + offset, + count, + dim, + degree, + metric); + } else { + rank_graph_rows_kernel + <<>>(dataset_rows.data_handle(), + dataset_stride, + nullptr, + device_input.data_handle(), + device_labels.data_handle(), + device_output.data_handle(), + offset, + count, + dim, + degree, + metric); + } + RAFT_CUDA_TRY(cudaPeekAtLastError()); + raft::copy(edge_order.data_handle() + offset * degree, + device_output.data_handle(), + static_cast(count) * degree, + stream); + } + raft::resource::sync_stream(res); + return edge_order; +} + +template +auto convert_uniform_ranked_graph( + raft::resources const& res, + index_params const& params, + raft::host_matrix_view graph, + DatasetViewT const& dataset) -> uniform_conversion_result +{ + using value_type = + std::remove_const_t().view())::value_type>; + constexpr bool host_dataset = cuvs::neighbors::is_host_dataset_view_v; + + auto const n_rows = static_cast(graph.extent(0)); + auto const degree = static_cast(graph.extent(1)); + auto const dim = dataset.dim(); + RAFT_EXPECTS(n_rows > 0 && degree > 0 && dim > 0, + "FlowANN integrated build requires a non-empty graph and dataset"); + RAFT_EXPECTS(dataset.n_rows() == graph.extent(0), + "FlowANN final graph and full-precision dataset row counts differ"); + RAFT_EXPECTS(params.node_per_cacheline > 0 && params.node_per_cacheline <= 32, + "FlowANN uniform build requires node_per_cacheline in [1, 32]"); + RAFT_EXPECTS(params.cagra_params.metric == cuvs::distance::DistanceType::L2Expanded || + params.cagra_params.metric == cuvs::distance::DistanceType::InnerProduct || + params.cagra_params.metric == cuvs::distance::DistanceType::CosineExpanded || + params.cagra_params.metric == cuvs::distance::DistanceType::L1 || + params.cagra_params.metric == cuvs::distance::DistanceType::BitwiseHamming, + "FlowANN uniform build does not support the requested CAGRA metric"); + + auto first_invalid = std::atomic{std::numeric_limits::max()}; +#ifdef _OPENMP +#pragma omp parallel for schedule(static) +#endif + for (int64_t node = 0; node < graph.extent(0); ++node) { + for (int64_t edge = 0; edge < graph.extent(1); ++edge) { + if (graph(node, edge) >= n_rows) { + auto const position = static_cast(node) * degree + edge; + auto expected = first_invalid.load(std::memory_order_relaxed); + while (position < expected && + !first_invalid.compare_exchange_weak( + expected, position, std::memory_order_relaxed, std::memory_order_relaxed)) {} + } + } + } + auto const invalid = first_invalid.load(std::memory_order_relaxed); + RAFT_EXPECTS(invalid == std::numeric_limits::max(), + "FlowANN final graph contains an out-of-range neighbor at row %llu, column %llu", + static_cast(invalid / degree), + static_cast(invalid % degree)); + + auto const layout = plan_uniform_layout( + n_rows, degree, params.node_per_cacheline, params.device_graph_budget_bytes); + auto inner_graph = layout.resident_degree == 0 + ? raft::make_device_matrix(res, 0, 0) + : raft::make_device_matrix( + res, + raft::div_rounding_up_safe(static_cast(n_rows), + params.node_per_cacheline), + layout.inner_row_bytes); + auto cross_graph = raft::make_host_matrix( + static_cast(n_rows), static_cast(layout.cross_degree) + 1); + + if (layout.resident_degree == 0) { +#ifdef _OPENMP +#pragma omp parallel for schedule(static) +#endif + for (int64_t node = 0; node < graph.extent(0); ++node) { + cross_graph(node, 0) = degree; + std::copy_n(graph.data_handle() + static_cast(node) * degree, + degree, + cross_graph.data_handle() + + static_cast(node) * (static_cast(degree) + 1) + 1); + } + return {std::move(inner_graph), std::move(cross_graph), layout}; + } + + constexpr auto target_staging_bytes = std::uint64_t{256} << 20; + auto bytes_per_staged_row = static_cast(degree + 1) * dim * sizeof(value_type); + auto batch_rows = host_dataset + ? std::max(params.node_per_cacheline, + target_staging_bytes / bytes_per_staged_row) + : std::uint64_t{1} << 20; + batch_rows -= batch_rows % params.node_per_cacheline; + batch_rows = + std::max(params.node_per_cacheline, std::min(batch_rows, n_rows)); + if (batch_rows < n_rows && batch_rows % params.node_per_cacheline != 0) { + batch_rows -= batch_rows % params.node_per_cacheline; + } + RAFT_EXPECTS(batch_rows > 0 && batch_rows <= std::numeric_limits::max(), + "FlowANN could not derive a valid rank batch size"); + + auto device_graph = raft::make_device_matrix( + res, static_cast(batch_rows), degree); + auto device_cross = raft::make_device_matrix( + res, static_cast(batch_rows), static_cast(layout.cross_degree) + 1); + + std::optional> device_staging; + std::optional> host_staging; + if constexpr (host_dataset) { + RAFT_EXPECTS(batch_rows <= std::numeric_limits::max() / + (static_cast(degree) + 1) / dim, + "FlowANN host rank staging size overflows size_t"); + auto const elements = static_cast(batch_rows) * (degree + 1) * dim; + host_staging.emplace(raft::make_pinned_vector(res, elements)); + device_staging.emplace(raft::make_device_matrix( + res, static_cast(batch_rows) * (degree + 1), dim)); + } + + auto const stream = raft::resource::get_cuda_stream(res); + auto const raw_stream = cuda_stream_handle(stream); + if (inner_graph.size() > 0) { + RAFT_CUDA_TRY(cudaMemsetAsync( + inner_graph.data_handle(), 0, inner_graph.size() * sizeof(std::uint8_t), raw_stream)); + } + auto const block_size = static_cast(params.node_per_cacheline) * raft::WarpSize; + auto shared_bytes = static_cast(params.node_per_cacheline) * degree * + (sizeof(float) + sizeof(std::uint8_t)); + shared_bytes = (shared_bytes + alignof(graph_index_type) - 1) & ~(alignof(graph_index_type) - 1); + shared_bytes += static_cast(params.node_per_cacheline) * layout.resident_degree * + sizeof(graph_index_type); + auto device_id = 0; + RAFT_CUDA_TRY(cudaGetDevice(&device_id)); + auto shared_memory_limit = 0; + RAFT_CUDA_TRY( + cudaDeviceGetAttribute(&shared_memory_limit, cudaDevAttrMaxSharedMemoryPerBlock, device_id)); + RAFT_EXPECTS(shared_bytes <= static_cast(shared_memory_limit), + "FlowANN rank kernel requires %zu bytes of shared memory, but the device limit is " + "%d bytes", + shared_bytes, + shared_memory_limit); + + auto const dataset_rows = dataset.view(); + auto const dataset_stride = dataset.stride(); + for (std::uint64_t offset = 0; offset < n_rows; offset += batch_rows) { + auto const count = + static_cast(std::min(batch_rows, n_rows - offset)); + raft::copy(device_graph.data_handle(), + graph.data_handle() + offset * degree, + static_cast(count) * degree, + stream); + + if constexpr (host_dataset) { + auto* staging = host_staging->data_handle(); +#ifdef _OPENMP +#pragma omp parallel for schedule(static) +#endif + for (std::uint32_t local = 0; local < count; ++local) { + auto const global = offset + local; + auto* output = staging + static_cast(local) * (degree + 1) * dim; + std::copy_n(dataset_rows.data_handle() + global * dataset_stride, dim, output); + auto const* node_graph = graph.data_handle() + global * degree; + for (std::uint32_t edge = 0; edge < degree; ++edge) { + std::copy_n(dataset_rows.data_handle() + + static_cast(node_graph[edge]) * dataset_stride, + dim, + output + static_cast(edge + 1) * dim); + } + } + raft::copy(device_staging->data_handle(), + host_staging->data_handle(), + static_cast(count) * (degree + 1) * dim, + stream); + auto const blocks = + raft::div_rounding_up_safe(count, params.node_per_cacheline); + rank_and_split_kernel + <<>>(nullptr, + 0, + device_staging->data_handle(), + device_graph.data_handle(), + offset, + count, + dim, + degree, + layout.resident_degree, + layout.cross_degree, + params.node_per_cacheline, + layout.n_bits, + layout.inner_row_bytes, + inner_graph.data_handle(), + device_cross.data_handle(), + params.cagra_params.metric); + } else { + auto const blocks = + raft::div_rounding_up_safe(count, params.node_per_cacheline); + rank_and_split_kernel + <<>>(dataset_rows.data_handle(), + dataset_stride, + nullptr, + device_graph.data_handle(), + offset, + count, + dim, + degree, + layout.resident_degree, + layout.cross_degree, + params.node_per_cacheline, + layout.n_bits, + layout.inner_row_bytes, + inner_graph.data_handle(), + device_cross.data_handle(), + params.cagra_params.metric); + } + RAFT_CUDA_TRY(cudaPeekAtLastError()); + raft::copy( + cross_graph.data_handle() + offset * (static_cast(layout.cross_degree) + 1), + device_cross.data_handle(), + static_cast(count) * (layout.cross_degree + 1), + stream); + raft::resource::sync_stream(res); + } + + return {std::move(inner_graph), std::move(cross_graph), layout}; +} + +template +auto gather_grouping_rows(raft::resources const& res, + DatasetViewT const& dataset, + std::vector const& ids) + -> raft::device_matrix +{ + auto const dim = dataset.dim(); + auto const stream = raft::resource::get_cuda_stream(res); + auto result = raft::make_device_matrix(res, ids.size(), dim); + if constexpr (cuvs::neighbors::is_host_dataset_view_v) { + auto staging = raft::make_pinned_vector(res, ids.size() * dim); +#ifdef _OPENMP +#pragma omp parallel for schedule(static) +#endif + for (std::size_t i = 0; i < ids.size(); ++i) { + for (int64_t col = 0; col < dim; ++col) { + staging(i * dim + col) = + static_cast(dataset.view().data_handle()[ids[i] * dataset.stride() + col]); + } + } + raft::copy(result.data_handle(), staging.data_handle(), result.size(), stream); + raft::resource::sync_stream(res); + } else { + auto device_ids = raft::make_device_vector(res, ids.size()); + raft::copy(device_ids.data_handle(), ids.data(), ids.size(), stream); + auto const blocks = + static_cast(std::min((result.size() + 255) / 256, 65535)); + gather_float_rows_kernel<<>>( + result.data_handle(), + dataset.view().data_handle(), + device_ids.data_handle(), + 0, + ids.size(), + dim, + dataset.stride()); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + raft::resource::sync_stream(res); + } + return result; +} + +// Compute nearest-child labels beside the distance matrix, before copying it to the host. +// This avoids another full CPU pass over N * fanout distances solely for argmin/validation. +__global__ void label_group_distances(float* costs, + std::uint8_t* labels, + std::uint64_t rows, + std::uint32_t branches, + bool inner_product, + int* invalid) +{ + for (std::uint64_t row = std::uint64_t{blockIdx.x} * blockDim.x + threadIdx.x; row < rows; + row += std::uint64_t{blockDim.x} * gridDim.x) { + auto* values = costs + row * branches; + std::uint8_t best = 0; + bool finite = true; + for (std::uint32_t group = 0; group < branches; ++group) { + auto value = values[group]; + if (inner_product) { + value = -value; + values[group] = value; + } + finite = finite && isfinite(value); + if (value < values[best]) { best = group; } + } + labels[row] = best; + if (!finite) { atomicExch(invalid, 1); } + } +} + +template +auto generate_group_labels(raft::resources const& res, + index_params const& index_config, + host_graph_view graph, + DatasetViewT const& dataset) -> grouping_result +{ + auto const& params = index_config.grouping; + auto const metric = index_config.cagra_params.metric; + RAFT_EXPECTS(graph.extent(0) == dataset.n_rows() && graph.extent(1) > 0 && graph.extent(1) <= 255, + "FlowANN grouping graph must match dataset rows and have degree in [1, 255]"); + RAFT_EXPECTS(dataset.n_rows() > 0 && static_cast(dataset.n_rows()) <= + std::numeric_limits::max(), + "FlowANN grouping dataset row count must fit uint32 graph IDs"); + RAFT_EXPECTS(dataset.dim() > 0 && params.kmeans_n_iters > 0, + "FlowANN grouping requires positive dimension and training iterations"); + RAFT_EXPECTS(metric == cuvs::distance::DistanceType::L2Expanded || + metric == cuvs::distance::DistanceType::InnerProduct || + metric == cuvs::distance::DistanceType::CosineExpanded, + "FlowANN grouping requires L2Expanded, InnerProduct, or CosineExpanded"); + auto const n_rows = static_cast(dataset.n_rows()); + auto const dim = dataset.dim(); + auto const plan = plan_grouping(n_rows, params.n_bits, params.n_groups, params.balance_tolerance); + RAFT_EXPECTS(plan.n_groups <= max_partitions, + "FlowANN group metadata supports at most %u groups", + max_partitions); + auto labels = raft::make_host_vector(n_rows); + if (plan.n_groups == 1) { + std::fill_n(labels.data_handle(), n_rows, std::uint32_t{0}); + return {std::move(labels), plan.n_bits, plan.n_groups}; + } + auto const stream = raft::resource::get_cuda_stream(res); + std::vector members(plan.n_groups > 16 ? n_rows : 0), scratch(members.size()); + std::iota(members.begin(), members.end(), std::uint32_t{0}); + auto source_id = [&](std::uint64_t position) { + return members.empty() ? static_cast(position) : members[position]; + }; + double training_seconds = 0, assignment_seconds = 0, rebalance_seconds = 0; + auto elapsed = [](auto begin) { + return std::chrono::duration(std::chrono::steady_clock::now() - begin).count(); + }; + auto partition = [&](auto&& self, + std::uint64_t begin, + std::uint64_t count, + std::uint32_t first_leaf, + std::uint32_t leaves) -> void { + RAFT_EXPECTS(count >= leaves * plan.lower && count <= leaves * plan.upper, + "Subtree violates descendant capacity bounds"); + if (leaves == 1) { + for (std::uint64_t i = begin; i < begin + count; ++i) { + labels(members[i]) = first_leaf; + } + return; + } + auto const children = grouping_children(leaves); + auto const branches = children.size(); + std::vector sizes; + std::vector starts(branches + 1, begin); + { + auto stage_begin = std::chrono::steady_clock::now(); + // Share the automatic sample budget across subtrees instead of multiplying a full + // budget by the number of nodes. Keep a minimum sample per centroid for deep trees. + auto const automatic_rows = std::max( + 256ull * branches, + std::min(10'000ull * branches, + (160'000ull * leaves + plan.n_groups - 1) / plan.n_groups)); + auto const training_rows = std::min( + count, params.training_rows == 0 ? automatic_rows : params.training_rows); + RAFT_EXPECTS(training_rows >= branches, + "Grouping training sample is smaller than node fanout"); + std::vector sample_ids(training_rows); + for (std::uint64_t i = 0; i < training_rows; ++i) { + sample_ids[i] = source_id(begin + i * count / training_rows); + } + auto sample = gather_grouping_rows(res, dataset, sample_ids); + auto centroids = raft::make_device_matrix(res, branches, dim); + cuvs::cluster::kmeans::balanced_params kmeans_params; + kmeans_params.n_iters = params.kmeans_n_iters; + kmeans_params.metric = metric; + kmeans_params.balance_lower_tolerance = 1.0f - params.balance_tolerance; + kmeans_params.balance_upper_tolerance = 1.0f + params.balance_tolerance; + cuvs::cluster::kmeans::fit( + res, kmeans_params, raft::make_const_mdspan(sample.view()), centroids.view()); + raft::resource::sync_stream(res); + training_seconds += elapsed(stage_begin); + stage_begin = std::chrono::steady_clock::now(); + // Every distance is written by the GPU copy; avoid an extra full-matrix zero fill. + auto costs = std::unique_ptr(new float[count * branches]); + std::vector node_labels(count); + auto const batch_rows = std::min( + count, + params.assignment_batch_rows == 0 + ? std::max(1, (256ull << 20) / (dim * sizeof(float))) + : params.assignment_batch_rows); + auto distances = raft::make_device_matrix(res, batch_rows, branches); + auto batch_labels = raft::make_device_vector(res, batch_rows); + auto invalid = raft::make_device_vector(res, 1); + RAFT_CUDA_TRY( + cudaMemsetAsync(invalid.data_handle(), 0, sizeof(int), cuda_stream_handle(stream))); + auto batch = raft::make_device_matrix(res, batch_rows, dim); + std::optional> staging; + std::optional> device_ids; + std::vector ids; + if constexpr (cuvs::neighbors::is_host_dataset_view_v) { + staging.emplace(raft::make_pinned_vector(res, batch_rows * dim)); + } else { + device_ids.emplace(raft::make_device_vector(res, batch_rows)); + ids.resize(batch_rows); + } + for (std::uint64_t offset = 0; offset < count; offset += batch_rows) { + auto const size = std::min(batch_rows, count - offset); + if constexpr (cuvs::neighbors::is_host_dataset_view_v) { +#ifdef _OPENMP +#pragma omp parallel for schedule(static) +#endif + for (std::uint64_t row = 0; row < size; ++row) { + auto const source = source_id(begin + offset + row); + for (int64_t col = 0; col < dim; ++col) { + (*staging)(row* dim + col) = static_cast( + dataset.view().data_handle()[std::uint64_t{source} * dataset.stride() + col]); + } + } + raft::copy(batch.data_handle(), staging->data_handle(), size * dim, stream); + } else { + for (std::uint64_t row = 0; row < size; ++row) { + ids[row] = source_id(begin + offset + row); + } + raft::copy(device_ids->data_handle(), ids.data(), size, stream); + auto const blocks = + static_cast(std::min((size * dim + 255) / 256, 65535)); + gather_float_rows_kernel<<>>( + batch.data_handle(), + dataset.view().data_handle(), + device_ids->data_handle(), + 0, + size, + dim, + dataset.stride()); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + } + auto batch_view = + raft::make_device_matrix_view(batch.data_handle(), size, dim); + auto distance_view = + raft::make_device_matrix_view(distances.data_handle(), size, branches); + cuvs::distance::pairwise_distance( + res, batch_view, raft::make_const_mdspan(centroids.view()), distance_view, metric); + auto const blocks = + static_cast(std::min((size + 255) / 256, 65535)); + label_group_distances<<>>( + distances.data_handle(), + batch_labels.data_handle(), + size, + branches, + metric == cuvs::distance::DistanceType::InnerProduct, + invalid.data_handle()); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + raft::copy(node_labels.data() + offset, batch_labels.data_handle(), size, stream); + raft::copy( + costs.get() + offset * branches, distances.data_handle(), size * branches, stream); + raft::resource::sync_stream(res); + } + int invalid_host = 0; + raft::copy(&invalid_host, invalid.data_handle(), 1, stream); + raft::resource::sync_stream(res); + RAFT_EXPECTS(invalid_host == 0, "Grouping distances must be finite"); + assignment_seconds += elapsed(stage_begin); + stage_begin = std::chrono::steady_clock::now(); + std::vector lower(branches), upper(branches); + for (std::size_t child = 0; child < branches; ++child) { + lower[child] = children[child] * plan.lower; + upper[child] = children[child] * plan.upper; + } + sizes = rebalance_group_distances(costs.get(), node_labels, lower, upper); + if (leaves <= 16) { + // The children are final leaves. Write labels directly, avoiding a full permutation + // and another pass over all IDs for the common single-level case. +#ifdef _OPENMP +#pragma omp parallel for schedule(static) +#endif + for (std::uint64_t row = 0; row < count; ++row) { + labels(source_id(begin + row)) = first_leaf + node_labels[row]; + } + rebalance_seconds += elapsed(stage_begin); + return; + } + for (std::size_t child = 0; child < branches; ++child) { + starts[child + 1] = starts[child] + sizes[child]; + } + auto next = starts; + for (std::uint64_t i = 0; i < count; ++i) { + scratch[next[node_labels[i]]++] = members[begin + i]; + } + std::copy_n(scratch.data() + begin, count, members.data() + begin); + // Do not retain per-row distances at ancestor levels while processing a child. + costs.reset(); + std::vector().swap(node_labels); + rebalance_seconds += elapsed(stage_begin); + } + for (std::size_t child = 0; child < branches; ++child) { + self(self, starts[child], sizes[child], first_leaf, children[child]); + first_leaf += children[child]; + } + }; + partition(partition, 0, n_rows, 0, plan.n_groups); + RAFT_LOG_INFO( + "FlowANN hierarchical grouping: bits=%u groups=%u max_fanout=16 " + "training=%.6f assignment=%.6f rebalance=%.6f seconds", + plan.n_bits, + plan.n_groups, + training_seconds, + assignment_seconds, + rebalance_seconds); + return {std::move(labels), plan.n_bits, plan.n_groups}; +} + +template +auto generate_medoid_seeds(raft::resources const& res, + index_params const& params, + DatasetViewT const& dataset) + -> raft::host_vector +{ + using source_type = + std::remove_const_t().view())::value_type>; + constexpr bool host_dataset = cuvs::neighbors::is_host_dataset_view_v; + + RAFT_EXPECTS(params.num_seeds > 0 && params.num_seeds <= dataset.n_rows(), + "FlowANN seed count must be in [1, dataset size]"); + RAFT_EXPECTS(params.cagra_params.metric == cuvs::distance::DistanceType::L2Expanded || + params.cagra_params.metric == cuvs::distance::DistanceType::InnerProduct || + params.cagra_params.metric == cuvs::distance::DistanceType::CosineExpanded, + "FlowANN balanced-kmeans seeds require L2Expanded, InnerProduct, or CosineExpanded"); + + auto const n_rows = static_cast(dataset.n_rows()); + auto const dim = dataset.dim(); + auto sample_rows = + params.seed_training_rows == 0 + ? std::min( + n_rows, + std::max(params.num_seeds, std::uint64_t{10'000} * params.num_seeds)) + : std::min(n_rows, params.seed_training_rows); + RAFT_EXPECTS(sample_rows >= params.num_seeds, + "FlowANN seed training sample must contain at least num_seeds rows"); + RAFT_EXPECTS( + dim > 0 && sample_rows <= static_cast(std::numeric_limits::max()) / dim, + "FlowANN seed training sample size exceeds int64"); + + auto sample_ids_host = raft::make_host_vector(sample_rows); + auto const rotation = params.seed % n_rows; + for (std::uint64_t row = 0; row < sample_rows; ++row) { + sample_ids_host(row) = + (rotation + (static_cast(row) * n_rows) / sample_rows) % n_rows; + } + + auto sample = raft::make_device_matrix(res, sample_rows, dim); + auto const rows = dataset.view(); + auto const stream = raft::resource::get_cuda_stream(res); + auto const raw_stream = cuda_stream_handle(stream); + if constexpr (host_dataset) { + auto sample_host = + raft::make_pinned_vector(res, static_cast(sample_rows * dim)); +#ifdef _OPENMP +#pragma omp parallel for schedule(static) +#endif + for (std::uint64_t row = 0; row < sample_rows; ++row) { + auto const* input = rows.data_handle() + sample_ids_host(row) * dataset.stride(); + auto* output = sample_host.data_handle() + row * dim; + for (std::uint32_t column = 0; column < dim; ++column) { + output[column] = static_cast(input[column]); + } + } + raft::copy(sample.data_handle(), + sample_host.data_handle(), + static_cast(sample_rows) * dim, + stream); + } else { + auto sample_ids_device = raft::make_device_vector(res, sample_rows); + raft::copy(sample_ids_device.data_handle(), sample_ids_host.data_handle(), sample_rows, stream); + constexpr auto threads = std::uint32_t{256}; + auto const blocks = static_cast( + std::min((sample_rows * dim + threads - 1) / threads, std::uint64_t{65535})); + gather_float_rows_kernel<<>>(sample.data_handle(), + rows.data_handle(), + sample_ids_device.data_handle(), + 0, + sample_rows, + dim, + dataset.stride()); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + } + + auto centroids = + raft::make_device_matrix(res, static_cast(params.num_seeds), dim); + cuvs::cluster::kmeans::balanced_params kmeans_params; + kmeans_params.n_iters = 20; + kmeans_params.metric = params.cagra_params.metric; + cuvs::cluster::kmeans::fit( + res, kmeans_params, raft::make_const_mdspan(sample.view()), centroids.view()); + raft::resource::sync_stream(res); + + constexpr auto target_batch_bytes = std::uint64_t{256} << 20; + auto batch_rows = std::max( + params.num_seeds, target_batch_bytes / (static_cast(dim) * sizeof(float))); + batch_rows = std::min(batch_rows, n_rows); + auto const batch_capacity = std::min(n_rows, batch_rows + params.num_seeds); + RAFT_EXPECTS( + batch_capacity <= static_cast(std::numeric_limits::max()) / dim, + "FlowANN seed scan batch size exceeds int64"); + auto batch = raft::make_device_matrix(res, batch_capacity, dim); + auto nearest = + raft::make_device_matrix(res, params.num_seeds, params.num_seeds); + auto nearest_distances = + raft::make_device_matrix(res, params.num_seeds, params.num_seeds); + auto nearest_host = raft::make_host_matrix(params.num_seeds, params.num_seeds); + auto distances_host = raft::make_host_matrix(params.num_seeds, params.num_seeds); + std::optional> batch_host; + if constexpr (host_dataset) { + batch_host.emplace( + raft::make_pinned_vector(res, static_cast(batch_capacity * dim))); + } + + using candidate = std::pair; + auto candidates = std::vector>(params.num_seeds); + for (auto& list : candidates) { + list.reserve(params.num_seeds * 2); + } + cuvs::neighbors::brute_force::index_params brute_params; + brute_params.metric = params.cagra_params.metric; + cuvs::neighbors::brute_force::search_params search_params; + for (std::uint64_t offset = 0; offset < n_rows;) { + auto count = std::min(batch_rows, n_rows - offset); + auto const remainder = n_rows - offset - count; + if (remainder > 0 && remainder < params.num_seeds) { count += remainder; } + + if constexpr (host_dataset) { +#ifdef _OPENMP +#pragma omp parallel for schedule(static) +#endif + for (std::uint64_t row = 0; row < count; ++row) { + auto const* input = rows.data_handle() + (offset + row) * dataset.stride(); + auto* output = batch_host->data_handle() + row * dim; + for (std::uint32_t column = 0; column < dim; ++column) { + output[column] = static_cast(input[column]); + } + } + raft::copy(batch.data_handle(), + batch_host->data_handle(), + static_cast(count) * dim, + stream); + } else { + constexpr auto threads = std::uint32_t{256}; + auto const blocks = static_cast( + std::min((count * dim + threads - 1) / threads, std::uint64_t{65535})); + gather_float_rows_kernel<<>>( + batch.data_handle(), rows.data_handle(), nullptr, offset, count, dim, dataset.stride()); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + } + + auto batch_view = raft::make_device_matrix_view( + batch.data_handle(), static_cast(count), dim); + auto brute_index = cuvs::neighbors::brute_force::build(res, brute_params, batch_view); + cuvs::neighbors::brute_force::search(res, + search_params, + brute_index, + raft::make_const_mdspan(centroids.view()), + nearest.view(), + nearest_distances.view()); + raft::copy(nearest_host.data_handle(), nearest.data_handle(), nearest.size(), stream); + raft::copy(distances_host.data_handle(), + nearest_distances.data_handle(), + nearest_distances.size(), + stream); + raft::resource::sync_stream(res); + + for (std::uint32_t center = 0; center < params.num_seeds; ++center) { + auto& list = candidates[center]; + for (std::uint32_t rank = 0; rank < params.num_seeds; ++rank) { + list.emplace_back(distances_host(center, rank), + static_cast(offset + nearest_host(center, rank))); + } + std::sort(list.begin(), list.end()); + if (list.size() > params.num_seeds) { list.resize(params.num_seeds); } + } + offset += count; + } + + struct assignment_candidate { + float distance; + graph_index_type row; + std::uint32_t center; + }; + auto assignments = std::vector(); + assignments.reserve(static_cast(params.num_seeds) * params.num_seeds); + for (std::uint32_t center = 0; center < params.num_seeds; ++center) { + for (auto const& [distance, row] : candidates[center]) { + assignments.push_back({distance, row, center}); + } + } + std::sort(assignments.begin(), assignments.end(), [](auto const& lhs, auto const& rhs) { + if (lhs.distance != rhs.distance) { return lhs.distance < rhs.distance; } + if (lhs.row != rhs.row) { return lhs.row < rhs.row; } + return lhs.center < rhs.center; + }); + + auto result = raft::make_host_vector(params.num_seeds); + auto assigned = std::vector(params.num_seeds, false); + auto used = std::set(); + auto assigned_count = std::uint32_t{0}; + for (auto const& candidate : assignments) { + if (assigned[candidate.center] || used.contains(candidate.row)) { continue; } + result(candidate.center) = candidate.row; + assigned[candidate.center] = true; + used.insert(candidate.row); + if (++assigned_count == params.num_seeds) { break; } + } + RAFT_EXPECTS(assigned_count == params.num_seeds, + "FlowANN could not select a unique medoid for every seed centroid"); + return result; +} + +void validate_build_params(build_params const& params, std::uint32_t graph_degree) +{ + RAFT_EXPECTS( + params.node_per_cacheline > 0 && params.node_per_cacheline <= max_nodes_per_packed_row, + "FlowANN node_per_cacheline must be in [1, %u]", + static_cast(max_nodes_per_packed_row)); + RAFT_EXPECTS(params.n_bits > 0 && params.n_bits <= 32, "FlowANN n_bits must be in [1, 32]"); + RAFT_EXPECTS(graph_degree > 0, "FlowANN cannot build from an empty CAGRA graph"); + RAFT_EXPECTS(graph_degree <= std::numeric_limits::max(), + "FlowANN graph degree must fit in the packed uint8 neighbor count"); + if (params.inner_graph_row_bytes != 0) { + RAFT_EXPECTS(params.inner_graph_row_bytes >= params.node_per_cacheline, + "FlowANN inner_graph_row_bytes is too small for its count header"); + } +} + +auto copy_graph_to_host( + raft::resources const& res, + raft::device_matrix_view graph) + -> raft::host_matrix +{ + auto result = raft::make_host_matrix(graph.extent(0), graph.extent(1)); + raft::copy( + result.data_handle(), graph.data_handle(), graph.size(), raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + return result; +} + +auto validate_and_count_partitions( + raft::host_vector_view partition_labels, std::uint16_t n_bits) + -> std::vector +{ + RAFT_EXPECTS(partition_labels.extent(0) > 0, "FlowANN partition labels must not be empty"); + + auto max_label = std::uint32_t{0}; + for (int64_t node = 0; node < partition_labels.extent(0); ++node) { + max_label = std::max(max_label, partition_labels(node)); + } + RAFT_EXPECTS(max_label < max_partitions, + "FlowANN supports at most %u compact partition identifiers", + max_partitions); + + auto partition_sizes = std::vector(static_cast(max_label) + 1, 0); + for (int64_t node = 0; node < partition_labels.extent(0); ++node) { + ++partition_sizes[partition_labels(node)]; + } + + auto const max_partition_size = + n_bits == 32 ? (std::uint64_t{1} << 32) : (std::uint64_t{1} << n_bits); + for (std::size_t partition = 0; partition < partition_sizes.size(); ++partition) { + RAFT_EXPECTS(partition_sizes[partition] > 0, + "FlowANN partition identifiers must be compact; partition %zu is empty", + partition); + RAFT_EXPECTS(partition_sizes[partition] <= max_partition_size, + "FlowANN partition %zu contains too many nodes for %u-bit local IDs", + partition, + static_cast(n_bits)); + } + return partition_sizes; +} + +auto count_local_neighbors(host_graph_view graph, + raft::host_vector_view partition_labels) + -> std::pair, std::uint64_t> +{ + auto const n_rows = graph.extent(0); + auto counts = std::vector(static_cast(n_rows)); + auto first_invalid = std::atomic{std::numeric_limits::max()}; + auto total_local = std::uint64_t{0}; + +#ifdef _OPENMP +#pragma omp parallel for schedule(static) reduction(+ : total_local) +#endif + for (int64_t node = 0; node < n_rows; ++node) { + auto local_count = std::uint32_t{0}; + for (int64_t edge = 0; edge < graph.extent(1); ++edge) { + auto const neighbor = graph(node, edge); + if (neighbor >= static_cast(n_rows)) { + auto const position = static_cast(node) * graph.extent(1) + edge; + auto expected = first_invalid.load(std::memory_order_relaxed); + while (position < expected && + !first_invalid.compare_exchange_weak( + expected, position, std::memory_order_relaxed, std::memory_order_relaxed)) {} + continue; + } + local_count += partition_labels(neighbor) == partition_labels(node); + } + counts[static_cast(node)] = static_cast(local_count); + total_local += local_count; + } + + auto const invalid = first_invalid.load(std::memory_order_relaxed); + RAFT_EXPECTS(invalid == std::numeric_limits::max(), + "CAGRA graph contains an out-of-range neighbor at row %llu, column %llu", + static_cast(invalid / graph.extent(1)), + static_cast(invalid % graph.extent(1))); + return {std::move(counts), total_local}; +} + +auto make_permutation(std::vector const& local_counts, + raft::host_vector_view partition_labels, + std::vector const& partition_sizes, + std::uint16_t node_per_cacheline) + -> std::pair, std::vector> +{ + auto groups = std::vector>(partition_sizes.size()); + for (std::size_t partition = 0; partition < groups.size(); ++partition) { + groups[partition].reserve(static_cast(partition_sizes[partition])); + } + for (int64_t node = 0; node < partition_labels.extent(0); ++node) { + groups[partition_labels(node)].push_back(static_cast(node)); + } + +#ifdef _OPENMP +#pragma omp parallel for schedule(dynamic) +#endif + for (int64_t partition = 0; partition < static_cast(groups.size()); ++partition) { + auto& group = groups[static_cast(partition)]; + std::sort(group.begin(), group.end(), [&](auto lhs, auto rhs) { + auto const lhs_count = local_counts[lhs]; + auto const rhs_count = local_counts[rhs]; + return lhs_count != rhs_count ? lhs_count > rhs_count : lhs < rhs; + }); + } + + auto const n_rows = static_cast(partition_labels.extent(0)); + auto old_to_new = std::vector(n_rows); + auto new_to_old = std::vector(); + new_to_old.reserve(n_rows); + + for (auto& group : groups) { + if (node_per_cacheline == 2 && new_to_old.size() % 2 != 0 && group.size() > 1) { + auto const middle = group.begin() + static_cast(group.size() / 2); + std::rotate(group.begin(), middle, middle + 1); + } + + std::size_t left = 0; + std::size_t right = group.size(); + bool take_left = true; + while (left < right) { + auto const old_node = take_left ? group[left++] : group[--right]; + auto const new_node = static_cast(new_to_old.size()); + new_to_old.push_back(old_node); + old_to_new[old_node] = new_node; + take_left = !take_left; + } + } + + RAFT_EXPECTS(new_to_old.size() == n_rows, "FlowANN node permutation is incomplete"); + return {std::move(old_to_new), std::move(new_to_old)}; +} + +auto compute_inner_row_bytes(build_params const& params, + std::uint64_t n_rows, + std::uint64_t total_local_neighbors) -> std::uint32_t +{ + if (params.inner_graph_row_bytes != 0) { return params.inner_graph_row_bytes; } + + auto const packed_rows = (n_rows + params.node_per_cacheline - 1) / params.node_per_cacheline; + auto const header_bits = static_cast(params.node_per_cacheline) * 8; + auto const total_bits = static_cast(header_bits) * packed_rows + + static_cast(total_local_neighbors) * params.n_bits; + auto const denominator = static_cast(packed_rows) * 8; + auto const row_bytes = (total_bits + denominator - 1) / denominator; + RAFT_EXPECTS(row_bytes <= std::numeric_limits::max(), + "FlowANN inner graph row length exceeds uint32"); + return std::max(params.node_per_cacheline, static_cast(row_bytes)); +} + +void set_bits(std::uint8_t* destination, + std::uint64_t bit_offset, + std::uint32_t value, + std::uint16_t n_bits) +{ + auto const byte_offset = bit_offset / 8; + auto const shift = static_cast(bit_offset % 8); + auto const bytes = static_cast((shift + n_bits + 7) / 8); + auto current = std::uint64_t{0}; + std::memcpy(¤t, destination + byte_offset, bytes); + auto const mask = n_bits == 32 ? std::uint64_t{0xffffffffu} : ((std::uint64_t{1} << n_bits) - 1); + current &= ~(mask << shift); + current |= (static_cast(value) & mask) << shift; + std::memcpy(destination + byte_offset, ¤t, bytes); +} + +auto get_bits(std::uint8_t const* source, std::uint64_t bit_offset, std::uint16_t n_bits) + -> std::uint32_t +{ + auto const byte_offset = bit_offset / 8; + auto const shift = static_cast(bit_offset % 8); + auto const bytes = static_cast((shift + n_bits + 7) / 8); + auto current = std::uint64_t{0}; + std::memcpy(¤t, source + byte_offset, bytes); + auto const mask = n_bits == 32 ? std::uint64_t{0xffffffffu} : ((std::uint64_t{1} << n_bits) - 1); + return static_cast((current >> shift) & mask); +} + +void assign_local_quotas(std::array& assigned, + std::array const& available, + std::uint16_t group_size, + std::uint32_t row_capacity, + std::uint16_t node_per_cacheline) +{ + auto remaining = row_capacity; + auto const base = row_capacity / node_per_cacheline; + for (std::uint16_t i = 0; i < group_size; ++i) { + assigned[i] = static_cast(std::min(base, available[i])); + remaining -= assigned[i]; + } + + while (remaining > 0) { + auto selected = group_size; + for (std::uint16_t i = 0; i < group_size; ++i) { + if (assigned[i] >= available[i]) { continue; } + if (selected == group_size || assigned[i] < assigned[selected]) { selected = i; } + } + if (selected == group_size) { break; } + ++assigned[selected]; + --remaining; + } +} + +__global__ void make_group_sort_keys_kernel(std::uint64_t* keys, + graph_index_type* values, + std::uint8_t* local_counts, + std::uint64_t* total_local, + graph_index_type const* graph, + std::uint32_t const* labels, + std::uint64_t n_rows, + std::uint32_t graph_degree) +{ + auto node = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + auto stride = static_cast(blockDim.x) * gridDim.x; + for (; node < n_rows; node += stride) { + auto const label = labels[node]; + auto count = std::uint32_t{0}; + auto const* row = graph + node * graph_degree; + for (std::uint32_t edge = 0; edge < graph_degree; ++edge) { + auto const neighbor = row[edge]; + if (neighbor < n_rows) { count += labels[neighbor] == label; } + } + local_counts[node] = static_cast(count); + atomicAdd(reinterpret_cast(total_local), + static_cast(count)); + keys[node] = (static_cast(label) << 40) | + (static_cast(std::uint8_t{255} - count) << 32) | node; + values[node] = static_cast(node); + } +} + +__global__ void make_interleaved_permutation_kernel(graph_index_type* new_to_old, + graph_index_type* old_to_new, + std::uint64_t const* sorted_keys, + graph_index_type const* sorted_nodes, + std::uint32_t const* group_offsets, + std::uint32_t const* group_sizes, + std::uint64_t n_rows, + std::uint16_t node_per_cacheline) +{ + auto new_node = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + auto stride = static_cast(blockDim.x) * gridDim.x; + for (; new_node < n_rows; new_node += stride) { + auto const group = static_cast(sorted_keys[new_node] >> 40); + auto const begin = static_cast(group_offsets[group]); + auto const size = static_cast(group_sizes[group]); + auto const local = new_node - begin; + auto source_local = (local & 1u) == 0 ? local / 2 : size - 1 - local / 2; + + // Preserve the reference rearrangement's group-boundary alignment. If a group begins in the + // middle of a packed row, rotate a few median-degree nodes to the front before alternating + // high- and low-degree nodes. This avoids pairing two high-degree nodes at a group boundary. + auto const misaligned = begin % node_per_cacheline; + auto const middle = size / 2; + auto const rotated = min(misaligned, size - middle); + if (source_local < rotated) { + source_local += middle; + } else if (source_local < middle + rotated) { + source_local -= rotated; + } + auto const old_node = sorted_nodes[begin + source_local]; + new_to_old[new_node] = old_node; + old_to_new[old_node] = static_cast(new_node); + } +} + +__global__ void pack_grouped_graph_kernel(std::uint8_t* inner_graph, + graph_index_type* cross_graph, + graph_index_type const* graph, + std::uint8_t const* edge_order, + std::uint32_t const* labels, + std::uint8_t const* local_counts, + graph_index_type const* new_to_old, + graph_index_type const* old_to_new, + std::uint32_t const* group_offsets, + std::uint64_t n_rows, + std::uint32_t graph_degree, + std::uint16_t node_per_cacheline, + std::uint16_t n_bits, + std::uint32_t inner_row_bytes, + std::uint32_t row_capacity) +{ + if (threadIdx.x != 0) { return; } + auto const packed_row = static_cast(blockIdx.x); + auto const first_node = packed_row * node_per_cacheline; + if (first_node >= n_rows) { return; } + auto const group_size = static_cast( + min(static_cast(node_per_cacheline), n_rows - first_node)); + std::uint16_t available[max_nodes_per_packed_row]{}; + std::uint16_t assigned[max_nodes_per_packed_row]{}; + for (std::uint16_t col = 0; col < group_size; ++col) { + available[col] = local_counts[new_to_old[first_node + col]]; + } + + auto remaining = row_capacity; + auto const base = row_capacity / node_per_cacheline; + for (std::uint16_t col = 0; col < group_size; ++col) { + assigned[col] = + static_cast(min(base, static_cast(available[col]))); + remaining -= assigned[col]; + } + while (remaining > 0) { + auto selected = group_size; + for (std::uint16_t col = 0; col < group_size; ++col) { + if (assigned[col] >= available[col]) { continue; } + if (selected == group_size || assigned[col] < assigned[selected]) { selected = col; } + } + if (selected == group_size) { break; } + ++assigned[selected]; + --remaining; + } + + auto* inner_row = inner_graph == nullptr ? nullptr : inner_graph + packed_row * inner_row_bytes; + auto bit_offset = static_cast(node_per_cacheline) * 8; + for (std::uint16_t col = 0; col < group_size; ++col) { + auto const new_node = first_node + col; + auto const old_node = new_to_old[new_node]; + auto const group = labels[old_node]; + auto const* old_row = graph + static_cast(old_node) * graph_degree; + auto* cross_row = cross_graph + new_node * (static_cast(graph_degree) + 1); + auto local_position = std::uint32_t{0}; + auto cross_position = std::uint32_t{0}; + if (inner_row != nullptr) { inner_row[col] = static_cast(assigned[col]); } + + for (std::uint32_t edge = 0; edge < graph_degree; ++edge) { + auto const old_neighbor = old_row[edge]; + auto const new_neighbor = old_to_new[old_neighbor]; + auto const local_rank = + edge_order[static_cast(old_node) * graph_degree + edge]; + if (labels[old_neighbor] == group && local_rank < assigned[col]) { + auto const local_offset = new_neighbor - group_offsets[group]; + set_packed_bits(inner_row, + bit_offset + static_cast(local_rank) * n_bits, + local_offset, + n_bits); + ++local_position; + } else { + cross_row[++cross_position] = new_neighbor; + } + } + cross_row[0] = cross_position; + for (auto edge = cross_position + 1; edge <= graph_degree; ++edge) { + cross_row[edge] = 0; + } + bit_offset += static_cast(assigned[col]) * n_bits; + } +} + +void verify_conversion( + host_graph_view graph, + std::uint8_t const* edge_order, + raft::host_vector_view partition_labels, + build_params const& params, + raft::host_matrix_view inner_graph, + raft::host_matrix_view cross_graph, + std::vector const& old_to_new, + std::vector const& new_to_old, + std::vector const& subgraph_offsets) +{ + auto errors = std::uint64_t{0}; + +#ifdef _OPENMP +#pragma omp parallel for schedule(static) reduction(+ : errors) +#endif + for (int64_t new_node = 0; new_node < graph.extent(0); ++new_node) { + auto const packed_row = new_node / params.node_per_cacheline; + auto const packed_col = new_node % params.node_per_cacheline; + auto const* inner_row = + inner_graph.size() == 0 + ? nullptr + : inner_graph.data_handle() + static_cast(packed_row) * inner_graph.extent(1); + auto const inner_count = + inner_row == nullptr ? 0u : static_cast(inner_row[packed_col]); + auto inner_bit_offset = static_cast(params.node_per_cacheline) * 8; + for (int64_t col = 0; inner_row != nullptr && col < packed_col; ++col) { + inner_bit_offset += static_cast(inner_row[col]) * params.n_bits; + } + + auto const old_node = new_to_old[static_cast(new_node)]; + auto const partition = partition_labels(old_node); + auto const* cross_row = + cross_graph.data_handle() + static_cast(new_node) * cross_graph.extent(1); + auto const cross_count = cross_row[0]; + auto local_position = std::uint32_t{0}; + auto cross_position = std::uint32_t{0}; + + for (int64_t edge = 0; edge < graph.extent(1); ++edge) { + auto const old_neighbor = graph(old_node, edge); + auto const new_neighbor = old_to_new[old_neighbor]; + auto local_rank = std::uint32_t{255}; + if (partition_labels(old_neighbor) == partition) { + if (edge_order != nullptr) { + local_rank = edge_order[static_cast(old_node) * graph.extent(1) + edge]; + } else { + local_rank = 0; + for (int64_t candidate = 0; candidate < edge; ++candidate) { + local_rank += partition_labels(graph(old_node, candidate)) == partition; + } + } + } + if (partition_labels(old_neighbor) == partition && local_rank < inner_count) { + auto const local_offset = + get_bits(inner_row, + inner_bit_offset + static_cast(local_rank) * params.n_bits, + params.n_bits); + auto const decoded = static_cast(subgraph_offsets[partition]) + local_offset; + errors += decoded != new_neighbor; + ++local_position; + } else { + errors += cross_position >= cross_count || cross_row[cross_position + 1] != new_neighbor; + ++cross_position; + } + } + errors += local_position != inner_count || cross_position != cross_count; + } + + RAFT_EXPECTS(errors == 0, + "FlowANN split verification found %llu mismatched rows or edges", + static_cast(errors)); +} + +auto convert_graph(host_graph_view graph, + raft::host_vector_view partition_labels, + build_params const& params) -> conversion_result +{ + auto const n_rows = graph.extent(0); + RAFT_EXPECTS(graph.extent(1) > 0 && static_cast(graph.extent(1)) <= + std::numeric_limits::max(), + "FlowANN graph degree must be in [1, 255]"); + auto const graph_degree = static_cast(graph.extent(1)); + validate_build_params(params, graph_degree); + RAFT_EXPECTS(n_rows > 0 && + static_cast(n_rows) <= std::numeric_limits::max(), + "FlowANN graph size must fit uint32"); + RAFT_EXPECTS(partition_labels.extent(0) == n_rows, + "FlowANN requires one partition label per CAGRA graph row"); + + auto partition_sizes = validate_and_count_partitions(partition_labels, params.n_bits); + auto [local_counts, total_local] = count_local_neighbors(graph, partition_labels); + auto [old_to_new, new_to_old] = + make_permutation(local_counts, partition_labels, partition_sizes, params.node_per_cacheline); + + auto subgraph_offsets = std::vector(partition_sizes.size()); + auto offset = std::uint64_t{0}; + for (std::size_t partition = 0; partition < partition_sizes.size(); ++partition) { + subgraph_offsets[partition] = static_cast(offset); + offset += partition_sizes[partition]; + } + RAFT_EXPECTS(offset == static_cast(n_rows), + "FlowANN partition sizes do not cover the graph"); + + auto subgraph_ids = std::vector(static_cast(n_rows)); + for (int64_t new_node = 0; new_node < n_rows; ++new_node) { + subgraph_ids[static_cast(new_node)] = partition_labels(new_to_old[new_node]); + } + + auto const inner_row_bytes = + compute_inner_row_bytes(params, static_cast(n_rows), total_local); + auto const header_bits = static_cast(params.node_per_cacheline) * 8; + RAFT_EXPECTS(static_cast(inner_row_bytes) * 8 >= header_bits, + "FlowANN inner graph row cannot hold its count header"); + auto const encoded_capacity = + (static_cast(inner_row_bytes) * 8 - header_bits) / params.n_bits; + auto const row_capacity = static_cast(std::min( + encoded_capacity, static_cast(params.node_per_cacheline) * graph_degree)); + + auto const packed_rows = raft::div_rounding_up_safe(n_rows, params.node_per_cacheline); + auto inner_graph = raft::make_host_matrix(packed_rows, inner_row_bytes); + auto cross_graph = raft::make_host_matrix(n_rows, graph_degree + 1); + std::fill_n(inner_graph.data_handle(), inner_graph.size(), std::uint8_t{0}); + +#ifdef _OPENMP +#pragma omp parallel for schedule(static) +#endif + for (int64_t packed_row = 0; packed_row < packed_rows; ++packed_row) { + auto const first_node = packed_row * params.node_per_cacheline; + auto const group_size = + static_cast(std::min(params.node_per_cacheline, n_rows - first_node)); + auto available = std::array{}; + auto assigned = std::array{}; + for (std::uint16_t col = 0; col < group_size; ++col) { + available[col] = local_counts[new_to_old[static_cast(first_node + col)]]; + } + assign_local_quotas(assigned, available, group_size, row_capacity, params.node_per_cacheline); + + auto* inner_row = + inner_graph.data_handle() + static_cast(packed_row) * inner_graph.extent(1); + auto bit_offset = header_bits; + for (std::uint16_t col = 0; col < group_size; ++col) { + auto const new_node = first_node + col; + auto const old_node = new_to_old[static_cast(new_node)]; + auto const partition = partition_labels(old_node); + auto* cross_row = + cross_graph.data_handle() + static_cast(new_node) * cross_graph.extent(1); + auto local_position = std::uint32_t{0}; + auto cross_position = std::uint32_t{0}; + + inner_row[col] = static_cast(assigned[col]); + for (std::uint32_t edge = 0; edge < graph_degree; ++edge) { + auto const old_neighbor = graph(old_node, edge); + auto const new_neighbor = old_to_new[old_neighbor]; + if (partition_labels(old_neighbor) == partition && local_position < assigned[col]) { + auto const local_offset = + static_cast(new_neighbor) - subgraph_offsets[partition]; + set_bits(inner_row, + bit_offset + static_cast(local_position) * params.n_bits, + static_cast(local_offset), + params.n_bits); + ++local_position; + } else { + cross_row[++cross_position] = new_neighbor; + } + } + cross_row[0] = cross_position; + std::fill(cross_row + cross_position + 1, + cross_row + static_cast(graph_degree) + 1, + graph_index_type{0}); + bit_offset += static_cast(assigned[col]) * params.n_bits; + } + } + + conversion_result result{std::move(inner_graph), + std::move(cross_graph), + std::move(old_to_new), + std::move(new_to_old), + std::move(subgraph_offsets), + std::move(subgraph_ids)}; + if (params.validate) { + verify_conversion(graph, + nullptr, + partition_labels, + params, + raft::make_const_mdspan(result.inner_graph.view()), + raft::make_const_mdspan(result.cross_graph.view()), + result.old_to_new, + result.new_to_old, + result.subgraph_offsets); + } + return result; +} + +auto convert_graph_on_device(raft::resources const& res, + host_graph_view graph, + host_edge_order_view edge_order, + raft::host_vector_view partition_labels, + build_params const& params, + std::size_t device_graph_budget_bytes) -> device_conversion_result +{ + auto const n_rows = static_cast(graph.extent(0)); + RAFT_EXPECTS(n_rows > 0 && n_rows <= std::numeric_limits::max(), + "FlowANN graph size must fit uint32"); + RAFT_EXPECTS(graph.extent(1) > 0 && static_cast(graph.extent(1)) <= + std::numeric_limits::max(), + "FlowANN graph degree must be in [1, 255]"); + auto const graph_degree = static_cast(graph.extent(1)); + validate_build_params(params, graph_degree); + RAFT_EXPECTS(partition_labels.extent(0) == graph.extent(0), + "FlowANN requires one partition label per graph row"); + RAFT_EXPECTS(edge_order.extent(0) == graph.extent(0) && edge_order.extent(1) == graph.extent(1), + "FlowANN edge-order matrix must match the graph shape"); + + auto const partition_sizes = validate_and_count_partitions(partition_labels, params.n_bits); + auto subgraph_offsets = std::vector(partition_sizes.size()); + auto group_sizes = std::vector(partition_sizes.size()); + auto offset = std::uint64_t{0}; + for (std::size_t group = 0; group < partition_sizes.size(); ++group) { + subgraph_offsets[group] = static_cast(offset); + group_sizes[group] = static_cast(partition_sizes[group]); + offset += partition_sizes[group]; + } + RAFT_EXPECTS(offset == n_rows, "FlowANN partition sizes do not cover the graph"); + + auto const stream = raft::resource::get_cuda_stream(res); + auto const raw_stream = cuda_stream_handle(stream); + auto device_graph = + raft::make_device_matrix(res, graph.extent(0), graph.extent(1)); + auto device_edge_order = + raft::make_device_matrix(res, graph.extent(0), graph.extent(1)); + auto device_labels = raft::make_device_vector(res, graph.extent(0)); + raft::copy(device_graph.data_handle(), graph.data_handle(), graph.size(), stream); + raft::copy(device_edge_order.data_handle(), edge_order.data_handle(), edge_order.size(), stream); + raft::copy(device_labels.data_handle(), + partition_labels.data_handle(), + partition_labels.extent(0), + stream); + + auto keys_in = raft::make_device_vector(res, graph.extent(0)); + auto keys_out = raft::make_device_vector(res, graph.extent(0)); + auto nodes_in = raft::make_device_vector(res, graph.extent(0)); + auto nodes_sorted = raft::make_device_vector(res, graph.extent(0)); + auto local_counts = raft::make_device_vector(res, graph.extent(0)); + auto total_local = raft::make_device_vector(res, 1); + RAFT_CUDA_TRY(cudaMemsetAsync(total_local.data_handle(), 0, sizeof(std::uint64_t), raw_stream)); + constexpr auto threads = std::uint32_t{256}; + auto const blocks = + static_cast(std::min((n_rows + threads - 1) / threads, 65535)); + make_group_sort_keys_kernel<<>>(keys_in.data_handle(), + nodes_in.data_handle(), + local_counts.data_handle(), + total_local.data_handle(), + device_graph.data_handle(), + device_labels.data_handle(), + n_rows, + graph_degree); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + + auto sort_workspace_bytes = std::size_t{0}; + RAFT_CUDA_TRY(cub::DeviceRadixSort::SortPairs(nullptr, + sort_workspace_bytes, + keys_in.data_handle(), + keys_out.data_handle(), + nodes_in.data_handle(), + nodes_sorted.data_handle(), + graph.extent(0), + 0, + 56, + raw_stream)); + rmm::device_buffer sort_workspace(sort_workspace_bytes, stream); + RAFT_CUDA_TRY(cub::DeviceRadixSort::SortPairs(sort_workspace.data(), + sort_workspace_bytes, + keys_in.data_handle(), + keys_out.data_handle(), + nodes_in.data_handle(), + nodes_sorted.data_handle(), + graph.extent(0), + 0, + 56, + raw_stream)); + + auto device_offsets = + raft::make_device_vector(res, subgraph_offsets.size()); + auto device_sizes = raft::make_device_vector(res, group_sizes.size()); + auto device_new_to_old = + raft::make_device_vector(res, graph.extent(0)); + auto device_old_to_new = + raft::make_device_vector(res, graph.extent(0)); + raft::copy( + device_offsets.data_handle(), subgraph_offsets.data(), subgraph_offsets.size(), stream); + raft::copy(device_sizes.data_handle(), group_sizes.data(), group_sizes.size(), stream); + make_interleaved_permutation_kernel<<>>( + device_new_to_old.data_handle(), + device_old_to_new.data_handle(), + keys_out.data_handle(), + nodes_sorted.data_handle(), + device_offsets.data_handle(), + device_sizes.data_handle(), + n_rows, + params.node_per_cacheline); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + + auto old_to_new = std::vector(n_rows); + auto new_to_old = std::vector(n_rows); + auto total_local_host = std::uint64_t{0}; + raft::copy(old_to_new.data(), device_old_to_new.data_handle(), n_rows, stream); + raft::copy(new_to_old.data(), device_new_to_old.data_handle(), n_rows, stream); + raft::copy(&total_local_host, total_local.data_handle(), 1, stream); + raft::resource::sync_stream(res); + + auto const packed_rows = raft::div_rounding_up_safe( + n_rows, static_cast(params.node_per_cacheline)); + auto inner_row_bytes = budgeted_grouped_row_bytes( + n_rows, graph_degree, params.node_per_cacheline, params.n_bits, device_graph_budget_bytes); + if (params.inner_graph_row_bytes != 0) { + inner_row_bytes = std::min(inner_row_bytes, params.inner_graph_row_bytes); + } + auto const header_bits = static_cast(params.node_per_cacheline) * 8; + auto row_capacity = std::uint32_t{0}; + if (inner_row_bytes >= params.node_per_cacheline) { + row_capacity = static_cast(std::min( + (static_cast(inner_row_bytes) * 8 - header_bits) / params.n_bits, + static_cast(params.node_per_cacheline) * graph_degree)); + } + if (row_capacity == 0) { inner_row_bytes = 0; } + + auto inner_graph = inner_row_bytes == 0 + ? raft::make_device_matrix(res, 0, 0) + : raft::make_device_matrix( + res, static_cast(packed_rows), inner_row_bytes); + auto device_cross = raft::make_device_matrix( + res, graph.extent(0), static_cast(graph_degree) + 1); + if (inner_graph.size() > 0) { + RAFT_CUDA_TRY(cudaMemsetAsync( + inner_graph.data_handle(), 0, inner_graph.size() * sizeof(std::uint8_t), raw_stream)); + } + RAFT_EXPECTS(packed_rows <= std::numeric_limits::max(), + "FlowANN packed row count exceeds the CUDA grid limit"); + pack_grouped_graph_kernel<<(packed_rows), 32, 0, raw_stream>>>( + inner_graph.size() == 0 ? nullptr : inner_graph.data_handle(), + device_cross.data_handle(), + device_graph.data_handle(), + device_edge_order.data_handle(), + device_labels.data_handle(), + local_counts.data_handle(), + device_new_to_old.data_handle(), + device_old_to_new.data_handle(), + device_offsets.data_handle(), + n_rows, + graph_degree, + params.node_per_cacheline, + params.n_bits, + inner_row_bytes, + row_capacity); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + + auto cross_graph = raft::make_host_matrix( + graph.extent(0), static_cast(graph_degree) + 1); + raft::copy(cross_graph.data_handle(), device_cross.data_handle(), device_cross.size(), stream); + raft::resource::sync_stream(res); + + auto subgraph_ids = std::vector(n_rows); + for (std::uint64_t new_node = 0; new_node < n_rows; ++new_node) { + subgraph_ids[new_node] = partition_labels(new_to_old[new_node]); + } + + if (params.validate) { + auto host_inner = + raft::make_host_matrix(inner_graph.extent(0), inner_graph.extent(1)); + if (inner_graph.size() > 0) { + raft::copy(host_inner.data_handle(), inner_graph.data_handle(), inner_graph.size(), stream); + raft::resource::sync_stream(res); + } + verify_conversion(graph, + edge_order.data_handle(), + partition_labels, + params, + raft::make_const_mdspan(host_inner.view()), + raft::make_const_mdspan(cross_graph.view()), + old_to_new, + new_to_old, + subgraph_offsets); + } + + RAFT_LOG_INFO("FlowANN grouped layout: groups=%zu local_edges=%llu row_bytes=%u budget_bytes=%zu", + partition_sizes.size(), + static_cast(total_local_host), + inner_row_bytes, + device_graph_budget_bytes); + return {std::move(inner_graph), + std::move(cross_graph), + std::move(old_to_new), + std::move(new_to_old), + std::move(subgraph_offsets), + std::move(subgraph_ids), + params.validate}; +} + +template +__global__ void gather_rows_kernel(T* destination, + T const* source, + graph_index_type const* new_to_old, + std::uint64_t elements, + std::uint32_t row_length) +{ + auto offset = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + auto stride = static_cast(blockDim.x) * gridDim.x; + for (; offset < elements; offset += stride) { + auto const new_row = offset / row_length; + auto const column = offset % row_length; + destination[offset] = + source[static_cast(new_to_old[new_row]) * row_length + column]; + } +} + +template +void gather_device_rows(raft::resources const& res, + T* destination, + T const* source, + std::uint64_t n_rows, + std::uint32_t row_length, + std::vector const& new_to_old) +{ + auto device_mapping = + raft::make_device_vector(res, static_cast(n_rows)); + auto const stream = raft::resource::get_cuda_stream(res); + auto const raw_stream = cuda_stream_handle(stream); + raft::copy(device_mapping.data_handle(), new_to_old.data(), n_rows, stream); + + auto const elements = n_rows * row_length; + constexpr auto block_size = std::uint32_t{256}; + auto const grid_size = static_cast( + std::min((elements + block_size - 1) / block_size, 65535)); + gather_rows_kernel<<>>( + destination, source, device_mapping.data_handle(), elements, row_length); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + raft::resource::sync_stream(res); +} + +template +auto reorder_dense_dataset(raft::resources const& res, + cuvs::neighbors::device_padded_dataset_view const& source, + std::vector const& new_to_old) + -> std::unique_ptr> +{ + auto const source_data = source.view(); + auto reordered = + raft::make_device_matrix(res, source_data.extent(0), source_data.extent(1)); + gather_device_rows(res, + reordered.data_handle(), + source_data.data_handle(), + source_data.extent(0), + static_cast(source_data.extent(1)), + new_to_old); + return std::make_unique>(std::move(reordered), + source.dim()); +} + +template +auto make_source_indices(raft::resources const& res, + CagraIndexT const& source, + std::vector const& new_to_old) + -> std::vector +{ + auto result = std::vector(new_to_old.size()); + if (source.source_indices().has_value()) { + auto const source_view = source.source_indices().value(); + RAFT_EXPECTS(source_view.extent(0) == static_cast(new_to_old.size()), + "FlowANN source-index mapping and graph row counts differ"); + auto source_ids = std::vector(new_to_old.size()); + raft::copy(source_ids.data(), + source_view.data_handle(), + source_view.extent(0), + raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + for (std::size_t new_node = 0; new_node < new_to_old.size(); ++new_node) { + result[new_node] = source_ids[new_to_old[new_node]]; + } + } else { + result = new_to_old; + } + return result; +} + +template +void populate_index(raft::resources const& res, + FlowIndexT& index, + conversion_result&& conversion, + std::vector const& source_indices, + std::optional> seeds) +{ + auto const n_rows = conversion.old_to_new.size(); + index.update_graph(res, std::move(conversion.inner_graph), std::move(conversion.cross_graph)); + index.update_source_indices( + res, + raft::make_host_vector_view(source_indices.data(), n_rows)); + index.update_subgraph_layout( + res, + raft::make_host_vector_view(conversion.subgraph_offsets.data(), + conversion.subgraph_offsets.size()), + raft::make_host_vector_view(conversion.subgraph_ids.data(), + conversion.subgraph_ids.size())); + + if (seeds.has_value()) { + RAFT_EXPECTS(seeds->extent(0) > 0, "FlowANN seed list must not be empty"); + auto remapped = std::vector(static_cast(seeds->extent(0))); + for (int64_t seed = 0; seed < seeds->extent(0); ++seed) { + RAFT_EXPECTS((*seeds)(seed) < n_rows, + "FlowANN input seed at position %ld is out of range", + static_cast(seed)); + remapped[static_cast(seed)] = conversion.old_to_new[(*seeds)(seed)]; + } + index.update_seeds(res, + raft::make_host_vector_view( + remapped.data(), remapped.size())); + } + raft::resource::sync_stream(res); +} + +template +void populate_index(raft::resources const& res, + FlowIndexT& index, + device_conversion_result&& conversion, + std::vector const& source_indices, + std::optional> seeds) +{ + auto const n_rows = conversion.old_to_new.size(); + auto const validated = conversion.validated; + index.update_graph( + res, std::move(conversion.inner_graph), std::move(conversion.cross_graph), validated); + index.update_source_indices( + res, + raft::make_host_vector_view(source_indices.data(), n_rows)); + index.update_subgraph_layout( + res, + raft::make_host_vector_view(conversion.subgraph_offsets.data(), + conversion.subgraph_offsets.size()), + raft::make_host_vector_view(conversion.subgraph_ids.data(), + conversion.subgraph_ids.size())); + + if (seeds.has_value()) { + auto remapped = std::vector(static_cast(seeds->extent(0))); + for (int64_t seed = 0; seed < seeds->extent(0); ++seed) { + RAFT_EXPECTS((*seeds)(seed) < n_rows, + "FlowANN input seed at position %ld is out of range", + static_cast(seed)); + remapped[static_cast(seed)] = conversion.old_to_new[(*seeds)(seed)]; + } + index.update_seeds(res, + raft::make_host_vector_view( + remapped.data(), remapped.size())); + } + raft::resource::sync_stream(res); +} + +template +void populate_uniform_index(raft::resources const& res, + FlowIndexT& index, + uniform_conversion_result&& conversion) +{ + auto const n_rows = conversion.cross_graph.extent(0); + auto const layout = conversion.layout; + index.update_graph(res, + std::move(conversion.inner_graph), + std::move(conversion.cross_graph), + layout.resident_degree + layout.cross_degree, + layout.resident_degree); + + auto offsets = raft::make_host_vector(1); + offsets(0) = 0; + auto subgraph_ids = raft::make_host_vector(n_rows); + std::fill_n(subgraph_ids.data_handle(), subgraph_ids.size(), std::uint32_t{0}); + index.update_subgraph_layout( + res, raft::make_const_mdspan(offsets.view()), raft::make_const_mdspan(subgraph_ids.view())); + raft::resource::sync_stream(res); +} + +inline void reject_ace(index_params const& params) +{ + RAFT_EXPECTS(!std::holds_alternative( + params.cagra_params.graph_build_params), + "FlowANN integrated rank/budget build does not support ACE graph construction"); +} + +template +auto convert_cagra_graph(raft::resources const& res, + build_params const& params, + CagraIndexT const& source, + raft::host_vector_view partition_labels) + -> conversion_result +{ + RAFT_EXPECTS(source.graph().extent(0) > 0 && source.graph().extent(1) > 0, + "FlowANN build requires an in-memory CAGRA graph"); + RAFT_EXPECTS(source.dataset().n_rows() == source.graph().extent(0), + "FlowANN build requires an in-memory dataset matching the CAGRA graph"); + auto host_graph = copy_graph_to_host(res, source.graph()); + return convert_graph(raft::make_const_mdspan(host_graph.view()), partition_labels, params); +} + +} // namespace + +} // namespace cuvs::neighbors::cagra::experimental::flowann::detail + +namespace cuvs::neighbors::cagra::experimental::flowann { + +template +auto build_dense_integrated(raft::resources const& res, + index_params const& params, + cuvs::neighbors::device_padded_dataset_view const& dataset) + -> device_padded_index_bundle +{ + detail::reject_ace(params); + + auto const build_begin = std::chrono::steady_clock::now(); + auto graph = cuvs::neighbors::cagra::detail::build_graph( + res, params.cagra_params, dataset); + auto const graph_end = std::chrono::steady_clock::now(); + if (params.grouping.enabled) { + auto grouping = + detail::generate_group_labels(res, params, raft::make_const_mdspan(graph.view()), dataset); + auto& labels = grouping.labels; + auto const grouping_end = std::chrono::steady_clock::now(); + auto edge_order = detail::rank_graph_edge_order(res, + params.cagra_params.metric, + raft::make_const_mdspan(graph.view()), + raft::make_const_mdspan(labels.view()), + dataset); + auto const rank_end = std::chrono::steady_clock::now(); + build_params grouped_params; + grouped_params.node_per_cacheline = params.node_per_cacheline; + grouped_params.n_bits = grouping.n_bits; + grouped_params.validate = params.grouping.validate; + auto conversion = detail::convert_graph_on_device(res, + raft::make_const_mdspan(graph.view()), + raft::make_const_mdspan(edge_order.view()), + raft::make_const_mdspan(labels.view()), + grouped_params, + params.device_graph_budget_bytes); + auto const split_end = std::chrono::steady_clock::now(); + graph = raft::make_host_matrix(0, 0); + std::optional> seeds; + if (params.num_seeds > 0) { + seeds.emplace(detail::generate_medoid_seeds(res, params, dataset)); + } + auto const seeds_end = std::chrono::steady_clock::now(); + + auto source_indices = conversion.new_to_old; + auto dataset_owner = detail::reorder_dense_dataset(res, dataset, conversion.new_to_old); + flowann::device_padded_index index(res, + params.cagra_params.metric, + dataset_owner->as_dataset_view(), + params.node_per_cacheline, + grouping.n_bits); + std::optional> seed_view; + if (seeds.has_value()) { seed_view.emplace(raft::make_const_mdspan(seeds->view())); } + detail::populate_index(res, index, std::move(conversion), source_indices, seed_view); + auto const finish = std::chrono::steady_clock::now(); + RAFT_LOG_INFO( + "FlowANN build stages: cagra=%.3f s, grouping=%.3f s, rank=%.3f s, rearrange=%.3f s, " + "seeds=%.3f s, dataset=%.3f s", + std::chrono::duration(graph_end - build_begin).count(), + std::chrono::duration(grouping_end - graph_end).count(), + std::chrono::duration(rank_end - grouping_end).count(), + std::chrono::duration(split_end - rank_end).count(), + std::chrono::duration(seeds_end - split_end).count(), + std::chrono::duration(finish - seeds_end).count()); + return {std::move(dataset_owner), std::move(index)}; + } + auto conversion = detail::convert_uniform_ranked_graph( + res, params, raft::make_const_mdspan(graph.view()), dataset); + auto const split_end = std::chrono::steady_clock::now(); + graph = raft::make_host_matrix(0, 0); + std::optional> seeds; + if (params.num_seeds > 0) { seeds.emplace(detail::generate_medoid_seeds(res, params, dataset)); } + auto const seeds_end = std::chrono::steady_clock::now(); + + auto source_rows = dataset.view(); + auto owned_rows = + raft::make_device_matrix(res, source_rows.extent(0), source_rows.extent(1)); + raft::copy(owned_rows.data_handle(), + source_rows.data_handle(), + source_rows.size(), + raft::resource::get_cuda_stream(res)); + auto dataset_owner = std::make_unique>( + std::move(owned_rows), dataset.dim()); + flowann::device_padded_index index(res, + params.cagra_params.metric, + dataset_owner->as_dataset_view(), + params.node_per_cacheline, + conversion.layout.n_bits); + detail::populate_uniform_index(res, index, std::move(conversion)); + if (seeds.has_value()) { index.update_seeds(res, raft::make_const_mdspan(seeds->view())); } + auto const finish = std::chrono::steady_clock::now(); + RAFT_LOG_INFO( + "FlowANN build stages: cagra=%.3f s, rank_split=%.3f s, seeds=%.3f s, dataset=%.3f s", + std::chrono::duration(graph_end - build_begin).count(), + std::chrono::duration(split_end - graph_end).count(), + std::chrono::duration(seeds_end - split_end).count(), + std::chrono::duration(finish - seeds_end).count()); + return {std::move(dataset_owner), std::move(index)}; +} + +#define CUVS_FLOWANN_DEFINE_INTEGRATED_BUILD(T) \ + auto build(raft::resources const& res, \ + index_params const& params, \ + cuvs::neighbors::device_padded_dataset_view const& dataset) \ + -> device_padded_index_bundle \ + { \ + return build_dense_integrated(res, params, dataset); \ + } + +CUVS_FLOWANN_DEFINE_INTEGRATED_BUILD(float) +CUVS_FLOWANN_DEFINE_INTEGRATED_BUILD(half) +CUVS_FLOWANN_DEFINE_INTEGRATED_BUILD(std::int8_t) +CUVS_FLOWANN_DEFINE_INTEGRATED_BUILD(std::uint8_t) + +#undef CUVS_FLOWANN_DEFINE_INTEGRATED_BUILD + +template +auto build_device_padded_from_cagra( + raft::resources const& res, + build_params const& params, + cuvs::neighbors::cagra::device_padded_index const& cagra_index, + raft::host_vector_view partition_labels, + std::optional> seeds) + -> device_padded_index_bundle +{ + auto conversion = detail::convert_cagra_graph(res, params, cagra_index, partition_labels); + auto source_indices = detail::make_source_indices(res, cagra_index, conversion.new_to_old); + auto dataset = detail::reorder_dense_dataset(res, cagra_index.dataset(), conversion.new_to_old); + flowann::device_padded_index index(res, + cagra_index.metric(), + dataset->as_dataset_view(), + params.node_per_cacheline, + params.n_bits); + detail::populate_index(res, index, std::move(conversion), source_indices, seeds); + return {std::move(dataset), std::move(index)}; +} + +#define CUVS_FLOWANN_DEFINE_GROUP_GRAPH(T) \ + auto group_graph(raft::resources const& res, \ + index_params const& params, \ + raft::host_matrix_view graph, \ + cuvs::neighbors::host_standard_dataset_view const& dataset) \ + -> grouping_result \ + { \ + return detail::generate_group_labels(res, params, graph, dataset); \ + } +CUVS_FLOWANN_DEFINE_GROUP_GRAPH(float) +CUVS_FLOWANN_DEFINE_GROUP_GRAPH(half) +CUVS_FLOWANN_DEFINE_GROUP_GRAPH(std::int8_t) +CUVS_FLOWANN_DEFINE_GROUP_GRAPH(std::uint8_t) +#undef CUVS_FLOWANN_DEFINE_GROUP_GRAPH + +#define CUVS_FLOWANN_DEFINE_BUILD(T) \ + auto build_from_cagra( \ + raft::resources const& res, \ + build_params const& params, \ + cuvs::neighbors::cagra::device_padded_index const& cagra_index, \ + raft::host_vector_view partition_labels, \ + std::optional> seeds) \ + -> device_padded_index_bundle \ + { \ + return build_device_padded_from_cagra(res, params, cagra_index, partition_labels, seeds); \ + } + +CUVS_FLOWANN_DEFINE_BUILD(float) +CUVS_FLOWANN_DEFINE_BUILD(half) +CUVS_FLOWANN_DEFINE_BUILD(std::int8_t) +CUVS_FLOWANN_DEFINE_BUILD(std::uint8_t) + +#undef CUVS_FLOWANN_DEFINE_BUILD + +} // namespace cuvs::neighbors::cagra::experimental::flowann diff --git a/cpp/src/neighbors/flowann_queue.cpp b/cpp/src/neighbors/flowann_queue.cpp new file mode 100644 index 0000000000..87f212b4d3 --- /dev/null +++ b/cpp/src/neighbors/flowann_queue.cpp @@ -0,0 +1,918 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * FlowANN is described in "Disentangling Graph Dependencies for Efficient + * Billion-Scale GPU Vector Search" (OSDI 2026). + */ + +#include + +#include "detail/flowann/queue.cuh" + +#include +#include +#include + +#ifdef CUVS_FLOWANN_USE_GDRCOPY +#include +#include +#endif + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#if defined(__linux__) +#include +#include +#include +#endif + +namespace cuvs::neighbors::cagra::experimental::flowann { +namespace { + +constexpr auto align_up(std::size_t value, std::size_t alignment) -> std::size_t +{ + return (value + alignment - 1) & ~(alignment - 1); +} + +struct queue_layout { + std::size_t state; + std::size_t rows; + std::size_t command_words; + std::size_t result_sequence; + std::size_t released_sequence; + std::size_t total_size; + + explicit queue_layout(std::uint32_t row_length) + { + constexpr std::size_t alignment = 128; + state = 0; + rows = align_up(sizeof(detail::queue_state), alignment); + + RAFT_EXPECTS(row_length <= std::numeric_limits::max() / detail::queue_capacity / + sizeof(std::uint32_t), + "FlowANN queue row storage size overflows"); + auto const row_bytes = + static_cast(detail::queue_capacity) * row_length * sizeof(std::uint32_t); + command_words = align_up(rows + row_bytes, alignment); + result_sequence = + align_up(command_words + detail::queue_capacity * sizeof(unsigned long long), alignment); + released_sequence = + align_up(result_sequence + detail::queue_capacity * sizeof(unsigned long long), alignment); + total_size = + align_up(released_sequence + detail::queue_capacity * sizeof(unsigned long long), alignment); + } +}; + +template +auto add_bytes(void* base, std::size_t offset) -> T* +{ + return reinterpret_cast(static_cast(base) + offset); +} + +class cuda_allocation { + public: + cuda_allocation() = default; + cuda_allocation(cuda_allocation const&) = delete; + auto operator=(cuda_allocation const&) -> cuda_allocation& = delete; + + ~cuda_allocation() + { + if (allocation_ != nullptr) { RAFT_CUDA_TRY_NO_THROW(cudaFree(allocation_)); } + } + + void allocate(std::size_t size, std::size_t alignment = 1) + { + RAFT_EXPECTS(allocation_ == nullptr, "FlowANN queue allocation is already initialized"); + RAFT_EXPECTS(alignment > 0 && (alignment & (alignment - 1)) == 0, + "FlowANN queue allocation alignment must be a power of two"); + RAFT_EXPECTS(size <= std::numeric_limits::max() - (alignment - 1), + "FlowANN queue allocation size overflows"); + RAFT_CUDA_TRY(cudaMalloc(&allocation_, size + alignment - 1)); + auto const address = reinterpret_cast(allocation_); + ptr_ = reinterpret_cast(align_up(address, alignment)); + } + + [[nodiscard]] auto get() const noexcept -> void* { return ptr_; } + + private: + void* allocation_{}; + void* ptr_{}; +}; + +class cuda_stream { + public: + cuda_stream() = default; + cuda_stream(cuda_stream const&) = delete; + auto operator=(cuda_stream const&) -> cuda_stream& = delete; + + ~cuda_stream() + { + if (stream_ != nullptr) { RAFT_CUDA_TRY_NO_THROW(cudaStreamDestroy(stream_)); } + } + + void create() + { + RAFT_EXPECTS(stream_ == nullptr, "FlowANN queue stream is already initialized"); + RAFT_CUDA_TRY(cudaStreamCreateWithFlags(&stream_, cudaStreamNonBlocking)); + } + + [[nodiscard]] auto get() const noexcept -> cudaStream_t { return stream_; } + + private: + cudaStream_t stream_{}; +}; + +#ifndef CUVS_FLOWANN_USE_GDRCOPY +class cuda_host_buffer { + public: + cuda_host_buffer() = default; + cuda_host_buffer(cuda_host_buffer const&) = delete; + auto operator=(cuda_host_buffer const&) -> cuda_host_buffer& = delete; + + ~cuda_host_buffer() + { + if (ptr_ != nullptr) { RAFT_CUDA_TRY_NO_THROW(cudaFreeHost(ptr_)); } + } + + void allocate(std::size_t size) + { + RAFT_EXPECTS(ptr_ == nullptr, "FlowANN host copy buffer is already initialized"); + RAFT_CUDA_TRY(cudaMallocHost(&ptr_, size)); + size_ = size; + } + + [[nodiscard]] auto get(std::size_t size) const -> void* + { + RAFT_EXPECTS(size <= size_, "FlowANN transfer exceeds the host copy buffer"); + return ptr_; + } + + private: + void* ptr_{}; + std::size_t size_{}; +}; +#endif + +#ifdef CUVS_FLOWANN_USE_GDRCOPY +class gdrcopy_mapping { + public: + gdrcopy_mapping(void* device_ptr, std::size_t size) : size_(size) + { + try { + unsigned int sync_memops = 1; + RAFT_EXPECTS(cuPointerSetAttribute(&sync_memops, + CU_POINTER_ATTRIBUTE_SYNC_MEMOPS, + reinterpret_cast(device_ptr)) == CUDA_SUCCESS, + "Failed to enable synchronous memory operations for a FlowANN GDRCopy queue"); + + handle_ = gdr_open(); + RAFT_EXPECTS(handle_ != nullptr, + "Failed to open GDRCopy; ensure the gdrdrv kernel module is loaded"); + RAFT_EXPECTS( + gdr_pin_buffer( + handle_, reinterpret_cast(device_ptr), size_, 0, 0, &memory_handle_) == 0, + "Failed to pin a FlowANN queue with GDRCopy"); + pinned_ = true; + RAFT_EXPECTS(gdr_map(handle_, memory_handle_, &mapped_base_, size_) == 0, + "Failed to map a FlowANN queue with GDRCopy"); + mapped_ = true; + + gdr_info_t info{}; + RAFT_EXPECTS(gdr_get_info(handle_, memory_handle_, &info) == 0, + "Failed to query a FlowANN GDRCopy mapping"); + auto const device_address = reinterpret_cast(device_ptr); + RAFT_EXPECTS(device_address >= info.va, + "GDRCopy returned an invalid FlowANN mapping address"); + auto const offset = device_address - info.va; + RAFT_EXPECTS(offset <= size_, "GDRCopy returned an invalid FlowANN mapping offset"); + mapped_device_ptr_ = static_cast(mapped_base_) + offset; + } catch (...) { + reset(); + throw; + } + } + + gdrcopy_mapping(gdrcopy_mapping const&) = delete; + auto operator=(gdrcopy_mapping const&) -> gdrcopy_mapping& = delete; + + ~gdrcopy_mapping() { reset(); } + + void read(void* dst, std::size_t offset, std::size_t size) + { + RAFT_EXPECTS(offset <= size_ && size <= size_ - offset, + "FlowANN GDRCopy read exceeds its queue mapping"); + RAFT_EXPECTS(gdr_copy_from_mapping(memory_handle_, dst, mapped_device_ptr_ + offset, size) == 0, + "Failed to read a FlowANN GDRCopy queue"); + } + + void write(std::size_t offset, void const* src, std::size_t size) + { + RAFT_EXPECTS(offset <= size_ && size <= size_ - offset, + "FlowANN GDRCopy write exceeds its queue mapping"); + RAFT_EXPECTS(gdr_copy_to_mapping(memory_handle_, mapped_device_ptr_ + offset, src, size) == 0, + "Failed to write a FlowANN GDRCopy queue"); + } + + private: + void reset() noexcept + { + if (mapped_) { + static_cast(gdr_unmap(handle_, memory_handle_, mapped_base_, size_)); + mapped_ = false; + } + if (pinned_) { + static_cast(gdr_unpin_buffer(handle_, memory_handle_)); + pinned_ = false; + } + if (handle_ != nullptr) { + static_cast(gdr_close(handle_)); + handle_ = nullptr; + } + } + + gdr_t handle_{}; + gdr_mh_t memory_handle_{}; + void* mapped_base_{}; + std::byte* mapped_device_ptr_{}; + std::size_t size_{}; + bool pinned_{}; + bool mapped_{}; +}; +#endif + +#if defined(__linux__) +auto read_topology_value(int cpu, char const* name) -> int +{ + std::ifstream input("/sys/devices/system/cpu/cpu" + std::to_string(cpu) + "/topology/" + name); + int value = -1; + input >> value; + return input ? value : -1; +} + +auto physical_cpus() -> std::vector +{ + cpu_set_t allowed; + CPU_ZERO(&allowed); + if (sched_getaffinity(0, sizeof(allowed), &allowed) != 0) { return {}; } + + std::vector> cores; + std::vector cpus; + for (int cpu = 0; cpu < CPU_SETSIZE; ++cpu) { + if (!CPU_ISSET(cpu, &allowed)) { continue; } + auto package = read_topology_value(cpu, "physical_package_id"); + auto core = read_topology_value(cpu, "core_id"); + if (package < 0 || core < 0) { + package = 0; + core = cpu; + } + auto const identity = std::tuple{package, core}; + if (std::find(cores.begin(), cores.end(), identity) == cores.end()) { + cores.push_back(identity); + cpus.push_back(cpu); + } + } + return cpus; +} + +void bind_to_physical_cpu(std::uint32_t queue_id) +{ + static auto const cpus = physical_cpus(); + if (cpus.empty()) { return; } + auto const cpu = cpus[queue_id % cpus.size()]; + cpu_set_t selected; + CPU_ZERO(&selected); + CPU_SET(cpu, &selected); + auto const error = pthread_setaffinity_np(pthread_self(), sizeof(selected), &selected); + if (error != 0) { + RAFT_LOG_WARN( + "FlowANN queue %u could not bind its poller to CPU %d (error=%d)", queue_id, cpu, error); + } + if (queue_id >= cpus.size()) { + RAFT_LOG_WARN("FlowANN queue %u reuses physical CPU %d because only %zu are available", + queue_id, + cpu, + cpus.size()); + } +} +#else +void bind_to_physical_cpu(std::uint32_t) {} +#endif + +class host_queue { + public: + static constexpr std::uint32_t max_batch_size = 128; + + host_queue(std::uint32_t const* cross_graph, + std::uint64_t row_count, + std::uint32_t row_length, + bool collect_statistics) + : cross_graph_(cross_graph), + row_count_(row_count), + row_length_(row_length), + collect_statistics_(collect_statistics), + layout_(row_length), + empty_commands_(max_batch_size, detail::empty_command), + command_staging_(max_batch_size), + result_sequence_staging_(max_batch_size), + row_staging_(static_cast(max_batch_size) * row_length) + { +#ifdef CUVS_FLOWANN_USE_GDRCOPY + // gdrdrv rejects mappings whose device address is not aligned to its 64 KiB GPU page. + // cudaMalloc can return a suballocation after the graph build, so align within an owning + // allocation instead of relying on cudaMalloc's current address. + storage_.allocate(layout_.total_size, 64u * 1024u); +#else + storage_.allocate(layout_.total_size); +#endif + RAFT_CUDA_TRY(cudaMemset(storage_.get(), 0, layout_.total_size)); + + detail::queue_state state{}; + state.capacity = detail::queue_capacity; + state.row_length = row_length_; + RAFT_CUDA_TRY(cudaMemcpy(add_bytes(storage_.get(), layout_.state), + &state, + sizeof(state), + cudaMemcpyHostToDevice)); + + view_.state = add_bytes(storage_.get(), layout_.state); + view_.rows = add_bytes(storage_.get(), layout_.rows); + view_.command_words = add_bytes(storage_.get(), layout_.command_words); + view_.result_sequence = add_bytes(storage_.get(), layout_.result_sequence); + view_.released_sequence = + add_bytes(storage_.get(), layout_.released_sequence); + +#ifdef CUVS_FLOWANN_USE_GDRCOPY + gdrcopy_.emplace(storage_.get(), layout_.total_size); +#else + copy_stream_.create(); + // Pageable copies may synchronize inside the driver and contend with lazy + // kernel initialization. Allocate once, before starting any pollers, and + // use pinned memory for every transfer on the nonblocking copy stream. + copy_staging_.allocate(std::max(row_staging_.size() * sizeof(std::uint32_t), + command_staging_.size() * sizeof(unsigned long long))); +#endif + } + + host_queue(host_queue const&) = delete; + auto operator=(host_queue const&) -> host_queue& = delete; + + [[nodiscard]] auto view() const noexcept -> detail::queue_view { return view_; } + + [[nodiscard]] auto poll_batch() -> bool + { + if (collect_statistics_) { polls_.fetch_add(1, std::memory_order_relaxed); } + auto const first_slot = static_cast(cpu_tail_) & detail::queue_mask; + // Sparse traffic needs only a short prefix. Read the rest only when the + // whole prefix is ready, retaining full service batches under load. + constexpr std::uint32_t prefix_size = 8; + read_ring(command_staging_.data(), layout_.command_words, first_slot, prefix_size); + if (!detail::command_ready(command_staging_[0])) { + if (collect_statistics_) { empty_polls_.fetch_add(1, std::memory_order_relaxed); } + return false; + } + + std::uint32_t count = 1; + while (count < prefix_size && detail::command_ready(command_staging_[count])) { + ++count; + } + if (count == prefix_size) { + auto const next_slot = (first_slot + prefix_size) & detail::queue_mask; + read_ring(command_staging_.data() + prefix_size, + layout_.command_words, + next_slot, + max_batch_size - prefix_size); + while (count < max_batch_size && detail::command_ready(command_staging_[count])) { + ++count; + } + } + + for (std::uint32_t i = 0; i < count; ++i) { + auto const row_id = static_cast(command_staging_[i]); + RAFT_EXPECTS( + row_id < row_count_, "FlowANN queue received out-of-range cross-graph row %u", row_id); + std::copy_n(cross_graph_ + static_cast(row_id) * row_length_, + row_length_, + row_staging_.data() + static_cast(i) * row_length_); + result_sequence_staging_[i] = cpu_tail_ + i + 1; + } + write_ring(row_staging_.data(), layout_.rows, first_slot, count, row_length_); + // Clearing the request is ordered before result publication. A producer + // cannot reuse this slot until it acquires and releases that result. + write_ring(empty_commands_.data(), layout_.command_words, first_slot, count); + std::atomic_thread_fence(std::memory_order_release); + write_ring(result_sequence_staging_.data(), layout_.result_sequence, first_slot, count); + + cpu_tail_ += count; + if (collect_statistics_) { processed_commands_.fetch_add(count, std::memory_order_relaxed); } + return true; + } + + void publish_progress() + { + // cpu_tail is an observational snapshot, not a GPU synchronization input: + // consumers use result_sequence and producers reclaim capacity via gpu_tail. + // Publish before acknowledging a pause instead of writing on every batch. + write(layout_.state + offsetof(detail::queue_state, cpu_tail), &cpu_tail_, sizeof(cpu_tail_)); + } + + [[nodiscard]] auto statistics() const noexcept -> queue_statistics + { + return {polls_.load(std::memory_order_relaxed), + empty_polls_.load(std::memory_order_relaxed), + processed_commands_.load(std::memory_order_relaxed)}; + } + + private: + template + void read_ring(T* destination, + std::size_t base_offset, + std::uint32_t first_slot, + std::uint32_t count) + { + auto const first_count = std::min(count, detail::queue_capacity - first_slot); + read(destination, + base_offset + static_cast(first_slot) * sizeof(T), + static_cast(first_count) * sizeof(T)); + if (first_count < count) { + read(destination + first_count, + base_offset, + static_cast(count - first_count) * sizeof(T)); + } + } + + template + void write_ring(T const* source, + std::size_t base_offset, + std::uint32_t first_slot, + std::uint32_t count, + std::uint32_t elements_per_slot = 1) + { + auto const first_count = std::min(count, detail::queue_capacity - first_slot); + auto const slot_bytes = static_cast(elements_per_slot) * sizeof(T); + write(base_offset + static_cast(first_slot) * slot_bytes, + source, + static_cast(first_count) * slot_bytes); + if (first_count < count) { + write(base_offset, + source + static_cast(first_count) * elements_per_slot, + static_cast(count - first_count) * slot_bytes); + } + } + + void read(void* dst, std::size_t offset, std::size_t size) + { +#ifdef CUVS_FLOWANN_USE_GDRCOPY + gdrcopy_->read(dst, offset, size); +#else + auto* staging = copy_staging_.get(size); + RAFT_CUDA_TRY(cudaMemcpyAsync(staging, + add_bytes(storage_.get(), offset), + size, + cudaMemcpyDeviceToHost, + copy_stream_.get())); + RAFT_CUDA_TRY(cudaStreamSynchronize(copy_stream_.get())); + std::memcpy(dst, staging, size); +#endif + } + + void write(std::size_t offset, void const* src, std::size_t size) + { +#ifdef CUVS_FLOWANN_USE_GDRCOPY + gdrcopy_->write(offset, src, size); +#else + auto* staging = copy_staging_.get(size); + std::memcpy(staging, src, size); + RAFT_CUDA_TRY(cudaMemcpyAsync(add_bytes(storage_.get(), offset), + staging, + size, + cudaMemcpyHostToDevice, + copy_stream_.get())); + RAFT_CUDA_TRY(cudaStreamSynchronize(copy_stream_.get())); +#endif + } + + std::uint32_t const* cross_graph_; + std::uint64_t row_count_; + std::uint32_t row_length_; + bool collect_statistics_; + queue_layout layout_; + cuda_allocation storage_; + detail::queue_view view_{}; + unsigned long long cpu_tail_{}; + std::atomic polls_{}; + std::atomic empty_polls_{}; + std::atomic processed_commands_{}; + std::vector empty_commands_; + std::vector command_staging_; + std::vector result_sequence_staging_; + std::vector row_staging_; +#ifdef CUVS_FLOWANN_USE_GDRCOPY + std::optional gdrcopy_; +#else + cuda_stream copy_stream_; + cuda_host_buffer copy_staging_; +#endif +}; + +} // namespace + +struct search_context::impl { + explicit impl(raft::host_matrix_view cross_graph, + queue_params params) + : params_(params), + row_count_(static_cast(cross_graph.extent(0))), + row_length_(static_cast(cross_graph.extent(1))) + { + RAFT_CUDA_TRY(cudaGetDevice(&device_id_)); + RAFT_EXPECTS(cross_graph.extent(0) > 0, "FlowANN cross graph must not be empty"); + RAFT_EXPECTS(cross_graph.extent(1) > 0, + "FlowANN cross graph must contain a neighbor-count column"); + RAFT_EXPECTS(params.num_queues > 0, "FlowANN must use at least one command queue"); + RAFT_EXPECTS(cross_graph.extent(1) <= std::numeric_limits::max(), + "FlowANN cross-graph row length exceeds uint32"); + + queues_.reserve(params.num_queues); + std::vector views; + views.reserve(params.num_queues); + for (std::uint32_t i = 0; i < params.num_queues; ++i) { + queues_.push_back( + std::make_unique(cross_graph.data_handle(), + static_cast(cross_graph.extent(0)), + static_cast(cross_graph.extent(1)), + params.collect_statistics)); + views.push_back(queues_.back()->view()); + } + + RAFT_CUDA_TRY(cudaMalloc(&device_views_, views.size() * sizeof(detail::queue_view))); + try { + RAFT_CUDA_TRY(cudaMemcpy(device_views_, + views.data(), + views.size() * sizeof(detail::queue_view), + cudaMemcpyHostToDevice)); + + pollers_.reserve(queues_.size()); + for (std::uint32_t i = 0; i < queues_.size(); ++i) { + pollers_.emplace_back([this, queue = queues_[i].get(), i](std::stop_token stop_token) { + poll(stop_token, *queue, i); + }); + } + + std::unique_lock lock(pause_mutex_); + pause_cv_.wait(lock, [this] { return paused_pollers_ + failed_pollers_ == pollers_.size(); }); + auto error = poller_error_; + lock.unlock(); + if (error != nullptr) { std::rethrow_exception(error); } + } catch (...) { + stop(); + throw; + } + } + + ~impl() { stop(); } + + void throw_if_failed() + { + if (poller_error_ != nullptr) { std::rethrow_exception(poller_error_); } + } + + [[nodiscard]] auto should_run() const noexcept -> bool + { + return manually_resumed_ || active_sessions_ > 0 || active_searches_ > 0; + } + + void update_pollers(std::unique_lock& lock) + { + while (true) { + throw_if_failed(); + if (should_run()) { + if (paused_.load(std::memory_order_acquire)) { + paused_.store(false, std::memory_order_release); + pause_cv_.notify_all(); + } + pause_cv_.wait(lock, [this] { + return !should_run() || paused_pollers_ == 0 || poller_error_ != nullptr; + }); + throw_if_failed(); + if (should_run() && paused_pollers_ == 0) { return; } + } else { + if (!paused_.load(std::memory_order_acquire)) { + paused_.store(true, std::memory_order_release); + pause_cv_.notify_all(); + } + pause_cv_.wait(lock, [this] { + return should_run() || paused_pollers_ + failed_pollers_ == pollers_.size() || + poller_error_ != nullptr; + }); + throw_if_failed(); + if (!should_run() && paused_pollers_ + failed_pollers_ == pollers_.size()) { return; } + } + } + } + + void stop() noexcept + { + for (auto& poller : pollers_) { + poller.request_stop(); + } + { + std::lock_guard lock(pause_mutex_); + paused_.store(false, std::memory_order_release); + } + pause_cv_.notify_all(); + pollers_.clear(); + if (device_views_ != nullptr) { + RAFT_CUDA_TRY_NO_THROW(cudaFree(device_views_)); + device_views_ = nullptr; + } + } + + void pause() + { + std::unique_lock lock(pause_mutex_); + RAFT_EXPECTS(active_sessions_ == 0, "Cannot pause while a FlowANN search session is active"); + manually_resumed_ = false; + ++pending_pause_requests_; + try { + pause_cv_.wait(lock, [this] { return active_searches_ == 0 || poller_error_ != nullptr; }); + throw_if_failed(); + update_pollers(lock); + } catch (...) { + --pending_pause_requests_; + pause_cv_.notify_all(); + throw; + } + --pending_pause_requests_; + pause_cv_.notify_all(); + } + + void resume() + { + std::unique_lock lock(pause_mutex_); + pause_cv_.wait(lock, + [this] { return pending_pause_requests_ == 0 || poller_error_ != nullptr; }); + throw_if_failed(); + manually_resumed_ = true; + update_pollers(lock); + } + + void acquire_search() + { + std::unique_lock lock(pause_mutex_); + pause_cv_.wait(lock, + [this] { return pending_pause_requests_ == 0 || poller_error_ != nullptr; }); + throw_if_failed(); + manually_resumed_ = false; + ++active_searches_; + try { + update_pollers(lock); + } catch (...) { + --active_searches_; + pause_cv_.notify_all(); + throw; + } + } + + void release_search() + { + std::unique_lock lock(pause_mutex_); + RAFT_EXPECTS(active_searches_ > 0, "Cannot release an inactive FlowANN search lease"); + --active_searches_; + if (active_searches_ == 0) { + manually_resumed_ = false; + pause_cv_.notify_all(); + update_pollers(lock); + } + } + + void begin_session() + { + std::unique_lock lock(pause_mutex_); + pause_cv_.wait(lock, + [this] { return pending_pause_requests_ == 0 || poller_error_ != nullptr; }); + throw_if_failed(); + ++active_sessions_; + try { + update_pollers(lock); + } catch (...) { + --active_sessions_; + throw; + } + } + + void end_session() + { + std::unique_lock lock(pause_mutex_); + RAFT_EXPECTS(active_sessions_ > 0, "Cannot end an inactive FlowANN search session"); + --active_sessions_; + update_pollers(lock); + } + + void poll(std::stop_token stop_token, host_queue& queue, std::uint32_t queue_id) noexcept + { + try { + RAFT_CUDA_TRY(cudaSetDevice(device_id_)); + bind_to_physical_cpu(queue_id); + while (!stop_token.stop_requested()) { + if (paused_.load(std::memory_order_acquire)) { + queue.publish_progress(); + std::unique_lock lock(pause_mutex_); + ++paused_pollers_; + pause_cv_.notify_all(); + pause_cv_.wait(lock, [this, &stop_token] { + return stop_token.stop_requested() || !paused_.load(std::memory_order_acquire); + }); + --paused_pollers_; + pause_cv_.notify_all(); + if (stop_token.stop_requested()) { break; } + } + + auto const processed = queue.poll_batch(); + if (!processed) { +#if defined(__x86_64__) || defined(__i386__) + for (std::uint32_t i = 0; i < params_.empty_pause; ++i) { + __asm__ __volatile__("pause"); + } +#else + std::this_thread::yield(); +#endif + } + } + } catch (...) { + std::lock_guard lock(pause_mutex_); + if (poller_error_ == nullptr) { poller_error_ = std::current_exception(); } + ++failed_pollers_; + pause_cv_.notify_all(); + } + } + + [[nodiscard]] auto statistics() const noexcept -> queue_statistics + { + queue_statistics total{}; + for (auto const& queue : queues_) { + auto const value = queue->statistics(); + total.polls += value.polls; + total.empty_polls += value.empty_polls; + total.processed_commands += value.processed_commands; + } + return total; + } + + queue_params params_; + int device_id_{}; + std::uint64_t row_count_; + std::uint32_t row_length_; + std::vector> queues_; + std::vector pollers_; + detail::queue_view* device_views_{}; + std::atomic paused_{true}; + mutable std::mutex pause_mutex_; + std::condition_variable pause_cv_; + std::size_t paused_pollers_{}; + std::size_t failed_pollers_{}; + std::size_t active_searches_{}; + std::size_t active_sessions_{}; + bool manually_resumed_{}; + std::size_t pending_pause_requests_{}; + std::exception_ptr poller_error_{}; +}; + +search_context::search_context( + raft::host_matrix_view cross_graph, + queue_params params) + : impl_(std::make_unique(cross_graph, params)) +{ +} + +search_context::~search_context() = default; + +search_context::search_context(search_context&&) noexcept = default; + +auto search_context::operator=(search_context&&) noexcept -> search_context& = default; + +search_session::search_session(search_context& context) : context_(&context) +{ + RAFT_EXPECTS(context.impl_ != nullptr, "Cannot start a session on a moved-from search context"); + context.impl_->begin_session(); +} + +search_session::~search_session() { reset(); } + +search_session::search_session(search_session&& other) noexcept + : context_(std::exchange(other.context_, nullptr)) +{ +} + +auto search_session::operator=(search_session&& other) noexcept -> search_session& +{ + if (this != &other) { + reset(); + context_ = std::exchange(other.context_, nullptr); + } + return *this; +} + +void search_session::reset() noexcept +{ + if (context_ == nullptr) { return; } + auto* context = std::exchange(context_, nullptr); + try { + context->impl_->end_session(); + } catch (...) { + } +} + +void search_context::pause() +{ + RAFT_EXPECTS(impl_ != nullptr, "Cannot pause a moved-from FlowANN search context"); + impl_->pause(); +} + +void search_context::resume() +{ + RAFT_EXPECTS(impl_ != nullptr, "Cannot resume a moved-from FlowANN search context"); + impl_->resume(); +} + +auto search_context::num_queues() const noexcept -> std::uint32_t +{ + return impl_ == nullptr ? 0 : static_cast(impl_->queues_.size()); +} + +auto search_context::statistics() const noexcept -> queue_statistics +{ + return impl_ == nullptr ? queue_statistics{} : impl_->statistics(); +} + +auto detail::search_context_access::device_queues(search_context const& context) noexcept + -> queue_view* +{ + return context.impl_ == nullptr ? nullptr : context.impl_->device_views_; +} + +auto detail::search_context_access::device_id(search_context const& context) noexcept -> int +{ + return context.impl_ == nullptr ? -1 : context.impl_->device_id_; +} + +auto detail::search_context_access::num_queues(search_context const& context) noexcept + -> std::uint32_t +{ + return context.num_queues(); +} + +auto detail::search_context_access::row_count(search_context const& context) noexcept + -> std::uint64_t +{ + return context.impl_ == nullptr ? 0 : context.impl_->row_count_; +} + +auto detail::search_context_access::row_length(search_context const& context) noexcept + -> std::uint32_t +{ + return context.impl_ == nullptr ? 0 : context.impl_->row_length_; +} + +void detail::search_context_access::acquire_search(search_context& context) +{ + RAFT_EXPECTS(context.impl_ != nullptr, "Cannot search with a moved-from FlowANN search context"); + context.impl_->acquire_search(); +} + +void detail::search_context_access::release_search(search_context& context) +{ + RAFT_EXPECTS(context.impl_ != nullptr, "Cannot search with a moved-from FlowANN search context"); + context.impl_->release_search(); +} + +auto detail::search_context_access::is_paused(search_context const& context) noexcept -> bool +{ + return context.impl_ == nullptr || context.impl_->paused_.load(std::memory_order_acquire); +} + +auto detail::search_context_access::pending_pause_requests(search_context const& context) noexcept + -> std::size_t +{ + if (context.impl_ == nullptr) { return 0; } + std::lock_guard lock(context.impl_->pause_mutex_); + return context.impl_->pending_pause_requests_; +} + +} // namespace cuvs::neighbors::cagra::experimental::flowann diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index 81a499419c..b115ca4a94 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -266,6 +266,39 @@ 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 + ) + if(CUVS_ENABLE_FLOWANN_BUILD) + list(APPEND CUVS_FLOWANN_TEST_SOURCES neighbors/ann_cagra/flowann_build.cu + ) + endif() + 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_build.cu b/cpp/tests/neighbors/ann_cagra/flowann_build.cu new file mode 100644 index 0000000000..0eb9778c59 --- /dev/null +++ b/cpp/tests/neighbors/ann_cagra/flowann_build.cu @@ -0,0 +1,744 @@ +/* + * 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 + +namespace flowann = cuvs::neighbors::cagra::experimental::flowann; + +TEST(FlowannBuild, ComputesMinimalGlobalIdWidth) +{ + EXPECT_EQ(flowann::detail::required_id_bits(1), 1); + EXPECT_EQ(flowann::detail::required_id_bits(2), 1); + EXPECT_EQ(flowann::detail::required_id_bits(3), 2); + EXPECT_EQ(flowann::detail::required_id_bits(87'555'327), 27); + EXPECT_EQ(flowann::detail::required_id_bits(1'000'000'000), 30); +} + +TEST(FlowannBuild, SelectsLargestUniformDegreeWithinBudget) +{ + constexpr std::uint64_t rows = 100; + constexpr std::uint16_t nodes_per_row = 2; + constexpr std::uint32_t degree = 32; + auto const bits = flowann::detail::required_id_bits(rows); + + auto const row_bytes_10 = flowann::detail::packed_row_bytes(nodes_per_row, bits, 10); + auto const budget = (rows / nodes_per_row) * row_bytes_10; + auto const plan = flowann::detail::plan_uniform_layout(rows, degree, nodes_per_row, budget); + + EXPECT_EQ(plan.n_bits, bits); + EXPECT_EQ(plan.resident_degree, 10); + EXPECT_EQ(plan.cross_degree, degree - 10); + EXPECT_EQ(plan.inner_row_bytes, row_bytes_10); + EXPECT_EQ(plan.inner_graph_bytes, budget); +} + +TEST(FlowannBuild, SupportsZeroAndFullResidentDegree) +{ + auto const zero = flowann::detail::plan_uniform_layout(101, 8, 2, 0); + EXPECT_EQ(zero.resident_degree, 0); + EXPECT_EQ(zero.cross_degree, 8); + EXPECT_EQ(zero.inner_row_bytes, 0); + EXPECT_EQ(zero.inner_graph_bytes, 0); + + auto const full_row = flowann::detail::packed_row_bytes(2, 7, 8); + auto const full = flowann::detail::plan_uniform_layout(101, 8, 2, 51 * full_row); + EXPECT_EQ(full.resident_degree, 8); + EXPECT_EQ(full.cross_degree, 0); + EXPECT_EQ(full.inner_graph_bytes, 51 * full_row); +} + +TEST(FlowannBuild, RejectsGraphSizesOutsideUint32IdSpace) +{ + EXPECT_ANY_THROW(flowann::detail::plan_uniform_layout( + std::numeric_limits::max(), 32, 2, std::numeric_limits::max())); +} + +TEST(FlowannBuild, GroupedCapacityUsesBudgetRatherThanAverageLocalDegree) +{ + auto row_bytes = flowann::detail::budgeted_grouped_row_bytes; + EXPECT_EQ(row_bytes(87'555'327, 32, 2, 24, 7'529'758'208), 172); + EXPECT_EQ(row_bytes(101, 32, 2, 24, 51 * 139), 139); + EXPECT_EQ(row_bytes(101, 32, 2, 24, 51 * 172 + 50), 172); + EXPECT_EQ(row_bytes(101, 32, 2, 24, 51 * 194), 194); + EXPECT_EQ(row_bytes(101, 32, 2, 24, 51 * 1000), 194); + EXPECT_EQ(row_bytes(101, 32, 2, 24, 0), 0); + EXPECT_EQ(row_bytes(101, 32, 2, 24, 51 * 4), 0); + EXPECT_EQ(row_bytes(101, 32, 2, 24, 51 * 5), 5); + EXPECT_EQ(row_bytes(101, 8, 2, 6, 51 * 14), 14); +} + +namespace { + +auto read_packed(std::uint8_t const* row, std::uint64_t bit_offset, std::uint16_t bits) + -> std::uint32_t +{ + auto result = std::uint32_t{0}; + for (std::uint16_t bit = 0; bit < bits; ++bit) { + result |= + static_cast((row[(bit_offset + bit) >> 3] >> ((bit_offset + bit) & 7)) & 1u) + << bit; + } + return result; +} + +auto l2(float const* lhs, float const* rhs, std::uint32_t dim) -> float +{ + auto result = 0.0f; + for (std::uint32_t column = 0; column < dim; ++column) { + auto const diff = lhs[column] - rhs[column]; + result += diff * diff; + } + return result; +} + +auto metric_score(float const* lhs, + float const* rhs, + std::uint32_t dim, + cuvs::distance::DistanceType metric) -> float +{ + auto dot = 0.0f; + auto lhs_norm = 0.0f; + auto rhs_norm = 0.0f; + for (std::uint32_t column = 0; column < dim; ++column) { + dot += lhs[column] * rhs[column]; + lhs_norm += lhs[column] * lhs[column]; + rhs_norm += rhs[column] * rhs[column]; + } + if (metric == cuvs::distance::DistanceType::InnerProduct) { return -dot; } + return 1.0f - dot / std::sqrt(lhs_norm * rhs_norm); +} + +} // namespace + +TEST(FlowannBuild, IntegratedBuildKeepsShortestNeighborsResident) +{ + constexpr int64_t rows = 128; + constexpr int64_t dim = 8; + constexpr std::uint32_t full_degree = 8; + constexpr std::uint32_t resident_degree = 4; + constexpr std::uint16_t nodes_per_row = 2; + + raft::resources res; + auto host = raft::make_host_matrix(rows, dim); + for (int64_t row = 0; row < rows; ++row) { + for (int64_t column = 0; column < dim; ++column) { + host(row, column) = + static_cast((row * 17 + column * 13) % 101) + static_cast(row) * 0.001f; + } + } + auto dataset = cuvs::neighbors::make_device_padded_dataset(res, host.view()); + + flowann::index_params params; + params.grouping.enabled = false; + params.cagra_params.metric = cuvs::distance::DistanceType::L2Expanded; + params.cagra_params.graph_degree = full_degree; + params.cagra_params.intermediate_graph_degree = 16; + params.cagra_params.graph_build_params = + cuvs::neighbors::cagra::graph_build_params::nn_descent_params( + 16, cuvs::distance::DistanceType::L2Expanded); + params.node_per_cacheline = nodes_per_row; + auto const bits = flowann::detail::required_id_bits(rows); + auto const row_bytes = flowann::detail::packed_row_bytes(nodes_per_row, bits, resident_degree); + params.device_graph_budget_bytes = (rows / nodes_per_row) * row_bytes; + + auto built = flowann::build(res, params, dataset->as_dataset_view()); + ASSERT_EQ(built.index.graph_degree(), full_degree); + ASSERT_EQ(built.index.resident_degree(), resident_degree); + ASSERT_EQ(built.index.cross_degree(), full_degree - resident_degree); + + auto inner = raft::make_host_matrix(built.index.inner_graph().extent(0), + built.index.inner_graph().extent(1)); + raft::copy(inner.data_handle(), + built.index.inner_graph().data_handle(), + inner.size(), + raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + auto const cross = built.index.cross_graph(); + + for (std::uint32_t node = 0; node < rows; ++node) { + auto const packed_row = node / nodes_per_row; + auto const packed_col = node % nodes_per_row; + auto const* row = inner.data_handle() + packed_row * inner.extent(1); + ASSERT_EQ(row[packed_col], resident_degree); + auto bit_offset = static_cast(nodes_per_row) * 8; + for (std::uint16_t column = 0; column < packed_col; ++column) { + bit_offset += static_cast(row[column]) * bits; + } + + std::vector> resident; + std::vector> spilled; + std::set all_neighbors; + for (std::uint32_t rank = 0; rank < resident_degree; ++rank) { + auto const neighbor = read_packed(row, bit_offset + rank * bits, bits); + resident.emplace_back(l2(&host(node, 0), &host(neighbor, 0), dim), neighbor); + all_neighbors.insert(neighbor); + } + ASSERT_EQ(cross(node, 0), full_degree - resident_degree); + for (std::uint32_t edge = 0; edge < full_degree - resident_degree; ++edge) { + auto const neighbor = cross(node, edge + 1); + spilled.emplace_back(l2(&host(node, 0), &host(neighbor, 0), dim), neighbor); + all_neighbors.insert(neighbor); + } + for (std::size_t rank = 1; rank < resident.size(); ++rank) { + EXPECT_LE(resident[rank - 1].first, resident[rank].first + 1e-3f) << "node=" << node; + } + auto const closest_spilled = *std::min_element(spilled.begin(), spilled.end()); + EXPECT_LE(resident.back().first, closest_spilled.first + 1e-3f) << "node=" << node; + EXPECT_EQ(all_neighbors.size(), full_degree); + } + + auto query_host = raft::make_host_matrix(2, dim); + std::copy_n(host.data_handle(), query_host.size(), query_host.data_handle()); + auto queries = raft::make_device_matrix(res, 2, dim); + raft::copy(queries.data_handle(), + query_host.data_handle(), + query_host.size(), + raft::resource::get_cuda_stream(res)); + auto neighbors = raft::make_device_matrix(res, 2, 2); + auto distances = raft::make_device_matrix(res, 2, 2); + for (auto algorithm : {cuvs::neighbors::cagra::search_algo::SINGLE_CTA, + cuvs::neighbors::cagra::search_algo::MULTI_CTA}) { + flowann::search_context context(built.index.cross_graph()); + flowann::search_params search_params; + search_params.algo = algorithm; + search_params.itopk_size = 32; + search_params.search_width = 1; + search_params.max_iterations = 16; + search_params.max_queries = 2; + flowann::search(res, + search_params, + built.index, + context, + raft::make_const_mdspan(queries.view()), + neighbors.view(), + distances.view()); + auto neighbor_host = raft::make_host_matrix(2, 2); + raft::copy(neighbor_host.data_handle(), + neighbors.data_handle(), + neighbors.size(), + raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + for (std::size_t i = 0; i < neighbor_host.size(); ++i) { + EXPECT_LT(neighbor_host.data_handle()[i], rows); + } + } +} + +TEST(FlowannBuild, IntegratedBuildGroupsRearrangesAndPreservesSourceIds) +{ + constexpr int64_t rows = 265; + constexpr int64_t dim = 8; + constexpr std::uint32_t group_count = 5; + raft::resources res; + auto host = raft::make_host_matrix(rows, dim); + for (int64_t row = 0; row < rows; ++row) { + auto const group = row / (rows / group_count); + for (int64_t column = 0; column < dim; ++column) { + host(row, column) = static_cast(group * 100 + column) + + static_cast(row % (rows / group_count)) * 0.001f; + } + } + auto dataset = cuvs::neighbors::make_device_padded_dataset(res, host.view()); + + flowann::index_params params; + params.cagra_params.metric = cuvs::distance::DistanceType::L2Expanded; + params.cagra_params.graph_degree = 8; + params.cagra_params.intermediate_graph_degree = 16; + params.cagra_params.graph_build_params = + cuvs::neighbors::cagra::graph_build_params::nn_descent_params( + 16, cuvs::distance::DistanceType::L2Expanded); + params.node_per_cacheline = 2; + params.device_graph_budget_bytes = ((rows + 1) / 2) * 64; + params.grouping.n_groups = 0; + params.grouping.n_bits = 6; + params.grouping.balance_tolerance = 0.10f; + params.grouping.training_rows = 128; + params.grouping.assignment_batch_rows = 64; + params.grouping.kmeans_n_iters = 4; + params.grouping.validate = true; + + auto built = flowann::build(res, params, dataset->as_dataset_view()); + EXPECT_EQ(built.index.graph_degree(), 8); + EXPECT_EQ(built.index.resident_degree(), 0); + EXPECT_GT(built.index.inner_graph().size(), 0); + EXPECT_EQ(built.index.inner_graph().extent(1), 14); // Two headers plus 16 six-bit IDs. + + auto source_view = built.index.source_indices(); + ASSERT_TRUE(source_view.has_value()); + auto source = raft::make_host_vector(rows); + raft::copy( + source.data_handle(), source_view->data_handle(), rows, raft::resource::get_cuda_stream(res)); + auto ids = std::set(); + for (int64_t row = 0; row < rows; ++row) { + ids.insert(source(row)); + } + EXPECT_EQ(ids.size(), rows); + EXPECT_EQ(*ids.begin(), 0); + EXPECT_EQ(*ids.rbegin(), rows - 1); + + auto offsets_view = built.index.subgraph_offsets(); + auto groups_view = built.index.subgraph_ids(); + ASSERT_TRUE(offsets_view.has_value()); + ASSERT_TRUE(groups_view.has_value()); + ASSERT_EQ(offsets_view->extent(0), group_count); + auto groups = raft::make_host_vector(rows); + raft::copy( + groups.data_handle(), groups_view->data_handle(), rows, raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + auto sizes = std::array{}; + for (int64_t row = 0; row < rows; ++row) { + ASSERT_LT(groups(row), group_count); + ++sizes[groups(row)]; + } + for (auto size : sizes) { + EXPECT_GE(size, 47); + EXPECT_LE(size, 59); + EXPECT_LE(size, 64); + } +} + +TEST(FlowannBuild, IntegratedBuildGeneratesUniqueMedoidSeeds) +{ + constexpr int64_t rows = 128; + constexpr int64_t dim = 8; + raft::resources res; + auto host = raft::make_host_matrix(rows, dim); + for (int64_t row = 0; row < rows; ++row) { + for (int64_t column = 0; column < dim; ++column) { + host(row, column) = static_cast((row * 19 + column * 7) % 113); + } + } + auto dataset = cuvs::neighbors::make_device_padded_dataset(res, host.view()); + + flowann::index_params params; + params.grouping.enabled = false; + params.cagra_params.metric = cuvs::distance::DistanceType::L2Expanded; + params.cagra_params.graph_degree = 8; + params.cagra_params.intermediate_graph_degree = 16; + params.cagra_params.graph_build_params = + cuvs::neighbors::cagra::graph_build_params::nn_descent_params( + 16, cuvs::distance::DistanceType::L2Expanded); + params.device_graph_budget_bytes = 0; + params.num_seeds = 4; + params.seed_training_rows = 64; + + auto built = flowann::build(res, params, dataset->as_dataset_view()); + auto seeds = built.index.seeds(); + ASSERT_TRUE(seeds.has_value()); + ASSERT_EQ(seeds->extent(0), params.num_seeds); + auto host_seeds = raft::make_host_vector(params.num_seeds); + raft::copy(host_seeds.data_handle(), + seeds->data_handle(), + params.num_seeds, + raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + auto unique = std::set(); + for (std::uint32_t i = 0; i < params.num_seeds; ++i) { + EXPECT_LT(host_seeds(i), rows); + unique.insert(host_seeds(i)); + } + EXPECT_EQ(unique.size(), params.num_seeds); +} + +TEST(FlowannBuild, IntegratedBuildUsesCagraMetricForRanking) +{ + constexpr int64_t rows = 128; + constexpr int64_t dim = 8; + constexpr std::uint32_t resident_degree = 4; + raft::resources res; + auto host = raft::make_host_matrix(rows, dim); + for (int64_t row = 0; row < rows; ++row) { + for (int64_t column = 0; column < dim; ++column) { + host(row, column) = 0.1f + static_cast((row * 29 + column * 17) % 137) * 0.01f; + } + } + auto dataset = cuvs::neighbors::make_device_padded_dataset(res, host.view()); + + for (auto metric : + {cuvs::distance::DistanceType::InnerProduct, cuvs::distance::DistanceType::CosineExpanded}) { + flowann::index_params params; + params.grouping.enabled = false; + params.cagra_params.metric = metric; + params.cagra_params.graph_degree = 8; + params.cagra_params.intermediate_graph_degree = 16; + params.cagra_params.graph_build_params = + cuvs::neighbors::cagra::graph_build_params::nn_descent_params(16, metric); + auto const bits = flowann::detail::required_id_bits(rows); + auto const row_bytes = flowann::detail::packed_row_bytes(2, bits, resident_degree); + params.device_graph_budget_bytes = (rows / 2) * row_bytes; + + auto built = flowann::build(res, params, dataset->as_dataset_view()); + auto inner = raft::make_host_matrix(built.index.inner_graph().extent(0), + built.index.inner_graph().extent(1)); + raft::copy(inner.data_handle(), + built.index.inner_graph().data_handle(), + inner.size(), + raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + auto const cross = built.index.cross_graph(); + + for (std::uint32_t node = 0; node < rows; ++node) { + auto const packed_row = node / 2; + auto const packed_col = node % 2; + auto const* row = inner.data_handle() + packed_row * inner.extent(1); + auto bit_offset = std::uint64_t{16}; + if (packed_col != 0) { bit_offset += static_cast(row[0]) * bits; } + auto worst_resident = -std::numeric_limits::infinity(); + for (std::uint32_t rank = 0; rank < resident_degree; ++rank) { + auto const neighbor = read_packed(row, bit_offset + rank * bits, bits); + worst_resident = + std::max(worst_resident, metric_score(&host(node, 0), &host(neighbor, 0), dim, metric)); + } + auto best_spilled = std::numeric_limits::infinity(); + for (std::uint32_t edge = 0; edge < cross(node, 0); ++edge) { + auto const neighbor = cross(node, edge + 1); + best_spilled = + std::min(best_spilled, metric_score(&host(node, 0), &host(neighbor, 0), dim, metric)); + } + EXPECT_LE(worst_resident, best_spilled + 1e-4f) << "node=" << node; + } + } +} + +TEST(FlowannBuild, SupportsIvfPqAndIterativeCagraGraphBuilders) +{ + constexpr int64_t rows = 512; + constexpr int64_t dim = 16; + raft::resources res; + auto host = raft::make_host_matrix(rows, dim); + for (int64_t row = 0; row < rows; ++row) { + for (int64_t column = 0; column < dim; ++column) { + host(row, column) = static_cast((row * 37 + column * 13) % 211) * 0.01f; + } + } + auto dataset = cuvs::neighbors::make_device_padded_dataset(res, host.view()); + + for (bool iterative : {false, true}) { + flowann::index_params params; + params.grouping.enabled = false; + params.cagra_params.metric = cuvs::distance::DistanceType::L2Expanded; + params.cagra_params.graph_degree = 8; + params.cagra_params.intermediate_graph_degree = 16; + params.cagra_params.graph_build_params = + iterative + ? cuvs::neighbors::cagra::graph_build_params_t{cuvs::neighbors::cagra::graph_build_params:: + iterative_search_params{}} + : cuvs::neighbors::cagra::graph_build_params_t{ + cuvs::neighbors::cagra::graph_build_params::ivf_pq_params( + raft::matrix_extent(rows, dim), params.cagra_params.metric)}; + + auto built = flowann::build(res, params, dataset->as_dataset_view()); + EXPECT_EQ(built.index.size(), rows); + EXPECT_EQ(built.index.graph_degree(), 8); + EXPECT_EQ(built.index.resident_degree(), 0); + EXPECT_EQ(built.index.cross_degree(), 8); + } +} + +TEST(FlowannBuild, SearchesZeroAndFullResidentLayouts) +{ + constexpr int64_t rows = 128; + constexpr int64_t dim = 8; + constexpr std::uint32_t degree = 8; + raft::resources res; + auto host = raft::make_host_matrix(rows, dim); + for (int64_t row = 0; row < rows; ++row) { + for (int64_t column = 0; column < dim; ++column) { + host(row, column) = static_cast((row * 31 + column * 5) % 149) * 0.01f; + } + } + auto dataset = cuvs::neighbors::make_device_padded_dataset(res, host.view()); + + for (bool all_resident : {false, true}) { + flowann::index_params params; + params.grouping.enabled = false; + params.cagra_params.metric = cuvs::distance::DistanceType::L2Expanded; + params.cagra_params.graph_degree = degree; + params.cagra_params.intermediate_graph_degree = 16; + params.cagra_params.graph_build_params = + cuvs::neighbors::cagra::graph_build_params::nn_descent_params( + 16, cuvs::distance::DistanceType::L2Expanded); + if (all_resident) { + auto const bits = flowann::detail::required_id_bits(rows); + auto const row_bytes = flowann::detail::packed_row_bytes(2, bits, degree); + params.device_graph_budget_bytes = (rows / 2) * row_bytes; + } + + auto built = flowann::build(res, params, dataset->as_dataset_view()); + EXPECT_EQ(built.index.resident_degree(), all_resident ? degree : 0); + EXPECT_EQ(built.index.cross_degree(), all_resident ? 0 : degree); + + auto query_host = raft::make_host_matrix(2, dim); + std::copy_n(host.data_handle(), query_host.size(), query_host.data_handle()); + auto queries = raft::make_device_matrix(res, 2, dim); + raft::copy(queries.data_handle(), + query_host.data_handle(), + query_host.size(), + raft::resource::get_cuda_stream(res)); + auto neighbors = raft::make_device_matrix(res, 2, 2); + auto distances = raft::make_device_matrix(res, 2, 2); + flowann::search_context context(built.index.cross_graph()); + flowann::search_params search_params; + search_params.algo = cuvs::neighbors::cagra::search_algo::SINGLE_CTA; + search_params.itopk_size = 32; + search_params.search_width = 1; + search_params.max_iterations = 16; + search_params.max_queries = 2; + flowann::search(res, + search_params, + built.index, + context, + raft::make_const_mdspan(queries.view()), + neighbors.view(), + distances.view()); + auto neighbor_host = raft::make_host_matrix(2, 2); + raft::copy(neighbor_host.data_handle(), + neighbors.data_handle(), + neighbors.size(), + raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + for (std::size_t i = 0; i < neighbor_host.size(); ++i) { + EXPECT_LT(neighbor_host.data_handle()[i], rows); + } + } +} + +TEST(FlowannBuild, AutomaticGroupingCapacityPlan) +{ + using flowann::detail::plan_grouping; + for (auto const& [rows, bits] : std::vector>{ + {1, 4}, {16, 4}, {17, 8}, {256, 8}, {257, 12}, {1u << 20, 20}, {1u << 24, 24}}) { + EXPECT_EQ(plan_grouping(rows).n_bits, bits); + EXPECT_EQ(plan_grouping(rows).n_groups, 1); + } + EXPECT_EQ(plan_grouping((1u << 24) + 1).n_groups, 2); + EXPECT_EQ(plan_grouping(87'555'327).n_groups, 6); + EXPECT_EQ(plan_grouping(10ull * (1u << 24)).n_groups, 11); + EXPECT_EQ(plan_grouping(1'000'000'000).n_groups, 66); + // Arithmetic scalability only: public graph construction still uses uint32 IDs. + EXPECT_EQ(plan_grouping(10'000'000'000ull).n_groups, 656); + EXPECT_EQ(plan_grouping(16, 4).upper, 16); + EXPECT_THROW(plan_grouping(17, 4, 1), raft::exception); + EXPECT_THROW(plan_grouping(0), raft::exception); + EXPECT_THROW(plan_grouping(256, 33), raft::exception); + EXPECT_THROW(plan_grouping(256, 0, 0, 0), raft::exception); + for (std::uint32_t leaves : {2, 16, 17, 32, 66, 656, 65536}) { + auto check = [&](auto&& self, std::uint32_t n) -> std::uint32_t { + if (n == 1) { return 1; } + auto children = flowann::detail::grouping_children(n); + EXPECT_LE(children.size(), 16); + std::uint32_t total = 0; + for (auto child : children) { + total += self(self, child); + } + EXPECT_EQ(total, n); + return total; + }; + EXPECT_EQ(check(check, leaves), leaves); + } +} + +TEST(FlowannBuild, AutomaticWidthIsIndependentOfGraphBudget) +{ + raft::resources res; + constexpr int64_t n = 256, dim = 4, degree = 32; + auto rows = raft::make_host_matrix(n, dim); + auto graph = raft::make_host_matrix(n, degree); + auto dataset = + cuvs::neighbors::make_host_standard_dataset_view(raft::make_const_mdspan(rows.view())); + flowann::index_params params; + // One group needs neither training nor graph inspection. + for (auto budget : {0, 172, 194}) { + params.device_graph_budget_bytes = (n / 2) * budget; + auto result = flowann::group_graph(res, params, raft::make_const_mdspan(graph.view()), dataset); + EXPECT_EQ(result.n_bits, 8); + EXPECT_EQ(result.n_groups, 1); + EXPECT_TRUE(std::all_of(result.labels.data_handle(), + result.labels.data_handle() + n, + [](auto label) { return label == 0; })); + } + params.grouping.n_bits = 24; + EXPECT_EQ( + flowann::group_graph(res, params, raft::make_const_mdspan(graph.view()), dataset).n_bits, 24); +} + +TEST(FlowannBuild, HierarchicalGroupingEnforcesEverySubtreeCapacity) +{ + raft::resources res; + constexpr int64_t n = 15'200, dim = 2, degree = 4; + auto rows = raft::make_host_matrix(n, dim); + auto graph = raft::make_host_matrix(n, degree); + for (int64_t i = 0; i < n; ++i) { + rows(i, 0) = float(i / 230); + rows(i, 1) = float(i % 230) * .001f; + for (int64_t j = 0; j < degree; ++j) { + graph(i, j) = (i + j + 1) % n; + } + } + auto dataset = + cuvs::neighbors::make_host_standard_dataset_view(raft::make_const_mdspan(rows.view())); + flowann::index_params params; + params.grouping.n_bits = 8; + params.grouping.kmeans_n_iters = 4; + params.grouping.assignment_batch_rows = 97; + auto grouped = flowann::group_graph(res, params, raft::make_const_mdspan(graph.view()), dataset); + ASSERT_EQ(grouped.n_bits, 8); + ASSERT_EQ(grouped.n_groups, 66); + auto plan = flowann::detail::plan_grouping(n, 8); + std::vector sizes(66); + std::uint64_t local = 0; + for (int64_t i = 0; i < n; ++i) { + ASSERT_LT(grouped.labels(i), 66); + ++sizes[grouped.labels(i)]; + for (int64_t j = 0; j < degree; ++j) { + local += grouped.labels(i) == grouped.labels(graph(i, j)); + } + } + auto check = [&](auto&& self, std::uint32_t first, std::uint32_t leaves) -> void { + auto count = + std::accumulate(sizes.begin() + first, sizes.begin() + first + leaves, std::uint64_t{0}); + EXPECT_GE(count, leaves * plan.lower); + EXPECT_LE(count, leaves * plan.upper); + EXPECT_LE(count, leaves * plan.capacity); + if (leaves == 1) { return; } + for (auto child : flowann::detail::grouping_children(leaves)) { + self(self, first, child); + first += child; + } + }; + check(check, 0, 66); + EXPECT_GT(double(local) / (n * degree), .8); +} + +TEST(FlowannBuild, CapacityRepairMovesCheapestPointsAndHandlesTies) +{ + // The first row is far from group 1; row-order repair would move the wrong point. + float costs[] = {0, 100, 0, 1, 0, 20, 20, 0}; + std::vector labels{0, 0, 0, 1}; + auto sizes = flowann::detail::rebalance_group_distances(costs, labels, {2, 2}, {2, 2}); + EXPECT_EQ(labels, (std::vector{0, 1, 0, 1})); + EXPECT_EQ(sizes, (std::vector{2, 2})); + // All points initially choose the same centroid, with unequal subtree capacities. + std::vector tied_costs(100 * 3, 0); + labels.assign(100, 0); + sizes = flowann::detail::rebalance_group_distances( + tied_costs.data(), labels, {40, 20, 20}, {50, 25, 25}); + EXPECT_EQ(sizes, (std::vector{50, 25, 25})); + labels = {0, 0, 0, 1}; + // No overflow, but the lower-bound repair must still choose the nearest eligible point. + sizes = flowann::detail::rebalance_group_distances(costs, labels, {1, 2}, {4, 4}); + EXPECT_EQ(labels, (std::vector{0, 1, 0, 1})); +} + +TEST(FlowannBuild, LargeCapacityRepairPreservesDistanceAndIdPriority) +{ + // Exceed the parallel radix-selection threshold; include negative costs and a tie at zero. + constexpr std::size_t rows = 1'050'000; + std::vector costs(rows * 2); + std::vector labels(rows, 0); + for (std::size_t row = 0; row < rows; ++row) { + costs[row * 2 + 1] = static_cast(row % 1000) - 500.0f; + } + auto const requested = rows / 2 + 1; + auto sizes = flowann::detail::rebalance_group_distances( + costs.data(), labels, {rows - requested, requested}, {rows - requested, requested}); + EXPECT_EQ(sizes[1], requested); + for (std::size_t row = 0; row < rows; ++row) { + ASSERT_EQ(labels[row], (row % 1000 < 500 || row == 500) ? 1 : 0) << row; + } +} + +TEST(FlowannBuild, SearchesAutomaticAndExplicitWidths) +{ + constexpr int64_t rows = 128; + constexpr int64_t dim = 8; + constexpr std::uint32_t degree = 8; + raft::resources res; + auto host = raft::make_host_matrix(rows, dim); + for (int64_t row = 0; row < rows; ++row) { + for (int64_t column = 0; column < dim; ++column) { + host(row, column) = static_cast((row * 31 + column * 5) % 149) * 0.01f; + } + } + auto dataset = cuvs::neighbors::make_device_padded_dataset(res, host.view()); + + for (std::uint16_t requested_bits : {0, 20}) { + flowann::index_params params; + params.grouping.enabled = true; + params.grouping.n_bits = requested_bits; + params.grouping.validate = true; + params.cagra_params.metric = cuvs::distance::DistanceType::L2Expanded; + params.cagra_params.graph_degree = degree; + params.cagra_params.intermediate_graph_degree = 16; + params.cagra_params.graph_build_params = + cuvs::neighbors::cagra::graph_build_params::nn_descent_params( + 16, cuvs::distance::DistanceType::L2Expanded); + { + auto const bits = requested_bits == 0 ? 8 : requested_bits; + auto const row_bytes = flowann::detail::packed_row_bytes(2, bits, degree); + params.device_graph_budget_bytes = (rows / 2) * row_bytes; + } + + auto built = flowann::build(res, params, dataset->as_dataset_view()); + ASSERT_EQ(built.index.n_bits(), requested_bits == 0 ? 8 : requested_bits); + EXPECT_EQ(built.index.resident_degree(), 0); + EXPECT_EQ(built.index.cross_degree(), degree); + + auto query_host = raft::make_host_matrix(2, dim); + std::copy_n(host.data_handle(), query_host.size(), query_host.data_handle()); + auto queries = raft::make_device_matrix(res, 2, dim); + raft::copy(queries.data_handle(), + query_host.data_handle(), + query_host.size(), + raft::resource::get_cuda_stream(res)); + auto neighbors = raft::make_device_matrix(res, 2, 2); + auto distances = raft::make_device_matrix(res, 2, 2); + flowann::search_context context(built.index.cross_graph()); + flowann::search_params search_params; + search_params.algo = cuvs::neighbors::cagra::search_algo::SINGLE_CTA; + search_params.itopk_size = 32; + search_params.search_width = 1; + search_params.max_iterations = 16; + search_params.max_queries = 2; + flowann::search(res, + search_params, + built.index, + context, + raft::make_const_mdspan(queries.view()), + neighbors.view(), + distances.view()); + auto neighbor_host = raft::make_host_matrix(2, 2); + raft::copy(neighbor_host.data_handle(), + neighbors.data_handle(), + neighbors.size(), + raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + EXPECT_EQ(neighbor_host(0, 0), 0); + EXPECT_EQ(neighbor_host(1, 0), 1); + for (std::size_t i = 0; i < neighbor_host.size(); ++i) { + EXPECT_LT(neighbor_host.data_handle()[i], rows); + } + } +} diff --git a/cpp/tests/neighbors/ann_cagra/flowann_queue.cu b/cpp/tests/neighbors/ann_cagra/flowann_queue.cu new file mode 100644 index 0000000000..4150516828 --- /dev/null +++ b/cpp/tests/neighbors/ann_cagra/flowann_queue.cu @@ -0,0 +1,716 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * FlowANN is described in "Disentangling Graph Dependencies for Efficient + * Billion-Scale GPU Vector Search" (OSDI 2026). + */ + +#include + +#include + +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include + +namespace flowann = cuvs::neighbors::cagra::experimental::flowann; + +namespace { + +RAFT_KERNEL queue_round_trip(flowann::detail::queue_view* queues, + std::uint32_t num_queues, + std::uint32_t row_count, + std::uint32_t row_length, + std::uint32_t* output) +{ + if (threadIdx.x != 0) { return; } + + auto const command = blockIdx.x % row_count; + auto queue = queues[blockIdx.x % num_queues]; + flowann::detail::queue_reservation reservation{}; + while (!flowann::detail::try_reserve(queue, &reservation)) {} + flowann::detail::submit(queue, reservation, command); + while (!flowann::detail::result_ready(queue, reservation)) {} + auto const* row = flowann::detail::result_row(queue, reservation); + for (std::uint32_t i = 0; i < row_length; ++i) { + output[static_cast(blockIdx.x) * row_length + i] = row[i]; + } + flowann::detail::release(queue, reservation); +} + +RAFT_KERNEL mixed_release_queue_round_trip(flowann::detail::queue_view* queues, + std::uint32_t num_queues, + std::uint32_t row_count, + std::uint32_t row_length, + std::uint32_t round, + std::uint32_t batch, + std::uint32_t* output) +{ + auto queue = queues[blockIdx.x % num_queues]; + if (threadIdx.x == 0) { + std::uint32_t ids[7]; + for (std::uint32_t i = 0; i < batch; ++i) { + ids[i] = (blockIdx.x * batch + i + round) % row_count; + } + flowann::detail::queue_reservation first{}; + while (!flowann::detail::try_reserve_deferred(queue, batch, &first)) {} + flowann::detail::submit_batch(queue, first, ids, batch); + for (std::uint32_t i = 0; i < batch; ++i) { + auto const reservation = flowann::detail::reservation_at(first, i); + while (!flowann::detail::result_ready(queue, reservation)) {} + auto const* row = flowann::detail::result_row(queue, reservation); + for (std::uint32_t col = 0; col < row_length; ++col) { + output[(static_cast(blockIdx.x) * batch + i) * row_length + col] = row[col]; + } + if (((blockIdx.x + i) & 1u) == 0) { + flowann::detail::release_deferred(queue, reservation); + } else { + flowann::detail::release(queue, reservation); + } + } + } + __syncthreads(); + flowann::detail::reclaim_released_warp(queue); +} + +TEST(FlowannQueue, CommandEncodingPreservesRowIdsAndRejectsPartialPublication) +{ + EXPECT_FALSE(flowann::detail::command_ready(flowann::detail::empty_command)); + for (std::uint32_t row : {0u, 1u, 65535u, 65536u, 0xabcdefu, 0x7fffffffu, 0xffffffffu}) { + auto const word = flowann::detail::make_command(row); + ASSERT_TRUE(flowann::detail::command_ready(word)); + EXPECT_EQ(static_cast(word), row); + // Reading either 32-bit half before the other is published must not + // produce a ready word carrying a different request. + for (auto partial : {word & 0xffffffff00000000ull, word & 0xffffffffull}) { + if (flowann::detail::command_ready(partial)) { EXPECT_EQ(partial, word); } + } + for (unsigned bit = 0; bit < 64; ++bit) { + EXPECT_FALSE(flowann::detail::command_ready(word ^ (1ull << bit))); + } + } +} + +void check_round_trip_pause_resume(std::uint32_t num_queues) +{ + raft::resources res; + constexpr std::uint32_t row_count = 17; + constexpr std::uint32_t row_length = 9; + constexpr std::uint32_t commands = 257; + + auto cross_graph = raft::make_host_matrix(row_count, row_length); + for (std::uint32_t row = 0; row < row_count; ++row) { + for (std::uint32_t col = 0; col < row_length; ++col) { + cross_graph(row, col) = row * 1000 + col; + } + } + + flowann::queue_params params{}; + params.num_queues = num_queues; + params.empty_pause = 8; + params.collect_statistics = true; + flowann::search_context context(raft::make_const_mdspan(cross_graph.view()), params); + auto output = raft::make_device_matrix(res, commands, row_length); + + context.resume(); + queue_round_trip<<>>( + flowann::detail::search_context_access::device_queues(context), + context.num_queues(), + row_count, + row_length, + output.data_handle()); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + raft::resource::sync_stream(res); + context.pause(); + + auto host_output = raft::make_host_matrix(commands, row_length); + raft::copy(host_output.data_handle(), + output.data_handle(), + output.size(), + raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + for (std::uint32_t command = 0; command < commands; ++command) { + auto const row = command % row_count; + for (std::uint32_t col = 0; col < row_length; ++col) { + EXPECT_EQ(host_output(command, col), cross_graph(row, col)); + } + } + EXPECT_EQ(context.statistics().processed_commands, commands); + + context.resume(); + context.pause(); +} + +TEST(FlowannQueue, RoundTripPauseResume) { check_round_trip_pause_resume(2); } + +// Also run this test alone in a fresh process to exercise lazy kernel loading +// while many idle CUDA-copy pollers are already issuing transfers. +TEST(FlowannQueue, ColdLaunchWithManyPollers) { check_round_trip_pause_resume(44); } + +TEST(FlowannQueue, PollersUseContextCudaDevice) +{ + if (raft::device_setter::get_device_count() < 2) { + GTEST_SKIP() << "Requires at least two CUDA devices"; + } + + raft::device_setter device_scope(1); + raft::resources res; + constexpr std::uint32_t row_count = 4; + constexpr std::uint32_t row_length = 3; + + auto cross_graph = raft::make_host_matrix(row_count, row_length); + for (std::uint32_t row = 0; row < row_count; ++row) { + for (std::uint32_t col = 0; col < row_length; ++col) { + cross_graph(row, col) = row * 10 + col; + } + } + + flowann::search_context context(raft::make_const_mdspan(cross_graph.view())); + auto output = raft::make_device_matrix(res, row_count, row_length); + context.resume(); + queue_round_trip<<>>( + flowann::detail::search_context_access::device_queues(context), + context.num_queues(), + row_count, + row_length, + output.data_handle()); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + raft::resource::sync_stream(res); + context.pause(); + + auto host_output = raft::make_host_matrix(row_count, row_length); + raft::copy(host_output.data_handle(), + output.data_handle(), + output.size(), + raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + for (std::uint32_t row = 0; row < row_count; ++row) { + for (std::uint32_t col = 0; col < row_length; ++col) { + EXPECT_EQ(host_output(row, col), cross_graph(row, col)); + } + } +} + +TEST(FlowannQueue, RealPollersPreservePayloadsAcrossWrapWithMixedReleases) +{ + raft::resources res; + auto const stream = raft::resource::get_cuda_stream(res); + constexpr std::uint32_t row_count = 17; + constexpr std::uint32_t row_length = 9; + constexpr std::uint32_t blocks = 4099; + constexpr std::uint32_t max_batch = 7; + constexpr std::uint32_t rounds = 65; + static_assert(static_cast(blocks) * rounds > + static_cast(flowann::detail::queue_capacity) * 4); + + auto cross_graph = raft::make_host_matrix(row_count, row_length); + for (std::uint32_t row = 0; row < row_count; ++row) { + for (std::uint32_t col = 0; col < row_length; ++col) { + cross_graph(row, col) = row * 1000 + col; + } + } + auto output = + raft::make_device_matrix(res, blocks * max_batch, row_length); + auto host_output = raft::make_host_matrix(blocks * max_batch, row_length); + + for (std::uint32_t num_queues : {1u, 4u}) { + for (std::uint32_t batch : {1u, max_batch}) { + auto const commands = blocks * batch; + flowann::queue_params params{}; + params.num_queues = num_queues; + params.empty_pause = 0; + params.collect_statistics = true; + flowann::search_context context(raft::make_const_mdspan(cross_graph.view()), params); + { + flowann::search_session session(context); + for (std::uint32_t round = 0; round < rounds; ++round) { + mixed_release_queue_round_trip<<>>( + flowann::detail::search_context_access::device_queues(context), + num_queues, + row_count, + row_length, + round, + batch, + output.data_handle()); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + // Finish the host-serviced kernel before copying into pageable host memory. + // A pageable D2H copy can block inside the driver while waiting for this + // kernel, preventing the CUDA-copy poller from submitting its transfers. + raft::resource::sync_stream(res); + raft::copy(host_output.data_handle(), + output.data_handle(), + static_cast(commands) * row_length, + stream); + raft::resource::sync_stream(res); + for (std::uint32_t command = 0; command < commands; ++command) { + auto const row = (command + round) % row_count; + for (std::uint32_t col = 0; col < row_length; ++col) { + if (host_output(command, col) != cross_graph(row, col)) { + FAIL() << "payload mismatch with " << num_queues << " queues at round " << round + << ", command " << command << ", column " << col; + } + } + } + } + } + + EXPECT_TRUE(flowann::detail::search_context_access::is_paused(context)); + EXPECT_EQ(context.statistics().processed_commands, + static_cast(commands) * rounds); + std::vector queues(num_queues); + RAFT_CUDA_TRY(cudaMemcpy(queues.data(), + flowann::detail::search_context_access::device_queues(context), + queues.size() * sizeof(flowann::detail::queue_view), + cudaMemcpyDeviceToHost)); + for (auto const& queue : queues) { + flowann::detail::queue_state state{}; + RAFT_CUDA_TRY(cudaMemcpy(&state, queue.state, sizeof(state), cudaMemcpyDeviceToHost)); + EXPECT_EQ(state.gpu_tail, state.reservation_head); + EXPECT_EQ(state.cpu_tail, state.reservation_head); + } + } + } +} + +TEST(FlowannQueue, PollerStopsAtUnpublishedCommandsAndResumesFullBursts) +{ + constexpr std::uint32_t row_count = 17; + constexpr std::uint32_t row_length = 33; + constexpr std::uint32_t commands = 140; + auto graph = raft::make_host_matrix(row_count, row_length); + for (std::uint32_t row = 0; row < row_count; ++row) { + for (std::uint32_t col = 0; col < row_length; ++col) { + graph(row, col) = row * 1000 + col; + } + } + + // Exercise holes within the short prefix, at its boundary, and in a full batch. + for (std::uint32_t hole : {5u, 8u, 127u}) { + SCOPED_TRACE(hole); + flowann::queue_params params{}; + params.num_queues = 1; + params.empty_pause = 0; + params.collect_statistics = true; + flowann::search_context context(raft::make_const_mdspan(graph.view()), params); + flowann::detail::queue_view queue{}; + RAFT_CUDA_TRY(cudaMemcpy(&queue, + flowann::detail::search_context_access::device_queues(context), + sizeof(queue), + cudaMemcpyDeviceToHost)); + unsigned long long const reserved = commands; + RAFT_CUDA_TRY(cudaMemcpy( + &queue.state->reservation_head, &reserved, sizeof(reserved), cudaMemcpyHostToDevice)); + std::vector ids(commands); + std::vector sequences(commands); + for (std::uint32_t i = 0; i < commands; ++i) { + ids[i] = i % row_count; + sequences[i] = + i == hole ? flowann::detail::empty_command : flowann::detail::make_command(ids[i]); + } + RAFT_CUDA_TRY(cudaMemcpy(queue.command_words, + sequences.data(), + sequences.size() * sizeof(sequences[0]), + cudaMemcpyHostToDevice)); + auto wait_for = [&](std::uint64_t expected) { + auto const deadline = std::chrono::steady_clock::now() + std::chrono::seconds(5); + while (context.statistics().processed_commands < expected && + std::chrono::steady_clock::now() < deadline) { + std::this_thread::yield(); + } + return context.statistics().processed_commands; + }; + + context.resume(); + auto const first_count = wait_for(hole); + context.pause(); + ASSERT_EQ(first_count, hole); + ASSERT_EQ(context.statistics().processed_commands, hole); + flowann::detail::queue_state state{}; + RAFT_CUDA_TRY(cudaMemcpy(&state, queue.state, sizeof(state), cudaMemcpyDeviceToHost)); + EXPECT_EQ(state.cpu_tail, hole); + std::vector published(commands); + RAFT_CUDA_TRY(cudaMemcpy(published.data(), + queue.result_sequence, + published.size() * sizeof(published[0]), + cudaMemcpyDeviceToHost)); + for (std::uint32_t i = 0; i < commands; ++i) { + EXPECT_EQ(published[i], i < hole ? i + 1ull : 0ull); + } + + auto const ready = flowann::detail::make_command(ids[hole]); + RAFT_CUDA_TRY( + cudaMemcpy(queue.command_words + hole, &ready, sizeof(ready), cudaMemcpyHostToDevice)); + context.resume(); + auto const final_count = wait_for(commands); + context.pause(); + ASSERT_EQ(final_count, commands); + RAFT_CUDA_TRY(cudaMemcpy(&state, queue.state, sizeof(state), cudaMemcpyDeviceToHost)); + EXPECT_EQ(state.cpu_tail, commands); + std::vector cleared(commands); + RAFT_CUDA_TRY(cudaMemcpy(cleared.data(), + queue.command_words, + cleared.size() * sizeof(cleared[0]), + cudaMemcpyDeviceToHost)); + for (auto word : cleared) { + EXPECT_EQ(word, flowann::detail::empty_command); + } + std::vector returned(commands * row_length); + RAFT_CUDA_TRY(cudaMemcpy( + returned.data(), queue.rows, returned.size() * sizeof(returned[0]), cudaMemcpyDeviceToHost)); + for (std::uint32_t i = 0; i < commands; ++i) { + for (std::uint32_t col = 0; col < row_length; ++col) { + EXPECT_EQ(returned[i * row_length + col], graph(ids[i], col)); + } + } + } +} + +RAFT_KERNEL deferred_capacity_reclaim(flowann::detail::queue_view q, unsigned* errors) +{ + if (threadIdx.x != 0 || blockIdx.x != 0) return; + for (unsigned i = 0; i < flowann::detail::queue_capacity * 3 + 17; i++) { + flowann::detail::queue_reservation r{}; + unsigned tries = 0; + while (!flowann::detail::try_reserve_deferred(q, 1, &r)) { + if (++tries > 4) { + *errors = 1; + return; + } + } + if (r.sequence != i) { + *errors = 2; + return; + } + flowann::detail::release_deferred(q, r); + } + flowann::detail::reclaim_released(q); + if (q.state->gpu_tail != q.state->reservation_head) *errors = 3; +} +RAFT_KERNEL deferred_release_hole(flowann::detail::queue_view q, unsigned* errors) +{ + if (threadIdx.x == 0) { + q.state->gpu_tail = 0; + q.state->reservation_head = 3; + q.released_sequence[0] = 0; + flowann::detail::release_deferred(q, {1, 1}); + flowann::detail::release_deferred(q, {2, 2}); + } + __syncthreads(); + flowann::detail::reclaim_released_warp(q); + if (threadIdx.x == 0 && q.state->gpu_tail != 0) *errors = 4; + if (threadIdx.x == 0) flowann::detail::release_deferred(q, {0, 0}); + __syncthreads(); + flowann::detail::reclaim_released_warp(q); + if (threadIdx.x == 0 && q.state->gpu_tail != 3) *errors = 5; +} + +RAFT_KERNEL reserve_deferred_concurrently(flowann::detail::queue_view queue, + std::uint64_t expected_base, + std::uint32_t reservations_per_block, + flowann::detail::queue_reservation* reservations, + std::uint32_t* valid, + std::uint32_t* sequence_owners, + std::uint32_t* errors) +{ + auto const reservation_index = blockIdx.x; + if (threadIdx.x == 0) { + flowann::detail::queue_reservation first{}; + bool reserved = false; + for (std::uint32_t attempt = 0; attempt < 4096 && !reserved; ++attempt) { + reserved = flowann::detail::try_reserve_deferred(queue, reservations_per_block, &first); + } + reservations[reservation_index] = first; + valid[reservation_index] = reserved; + if (!reserved) { atomicOr(errors, 1u); } + } + __syncthreads(); + + if (valid[reservation_index] == 0) { return; } + auto const first = reservations[reservation_index]; + for (std::uint32_t i = threadIdx.x; i < reservations_per_block; i += blockDim.x) { + auto const reservation = flowann::detail::reservation_at(first, i); + if (reservation.sequence < expected_base || + reservation.sequence >= + expected_base + static_cast(gridDim.x) * reservations_per_block) { + atomicOr(errors, 2u); + continue; + } + auto const offset = static_cast(reservation.sequence - expected_base); + if (atomicCAS(&sequence_owners[offset], ~std::uint32_t{0}, reservation_index) != + ~std::uint32_t{0}) { + atomicOr(errors, 4u); + } + } +} + +RAFT_KERNEL release_deferred_except_first(flowann::detail::queue_view queue, + std::uint64_t first_sequence, + std::uint32_t reservations_per_block, + flowann::detail::queue_reservation const* reservations, + std::uint32_t const* valid) +{ + if (valid[blockIdx.x] != 0) { + auto const first = reservations[blockIdx.x]; + for (std::uint32_t i = threadIdx.x; i < reservations_per_block; i += blockDim.x) { + auto const reservation = flowann::detail::reservation_at(first, i); + if (reservation.sequence != first_sequence) { + flowann::detail::release_deferred(queue, reservation); + } + } + } + __syncthreads(); + if (threadIdx.x < 32) { flowann::detail::reclaim_released_warp(queue); } +} + +RAFT_KERNEL release_first_and_reclaim(flowann::detail::queue_view queue, + std::uint64_t first_sequence, + std::uint32_t reservations_per_block, + flowann::detail::queue_reservation const* reservations, + std::uint32_t const* valid) +{ + if (valid[blockIdx.x] != 0) { + auto const first = reservations[blockIdx.x]; + if (first.sequence <= first_sequence && + first_sequence < first.sequence + reservations_per_block && threadIdx.x == 0) { + flowann::detail::release_deferred( + queue, + flowann::detail::reservation_at( + first, static_cast(first_sequence - first.sequence))); + } + } + __syncthreads(); + if (threadIdx.x < 32) { flowann::detail::reclaim_released_warp(queue); } +} + +RAFT_KERNEL check_deferred_queue_counters(flowann::detail::queue_view queue, + std::uint64_t expected_tail, + std::uint32_t error_bit, + std::uint32_t* errors) +{ + if (threadIdx.x == 0 && blockIdx.x == 0 && + (queue.state->gpu_tail != expected_tail || + (error_bit != 0 && queue.state->reservation_head != expected_tail))) { + atomicOr(errors, error_bit == 0 ? 8u : error_bit); + } +} + +TEST(FlowannQueue, DeferredReclamationHandlesCapacityWrapAndUnreleasedHoles) +{ + raft::resources res; + auto const stream = raft::resource::get_cuda_stream(res); + auto cross = raft::make_host_matrix(1, 2); + cross(0, 0) = 1; + cross(0, 1) = 0; + flowann::search_context context(raft::make_const_mdspan(cross.view())); + flowann::detail::queue_view queue; + RAFT_CUDA_TRY(cudaMemcpy(&queue, + flowann::detail::search_context_access::device_queues(context), + sizeof(queue), + cudaMemcpyDeviceToHost)); + auto errors = raft::make_device_matrix(res, 1, 1); + RAFT_CUDA_TRY(cudaMemsetAsync(errors.data_handle(), 0, sizeof(std::uint32_t), stream.get())); + deferred_capacity_reclaim<<<1, 1, 0, stream.get()>>>(queue, errors.data_handle()); + deferred_release_hole<<<1, 32, 0, stream.get()>>>(queue, errors.data_handle()); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + std::uint32_t host_error{}; + raft::copy(&host_error, errors.data_handle(), 1, stream); + raft::resource::sync_stream(res); + EXPECT_EQ(host_error, 0u); +} + +TEST(FlowannQueue, ConcurrentDeferredReclamationPreservesOwnershipAcrossWrap) +{ + raft::resources res; + auto const stream = raft::resource::get_cuda_stream(res); + auto cross = raft::make_host_matrix(1, 2); + cross(0, 0) = 1; + cross(0, 1) = 0; + flowann::search_context context(raft::make_const_mdspan(cross.view())); + flowann::detail::queue_view queue; + RAFT_CUDA_TRY(cudaMemcpy(&queue, + flowann::detail::search_context_access::device_queues(context), + sizeof(queue), + cudaMemcpyDeviceToHost)); + + constexpr std::uint32_t blocks = 255; + constexpr std::uint32_t reservations_per_block = 257; + constexpr std::uint32_t reservations_per_wave = blocks * reservations_per_block; + static_assert(reservations_per_wave == flowann::detail::queue_capacity - 1); + + auto reservations = + raft::make_device_matrix(res, blocks, 1); + auto valid = raft::make_device_matrix(res, blocks, 1); + auto owners = raft::make_device_matrix(res, reservations_per_wave, 1); + auto errors = raft::make_device_matrix(res, 1, 1); + RAFT_CUDA_TRY(cudaMemsetAsync(errors.data_handle(), 0, sizeof(std::uint32_t), stream.get())); + + std::uint64_t expected_base = 0; + for (std::uint32_t wave = 0; wave < 3; ++wave) { + RAFT_CUDA_TRY( + cudaMemsetAsync(valid.data_handle(), 0, valid.size() * sizeof(std::uint32_t), stream.get())); + RAFT_CUDA_TRY(cudaMemsetAsync( + owners.data_handle(), 0xff, owners.size() * sizeof(std::uint32_t), stream.get())); + reserve_deferred_concurrently<<>>(queue, + expected_base, + reservations_per_block, + reservations.data_handle(), + valid.data_handle(), + owners.data_handle(), + errors.data_handle()); + release_deferred_except_first<<>>(queue, + expected_base, + reservations_per_block, + reservations.data_handle(), + valid.data_handle()); + check_deferred_queue_counters<<<1, 1, 0, stream.get()>>>( + queue, expected_base, 0, errors.data_handle()); + release_first_and_reclaim<<>>(queue, + expected_base, + reservations_per_block, + reservations.data_handle(), + valid.data_handle()); + expected_base += reservations_per_wave; + check_deferred_queue_counters<<<1, 1, 0, stream.get()>>>( + queue, expected_base, 16u, errors.data_handle()); + } + RAFT_CUDA_TRY(cudaPeekAtLastError()); + + std::uint32_t host_error{}; + raft::copy(&host_error, errors.data_handle(), 1, stream); + raft::resource::sync_stream(res); + EXPECT_EQ(host_error, 0u); +} + +TEST(FlowannQueue, SearchLeaseConsumesRawResume) +{ + constexpr std::uint32_t row_count = 4; + constexpr std::uint32_t row_length = 3; + auto cross_graph = raft::make_host_matrix(row_count, row_length); + std::fill_n(cross_graph.data_handle(), cross_graph.size(), 0u); + + flowann::search_context context(raft::make_const_mdspan(cross_graph.view())); + EXPECT_TRUE(flowann::detail::search_context_access::is_paused(context)); + + context.resume(); + EXPECT_FALSE(flowann::detail::search_context_access::is_paused(context)); + flowann::detail::search_context_access::acquire_search(context); + flowann::detail::search_context_access::release_search(context); + EXPECT_TRUE(flowann::detail::search_context_access::is_paused(context)); + + context.resume(); + context.pause(); + EXPECT_TRUE(flowann::detail::search_context_access::is_paused(context)); + + flowann::detail::search_context_access::acquire_search(context); + context.resume(); + flowann::detail::search_context_access::release_search(context); + EXPECT_TRUE(flowann::detail::search_context_access::is_paused(context)); +} + +TEST(FlowannQueue, SearchSessionKeepsPollersRunningAcrossSearchLease) +{ + constexpr std::uint32_t row_count = 4; + constexpr std::uint32_t row_length = 3; + auto cross_graph = raft::make_host_matrix(row_count, row_length); + std::fill_n(cross_graph.data_handle(), cross_graph.size(), 0u); + + flowann::search_context context(raft::make_const_mdspan(cross_graph.view())); + { + flowann::search_session session(context); + EXPECT_FALSE(flowann::detail::search_context_access::is_paused(context)); + flowann::detail::search_context_access::acquire_search(context); + flowann::detail::search_context_access::release_search(context); + EXPECT_FALSE(flowann::detail::search_context_access::is_paused(context)); + } + EXPECT_TRUE(flowann::detail::search_context_access::is_paused(context)); +} + +TEST(FlowannQueue, PauseBlocksNewSearchesUntilPollersReachQuiescence) +{ + using namespace std::chrono_literals; + constexpr std::uint32_t row_count = 4; + constexpr std::uint32_t row_length = 3; + auto cross_graph = raft::make_host_matrix(row_count, row_length); + std::fill_n(cross_graph.data_handle(), cross_graph.size(), 0u); + + flowann::search_context context(raft::make_const_mdspan(cross_graph.view())); + flowann::detail::search_context_access::acquire_search(context); + std::atomic pauses_returned{}; + std::jthread first_pauser([&] { + context.pause(); + pauses_returned.fetch_add(1, std::memory_order_release); + }); + + auto const deadline = std::chrono::steady_clock::now() + 2s; + while (flowann::detail::search_context_access::pending_pause_requests(context) < 1 && + std::chrono::steady_clock::now() < deadline) { + std::this_thread::yield(); + } + std::jthread second_pauser([&] { + context.pause(); + pauses_returned.fetch_add(1, std::memory_order_release); + }); + auto const second_deadline = std::chrono::steady_clock::now() + 2s; + while (flowann::detail::search_context_access::pending_pause_requests(context) < 2 && + std::chrono::steady_clock::now() < second_deadline) { + std::this_thread::yield(); + } + auto const pending_pauses = + flowann::detail::search_context_access::pending_pause_requests(context); + + std::atomic later_search_started{}; + std::atomic later_search_acquired{}; + std::jthread later_search([&] { + later_search_started.store(true, std::memory_order_release); + flowann::detail::search_context_access::acquire_search(context); + later_search_acquired.store(true, std::memory_order_release); + flowann::detail::search_context_access::release_search(context); + }); + while (!later_search_started.load(std::memory_order_acquire)) { + std::this_thread::yield(); + } + std::this_thread::sleep_for(20ms); + if (pending_pauses == 2) { EXPECT_FALSE(later_search_acquired.load(std::memory_order_acquire)); } + + flowann::detail::search_context_access::release_search(context); + first_pauser.join(); + second_pauser.join(); + later_search.join(); + EXPECT_EQ(pending_pauses, 2); + EXPECT_EQ(pauses_returned.load(std::memory_order_acquire), 2); + EXPECT_TRUE(later_search_acquired.load(std::memory_order_acquire)); + EXPECT_TRUE(flowann::detail::search_context_access::is_paused(context)); +} + +TEST(FlowannQueue, SearchLeasePausesAfterLastSearch) +{ + constexpr std::uint32_t row_count = 4; + constexpr std::uint32_t row_length = 3; + auto cross_graph = raft::make_host_matrix(row_count, row_length); + std::fill_n(cross_graph.data_handle(), cross_graph.size(), 0u); + + flowann::search_context context(raft::make_const_mdspan(cross_graph.view())); + flowann::detail::search_context_access::acquire_search(context); + flowann::detail::search_context_access::acquire_search(context); + EXPECT_FALSE(flowann::detail::search_context_access::is_paused(context)); + + flowann::detail::search_context_access::release_search(context); + EXPECT_FALSE(flowann::detail::search_context_access::is_paused(context)); + flowann::detail::search_context_access::release_search(context); + EXPECT_TRUE(flowann::detail::search_context_access::is_paused(context)); +} + +} // namespace diff --git a/cpp/tests/neighbors/ann_cagra/flowann_search.cu b/cpp/tests/neighbors/ann_cagra/flowann_search.cu new file mode 100644 index 0000000000..da8973518a --- /dev/null +++ b/cpp/tests/neighbors/ann_cagra/flowann_search.cu @@ -0,0 +1,388 @@ +/* + * 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 + +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) -> search_result +{ + raft::resources res; + constexpr int64_t rows = 8; + constexpr int64_t dim = 4; + + 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) = 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)); + auto dataset = cuvs::neighbors::make_device_padded_dataset_view(res, device_dataset.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) = 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, + cuvs::neighbors::cagra::search_algo::MULTI_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(FlowannMultiCtaSearch, EvaluatesSeedsBeyondGraphDegree) +{ + // This graph has degree two. Stop before expansion so that node four can only be found + // by evaluating the third seed. Source IDs are reversed by run_dense_search. + auto result = + run_dense_search(cuvs::neighbors::cagra::search_algo::MULTI_CTA, {10.1f}, {0, 1, 4}, 1); + + 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, 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); + } + } +} + +TEST(FlowannMultiCtaSearch, UniversalProviderHandlesInternallyChunkedRequests) +{ + auto result = run_dense_search( + cuvs::neighbors::cagra::search_algo::MULTI_CTA, {0.1f, 12.2f}, {}, 16, true, 1, 256, 2, 1); + EXPECT_EQ(result.neighbors(0, 0), 7u); + EXPECT_EQ(result.neighbors(1, 0), 1u); +} + +TEST(FlowannSearchModes, AlternatesAlgorithmsWidthsAndShapesInOneSession) +{ + auto result = run_dense_search(cuvs::neighbors::cagra::search_algo::MULTI_CTA, + {0.1f, 12.2f, 1.1f, 11.2f, 2.1f, 10.2f, 3.1f, 13.2f}, + {}, + 16, + true, + 8, + 256, + 1, + 1, + true); + EXPECT_EQ(result.neighbors(0, 0), 7u); + EXPECT_EQ(result.neighbors(1, 0), 1u); + EXPECT_FALSE(result.pollers_paused); +} + +INSTANTIATE_TEST_SUITE_P(SingleAndMultiCta, + FlowannSearchModesTest, + ::testing::Values(cuvs::neighbors::cagra::search_algo::SINGLE_CTA, + cuvs::neighbors::cagra::search_algo::MULTI_CTA)); + +} // namespace diff --git a/cpp/tests/neighbors/ann_cagra/flowann_storage.cu b/cpp/tests/neighbors/ann_cagra/flowann_storage.cu new file mode 100644 index 0000000000..f2f85b7af9 --- /dev/null +++ b/cpp/tests/neighbors/ann_cagra/flowann_storage.cu @@ -0,0 +1,76 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * FlowANN is described in "Disentangling Graph Dependencies for Efficient + * Billion-Scale GPU Vector Search" (OSDI 2026). + */ + +#include + +#include +#include + +#include + +namespace flowann = cuvs::neighbors::cagra::experimental::flowann; + +TEST(FlowannStorage, OwnsCrossGraphAndPreservesItAcrossMove) +{ + raft::resources res; + flowann::device_padded_index index(res); + auto inner = raft::make_host_matrix(0, 0); + auto cross = raft::make_host_matrix(2, 2); + cross(0, 0) = 1; + cross(0, 1) = 1; + cross(1, 0) = 1; + cross(1, 1) = 0; + index.update_graph( + res, raft::make_const_mdspan(inner.view()), raft::make_const_mdspan(cross.view()), 1, 0); + cross(0, 1) = 0; + auto moved = std::move(index); + EXPECT_EQ(moved.size(), 2); + EXPECT_EQ(moved.graph_degree(), 1); + EXPECT_EQ(moved.resident_degree(), 0); + EXPECT_EQ(moved.cross_graph()(0, 1), 1); + raft::resource::sync_stream(res); +} + +TEST(FlowannStorage, RejectsInvalidCrossEdgesBeforeReplacingGraph) +{ + raft::resources res; + flowann::device_padded_index index(res); + auto inner = raft::make_host_matrix(0, 0); + auto cross = raft::make_host_matrix(2, 2); + cross(0, 0) = 1; + cross(0, 1) = 1; + cross(1, 0) = 1; + cross(1, 1) = 0; + index.update_graph( + res, raft::make_const_mdspan(inner.view()), raft::make_const_mdspan(cross.view()), 1, 0); + cross(0, 1) = 2; + EXPECT_ANY_THROW(index.update_graph( + res, raft::make_const_mdspan(inner.view()), raft::make_const_mdspan(cross.view()), 1, 0)); + EXPECT_EQ(index.cross_graph()(0, 1), 1); + EXPECT_EQ(index.size(), 2); + raft::resource::sync_stream(res); +} + +TEST(FlowannStorage, ComputesCosineNormsWithoutSearchImplementation) +{ + raft::resources res; + auto host = raft::make_host_matrix(2, 2); + host(0, 0) = 3.0f; + host(0, 1) = 4.0f; + host(1, 0) = 0.0f; + host(1, 1) = 2.0f; + auto dataset = cuvs::neighbors::make_device_padded_dataset(res, host.view()); + flowann::device_padded_index index( + res, cuvs::distance::DistanceType::CosineExpanded, dataset->as_dataset_view()); + ASSERT_TRUE(index.dataset_norms().has_value()); + float norms[2]{}; + raft::copy(norms, index.dataset_norms()->data_handle(), 2, raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + EXPECT_FLOAT_EQ(norms[0], 5.0f); + EXPECT_FLOAT_EQ(norms[1], 2.0f); +}