diff --git a/c/CMakeLists.txt b/c/CMakeLists.txt index 59932c32ee..602288d2b1 100644 --- a/c/CMakeLists.txt +++ b/c/CMakeLists.txt @@ -123,6 +123,12 @@ set_target_properties( VISIBILITY_INLINES_HIDDEN ON ) +if(CUVSC_STATIC_CUVS_LIBRARY AND CUVS_ENABLE_FLOWANN_SEARCH AND CUVS_FLOWANN_USE_GDRCOPY) + # This shared library carries the GDRCopy dependency when cuvs itself is linked statically. + set_property(TARGET cuvs_c APPEND PROPERTY INSTALL_RPATH + "${CUVS_FLOWANN_GDRCOPY_INSTALL_RPATH}") +endif() + # Setup the VERSION and SOVERSION of our library # # Long term goal is that every release that breaks the C ABI will update `cuvs_abi_version` allowing diff --git a/c/include/cuvs/neighbors/cagra.h b/c/include/cuvs/neighbors/cagra.h index 602f52dbfc..b5430f92bf 100644 --- a/c/include/cuvs/neighbors/cagra.h +++ b/c/include/cuvs/neighbors/cagra.h @@ -203,9 +203,77 @@ struct cuvsAceParams { typedef struct cuvsAceParams* cuvsAceParams_t; -/** - * @brief Supplemental parameters to build CAGRA Index - * +/** Experimental CAGRA graph storage mode. */ +enum cuvsCagraGraphStorage { + CUVS_CAGRA_GRAPH_DEVICE = 0, + CUVS_CAGRA_GRAPH_TIERED = 1 +}; +/** Tiered graph settings. The GPU budget covers packed edges only. */ +struct cuvsCagraTieredGraphParams { + /** Maximum packed resident graph bytes, excluding datasets and queues. Zero + * stores all edges on host. */ + size_t device_graph_budget_bytes; + /** Consecutive nodes per packed row: [1, 64] with grouping, [1, 32] without + * grouping. */ + uint16_t node_per_cacheline; + /** Group and reorder nodes for locality; false retains the rank-only layout. + */ + bool grouping_enabled; + /** Final group count, at most 65,536. Zero uses one group if IDs fit; + * otherwise ceil(N * (1 + balance_tolerance) / 2^b). Each k-means node has at + * most 16 children. */ + uint32_t n_groups; + /** Local-ID width. Zero selects min(24, max(4, align_up_4(ceil(log2(N))))). + * Explicit widths must be in [1, 32]. */ + uint16_t n_bits; + /** Allowed relative deviation from average group size, in the open interval + * (0, 1). Hard ID capacity still applies. */ + double balance_tolerance; + /** Maximum sample rows per k-means node. Zero shares a 160,000-row budget by + * subtree leaf count, bounded by 256 to 10,000 rows per child and actual node + * size. */ + size_t training_rows; + /** Rows assigned per GPU batch. Zero targets a 256 MiB vector buffer; + * distance/label buffers and host capacity-repair storage are additional. */ + size_t assignment_batch_rows; + /** Positive number of balanced k-means training iterations. */ + uint32_t kmeans_n_iters; + /** Verify encoded edges after rearrangement; intended for tests and small + * builds. */ + bool validate; + /** Number of medoid seeds generated at build time; zero disables seed + * generation. */ + uint32_t num_seeds; + /** Rows sampled to train seed centroids. Zero selects an automatic limit. */ + size_t seed_training_rows; + /** Deterministic seed for rotating the uniform k-means training sample. */ + uint64_t seed; +}; +/** Tiered search settings; queue state is retained by the index. */ +struct cuvsCagraTieredSearchParams { + /** Maximum number of stored medoid seeds used per query; zero uses random + * initialization. */ + uint32_t num_seeds; + /** Scale for the adaptive deferred-cross-edge synchronization window. */ + float sync_window_scale; + /** Parent-position drop that forces synchronization of a submitted cross-edge + * request. */ + uint32_t sync_drop_threshold; + /** Positive number of independent GPU-to-CPU command queues. */ + uint32_t num_queues; + /** Number of CPU pause instructions after an empty poll; zero disables this + * pause. */ + uint32_t empty_pause; + /** Collect aggregate queue and polling counters. */ + bool collect_statistics; + /** Keep pollers active between calls for low latency. Consumes host CPU while + * idle. */ + bool keep_pollers_running; +}; + +/** Supplemental parameters to build a CAGRA index. + * Allocate with cuvsCagraIndexParamsCreate(); stack allocation and by-value + * copies are not covered by the C ABI stability guarantee. */ struct cuvsCagraIndexParams { /** Distance type. */ @@ -225,6 +293,12 @@ struct cuvsCagraIndexParams { * - Others: nullptr */ void* graph_build_params; + /** Graph representation selected for construction. The default is + * CUVS_CAGRA_GRAPH_DEVICE. */ + enum cuvsCagraGraphStorage graph_storage; + /** Tiered layout settings used when graph_storage is CUVS_CAGRA_GRAPH_TIERED. + */ + struct cuvsCagraTieredGraphParams tiered; }; typedef struct cuvsCagraIndexParams* cuvsCagraIndexParams_t; @@ -256,7 +330,8 @@ typedef struct cuvsCagraMergeParams* cuvsCagraMergeParams_t; * @param[in] params cuvsCagraIndexParams_t to allocate * @return cuvsError_t */ -CUVS_EXPORT cuvsError_t cuvsCagraIndexParamsCreate(cuvsCagraIndexParams_t* params); +CUVS_EXPORT cuvsError_t +cuvsCagraIndexParamsCreate(cuvsCagraIndexParams_t *params); /** * @brief De-allocate CAGRA Index params @@ -264,13 +339,16 @@ CUVS_EXPORT cuvsError_t cuvsCagraIndexParamsCreate(cuvsCagraIndexParams_t* param * @param[in] params * @return cuvsError_t */ -CUVS_EXPORT cuvsError_t cuvsCagraIndexParamsDestroy(cuvsCagraIndexParams_t params); +CUVS_EXPORT cuvsError_t +cuvsCagraIndexParamsDestroy(cuvsCagraIndexParams_t params); /** Allocate CAGRA merge params and populate them with AUTO defaults. */ -CUVS_EXPORT cuvsError_t cuvsCagraMergeParamsCreate(cuvsCagraMergeParams_t* params); +CUVS_EXPORT cuvsError_t +cuvsCagraMergeParamsCreate(cuvsCagraMergeParams_t *params); /** De-allocate CAGRA merge params. */ -CUVS_EXPORT cuvsError_t cuvsCagraMergeParamsDestroy(cuvsCagraMergeParams_t params); +CUVS_EXPORT cuvsError_t +cuvsCagraMergeParamsDestroy(cuvsCagraMergeParams_t params); /** * @brief Allocate CAGRA Compression params, and populate with default values. @@ -281,7 +359,8 @@ CUVS_EXPORT cuvsError_t cuvsCagraMergeParamsDestroy(cuvsCagraMergeParams_t param * @param[in] params cuvsCagraCompressionParams_t to allocate * @return cuvsError_t */ -CUVS_EXPORT cuvsError_t cuvsCagraCompressionParamsCreate(cuvsCagraCompressionParams_t* params); +CUVS_EXPORT cuvsError_t +cuvsCagraCompressionParamsCreate(cuvsCagraCompressionParams_t *params); /** * @brief De-allocate CAGRA Compression params. @@ -292,7 +371,8 @@ CUVS_EXPORT cuvsError_t cuvsCagraCompressionParamsCreate(cuvsCagraCompressionPar * @param[in] params * @return cuvsError_t */ -CUVS_EXPORT cuvsError_t cuvsCagraCompressionParamsDestroy(cuvsCagraCompressionParams_t params); +CUVS_EXPORT cuvsError_t +cuvsCagraCompressionParamsDestroy(cuvsCagraCompressionParams_t params); /** * @brief Allocate ACE params, and populate with default values @@ -325,12 +405,9 @@ CUVS_EXPORT cuvsError_t cuvsAceParamsDestroy(cuvsAceParams_t params); * @param[in] metric Distance metric to use * @return cuvsError_t */ -CUVS_EXPORT cuvsError_t cuvsCagraIndexParamsFromHnswParams(cuvsCagraIndexParams_t params, - int64_t n_rows, - int64_t dim, - int M, - int ef_construction, - enum cuvsCagraHnswHeuristicType heuristic, +CUVS_EXPORT cuvsError_t cuvsCagraIndexParamsFromHnswParams( + cuvsCagraIndexParams_t params, int64_t n_rows, int64_t dim, int M, + int ef_construction, enum cuvsCagraHnswHeuristicType heuristic, cuvsDistanceType metric); /** @@ -347,12 +424,9 @@ CUVS_EXPORT cuvsError_t cuvsCagraIndexParamsFromHnswParams(cuvsCagraIndexParams_ * @param[in] build_quality Higher values increase build quality (and cost) up to a point * @return cuvsError_t */ -CUVS_EXPORT cuvsError_t cuvsCagraIndexParamsFromDataset(cuvsCagraIndexParams_t params, - int64_t n_rows, - int64_t dim, - size_t graph_degree, - cuvsDistanceType metric, - size_t build_quality); +CUVS_EXPORT cuvsError_t cuvsCagraIndexParamsFromDataset( + cuvsCagraIndexParams_t params, int64_t n_rows, int64_t dim, + size_t graph_degree, cuvsDistanceType metric, size_t build_quality); /** * @} @@ -383,7 +457,8 @@ typedef struct cuvsCagraExtendParams* cuvsCagraExtendParams_t; * @param[in] params cuvsCagraExtendParams_t to allocate * @return cuvsError_t */ -CUVS_EXPORT cuvsError_t cuvsCagraExtendParamsCreate(cuvsCagraExtendParams_t* params); +CUVS_EXPORT cuvsError_t +cuvsCagraExtendParamsCreate(cuvsCagraExtendParams_t *params); /** * @brief De-allocate CAGRA Extend params @@ -391,7 +466,8 @@ CUVS_EXPORT cuvsError_t cuvsCagraExtendParamsCreate(cuvsCagraExtendParams_t* par * @param[in] params * @return cuvsError_t */ -CUVS_EXPORT cuvsError_t cuvsCagraExtendParamsDestroy(cuvsCagraExtendParams_t params); +CUVS_EXPORT cuvsError_t +cuvsCagraExtendParamsDestroy(cuvsCagraExtendParams_t params); /** * @} @@ -491,6 +567,12 @@ struct cuvsCagraSearchParams { * impact on the throughput. */ float persistent_device_usage; + /** Enable the tiered settings below. Allocate this struct with + * cuvsCagraSearchParamsCreate(). */ + bool use_tiered_params; + /** Tiered queue and search settings used only when use_tiered_params is true. + */ + struct cuvsCagraTieredSearchParams tiered; }; typedef struct cuvsCagraSearchParams* cuvsCagraSearchParams_t; @@ -501,7 +583,8 @@ typedef struct cuvsCagraSearchParams* cuvsCagraSearchParams_t; * @param[in] params cuvsCagraSearchParams_t to allocate * @return cuvsError_t */ -CUVS_EXPORT cuvsError_t cuvsCagraSearchParamsCreate(cuvsCagraSearchParams_t* params); +CUVS_EXPORT cuvsError_t +cuvsCagraSearchParamsCreate(cuvsCagraSearchParams_t *params); /** * @brief De-allocate CAGRA search params @@ -509,7 +592,8 @@ CUVS_EXPORT cuvsError_t cuvsCagraSearchParamsCreate(cuvsCagraSearchParams_t* par * @param[in] params * @return cuvsError_t */ -CUVS_EXPORT cuvsError_t cuvsCagraSearchParamsDestroy(cuvsCagraSearchParams_t params); +CUVS_EXPORT cuvsError_t +cuvsCagraSearchParamsDestroy(cuvsCagraSearchParams_t params); /** * @} @@ -556,7 +640,8 @@ CUVS_EXPORT cuvsError_t cuvsCagraIndexDestroy(cuvsCagraIndex_t index); * @param[out] dim return dimension of the index * @return cuvsError_t */ -CUVS_EXPORT cuvsError_t cuvsCagraIndexGetDims(cuvsCagraIndex_t index, int64_t* dim); +CUVS_EXPORT cuvsError_t cuvsCagraIndexGetDims(cuvsCagraIndex_t index, + int64_t *dim); /** * @brief Get size of the CAGRA index @@ -565,16 +650,32 @@ CUVS_EXPORT cuvsError_t cuvsCagraIndexGetDims(cuvsCagraIndex_t index, int64_t* d * @param[out] size return number of vectors in the index * @return cuvsError_t */ -CUVS_EXPORT cuvsError_t cuvsCagraIndexGetSize(cuvsCagraIndex_t index, int64_t* size); +CUVS_EXPORT cuvsError_t cuvsCagraIndexGetSize(cuvsCagraIndex_t index, + int64_t *size); /** - * @brief Get graph degree of the CAGRA index + * @brief Get the graph representation of a CAGRA index. + * + * Returns the selected storage kind without exposing the device and host edge + * buffers. + * @param[in] index Initialized CAGRA index handle. + * @param[out] storage Receives CUVS_CAGRA_GRAPH_DEVICE or + * CUVS_CAGRA_GRAPH_TIERED. + * @return CUVS_SUCCESS on success, or CUVS_ERROR when an argument is invalid. + */ +CUVS_EXPORT cuvsError_t cuvsCagraIndexGetGraphStorage( + cuvsCagraIndex_t index, enum cuvsCagraGraphStorage *storage); + +/** + * @brief Get the complete graph degree of a CAGRA index. * - * @param[in] index CAGRA index - * @param[out] graph_degree return graph degree - * @return cuvsError_t + * For a tiered index this includes both resident and host edges. + * @param[in] index Initialized CAGRA index handle. + * @param[out] graph_degree Receives the complete graph degree. + * @return CUVS_SUCCESS on success, or CUVS_ERROR when an argument is invalid. */ -CUVS_EXPORT cuvsError_t cuvsCagraIndexGetGraphDegree(cuvsCagraIndex_t index, int64_t* graph_degree); +CUVS_EXPORT cuvsError_t cuvsCagraIndexGetGraphDegree(cuvsCagraIndex_t index, + int64_t *graph_degree); /** * @brief Returns a view of the CAGRA dataset @@ -595,7 +696,8 @@ CUVS_EXPORT cuvsError_t cuvsCagraIndexGetGraphDegree(cuvsCagraIndex_t index, int * @param[out] dataset the dataset used in cagra * @return cuvsError_t */ -CUVS_EXPORT cuvsError_t cuvsCagraIndexGetDataset(cuvsCagraIndex_t index, DLManagedTensor* dataset); +CUVS_EXPORT cuvsError_t cuvsCagraIndexGetDataset(cuvsCagraIndex_t index, + DLManagedTensor *dataset); /** * @brief Returns a view of the CAGRA graph @@ -616,7 +718,8 @@ CUVS_EXPORT cuvsError_t cuvsCagraIndexGetDataset(cuvsCagraIndex_t index, DLManag * @param[out] graph the output knn graph. * @return cuvsError_t */ -CUVS_EXPORT cuvsError_t cuvsCagraIndexGetGraph(cuvsCagraIndex_t index, DLManagedTensor* graph); +CUVS_EXPORT cuvsError_t cuvsCagraIndexGetGraph(cuvsCagraIndex_t index, + DLManagedTensor *graph); /** * @brief Update a CAGRA index with a device-padded or device-PQ dataset. @@ -712,6 +815,32 @@ CUVS_EXPORT cuvsError_t cuvsCagraBuild(cuvsResources_t res, cuvsDataset_t dataset, cuvsCagraIndex_t index); +/** + * @brief Build an owning tiered graph and VPQ-F16 dataset from host standard + * vectors. + * + * Requires graph_storage=CUVS_CAGRA_GRAPH_TIERED and enabled FlowANN building. + * The caller retains the input vectors and may release them after this call. + * The resulting index owns its reordered compressed dataset. + * + * @param[in] res Resources used for construction, compression, and device + * transfers. + * @param[in] params CAGRA graph construction parameters selecting tiered + * storage. + * @param[in] compression VPQ compression parameters using 8-bit PQ codes and + * FP16 codebooks. + * @param[in] dataset Host standard dataset with float32, float16, int8, or + * uint8 values. + * @param[out] index Output handle initialized with cuvsCagraIndexCreate. + * @return CUVS_SUCCESS on success, or CUVS_ERROR for unsupported parameters or + * build failure. + */ +CUVS_EXPORT cuvsError_t cuvsCagraBuildCompressed(cuvsResources_t res, + cuvsCagraIndexParams_t params, + cuvsPqParams_t compression, + cuvsDataset_t dataset, + cuvsCagraIndex_t index); + /** * @} */ @@ -801,13 +930,10 @@ CUVS_EXPORT cuvsError_t cuvsCagraExtend(cuvsResources_t res, * @param[in] filter cuvsFilter input filter that can be used to filter queries and neighbors based on the given bitset. */ -CUVS_EXPORT cuvsError_t cuvsCagraSearch(cuvsResources_t res, - cuvsCagraSearchParams_t params, - cuvsCagraIndex_t index, - DLManagedTensor* queries, - DLManagedTensor* neighbors, - DLManagedTensor* distances, - cuvsFilter filter); +CUVS_EXPORT cuvsError_t cuvsCagraSearch( + cuvsResources_t res, cuvsCagraSearchParams_t params, cuvsCagraIndex_t index, + DLManagedTensor *queries, DLManagedTensor *neighbors, + DLManagedTensor *distances, cuvsFilter filter); /** * @brief Search multiple CAGRA index partitions concurrently and return the global top-k per @@ -842,14 +968,11 @@ CUVS_EXPORT cuvsError_t cuvsCagraSearch(cuvsResources_t res, * DLManagedTensor* holding that partition's own bitset (one bit per * vector in that partition; standard 32-bit packing). */ -CUVS_EXPORT cuvsError_t cuvsCagraSearchMultiPartition(cuvsResources_t res, - cuvsCagraSearchParams_t params, - uint32_t num_partitions, - cuvsCagraIndex_t* indices, - DLManagedTensor* queries, - DLManagedTensor* partition_ids, - DLManagedTensor* neighbors, - DLManagedTensor* distances, +CUVS_EXPORT cuvsError_t cuvsCagraSearchMultiPartition( + cuvsResources_t res, cuvsCagraSearchParams_t params, + uint32_t num_partitions, cuvsCagraIndex_t *indices, + DLManagedTensor *queries, DLManagedTensor *partition_ids, + DLManagedTensor *neighbors, DLManagedTensor *distances, cuvsFilter* filters); /** @@ -891,9 +1014,8 @@ CUVS_EXPORT cuvsError_t cuvsCagraSerializeGraph(cuvsResources_t res, * @param[in] filename the file name for saving the graph and dataset * @param[in] index CAGRA index with an attached host or device dataset */ -CUVS_EXPORT cuvsError_t cuvsCagraSerializeGraphAndDataset(cuvsResources_t res, - const char* filename, - cuvsCagraIndex_t index); +CUVS_EXPORT cuvsError_t cuvsCagraSerializeGraphAndDataset( + cuvsResources_t res, const char *filename, cuvsCagraIndex_t index); /** * Save the CAGRA index to file in hnswlib format. @@ -942,25 +1064,33 @@ CUVS_EXPORT cuvsError_t cuvsCagraDeserializeGraph(cuvsResources_t res, /** * Load the CAGRA graph and dataset from file. * - * The returned owning dataset preserves the serialized host/device memory type and - * standard/padded layout. The index stores a non-owning view into it, so the caller must keep the - * dataset alive while the index uses it and destroy it separately with cuvsDatasetDestroy. Only a - * device-padded result is immediately searchable through the C API; attach a caller-owned - * device-padded view with cuvsCagraUpdateDataset for any other kind. The output pointer - * must point to a null handle on entry; deserialization acts as a factory and transfers ownership - * of the allocated dataset handle on success. Returns CUVS_ERROR when the file has no dataset; the - * index and output handle are unchanged on failure. + * The returned owning dataset preserves the serialized host/device memory type + * and standard/padded layout. The index stores a non-owning view into it, so + * the caller must keep the dataset alive while the index uses it and destroy it + * separately with cuvsDatasetDestroy. Only a device-padded result is + * immediately searchable through the C API; attach a caller-owned device-padded + * view with cuvsCagraUpdateDataset for any other kind. The output pointer must + * point to a null handle on entry; deserialization acts as a factory and + * transfers ownership of the allocated dataset handle on success. Returns + * CUVS_ERROR when the file has no dataset; the index and output handle are + * unchanged on failure. * - * Experimental, both the API and the serialization format are subject to change. + * For a tiered file the index owns the reordered dataset internally, and + * out_dataset remains null. The loaded index is immediately searchable and has + * no separately owned dataset handle. + * + * Experimental, both the API and the serialization format are subject to + * change. * * @param[in] res cuvsResources_t opaque C handle * @param[in] filename the name of the file that stores the graph and dataset - * @param[inout] index pre-created CAGRA index populated on success and unchanged on failure - * @param[out] out_dataset receives the allocated owning dataset handle; must point to null on entry + * @param[inout] index pre-created CAGRA index populated on success and + * unchanged on failure + * @param[out] out_dataset receives the allocated owning dataset handle; must + * point to null on entry */ -CUVS_EXPORT cuvsError_t cuvsCagraDeserializeGraphAndDataset(cuvsResources_t res, - const char* filename, - cuvsCagraIndex_t index, +CUVS_EXPORT cuvsError_t cuvsCagraDeserializeGraphAndDataset( + cuvsResources_t res, const char *filename, cuvsCagraIndex_t index, cuvsDataset_t* out_dataset); /** @@ -1079,8 +1209,7 @@ CUVS_EXPORT cuvsError_t cuvsCagraIndexFromArgs(cuvsResources_t res, CUVS_EXPORT cuvsError_t cuvsCagraMerge(cuvsResources_t res, cuvsCagraIndexParams_t params, cuvsCagraIndex_t* indices, - size_t num_indices, - cuvsFilter filter, + size_t num_indices, cuvsFilter filter, cuvsDataset_t merged_dataset, cuvsCagraIndex_t output_index); @@ -1103,13 +1232,10 @@ CUVS_EXPORT cuvsError_t cuvsCagraMerge(cuvsResources_t res, * `cuvsCagraUpdateDataset` before device search. * @param[out] output_index Output handle initialized with `cuvsCagraIndexCreate` */ -CUVS_EXPORT cuvsError_t cuvsCagraMergeWithParams(cuvsResources_t res, - cuvsCagraIndexParams_t params, - cuvsCagraMergeParams_t merge_params, - cuvsCagraIndex_t* indices, - size_t num_indices, - cuvsFilter filter, - cuvsDataset_t merged_dataset, +CUVS_EXPORT cuvsError_t cuvsCagraMergeWithParams( + cuvsResources_t res, cuvsCagraIndexParams_t params, + cuvsCagraMergeParams_t merge_params, cuvsCagraIndex_t *indices, + size_t num_indices, cuvsFilter filter, cuvsDataset_t merged_dataset, cuvsCagraIndex_t output_index); /** diff --git a/c/src/neighbors/cagra.cpp b/c/src/neighbors/cagra.cpp index 30d356101f..97484e1042 100644 --- a/c/src/neighbors/cagra.cpp +++ b/c/src/neighbors/cagra.cpp @@ -1113,6 +1113,10 @@ void _serialize(cuvsResources_t res, const char *filename, : "cuvsCagraSerializeGraph: null index handle"; with_index_by_layout(box, null_handle_err, "", [&](auto &idx) { + bool const tiered = idx.graph_storage() == + cuvs::neighbors::cagra::graph_storage_kind::tiered; + RAFT_EXPECTS(!tiered || include_dataset, + "Tiered serialization requires its dataset"); using index_dataset_view_t = std::remove_cvref_t; if constexpr (cuvs::neighbors::is_bbq_dataset_view_v) { RAFT_EXPECTS(!include_dataset, @@ -1120,7 +1124,7 @@ void _serialize(cuvsResources_t res, const char *filename, cuvs::neighbors::cagra::serialize(*res_ptr, std::string(filename), idx); } else if constexpr (cuvs::neighbors::is_vpq_dataset_view_v) { RAFT_EXPECTS( - !include_dataset, + !include_dataset || tiered, "cuvsCagraSerializeGraphAndDataset is not supported for PQ indices; serialize the PQ " "dataset separately"); cuvs::neighbors::cagra::serialize(*res_ptr, std::string(filename), idx); @@ -1139,6 +1143,7 @@ void _serialize(cuvsResources_t res, const char *filename, struct serialized_cagra_header { DLDataType dtype; cuvs::neighbors::cagra::serialized_dataset_kind dataset_kind; + bool tiered_pq = false; }; static auto read_serialized_header(cuvsResources_t res, const char *filename) @@ -1154,6 +1159,14 @@ static auto read_serialized_header(cuvsResources_t res, const char *filename) RAFT_FAIL("Invalid or truncated index header in file %s", filename); } + bool const tiered = std::memcmp(dtype_string, "CTIR", 4) == 0; + if (tiered) { + auto version = raft::deserialize_scalar(*res_ptr, is); + RAFT_EXPECTS(version == 1, "Unsupported CAGRA tiered envelope version"); + RAFT_EXPECTS(is.read(dtype_string, 4), + "Truncated CAGRA tiered dtype header"); + } + auto const dtype = raft::numpy_serializer::parse_descr( std::string(dtype_string, sizeof(dtype_string))); DLDataType output_dtype{ @@ -1173,6 +1186,14 @@ static auto read_serialized_header(cuvsResources_t res, const char *filename) auto const version = raft::deserialize_scalar(*res_ptr, is); auto const dataset_kind_raw = raft::deserialize_scalar(*res_ptr, is); + if (tiered) { + RAFT_EXPECTS(version == 2 && + (dataset_kind_raw == 1 || dataset_kind_raw == 2), + "Unsupported tiered serialization header"); + return {output_dtype, + cuvs::neighbors::cagra::serialized_dataset_kind::device_padded, + dataset_kind_raw == 2}; + } RAFT_EXPECTS( version == cuvs::neighbors::cagra::cagra_serialization_version, "serialization version mismatch, expected %d, got %d", @@ -1246,12 +1267,14 @@ void _deserialize(cuvsResources_t res, const char *filename, if (include_dataset) { RAFT_EXPECTS( - dataset_owner != nullptr, + dataset_owner != nullptr || + holder->idx.graph_storage() == + cuvs::neighbors::cagra::graph_storage_kind::tiered, "cuvsCagraDeserializeGraphAndDataset: serialized index has no dataset"); } std::unique_ptr dataset_handle{}; - if (include_dataset) { + if (include_dataset && dataset_owner) { dataset_handle = std::make_unique(); dataset_handle->addr = reinterpret_cast(dataset_owner.get()); dataset_handle->destroy_addr = &destroy_typed_addr; @@ -1273,13 +1296,31 @@ void _deserialize(cuvsResources_t res, const char *filename, output_index->dtype = dtype; holder.release(); - if (include_dataset) { + if (include_dataset && dataset_owner) { dataset_handle->addr = reinterpret_cast(dataset_owner.release()); *out_dataset = dataset_handle.release(); } destroy_sg_cagra_c_api_box(old_addr); } +template +void deserialize_tiered_pq(cuvsResources_t res, char const *filename, + cuvsCagraIndex_t output, DLDataType dtype) { + using view_t = cuvs::neighbors::device_vpq_dataset_view; + using holder_t = cuvs_cagra_c_api_index_lifetime_holder; + auto &resources = *reinterpret_cast(res); + auto holder = std::make_unique( + cuvs::neighbors::cagra::device_pq_index(resources)); + cuvs::neighbors::cagra::deserialize(resources, std::string(filename), + &holder->idx); + auto box = make_sg_cagra_c_api_index_box(holder.get()); + auto old_addr = output->addr; + output->addr = reinterpret_cast(box.release()); + output->dtype = dtype; + holder.release(); + destroy_sg_cagra_c_api_box(old_addr); +} + template void _serialize_to_hnswlib(cuvsResources_t res, const char *filename, cuvsCagraIndex_t index) { @@ -1473,6 +1514,29 @@ void convert_c_index_params(cuvsCagraIndexParams params, out->metric = static_cast((int)params.metric); out->intermediate_graph_degree = params.intermediate_graph_degree; out->graph_degree = params.graph_degree; + RAFT_EXPECTS(params.graph_storage == CUVS_CAGRA_GRAPH_DEVICE || + params.graph_storage == CUVS_CAGRA_GRAPH_TIERED, + "Invalid CAGRA graph storage mode"); + out->graph_storage = static_cast(params.graph_storage); + out->tiered.reset(); + if (out->graph_storage == graph_storage_kind::tiered) { + out->tiered.emplace(); + out->tiered->device_graph_budget_bytes = + params.tiered.device_graph_budget_bytes; + out->tiered->node_per_cacheline = params.tiered.node_per_cacheline; + out->tiered->grouping_enabled = params.tiered.grouping_enabled; + out->tiered->n_groups = params.tiered.n_groups; + out->tiered->n_bits = params.tiered.n_bits; + out->tiered->balance_tolerance = params.tiered.balance_tolerance; + out->tiered->training_rows = params.tiered.training_rows; + out->tiered->assignment_batch_rows = params.tiered.assignment_batch_rows; + out->tiered->kmeans_n_iters = params.tiered.kmeans_n_iters; + out->tiered->validate = params.tiered.validate; + out->tiered->num_seeds = params.tiered.num_seeds; + out->tiered->seed_training_rows = params.tiered.seed_training_rows; + out->tiered->seed = params.tiered.seed; + } + _set_graph_build_params(out->graph_build_params, params, params.build_algo, n_rows, dim); } @@ -1495,6 +1559,17 @@ void convert_c_search_params(cuvsCagraSearchParams params, out->persistent = params.persistent; out->persistent_lifetime = params.persistent_lifetime; out->persistent_device_usage = params.persistent_device_usage; + out->tiered.reset(); + if (params.use_tiered_params) { + out->tiered.emplace(); + out->tiered->num_seeds = params.tiered.num_seeds; + out->tiered->sync_window_scale = params.tiered.sync_window_scale; + out->tiered->sync_drop_threshold = params.tiered.sync_drop_threshold; + out->tiered->num_queues = params.tiered.num_queues; + out->tiered->empty_pause = params.tiered.empty_pause; + out->tiered->collect_statistics = params.tiered.collect_statistics; + out->tiered->keep_pollers_running = params.tiered.keep_pollers_running; + } } void* cagra_c_api_index_ptr(cuvsCagraIndex const* idx) @@ -1550,6 +1625,22 @@ extern "C" cuvsError_t cuvsCagraIndexGetSize(cuvsCagraIndex_t index, int64_t* si }); } +extern "C" cuvsError_t +cuvsCagraIndexGetGraphStorage(cuvsCagraIndex_t index, + cuvsCagraGraphStorage *storage) { + return cuvs::core::translate_exceptions([=] { + RAFT_EXPECTS(index && index->addr && storage, + "cuvsCagraIndexGetGraphStorage: null argument"); + auto *box = reinterpret_cast(index->addr); + dispatch_serialized_dtype(index->dtype, [&]() { + with_index_by_layout( + box, "null index", "", [&](auto const &idx) { + *storage = static_cast(idx.graph_storage()); + }); + }); + }); +} + extern "C" cuvsError_t cuvsCagraIndexGetGraphDegree(cuvsCagraIndex_t index, int64_t* graph_degree) { return cuvs::core::translate_exceptions([=] { @@ -2006,6 +2097,48 @@ extern "C" cuvsError_t cuvsCagraBuild(cuvsResources_t res, }); } +extern "C" cuvsError_t cuvsCagraBuildCompressed(cuvsResources_t res, + cuvsCagraIndexParams_t params, + cuvsPqParams_t compression, + cuvsDataset_t dataset, + cuvsCagraIndex_t index) { + return cuvs::core::translate_exceptions([=] { + RAFT_EXPECTS(res && params && compression && dataset && dataset->addr && + index, + "cuvsCagraBuildCompressed: null argument"); + RAFT_EXPECTS(dataset->mem_type == CUVS_DATASET_MEM_TYPE_HOST && + dataset->layout == CUVS_DATASET_LAYOUT_STANDARD, + "Compressed tiered build requires host standard vectors"); +#ifdef CUVS_ENABLE_FLOWANN_BUILD + auto *resources = reinterpret_cast(res); + dispatch_serialized_dtype(dataset->dtype, [&]() { + using view_t = cuvs::neighbors::host_standard_dataset_view; + using owner_t = cuvs::neighbors::host_standard_dataset; + using output_view_t = + cuvs::neighbors::device_vpq_dataset_view; + with_dataset_view(dataset, [&](auto const &view) { + cuvs::neighbors::cagra::index_params cpp_params; + convert_c_index_params(*params, view.n_rows(), view.dim(), &cpp_params); + auto result = cuvs::neighbors::cagra::build( + *resources, cpp_params, make_cpp_vpq_params(*compression), view); + using holder_t = + cuvs_cagra_c_api_index_lifetime_holder; + auto holder = std::make_unique(std::move(result)); + auto box = + make_sg_cagra_c_api_index_box(holder.get()); + auto old_addr = index->addr; + index->addr = reinterpret_cast(box.release()); + index->dtype = dataset->dtype; + holder.release(); + destroy_sg_cagra_c_api_box(old_addr); + }); + }); +#else + RAFT_FAIL("Tiered building is disabled in this cuVS library"); +#endif + }); +} + extern "C" cuvsError_t cuvsCagraIndexFromArgs(cuvsResources_t res, cuvsDistanceType metric, DLManagedTensor* graph_tensor, @@ -2250,6 +2383,22 @@ extern "C" cuvsError_t cuvsCagraIndexParamsCreate(cuvsCagraIndexParams_t* params .build_algo = IVF_PQ, .nn_descent_niter = 20}; (*params)->graph_build_params = new cuvsIvfPqParams{nullptr, nullptr, 1}; + auto const defaults = cuvs::neighbors::cagra::tiered_graph_params{}; + (*params)->tiered.device_graph_budget_bytes = + defaults.device_graph_budget_bytes; + (*params)->tiered.node_per_cacheline = defaults.node_per_cacheline; + (*params)->tiered.grouping_enabled = defaults.grouping_enabled; + (*params)->tiered.n_groups = defaults.n_groups; + (*params)->tiered.n_bits = defaults.n_bits; + (*params)->tiered.balance_tolerance = defaults.balance_tolerance; + (*params)->tiered.training_rows = defaults.training_rows; + (*params)->tiered.assignment_batch_rows = defaults.assignment_batch_rows; + (*params)->tiered.kmeans_n_iters = defaults.kmeans_n_iters; + (*params)->tiered.validate = defaults.validate; + (*params)->tiered.num_seeds = defaults.num_seeds; + (*params)->tiered.seed_training_rows = defaults.seed_training_rows; + (*params)->tiered.seed = defaults.seed; + }); } @@ -2415,6 +2564,15 @@ extern "C" cuvsError_t cuvsCagraSearchParamsCreate(cuvsCagraSearchParams_t* para .persistent_lifetime = 2, .persistent_device_usage = 1.0, }; + auto const defaults = cuvs::neighbors::cagra::tiered_search_params{}; + (*params)->tiered.num_seeds = defaults.num_seeds; + (*params)->tiered.sync_window_scale = defaults.sync_window_scale; + (*params)->tiered.sync_drop_threshold = defaults.sync_drop_threshold; + (*params)->tiered.num_queues = defaults.num_queues; + (*params)->tiered.empty_pause = defaults.empty_pause; + (*params)->tiered.collect_statistics = defaults.collect_statistics; + (*params)->tiered.keep_pollers_running = defaults.keep_pollers_running; + }); } @@ -2434,6 +2592,10 @@ extern "C" cuvsError_t cuvsCagraDeserializeGraph(cuvsResources_t res, "cuvsCagraDeserializeGraph: null index handle"); auto const header = read_serialized_header(res, filename); dispatch_serialized_dtype(header.dtype, [&]() { + if (header.tiered_pq) { + deserialize_tiered_pq(res, filename, index, header.dtype); + return; + } using view_t = cuvs::neighbors::device_padded_dataset_view; _deserialize( res, filename, index, header.dtype, false, nullptr); @@ -2460,6 +2622,10 @@ cuvsCagraDeserializeGraphAndDataset(cuvsResources_t res, const char *filename, "must be null"); auto const header = read_serialized_header(res, filename); dispatch_serialized_dtype(header.dtype, [&]() { + if (header.tiered_pq) { + deserialize_tiered_pq(res, filename, index, header.dtype); + return; + } dispatch_serialized_dataset_kind(header.dataset_kind, [&]() { _deserialize( res, filename, index, header.dtype, true, out_dataset); diff --git a/c/tests/CMakeLists.txt b/c/tests/CMakeLists.txt index cfceb58f33..848001b3e3 100644 --- a/c/tests/CMakeLists.txt +++ b/c/tests/CMakeLists.txt @@ -89,6 +89,9 @@ ConfigureTest(NAME IVF_FLAT_C_TEST PATH neighbors/run_ivf_flat_c.c neighbors/ann ConfigureTest(NAME IVF_PQ_C_TEST PATH neighbors/run_ivf_pq_c.c neighbors/ann_ivf_pq_c.cu) ConfigureTest(NAME IVF_SQ_C_TEST PATH neighbors/run_ivf_sq_c.c neighbors/ann_ivf_sq_c.cu) ConfigureTest(NAME CAGRA_C_TEST PATH neighbors/ann_cagra_c.cu) +target_compile_definitions( + CAGRA_C_TEST PRIVATE $<$:CUVS_ENABLE_FLOWANN_BUILD> +) ConfigureTest(NAME MG_C_TEST PATH neighbors/run_mg_c.c neighbors/ann_mg_c.cu) ConfigureTest( NAME ALL_NEIGHBORS_C_TEST PATH neighbors/run_all_neighbors_c.c neighbors/all_neighbors_c.cu diff --git a/c/tests/neighbors/ann_cagra_c.cu b/c/tests/neighbors/ann_cagra_c.cu index d08357b91e..d274a253bf 100644 --- a/c/tests/neighbors/ann_cagra_c.cu +++ b/c/tests/neighbors/ann_cagra_c.cu @@ -2208,3 +2208,90 @@ TEST(CagraC, SearchMultiPartitionMultiKernelRejected) } cuvsResourcesDestroy(res); } + +TEST(CagraC, TieredBuildSearchAndLoad) { + constexpr int64_t rows = 128, dim = 8; + cuvsResources_t res{}; + ASSERT_EQ(cuvsResourcesCreate(&res), CUVS_SUCCESS); + cudaStream_t stream{}; + ASSERT_EQ(cuvsStreamGet(res, &stream), CUVS_SUCCESS); + cuvsCagraIndexParams_t bp{}; + ASSERT_EQ(cuvsCagraIndexParamsCreate(&bp), CUVS_SUCCESS); + bp->build_algo = NN_DESCENT; + bp->graph_degree = 8; + bp->intermediate_graph_degree = 16; + bp->graph_storage = CUVS_CAGRA_GRAPH_TIERED; + bp->tiered.device_graph_budget_bytes = 512; + std::vector data(rows * dim); + for (int64_t i = 0; i < rows; ++i) { + for (int64_t j = 0; j < dim; ++j) + data[i * dim + j] = ((i * 31 + j * 5) % 149) * .01f; + } + int64_t data_shape[]{rows, dim}, query_shape[]{4, dim}, out_shape[]{4, 1}; + auto tensor = [](void *data, int64_t *shape, DLDeviceType device, + uint8_t dtype) { + DLManagedTensor result{}; + result.dl_tensor.data = data; + result.dl_tensor.device = {device, 0}; + result.dl_tensor.ndim = 2; + result.dl_tensor.dtype = {dtype, 32, 1}; + result.dl_tensor.shape = shape; + return result; + }; + auto input = tensor(data.data(), data_shape, kDLCPU, kDLFloat); + cuvsDataset_t padded{}; + ASSERT_EQ( + cuvsDatasetMakePadded(res, &input, CUVS_DATASET_MEM_TYPE_DEVICE, &padded), + CUVS_SUCCESS); + cuvsCagraIndex_t index{}; + ASSERT_EQ(cuvsCagraIndexCreate(&index), CUVS_SUCCESS); +#ifdef CUVS_ENABLE_FLOWANN_BUILD + ASSERT_EQ(cuvsCagraBuild(res, bp, padded, index), CUVS_SUCCESS); + cuvsCagraGraphStorage storage{}; + ASSERT_EQ(cuvsCagraIndexGetGraphStorage(index, &storage), CUVS_SUCCESS); + EXPECT_EQ(storage, CUVS_CAGRA_GRAPH_TIERED); + ASSERT_EQ(cuvsDatasetDestroy(padded), CUVS_SUCCESS); + padded = nullptr; + rmm::device_uvector q(4 * dim, stream), d(4, stream); + rmm::device_uvector n(4, stream); + raft::copy(q.data(), data.data(), 4 * dim, stream); + auto queries = tensor(q.data(), query_shape, kDLCUDA, kDLFloat); + auto neighbors = tensor(n.data(), out_shape, kDLCUDA, kDLUInt); + auto distances = tensor(d.data(), out_shape, kDLCUDA, kDLFloat); + cuvsCagraSearchParams_t sp{}; + ASSERT_EQ(cuvsCagraSearchParamsCreate(&sp), CUVS_SUCCESS); + sp->algo = SINGLE_CTA; + sp->itopk_size = 32; + sp->max_iterations = 32; + sp->use_tiered_params = true; + cuvsFilter no_filter{}; + no_filter.type = NO_FILTER; + auto search = [&] { + ASSERT_EQ(cuvsCagraSearch(res, sp, index, &queries, &neighbors, &distances, + no_filter), + CUVS_SUCCESS); + uint32_t expected[]{0, 1, 2, 3}; + EXPECT_TRUE(cuvs::devArrMatchHost(expected, n.data(), 4, + cuvs::Compare())); + }; + search(); + auto filename = std::filesystem::temp_directory_path() / + ("cagra-tiered-" + std::to_string(getpid()) + ".bin"); + ASSERT_EQ(cuvsCagraSerializeGraphAndDataset(res, filename.c_str(), index), + CUVS_SUCCESS); + ASSERT_EQ(cuvsCagraIndexDestroy(index), CUVS_SUCCESS); + ASSERT_EQ(cuvsCagraIndexCreate(&index), CUVS_SUCCESS); + ASSERT_EQ(cuvsCagraDeserializeGraph(res, filename.c_str(), index), + CUVS_SUCCESS); + search(); + std::filesystem::remove(filename); + ASSERT_EQ(cuvsCagraSearchParamsDestroy(sp), CUVS_SUCCESS); +#else + EXPECT_EQ(cuvsCagraBuild(res, bp, padded, index), CUVS_ERROR); +#endif + if (padded) + EXPECT_EQ(cuvsDatasetDestroy(padded), CUVS_SUCCESS); + EXPECT_EQ(cuvsCagraIndexDestroy(index), CUVS_SUCCESS); + EXPECT_EQ(cuvsCagraIndexParamsDestroy(bp), CUVS_SUCCESS); + EXPECT_EQ(cuvsResourcesDestroy(res), CUVS_SUCCESS); +} diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 78bde6f657..5d436f9c97 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -84,7 +84,22 @@ 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_BUILD_FLOWANN_BENCHMARKS "Build the standalone FlowANN benchmark programs" OFF) +option(CUVS_FLOWANN_USE_GDRCOPY "Use GDRCopy for FlowANN GPU-to-CPU command queues" ON) +set(CUVS_FLOWANN_GDRCOPY_INSTALL_RPATH "\$ORIGIN" CACHE STRING + "Runtime search paths for externally deployed GDRCopy; never inferred from the build host") + +if(CUVS_ENABLE_FLOWANN_SEARCH AND BUILD_CPU_ONLY) + message(FATAL_ERROR "CUVS_ENABLE_FLOWANN_SEARCH requires a CUDA-enabled build") +endif() +if(CUVS_ENABLE_FLOWANN_BUILD AND NOT CUVS_ENABLE_FLOWANN_SEARCH) + message(FATAL_ERROR "CUVS_ENABLE_FLOWANN_BUILD requires CUVS_ENABLE_FLOWANN_SEARCH") +endif() +if(CUVS_BUILD_FLOWANN_BENCHMARKS AND NOT CUVS_ENABLE_FLOWANN_SEARCH) + message(FATAL_ERROR "CUVS_BUILD_FLOWANN_BENCHMARKS requires CUVS_ENABLE_FLOWANN_SEARCH") +endif() if(BUILD_CPU_ONLY) set(BUILD_SHARED_LIBS OFF) set(BUILD_TESTS OFF) @@ -113,7 +128,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 +141,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 +186,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 +266,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 +996,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 +1561,10 @@ 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/cagra_tiered.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 +1671,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 +1693,7 @@ if(NOT BUILD_CPU_ONLY) nvidia::cutlass::cutlass rtcx::rtcx ${CUVS_CTK_MATH_DEPENDENCIES} + ${CUVS_FLOWANN_GDRCOPY_DEPENDENCIES} $ $ kvikio::kvikio @@ -1687,10 +1773,30 @@ if(NOT BUILD_CPU_ONLY) $> $> $<$:CUDA::nvtx3> - PRIVATE $ $ - $ CUDA::nvJitLink CUDA::nvrtc rtcx::rtcx kvikio::kvikio + PRIVATE $ + $ + $ + CUDA::nvJitLink + CUDA::nvrtc + rtcx::rtcx + kvikio::kvikio + ${CUVS_FLOWANN_GDRCOPY_DEPENDENCIES} INTERFACE $ ) + if(CUVS_ENABLE_FLOWANN_SEARCH) + target_link_libraries(cuvs INTERFACE $>) + endif() + if(CUVS_ENABLE_FLOWANN_SEARCH AND CUVS_FLOWANN_USE_GDRCOPY) + # GDRCopy is an external runtime dependency. Default to an install-relative path; + # packagers may supply an explicit deployment path, independent of the build prefix. + set_property(TARGET cuvs APPEND PROPERTY INSTALL_RPATH + "${CUVS_FLOWANN_GDRCOPY_INSTALL_RPATH}") + # Consumers must resolve this runtime dependency even when it is installed outside the + # linker's default search paths. Its headers remain private to the queue implementation. + target_link_libraries( + cuvs INTERFACE $ $ + ) + endif() endif() @@ -1735,6 +1841,7 @@ if(NOT BUILD_CPU_ONLY) CUDA::nvJitLink CUDA::nvrtc rtcx::rtcx + ${CUVS_FLOWANN_GDRCOPY_DEPENDENCIES} $<$:CUDA::nvtx3> $ $ @@ -1742,6 +1849,56 @@ if(NOT BUILD_CPU_ONLY) ) endif() + if(CUVS_ENABLE_FLOWANN_SEARCH AND CUVS_BUILD_FLOWANN_BENCHMARKS) + add_executable(FLOWANN_BENCHMARK ../examples/cpp/src/flowann_benchmark.cu) + target_include_directories(FLOWANN_BENCHMARK PRIVATE "${CUVS_SOURCE_DIR}/src") + set_target_properties( + FLOWANN_BENCHMARK + PROPERTIES CXX_STANDARD 20 + CXX_STANDARD_REQUIRED ON + CUDA_STANDARD 20 + CUDA_STANDARD_REQUIRED ON + ) + if(TARGET cuvs::cuvs) + target_link_libraries( + FLOWANN_BENCHMARK + PRIVATE cuvs::cuvs + ${CUVS_FLOWANN_GDRCOPY_DEPENDENCIES} + $ + $ + ) + else() + target_link_libraries( + FLOWANN_BENCHMARK + PRIVATE cuvs::cuvs_static + ${CUVS_FLOWANN_GDRCOPY_DEPENDENCIES} + $ + $ + ) + endif() + endif() + if(CUVS_ENABLE_FLOWANN_BUILD AND CUVS_BUILD_FLOWANN_BENCHMARKS) + add_executable(FLOWANN_BUILD_BENCHMARK ../examples/cpp/src/flowann_build_benchmark.cu) + target_include_directories(FLOWANN_BUILD_BENCHMARK PRIVATE "${CUVS_SOURCE_DIR}/src") + set_target_properties( + FLOWANN_BUILD_BENCHMARK + PROPERTIES CXX_STANDARD 20 + CXX_STANDARD_REQUIRED ON + CUDA_STANDARD 20 + CUDA_STANDARD_REQUIRED ON + ) + if(TARGET cuvs::cuvs) + target_link_libraries( + FLOWANN_BUILD_BENCHMARK PRIVATE cuvs::cuvs $ + ) + else() + target_link_libraries( + FLOWANN_BUILD_BENCHMARK PRIVATE cuvs::cuvs_static + $ + ) + endif() + endif() + # ################################################################################################ # * NVTX support in cuvs ----------------------------------------------------- if(NOT CUVS_NVTX) 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/include/cuvs/neighbors/cagra.hpp b/cpp/include/cuvs/neighbors/cagra.hpp index 98df8cdadf..fe8369c459 100644 --- a/cpp/include/cuvs/neighbors/cagra.hpp +++ b/cpp/include/cuvs/neighbors/cagra.hpp @@ -24,6 +24,7 @@ #include #include +#include #include #include #include @@ -53,6 +54,67 @@ enum class hash_mode { HASH = 0, SMALL = 1, AUTO = 100 }; enum class internal_dtype { F16 = 0, E5M2 = 1 }; +/** Graph representation used by a CAGRA index. */ +enum class graph_storage_kind { device = 0, tiered = 1 }; + +/** Experimental GPU/CPU graph layout. The budget excludes datasets and search workspaces. */ +struct tiered_graph_params { + /** Maximum packed resident graph bytes, excluding datasets and queues. Zero stores all edges on + * host. */ + std::size_t device_graph_budget_bytes = 0; + /** Consecutive nodes per packed row: [1, 64] with grouping, [1, 32] without grouping. */ + std::uint16_t node_per_cacheline = 2; + /** Group and reorder nodes for locality; false retains the rank-only layout. */ + bool grouping_enabled = true; + /** Final group count, at most 65,536. Zero uses one group if IDs fit; otherwise ceil(N * (1 + + * balance_tolerance) / 2^b). Each k-means node has at most 16 children. */ + std::uint32_t n_groups = 0; + /** Local-ID width. Zero selects min(24, max(4, align_up_4(ceil(log2(N))))). Explicit widths must + * be in [1, 32]. */ + std::uint16_t n_bits = 0; + /** Allowed relative deviation from average group size, in the open interval (0, 1). Hard ID + * capacity still applies. */ + double balance_tolerance = 0.10; + /** Maximum sample rows per k-means node. Zero shares a 160,000-row budget by subtree leaf count, + * bounded by 256 to 10,000 rows per child and actual node size. */ + std::size_t training_rows = 0; + /** Rows assigned per GPU batch. Zero targets a 256 MiB vector buffer; distance/label buffers and + * host capacity-repair storage are additional. */ + std::size_t assignment_batch_rows = 0; + /** Positive number of balanced k-means training iterations. */ + std::uint32_t kmeans_n_iters = 20; + /** Verify encoded edges after rearrangement; intended for tests and small builds. */ + bool validate = false; + /** Number of medoid seeds generated at build time; zero disables seed generation. */ + std::uint32_t num_seeds = 0; + /** Rows sampled to train seed centroids. Zero selects an automatic limit. */ + std::size_t seed_training_rows = 0; + /** Deterministic seed for rotating the uniform k-means training sample. */ + std::uint64_t seed = 0x9e3779b97f4a7c15ULL; +}; + +/** Experimental tiered search and reusable host queue settings. */ +struct tiered_search_params { + /** Maximum number of stored medoid seeds used per query; zero uses random initialization. */ + std::uint32_t num_seeds = 0; + /** Scale for the adaptive deferred-cross-edge synchronization window. */ + float sync_window_scale = 10.0f; + /** Parent-position drop that forces synchronization of a submitted cross-edge request. */ + std::uint32_t sync_drop_threshold = 51; + /** Positive number of independent GPU-to-CPU command queues. */ + std::uint32_t num_queues = 1; + /** Number of CPU pause instructions after an empty poll; zero disables this pause. */ + std::uint32_t empty_pause = 64; + /** Collect aggregate queue and polling counters. */ + bool collect_statistics = false; + /** Keep pollers active between calls for low latency. Consumes host CPU while idle. */ + bool keep_pollers_running = false; +}; + +/** CAGRA search parameters. + * @note The optional tiered settings change this struct's C++ ABI, including builds with + * FlowANN disabled. Rebuild consumers with headers matching the linked cuVS library. + */ struct search_params : cuvs::neighbors::search_params { /** Maximum number of queries to search at the same time (batch size). Auto select when 0.*/ size_t max_queries = 0; @@ -133,6 +195,9 @@ struct search_params : cuvs::neighbors::search_params { /** Data type of the query vector and codebook table on shared memory. Currently, only VPQ * supports FP8. **/ internal_dtype smem_dtype = internal_dtype::F16; + /** Settings for tiered indexes. Must be unset for ordinary device-graph indexes; + * search throws otherwise. Tiered indexes use default settings when this is unset. */ + std::optional tiered; }; /** @@ -290,6 +355,10 @@ enum class hnsw_heuristic_type : uint32_t { SAME_GRAPH_FOOTPRINT = 1 }; +/** CAGRA index construction parameters. + * @note The graph storage selector and optional tiered settings change this struct's C++ ABI, + * including builds with FlowANN disabled. Rebuild consumers against matching cuVS headers. + */ struct index_params : cuvs::neighbors::index_params { /** Degree of input graph for pruning. */ size_t intermediate_graph_degree = 128; @@ -356,6 +425,10 @@ struct index_params : cuvs::neighbors::index_params { * @endcode */ bool attach_dataset_on_build = true; + /** Ordinary device graph by default; tiered storage is experimental and opt-in. */ + graph_storage_kind graph_storage = graph_storage_kind::device; + /** A zero tiered graph budget stores all graph edges on the host. */ + std::optional tiered; /** * @brief Select the graph build algorithm and its parameters for a dataset. @@ -461,6 +534,34 @@ struct extend_params { * @} */ +namespace detail { +struct tiered_index_access; + +// The owning implementation lives out of line. Ordinary indexes allocate no tiered state. +template +struct tiered_index_state { + virtual ~tiered_index_state() = default; + virtual auto size() const noexcept -> IdxT = 0; + virtual auto graph_degree() const noexcept -> std::uint32_t = 0; + virtual auto dataset() const noexcept -> DatasetViewT = 0; + virtual auto source_indices() const noexcept + -> std::optional> = 0; + virtual void serialize(raft::resources const&, std::ostream&, bool) const = 0; + virtual void search(raft::resources const&, + search_params const&, + raft::device_matrix_view, + raft::device_matrix_view, + raft::device_matrix_view, + cuvs::neighbors::filtering::base_filter const&) const = 0; + virtual void search(raft::resources const&, + search_params const&, + raft::device_matrix_view, + raft::device_matrix_view, + raft::device_matrix_view, + cuvs::neighbors::filtering::base_filter const&) const = 0; +}; +} // namespace detail + static_assert(std::is_aggregate_v); static_assert(std::is_aggregate_v); @@ -504,6 +605,7 @@ struct CUVS_EXPORT index : cuvs::neighbors::index { /** Total length of the index (number of vectors). */ [[nodiscard]] constexpr inline auto size() const noexcept -> IdxT { + if (tiered_) { return tiered_->size(); } if (dataset_fd_.has_value() || graph_fd_.has_value()) { return n_rows_; } auto data_rows = dataset_.n_rows(); return data_rows > 0 ? data_rows : graph_view_.extent(0); @@ -517,23 +619,36 @@ struct CUVS_EXPORT index : cuvs::neighbors::index { /** Graph degree */ [[nodiscard]] constexpr inline auto graph_degree() const noexcept -> uint32_t { + if (tiered_) { return tiered_->graph_degree(); } return graph_fd_.has_value() ? graph_degree_ : graph_view_.extent(1); } /** Number of rows represented by the graph. */ [[nodiscard]] constexpr inline auto graph_size() const noexcept -> IdxT { + if (tiered_) { return tiered_->size(); } return graph_fd_.has_value() ? static_cast(n_rows_) : static_cast(graph_view_.extent(0)); } - /** Non-owning dataset binding stored by the index. */ + /** Return the actual graph representation. */ + [[nodiscard]] auto graph_storage() const noexcept -> graph_storage_kind + { + return tiered_ ? graph_storage_kind::tiered : graph_storage_kind::device; + } + + /** Dataset binding; a tiered index owns its reordered dataset internally. */ [[nodiscard]] inline auto dataset() const noexcept -> DatasetViewT const& { return dataset_; } - /** neighborhood graph [size, graph-degree] */ - [[nodiscard]] inline auto graph() const noexcept + /** Device neighborhood graph with shape [size, graph_degree]. + * + * Throws raft::logic_error when graph_storage() is tiered. Tiered graphs have separate + * device and host edge storage and cannot be represented by a single device adjacency matrix. + */ + [[nodiscard]] inline auto graph() const -> raft::device_matrix_view { + require_device_graph_(); return graph_view_; } @@ -541,6 +656,7 @@ struct CUVS_EXPORT index : cuvs::neighbors::index { [[nodiscard]] inline auto source_indices() const noexcept -> std::optional> { + if (tiered_) { return tiered_->source_indices(); } return source_indices_.has_value() ? std::optional>( source_indices_->view()) @@ -655,7 +771,7 @@ struct CUVS_EXPORT index : cuvs::neighbors::index { * dataset*/ template index(raft::resources const& res, index&& other, DatasetViewT dataset) - : metric_(other.metric_), + : metric_(other.require_device_graph_()), graph_(std::move(other.graph_)), graph_view_(other.graph_view_), dataset_(dataset), @@ -684,6 +800,7 @@ struct CUVS_EXPORT index : cuvs::neighbors::index { raft::resources const& res, raft::device_matrix_view knn_graph) { + require_device_graph_(); graph_view_ = knn_graph; } @@ -693,6 +810,7 @@ struct CUVS_EXPORT index : cuvs::neighbors::index { void update_graph(raft::resources const&, raft::device_matrix&& knn_graph) { + require_device_graph_(); graph_ = std::move(knn_graph); graph_view_ = graph_.view(); } @@ -706,6 +824,7 @@ struct CUVS_EXPORT index : cuvs::neighbors::index { raft::resources const& res, raft::host_matrix_view knn_graph) { + require_device_graph_(); RAFT_LOG_DEBUG("Copying CAGRA knn graph from host to device"); if ((graph_.extent(0) != knn_graph.extent(0)) || (graph_.extent(1) != knn_graph.extent(1))) { @@ -728,6 +847,7 @@ struct CUVS_EXPORT index : cuvs::neighbors::index { */ void update_source_indices(raft::device_vector&& source_indices) { + require_device_graph_(); RAFT_EXPECTS(source_indices.extent(0) == size(), "Source indices must have the same number of rows as the index"); source_indices_.emplace(std::move(source_indices)); @@ -742,6 +862,7 @@ struct CUVS_EXPORT index : cuvs::neighbors::index { raft::mdspan, raft::row_major, Accessor> source_indices) { + require_device_graph_(); RAFT_EXPECTS(source_indices.extent(0) == size(), "Source indices must have the same number of rows as the index"); // Reset the array if it's not compatible to avoid using more memory than necessary. @@ -774,6 +895,7 @@ struct CUVS_EXPORT index : cuvs::neighbors::index { */ void update_dataset(raft::resources const& res, cuvs::util::file_descriptor&& fd) { + require_device_graph_(); RAFT_EXPECTS(fd.is_valid(), "Invalid file descriptor provided for dataset"); auto stream = fd.make_istream(); @@ -814,6 +936,7 @@ struct CUVS_EXPORT index : cuvs::neighbors::index { */ void update_graph(raft::resources const& res, cuvs::util::file_descriptor&& fd) { + require_device_graph_(); RAFT_EXPECTS(fd.is_valid(), "Invalid file descriptor provided for graph"); auto stream = fd.make_istream(); @@ -854,6 +977,7 @@ struct CUVS_EXPORT index : cuvs::neighbors::index { */ void update_mapping(raft::resources const& res, cuvs::util::file_descriptor&& fd) { + require_device_graph_(); RAFT_EXPECTS(fd.is_valid(), "Invalid file descriptor provided for mapping"); // Read header from file using ifstream @@ -883,6 +1007,16 @@ struct CUVS_EXPORT index : cuvs::neighbors::index { friend struct index; friend struct detail::fd_transfer; + friend struct detail::tiered_index_access; + + auto require_device_graph_() const -> cuvs::distance::DistanceType + { + RAFT_EXPECTS(!tiered_, "This operation requires a device adjacency graph, not a tiered index"); + return metric_; + } + // This member changes index's C++ layout even when FlowANN is disabled. Applications must + // rebuild against the same cuVS headers as the library they link. + std::shared_ptr> tiered_; [[nodiscard]] inline auto steal_dataset_fd_() noexcept -> std::optional @@ -959,6 +1093,138 @@ using cagra_index_t = * @} */ +namespace detail { +struct tiered_index_access { + template + static auto get(Index const& idx) + { + return idx.tiered_.get(); + } + template + static void install(Index& idx, std::shared_ptr state) + { + RAFT_EXPECTS(idx.graph_size() == 0, "Tiered state must be installed on an empty index"); + idx.dataset_ = state->dataset(); + idx.tiered_ = std::move(state); + } +}; + +template +auto build_tiered(raft::resources const&, index_params const&, DatasetViewT const&) + -> cagra_index_t +{ + RAFT_FAIL("Tiered build requires enabled FlowANN building and a device-padded dense dataset"); +} + +template +void deserialize_tiered(raft::resources const&, std::istream&, index*) +{ + RAFT_FAIL("Tiered deserialization is disabled or the dataset/index type is unsupported"); +} + +#ifdef CUVS_ENABLE_FLOWANN_SEARCH +#define CUVS_DECLARE_TIERED_IO(T) \ + CUVS_EXPORT void deserialize_tiered( \ + raft::resources const&, std::istream&, device_padded_index*); \ + CUVS_EXPORT void deserialize_tiered( \ + raft::resources const&, std::istream&, device_pq_index*); +CUVS_DECLARE_TIERED_IO(float) +CUVS_DECLARE_TIERED_IO(half) +CUVS_DECLARE_TIERED_IO(std::int8_t) +CUVS_DECLARE_TIERED_IO(std::uint8_t) +#undef CUVS_DECLARE_TIERED_IO +#endif +#ifdef CUVS_ENABLE_FLOWANN_BUILD +#define CUVS_DECLARE_TIERED_BUILD(T) \ + CUVS_EXPORT auto build_tiered( \ + raft::resources const&, index_params const&, device_padded_dataset_view const&) \ + -> device_padded_index; +CUVS_DECLARE_TIERED_BUILD(float) +CUVS_DECLARE_TIERED_BUILD(half) +CUVS_DECLARE_TIERED_BUILD(std::int8_t) +CUVS_DECLARE_TIERED_BUILD(std::uint8_t) +#undef CUVS_DECLARE_TIERED_BUILD +#endif +} // namespace detail + +#ifdef CUVS_ENABLE_FLOWANN_BUILD +/** + * @brief Build an owning tiered VPQ index from host vectors of type `float`. + * + * Graph construction and grouping use full-precision input before compression. Requires + * `params.graph_storage = graph_storage_kind::tiered` and FlowANN build support. The returned + * index owns the reordered VPQ dataset. The input storage may be released after this call. + * + * @param[in] res RAFT resources used for construction and compression. + * @param[in] params Graph construction and tiered storage parameters. + * @param[in] compression VPQ compression parameters with FP16 codebooks and 8-bit PQ codes. + * @param[in] dataset Non-owning host view of the uncompressed input vectors. + * @return An owning index accepting query vectors of type `float`. + */ +CUVS_EXPORT auto build(raft::resources const& res, + index_params const& params, + cuvs::neighbors::vpq_params const& compression, + host_standard_dataset_view const& dataset) + -> device_pq_index; + +/** + * @brief Build an owning tiered VPQ index from host vectors of type `half`. + * + * Graph construction and grouping use full-precision input before compression. Requires + * `params.graph_storage = graph_storage_kind::tiered` and FlowANN build support. The returned + * index owns the reordered VPQ dataset. The input storage may be released after this call. + * + * @param[in] res RAFT resources used for construction and compression. + * @param[in] params Graph construction and tiered storage parameters. + * @param[in] compression VPQ compression parameters with FP16 codebooks and 8-bit PQ codes. + * @param[in] dataset Non-owning host view of the uncompressed input vectors. + * @return An owning index accepting query vectors of type `half`. + */ +CUVS_EXPORT auto build(raft::resources const& res, + index_params const& params, + cuvs::neighbors::vpq_params const& compression, + host_standard_dataset_view const& dataset) + -> device_pq_index; + +/** + * @brief Build an owning tiered VPQ index from host vectors of type `std::int8_t`. + * + * Graph construction and grouping use full-precision input before compression. Requires + * `params.graph_storage = graph_storage_kind::tiered` and FlowANN build support. The returned + * index owns the reordered VPQ dataset. The input storage may be released after this call. + * + * @param[in] res RAFT resources used for construction and compression. + * @param[in] params Graph construction and tiered storage parameters. + * @param[in] compression VPQ compression parameters with FP16 codebooks and 8-bit PQ codes. + * @param[in] dataset Non-owning host view of the uncompressed input vectors. + * @return An owning index accepting query vectors of type `std::int8_t`. + */ +CUVS_EXPORT auto build(raft::resources const& res, + index_params const& params, + cuvs::neighbors::vpq_params const& compression, + host_standard_dataset_view const& dataset) + -> device_pq_index; + +/** + * @brief Build an owning tiered VPQ index from host vectors of type `std::uint8_t`. + * + * Graph construction and grouping use full-precision input before compression. Requires + * `params.graph_storage = graph_storage_kind::tiered` and FlowANN build support. The returned + * index owns the reordered VPQ dataset. The input storage may be released after this call. + * + * @param[in] res RAFT resources used for construction and compression. + * @param[in] params Graph construction and tiered storage parameters. + * @param[in] compression VPQ compression parameters with FP16 codebooks and 8-bit PQ codes. + * @param[in] dataset Non-owning host view of the uncompressed input vectors. + * @return An owning index accepting query vectors of type `std::uint8_t`. + */ +CUVS_EXPORT auto build(raft::resources const& res, + index_params const& params, + cuvs::neighbors::vpq_params const& compression, + host_standard_dataset_view const& dataset) + -> device_pq_index; +#endif + /** * @defgroup cagra_cpp_index_build CAGRA index build functions * @{ @@ -1226,22 +1492,55 @@ auto build(raft::resources const& res, cuvs::neighbors::device_bbq_dataset_view dataset) -> cuvs::neighbors::cagra::device_bbq_index; -/** @copydoc build(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, - * cuvs::neighbors::device_bbq_dataset_view dataset) */ +/** + * @brief Build from a device BBQ-quantized dataset view of type `half`. + * + * Construction uses quantized codes and requires nn-descent. The metric must match the + * quantizer corrections and be L2Expanded, L2SqrtExpanded, CosineExpanded, or InnerProduct. + * The returned index has no BBQ search support; use the type-changing `update_dataset` + * with an uncompressed device-padded dataset before searching the resulting graph. + * + * @param[in] res RAFT resources used for graph construction. + * @param[in] params CAGRA index build parameters. + * @param[in] dataset Device BBQ dataset view with shape [n_rows, dim]. + * @return Built `device_bbq_index`. + */ auto build(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, cuvs::neighbors::device_bbq_dataset_view dataset) -> cuvs::neighbors::cagra::device_bbq_index; -/** @copydoc build(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, - * cuvs::neighbors::device_bbq_dataset_view dataset) */ +/** + * @brief Build from a device BBQ-quantized dataset view of type `int8_t`. + * + * Construction uses quantized codes and requires nn-descent. The metric must match the + * quantizer corrections and be L2Expanded, L2SqrtExpanded, CosineExpanded, or InnerProduct. + * The returned index has no BBQ search support; use the type-changing `update_dataset` + * with an uncompressed device-padded dataset before searching the resulting graph. + * + * @param[in] res RAFT resources used for graph construction. + * @param[in] params CAGRA index build parameters. + * @param[in] dataset Device BBQ dataset view with shape [n_rows, dim]. + * @return Built `device_bbq_index`. + */ auto build(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, cuvs::neighbors::device_bbq_dataset_view dataset) -> cuvs::neighbors::cagra::device_bbq_index; -/** @copydoc build(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, - * cuvs::neighbors::device_bbq_dataset_view dataset) */ +/** + * @brief Build from a device BBQ-quantized dataset view of type `uint8_t`. + * + * Construction uses quantized codes and requires nn-descent. The metric must match the + * quantizer corrections and be L2Expanded, L2SqrtExpanded, CosineExpanded, or InnerProduct. + * The returned index has no BBQ search support; use the type-changing `update_dataset` + * with an uncompressed device-padded dataset before searching the resulting graph. + * + * @param[in] res RAFT resources used for graph construction. + * @param[in] params CAGRA index build parameters. + * @param[in] dataset Device BBQ dataset view with shape [n_rows, dim]. + * @return Built `device_bbq_index`. + */ auto build(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, cuvs::neighbors::device_bbq_dataset_view dataset) diff --git a/cpp/src/neighbors/cagra.cuh b/cpp/src/neighbors/cagra.cuh index 550b317552..34320dc280 100644 --- a/cpp/src/neighbors/cagra.cuh +++ b/cpp/src/neighbors/cagra.cuh @@ -302,6 +302,11 @@ auto build(raft::resources const& res, const index_params& params, DatasetViewT using index_type = cuvs::neighbors::cagra::cagra_index_t; using T = typename index_type::value_type; using IdxT = uint32_t; + if (params.graph_storage == graph_storage_kind::tiered) { + RAFT_EXPECTS(params.attach_dataset_on_build, "Tiered builds must attach their dataset"); + return detail::build_tiered(res, params, dataset); + } + RAFT_EXPECTS(!params.tiered.has_value(), "Tiered options require graph_storage=tiered"); // Dense paths build the graph and optionally attach the input dataset view. Host indexes remain // non-searchable until the type-changing update_dataset(...) supplies a device-padded dataset. @@ -454,6 +459,10 @@ void search(raft::resources const& res, raft::device_matrix_view distances, const cuvs::neighbors::filtering::base_filter& sample_filter_ref) { + if (auto* tiered = detail::tiered_index_access::get(idx)) { + return tiered->search(res, params, queries, neighbors, distances, sample_filter_ref); + } + RAFT_EXPECTS(!params.tiered.has_value(), "Tiered search options require a tiered index"); try { using none_filter_type = cuvs::neighbors::filtering::none_sample_filter; auto& sample_filter = dynamic_cast(sample_filter_ref); @@ -586,6 +595,12 @@ void search( const std::vector>& partition_bitsets = {}) { RAFT_EXPECTS(!indices.empty(), "At least one index partition must be provided."); + RAFT_EXPECTS(!params.tiered.has_value(), + "Multi-partition search does not support tiered options"); + for (auto const* idx : indices) { + RAFT_EXPECTS(idx && idx->graph_storage() == graph_storage_kind::device, + "Multi-partition search does not support tiered graph inputs"); + } RAFT_EXPECTS(queries.extent(0) == partition_ids.extent(0) && queries.extent(0) == neighbors.extent(0) && diff --git a/cpp/src/neighbors/cagra_serialize.cuh b/cpp/src/neighbors/cagra_serialize.cuh index 18e807b90f..3aba47d1e1 100644 --- a/cpp/src/neighbors/cagra_serialize.cuh +++ b/cpp/src/neighbors/cagra_serialize.cuh @@ -159,7 +159,8 @@ namespace cuvs::neighbors::cagra { const std::string& filename, \ const cuvs::neighbors::cagra::device_pq_index& index) \ { \ - cuvs::neighbors::cagra::detail::serialize(handle, filename, index, false); \ + cuvs::neighbors::cagra::detail::serialize( \ + handle, filename, index, index.graph_storage() == graph_storage_kind::tiered); \ } \ \ void deserialize(raft::resources const& handle, \ @@ -174,7 +175,8 @@ namespace cuvs::neighbors::cagra { std::ostream& os, \ const cuvs::neighbors::cagra::device_pq_index& index) \ { \ - cuvs::neighbors::cagra::detail::serialize(handle, os, index, false); \ + cuvs::neighbors::cagra::detail::serialize( \ + handle, os, index, index.graph_storage() == graph_storage_kind::tiered); \ } \ \ void deserialize(raft::resources const& handle, \ diff --git a/cpp/src/neighbors/cagra_tiered.cu b/cpp/src/neighbors/cagra_tiered.cu new file mode 100644 index 0000000000..d8d4829990 --- /dev/null +++ b/cpp/src/neighbors/cagra_tiered.cu @@ -0,0 +1,243 @@ +/* + * 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 +#ifdef CUVS_ENABLE_FLOWANN_BUILD +#include +#endif +#include "detail/flowann/cagra_adapter.hpp" +#include + +namespace cuvs::neighbors::cagra::detail { +namespace flowann = cuvs::neighbors::cagra::experimental::flowann; + +// This owner is stable across moves of the public index. The context and session are destroyed +// before the graph/dataset bundle, so no polling thread can reference released graph storage. +template +class tiered_state final + : public tiered_index_state { + using view_type = typename decltype(Bundle::index)::dataset_view_type; + + public: + explicit tiered_state(Bundle&& bundle) : bundle_(std::move(bundle)) + { + RAFT_CUDA_TRY(cudaGetDevice(&device_)); + } + auto size() const noexcept -> uint32_t override { return bundle_.index.size(); } + auto graph_degree() const noexcept -> uint32_t override { return bundle_.index.graph_degree(); } + auto dataset() const noexcept -> view_type override { return bundle_.index.dataset(); } + auto source_indices() const noexcept + -> std::optional> override + { + return bundle_.index.source_indices(); + } + void serialize(raft::resources const& res, std::ostream& os, bool include_dataset) const override + { + RAFT_EXPECTS(include_dataset, "Tiered serialization currently requires include_dataset=true"); + std::lock_guard lock(mutex_); + int device; + RAFT_CUDA_TRY(cudaGetDevice(&device)); + RAFT_EXPECTS(device == device_, "Tiered serialization must use the index CUDA device"); + os.write("CTIR", 4); + raft::serialize_scalar(res, os, uint32_t{1}); + flowann::serialize(res, os, bundle_.index); + } + auto statistics() const -> flowann::queue_statistics + { + std::lock_guard lock(mutex_); + return context_ ? context_->statistics() : flowann::queue_statistics{}; + } + template + void run(raft::resources const& res, + search_params const& params, + raft::device_matrix_view queries, + raft::device_matrix_view neighbors, + raft::device_matrix_view distances, + cuvs::neighbors::filtering::base_filter const& filter) const + { + std::lock_guard lock(mutex_); + int device; + RAFT_CUDA_TRY(cudaGetDevice(&device)); + RAFT_EXPECTS(device == device_, "Tiered search must use the index's CUDA device"); + auto const options = params.tiered.value_or(tiered_search_params{}); + RAFT_EXPECTS(options.num_queues > 0, "Tiered search requires at least one queue"); + if (!context_ || options.num_queues != queue_.num_queues || + options.empty_pause != queue_.empty_pause || + options.collect_statistics != queue_.collect_statistics) { + session_.reset(); + context_.reset(); + flowann::queue_params queue{ + options.num_queues, options.empty_pause, options.collect_statistics}; + context_ = std::make_unique(bundle_.index.cross_graph(), queue); + queue_ = options; + } + if (!options.keep_pollers_running) { + session_.reset(); + } else if (!session_) { + session_ = std::make_unique(*context_); + } + flowann::search_params effective; + static_cast(effective) = params; + effective.tiered.reset(); + effective.num_seeds = options.num_seeds; + effective.sync_window_scale = options.sync_window_scale; + effective.sync_drop_threshold = options.sync_drop_threshold; + flowann::search( + res, effective, bundle_.index, *context_, queries, neighbors, distances, filter); + } +#define SEARCH_OVERRIDE(I) \ + void search(raft::resources const& res, \ + search_params const& params, \ + raft::device_matrix_view queries, \ + raft::device_matrix_view neighbors, \ + raft::device_matrix_view distances, \ + cuvs::neighbors::filtering::base_filter const& filter) const override \ + { \ + run(res, params, queries, neighbors, distances, filter); \ + } + SEARCH_OVERRIDE(uint32_t) + SEARCH_OVERRIDE(int64_t) +#undef SEARCH_OVERRIDE + private: + Bundle bundle_; + int device_{}; + mutable std::mutex mutex_; + mutable tiered_search_params queue_; + mutable std::unique_ptr context_; + mutable std::unique_ptr session_; +}; + +template +auto wrap_tiered(raft::resources const& res, Bundle&& bundle) +{ + using view_type = typename decltype(Bundle::index)::dataset_view_type; + index result(res, bundle.index.metric()); + auto state = std::make_shared>(std::move(bundle)); + tiered_index_access::install(result, std::move(state)); + return result; +} + +#define DEFINE_ADOPT_TIERED(T) \ + auto adopt_tiered(raft::resources const& res, flowann::device_padded_index_bundle&& bundle) \ + -> device_padded_index \ + { \ + return wrap_tiered(res, std::move(bundle)); \ + } \ + auto adopt_tiered(raft::resources const& res, flowann::vpq_f16_index_bundle&& bundle) \ + -> device_pq_index \ + { \ + return wrap_tiered(res, std::move(bundle)); \ + } +DEFINE_ADOPT_TIERED(float) +DEFINE_ADOPT_TIERED(half) +DEFINE_ADOPT_TIERED(int8_t) +DEFINE_ADOPT_TIERED(uint8_t) +#undef DEFINE_ADOPT_TIERED + +auto tiered_statistics(device_pq_index const& idx) -> flowann::queue_statistics +{ + auto* state = dynamic_cast>*>( + tiered_index_access::get(idx)); + RAFT_EXPECTS(state != nullptr, "Expected an owning tiered VPQ index"); + return state->statistics(); +} + +template +void read_tiered(raft::resources const& res, std::istream& is, index* idx) +{ + auto const version = raft::deserialize_scalar(res, is); + RAFT_EXPECTS(version == 1, "Unsupported CAGRA tiered serialization version"); + if constexpr (is_device_padded_dataset_view_v) { + auto bundle = flowann::deserialize_device_padded(res, is); + *idx = wrap_tiered(res, std::move(bundle)); + } else { + auto bundle = flowann::deserialize_vpq_f16(res, is); + *idx = wrap_tiered(res, std::move(bundle)); + } +} + +#define DEFINE_TIERED_IO(T) \ + void deserialize_tiered( \ + raft::resources const& res, std::istream& is, device_padded_index* idx) \ + { \ + read_tiered(res, is, idx); \ + } \ + void deserialize_tiered( \ + raft::resources const& res, std::istream& is, device_pq_index* idx) \ + { \ + read_tiered(res, is, idx); \ + } +DEFINE_TIERED_IO(float) +DEFINE_TIERED_IO(half) +DEFINE_TIERED_IO(int8_t) +DEFINE_TIERED_IO(uint8_t) +#undef DEFINE_TIERED_IO + +#ifdef CUVS_ENABLE_FLOWANN_BUILD +inline auto flowann_build_params(index_params const& params) +{ + auto const options = params.tiered.value_or(tiered_graph_params{}); + flowann::index_params out; + out.cagra_params = params; + out.cagra_params.graph_storage = graph_storage_kind::device; + out.cagra_params.tiered.reset(); + out.device_graph_budget_bytes = options.device_graph_budget_bytes; + out.node_per_cacheline = options.node_per_cacheline; + out.grouping.enabled = options.grouping_enabled; + out.grouping.n_groups = options.n_groups; + out.grouping.n_bits = options.n_bits; + out.grouping.balance_tolerance = options.balance_tolerance; + out.grouping.training_rows = options.training_rows; + out.grouping.assignment_batch_rows = options.assignment_batch_rows; + out.grouping.kmeans_n_iters = options.kmeans_n_iters; + out.grouping.validate = options.validate; + out.num_seeds = options.num_seeds; + out.seed_training_rows = options.seed_training_rows; + out.seed = options.seed; + return out; +} +#define DEFINE_TIERED_BUILD(T) \ + auto build_tiered(raft::resources const& res, \ + index_params const& params, \ + device_padded_dataset_view const& dataset) \ + -> device_padded_index \ + { \ + auto bundle = flowann::build(res, flowann_build_params(params), dataset); \ + return wrap_tiered(res, std::move(bundle)); \ + } +DEFINE_TIERED_BUILD(float) +DEFINE_TIERED_BUILD(half) +DEFINE_TIERED_BUILD(int8_t) +DEFINE_TIERED_BUILD(uint8_t) +#undef DEFINE_TIERED_BUILD +#endif +} // namespace cuvs::neighbors::cagra::detail + +#ifdef CUVS_ENABLE_FLOWANN_BUILD +namespace cuvs::neighbors::cagra { +#define DEFINE_TIERED_PQ_BUILD(T) \ + auto build(raft::resources const& res, \ + index_params const& params, \ + cuvs::neighbors::vpq_params const& compression, \ + host_standard_dataset_view const& dataset) -> device_pq_index \ + { \ + RAFT_EXPECTS(params.graph_storage == graph_storage_kind::tiered, \ + "This compressed build overload requires a tiered graph"); \ + RAFT_EXPECTS(params.attach_dataset_on_build, "Tiered builds must attach their dataset"); \ + auto bundle = \ + detail::flowann::build(res, detail::flowann_build_params(params), compression, dataset); \ + return detail::wrap_tiered(res, std::move(bundle)); \ + } +DEFINE_TIERED_PQ_BUILD(float) +DEFINE_TIERED_PQ_BUILD(half) +DEFINE_TIERED_PQ_BUILD(int8_t) +DEFINE_TIERED_PQ_BUILD(uint8_t) +#undef DEFINE_TIERED_PQ_BUILD +} // namespace cuvs::neighbors::cagra +#endif diff --git a/cpp/src/neighbors/detail/cagra/add_nodes.cuh b/cpp/src/neighbors/detail/cagra/add_nodes.cuh index 294aefd425..66e19ef889 100644 --- a/cpp/src/neighbors/detail/cagra/add_nodes.cuh +++ b/cpp/src/neighbors/detail/cagra/add_nodes.cuh @@ -356,6 +356,8 @@ void extend_core(raft::resources const& handle, { static_assert(cuvs::neighbors::is_padded_dataset_view_v, "cagra::extend requires a padded dataset view index type"); + RAFT_EXPECTS(index.graph_storage() == cagra::graph_storage_kind::device, + "Extending a tiered CAGRA graph is not supported"); RAFT_EXPECTS(!index.dataset_fd().has_value(), "Cannot extend a disk-backed CAGRA index. Convert it with " "cuvs::neighbors::hnsw::from_cagra() and load it into memory via " diff --git a/cpp/src/neighbors/detail/cagra/cagra_build.cuh b/cpp/src/neighbors/detail/cagra/cagra_build.cuh index 8c8766cbd5..074b071e4e 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_build.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_build.cuh @@ -719,6 +719,8 @@ auto iterative_build_graph(raft::resources const& res, { size_t intermediate_degree = params.intermediate_graph_degree; size_t graph_degree = params.graph_degree; + // Other callers also pass the original parameters after checking local degree copies. + check_graph_degree(intermediate_degree, graph_degree, dataset.n_rows()); const auto& iter_params = std::get(params.graph_build_params); @@ -755,11 +757,15 @@ auto iterative_build_graph(raft::resources const& res, final_graph_size = static_cast(dataset.n_rows()); } + // Every initial row must have enough distinct non-self neighbors for the + // first search-and-optimize iteration, including very small datasets. + auto small_graph_degree = std::max(graph_degree / 2, std::min(graph_degree, (uint64_t)24)); // Determine initial graph size. uint64_t initial_graph_size = (final_graph_size + 1) / 2; while (initial_graph_size > graph_degree * 64) { initial_graph_size = (initial_graph_size + 1) / 2; } + initial_graph_size = std::max(initial_graph_size, small_graph_degree + 1); RAFT_LOG_DEBUG("# initial graph size = %lu", (uint64_t)initial_graph_size); // Preallocate the kNN graph at last-iteration size (N × (intermediate_degree+1)) from the large @@ -793,7 +799,6 @@ auto iterative_build_graph(raft::resources const& res, // Determine graph degree and number of search results while increasing // graph size. - auto small_graph_degree = std::max(graph_degree / 2, std::min(graph_degree, (uint64_t)24)); RAFT_LOG_DEBUG("# small_graph_degree = %lu", (uint64_t)small_graph_degree); RAFT_LOG_DEBUG("# graph_degree = %lu", (uint64_t)graph_degree); RAFT_LOG_DEBUG("# topk = %lu", (uint64_t)topk); @@ -1062,6 +1067,61 @@ 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 adjusted_params = params; + adjusted_params.intermediate_graph_degree = intermediate_degree; + adjusted_params.graph_degree = graph_degree; + auto device_graph = [&]() { + if constexpr (cuvs::neighbors::is_host_dataset_view_v) { + auto padded = cuvs::neighbors::make_device_padded_dataset(res, dataset.data_view()); + auto graph = iterative_build_graph(res, adjusted_params, padded->view()); + raft::resource::sync_stream(res); + return graph; + } else { + return iterative_build_graph(res, adjusted_params, dataset); + } + }(); + auto host_graph = + raft::make_host_matrix(device_graph.extent(0), device_graph.extent(1)); + 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, + cuvs::neighbors::dense_row_major_view(dataset)); +} + /** * 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/cagra_merge.cuh b/cpp/src/neighbors/detail/cagra/cagra_merge.cuh index e33c13aa10..f8ec572660 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_merge.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_merge.cuh @@ -499,6 +499,13 @@ auto merge(raft::resources const& handle, raft::common::nvtx::range merge_scope( "cagra::merge(algo=%d,parts=%zu)", static_cast(merge_params.algo), indices.size()); + RAFT_EXPECTS(params.graph_storage == cagra::graph_storage_kind::device, + "Merging into a tiered graph is not supported"); + for (auto* idx : indices) { + RAFT_EXPECTS(idx && idx->graph_storage() == cagra::graph_storage_kind::device, + "Merging tiered graph inputs is not supported"); + } + RAFT_EXPECTS(merge_params.algo == cagra::merge_algo::AUTO || merge_params.algo == cagra::merge_algo::FASTENER || merge_params.algo == cagra::merge_algo::REBUILD, diff --git a/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh b/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh index 67947a9d2f..7a090c1a86 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh @@ -5,6 +5,8 @@ #pragma once +#include + #include #include #include @@ -114,6 +116,9 @@ void serialize(raft::resources const& res, { raft::common::nvtx::range fun_scope("cagra::serialize"); + if (auto* tiered = tiered_index_access::get(index_)) { + return tiered->serialize(res, os, include_dataset); + } RAFT_EXPECTS(!index_.dataset_fd().has_value(), "Cannot serialize a disk-backed CAGRA index. Convert it with " "cuvs::neighbors::hnsw::from_cagra() and load it into memory via " @@ -552,6 +557,11 @@ void deserialize_impl( char dtype_string[4]; RAFT_EXPECTS(is.read(dtype_string, 4), "cagra::deserialize: failed to read dtype prefix"); + if (std::memcmp(dtype_string, "CTIR", 4) == 0) { + deserialize_tiered(res, is, index_); + if (out_dataset != nullptr) { out_dataset->reset(); } + return; + } RAFT_EXPECTS(cuvs::util::validate_serialized_dtype(dtype_string, sizeof(dtype_string)), "cagra::deserialize: serialized dtype prefix does not match requested type"); 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..43da32e41f 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, @@ -73,7 +157,8 @@ RAFT_DEVICE_INLINE_FUNCTION void search_core( uint32_t* const num_executed_iterations, /* stats */ const IndexT graph_size, const uint32_t query_id_offset, // Offset to add to query_id when calling filter - cagra_sample_filter filter_payload) + cagra_sample_filter filter_payload, + const std::uint32_t query_stride = 0) { using DATA_T = DataT; using INDEX_T = IndexT; @@ -116,15 +201,21 @@ 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); // Get smem_ws_size_in_bytes directly from base descriptor uint32_t smem_ws_size_in_bytes = dataset_desc->smem_ws_size_in_bytes(); - auto smem_desc = - setup_workspace(dataset_desc, smem, queries_ptr, query_id); + auto smem_desc = setup_workspace( + dataset_desc, + smem, + query_stride == 0 ? queries_ptr + : queries_ptr + static_cast(query_id) * query_stride, + query_stride == 0 ? query_id : 0); auto* __restrict__ result_indices_buffer = reinterpret_cast(smem + smem_ws_size_in_bytes); @@ -147,20 +238,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 +341,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 +390,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 +508,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 +569,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..3e409d4b85 100644 --- a/cpp/src/neighbors/detail/cagra/jit_lto_kernels/search_single_cta_jit.cuh +++ b/cpp/src/neighbors/detail/cagra/jit_lto_kernels/search_single_cta_jit.cuh @@ -46,6 +46,74 @@ using cuvs::neighbors::detail::sample_filter; using cuvs::neighbors::cagra::detail::device::compute_distance_to_child_nodes_jit; using cuvs::neighbors::cagra::detail::device::compute_distance_to_random_nodes_jit; +template +struct device_graph_provider { + const IndexT* graph; + std::uint32_t degree; + + RAFT_DEVICE_INLINE_FUNCTION void initialize() const {} + + RAFT_DEVICE_INLINE_FUNCTION auto candidate_capacity(std::uint32_t search_width) const + -> std::uint32_t + { + return search_width * degree; + } + + RAFT_DEVICE_INLINE_FUNCTION auto result_buffer_size(std::uint32_t base_size, std::uint32_t) const + -> std::uint32_t + { + return base_size; + } + + RAFT_DEVICE_INLINE_FUNCTION auto seed_pointer(const IndexT* seeds, + std::uint32_t num_seeds, + std::uint32_t query_id) const -> const IndexT* + { + return seeds == nullptr ? nullptr : seeds + static_cast(num_seeds) * query_id; + } + + RAFT_DEVICE_INLINE_FUNCTION auto initial_candidate_count(std::uint32_t candidate_count, + std::uint32_t, + std::uint32_t, + std::uint32_t) const -> std::uint32_t + { + return candidate_count; + } + + RAFT_DEVICE_INLINE_FUNCTION auto initial_distillation_count(std::uint32_t count, + std::uint32_t) const -> std::uint32_t + { + return count; + } + + RAFT_DEVICE_INLINE_FUNCTION auto expand( + IndexT* result_indices, + DistanceT* result_distances, + const dataset_descriptor_base_t* dataset_desc, + IndexT* visited_hashmap, + std::uint32_t visited_hash_bitlen, + const IndexT* parent_indices, + const IndexT* internal_topk, + std::uint32_t search_width) const -> std::uint32_t + { + compute_distance_to_child_nodes_jit(result_indices, + result_distances, + dataset_desc, + graph, + degree, + visited_hashmap, + visited_hash_bitlen, + static_cast(nullptr), + 0u, + parent_indices, + internal_topk, + search_width); + return search_width * degree; + } + + RAFT_DEVICE_INLINE_FUNCTION void finalize() const {} +}; + // JIT search_core - setup_workspace/compute_distance via function pointers template > RAFT_DEVICE_INLINE_FUNCTION void search_core( uintptr_t result_indices_ptr, DistanceT* const result_distances_ptr, const std::uint32_t top_k, const DataT* const queries_ptr, - const IndexT* const knn_graph, + GraphProviderT graph_provider, const std::uint32_t graph_degree, const SourceIndexT* source_indices_ptr, const unsigned num_distilation, @@ -81,7 +150,8 @@ RAFT_DEVICE_INLINE_FUNCTION void search_core( const std::uint32_t query_id_offset, // Offset to add to query_id when calling filter const dataset_descriptor_base_t* dataset_desc, cagra_sample_filter filter_payload, - const IndexT graph_size = 0) // Original number of bits + const IndexT graph_size = 0, + const std::uint32_t query_stride = 0) // Zero retains contiguous query addressing. { using LOAD_T = device::LOAD_128BIT_T; @@ -111,7 +181,10 @@ RAFT_DEVICE_INLINE_FUNCTION void search_core( extern __shared__ uint8_t smem[]; // Layout of result_buffer - const auto result_buffer_size = internal_topk + (search_width * graph_degree); + const auto candidate_capacity = graph_provider.candidate_capacity(search_width); + const auto base_result_buffer_size = internal_topk + candidate_capacity; + const auto result_buffer_size = + graph_provider.result_buffer_size(base_result_buffer_size, num_seeds); const auto result_buffer_size_32 = raft::round_up_safe(result_buffer_size, 32); const auto small_hash_size = hashmap::get_size(small_hash_bitlen); @@ -119,8 +192,12 @@ RAFT_DEVICE_INLINE_FUNCTION void search_core( uint32_t dim = dataset_desc->args.dim; uint32_t smem_ws_size_in_bytes = dataset_desc->smem_ws_size_in_bytes(); - auto smem_desc = - setup_workspace(dataset_desc, smem, queries_ptr, query_id); + auto smem_desc = setup_workspace( + dataset_desc, + smem, + query_stride == 0 ? queries_ptr + : queries_ptr + static_cast(query_id) * query_stride, + query_stride == 0 ? query_id : 0); auto* __restrict__ result_indices_buffer = reinterpret_cast(smem + smem_ws_size_in_bytes); @@ -150,19 +227,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 +255,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 +313,7 @@ RAFT_DEVICE_INLINE_FUNCTION void search_core( result_distances_buffer + internal_topk, result_indices_buffer + internal_topk, max_candidates, - search_width * graph_degree, + computed_children_this_iteration, topk_ws, (iter == 0)); __syncthreads(); @@ -280,24 +368,30 @@ RAFT_DEVICE_INLINE_FUNCTION void search_core( __syncthreads(); // compute the norms between child nodes and query node using JIT version _CLK_START(); - compute_distance_to_child_nodes_jit( - result_indices_buffer + internal_topk, - result_distances_buffer + internal_topk, - smem_desc, - knn_graph, - graph_degree, - local_visited_hashmap_ptr, - hash_bitlen, - (IndexT*)nullptr, - 0u, - parent_list_buffer, - result_indices_buffer, - search_width); + computed_children_this_iteration = + graph_provider.expand(result_indices_buffer + internal_topk, + result_distances_buffer + internal_topk, + smem_desc, + local_visited_hashmap_ptr, + hash_bitlen, + parent_list_buffer, + result_indices_buffer, + search_width); // Critical: __syncthreads() must be reached by ALL threads // If any thread is stuck in compute_distance_to_child_nodes_jit, this will hang __syncthreads(); _CLK_REC(clk_compute_distance); + // Radix selection reads the entire result buffer. A provider may produce fewer + // children than its capacity, so old candidates must not survive in the tail. + for (std::uint32_t i = internal_topk + computed_children_this_iteration + threadIdx.x; + i < result_buffer_size; + i += blockDim.x) { + result_indices_buffer[i] = raft::upper_bound(); + result_distances_buffer[i] = raft::upper_bound(); + } + + // The filter-flag barrier below also orders these tail writes. // Filtering - use extern sample_filter function if (threadIdx.x == 0) { *filter_flag = 0; } __syncthreads(); @@ -322,11 +416,13 @@ RAFT_DEVICE_INLINE_FUNCTION void search_core( iter++; } + graph_provider.finalize(); + // Post process for filtering - use extern sample_filter function constexpr IndexT index_msb_1_mask = utils::gen_index_msb_1_mask::value; const IndexT invalid_index = utils::get_max_value(); - for (unsigned i = threadIdx.x; i < internal_topk + search_width * graph_degree; i += blockDim.x) { + for (unsigned i = threadIdx.x; i < result_buffer_size; i += blockDim.x) { const auto node_id = result_indices_buffer[i] & ~index_msb_1_mask; if (node_id != (invalid_index & ~index_msb_1_mask) && !sample_filter(query_id + query_id_offset, @@ -387,7 +483,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 +514,21 @@ 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) { + // Keep the historical resident-graph sentinel without indexing a source map. + if (source_indices_ptr == nullptr) { + write_indices( + result_indices_ptr, j, static_cast(invalid_index & ~index_msb_1_mask)); + } else if (index_element_tag == 3) { + reinterpret_cast(result_indices_ptr)[j] = ~uint64_t{0}; + } else { + write_indices(result_indices_ptr, j, raft::upper_bound()); + } + continue; + } + auto internal_index = result_index & ~index_msb_1_mask; // clear most significant bit + auto source_index = to_source_index(internal_index); write_indices(result_indices_ptr, j, source_index); } if (threadIdx.x == 0 && num_executed_iterations != nullptr) { @@ -495,33 +603,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 +714,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 +819,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..8b700d9ced --- /dev/null +++ b/cpp/src/neighbors/detail/flowann/README.md @@ -0,0 +1,44 @@ +# FlowANN implementation + +FlowANN is described in "Disentangling Graph Dependencies for Efficient Billion-Scale GPU Vector Search" (OSDI 2026). + +These are internal implementation headers. The integration is introduced in separate PRs; only the interfaces present in the current revision are available. + + +## Host service thread lifetime + +The experimental queue runtime uses `std::jthread` for persistent CPU pollers. They must +continue servicing device requests while the calling thread waits for CUDA completion, +and an explicit search session can keep them alive between calls. An OpenMP parallel +region scoped to one search call cannot provide that lifetime. Threads select the +context's CUDA device and are stopped and joined before queue storage is released. +This threading choice needs an explicit exception to the repository's OpenMP-only host +threading guideline during maintainer review; this rationale does not claim that approval. + +Pollers inherit the caller's CPU affinity by default. Internal callers may opt into +physical-core pinning with `queue_params::pin_pollers`; pinned contexts share a process-wide +allocation counter rather than all starting at CPU zero. + +## C++ compatibility + +The tiered fields change the layouts of `cagra::index`, `cagra::index_params`, and +`cagra::search_params`, including builds with FlowANN disabled. Rebuild consumers with +headers from the same cuVS version as the linked library. `index::graph()` throws for +tiered indexes because their device and host edges do not form one device adjacency matrix. +Use `graph_storage()` to inspect the representation before requesting a device graph. + + +## Deploying a GDRCopy-enabled library + +GDRCopy remains an external runtime dependency; installing cuVS does not bundle `libgdrapi`. +The runtime environment must provide a compatible `libgdrapi` and its SONAME links. Place them +beside the installed cuVS libraries, install them in the system loader paths, or set +`LD_LIBRARY_PATH` to the runtime GDRCopy directory before loading cuVS. The kernel driver is +also required for GDRCopy transfers. + +The installed RPATH defaults to `$ORIGIN`, so moving an installation does not retain the +build host's GDRCopy prefix. Packagers with a fixed runtime layout may set +`CUVS_FLOWANN_GDRCOPY_INSTALL_RPATH` to their deployment search paths. The same setting applies +to `libcuvs_c` when `CUVSC_STATIC_CUVS_LIBRARY=ON`, since that library then carries the GDRCopy +dependency directly. Compile-time package discovery still uses `GDRCopy_ROOT` or the usual +CMake search paths and does not establish the runtime layout. diff --git a/cpp/src/neighbors/detail/flowann/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/cagra_adapter.hpp b/cpp/src/neighbors/detail/flowann/cagra_adapter.hpp new file mode 100644 index 0000000000..10d6c9e2bc --- /dev/null +++ b/cpp/src/neighbors/detail/flowann/cagra_adapter.hpp @@ -0,0 +1,27 @@ +/* + * 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 + +namespace cuvs::neighbors::cagra::detail { +// Transitional internal bridge for legacy importers. Moves the owner without copying its buffers. +#define CUVS_DECLARE_ADOPT_TIERED(T) \ + CUVS_EXPORT auto adopt_tiered(raft::resources const&, \ + experimental::flowann::device_padded_index_bundle&&) \ + -> device_padded_index; \ + CUVS_EXPORT auto adopt_tiered(raft::resources const&, \ + experimental::flowann::vpq_f16_index_bundle&&) \ + -> device_pq_index; +CUVS_DECLARE_ADOPT_TIERED(float) +CUVS_DECLARE_ADOPT_TIERED(half) +CUVS_DECLARE_ADOPT_TIERED(std::int8_t) +CUVS_DECLARE_ADOPT_TIERED(std::uint8_t) +#undef CUVS_DECLARE_ADOPT_TIERED +CUVS_EXPORT auto tiered_statistics(device_pq_index const&) + -> experimental::flowann::queue_statistics; +} // namespace cuvs::neighbors::cagra::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..24acf63e9b --- /dev/null +++ b/cpp/src/neighbors/detail/flowann/flowann.hpp @@ -0,0 +1,836 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * FlowANN is described in "Disentangling Graph Dependencies for Efficient + * Billion-Scale GPU Vector Search" (OSDI 2026). + */ + +#pragma once + +#ifdef CUVS_ENABLE_FLOWANN_SEARCH + +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace cuvs::neighbors::cagra::experimental::flowann { + +namespace detail { +struct search_context_access; +} + +class search_session; + +/** Parameters controlling the host-serviced cross-graph queues. */ +struct queue_params { + /** Number of independent GPU-to-CPU command queues. */ + std::uint32_t num_queues = 1; + /** Number of CPU pause instructions issued after an empty poll. */ + std::uint32_t empty_pause = 64; + /** Collect aggregate queue and polling statistics. */ + bool collect_statistics = false; + /** Opt in to pinning pollers; otherwise inherit caller CPU affinity. */ + bool pin_pollers = false; +}; + +/** Aggregate counters collected by the host-serviced cross-graph queues. */ +struct queue_statistics { + std::uint64_t polls{}; + std::uint64_t empty_polls{}; + std::uint64_t processed_commands{}; +}; + +/** FlowANN search parameters. */ +struct search_params : cuvs::neighbors::cagra::search_params { + /** Maximum number of preloaded medoid seed nodes used to initialize each query. */ + std::uint32_t num_seeds = 0; + /** Scale used by the adaptive deferred-cross-edge synchronization window. */ + float sync_window_scale = 10.0f; + /** Parent-position drop that forces synchronization for a submitted cross-edge request. */ + std::uint32_t sync_drop_threshold = 51; +}; + +/** + * @brief Owns FlowANN command queues and their CPU polling threads. + * + * The cross-graph storage passed to the constructor must outlive this context. A context can be + * reused across search calls for the same index. Construction and search must use the same current + * CUDA device; CPU pollers select that device before servicing requests. + */ +class CUVS_EXPORT search_context { + public: + search_context(raft::host_matrix_view cross_graph, + queue_params params = {}); + ~search_context(); + + search_context(search_context const&) = delete; + auto operator=(search_context const&) -> search_context& = delete; + search_context(search_context&&) noexcept; + auto operator=(search_context&&) noexcept -> search_context&; + + /** + * Pause all CPU pollers after in-flight work has completed. + * + * @throws raft::logic_error if a search_session is active. + */ + void pause(); + /** Resume all CPU pollers before submitting search work. */ + void resume(); + [[nodiscard]] auto num_queues() const noexcept -> std::uint32_t; + /** Return aggregate counters. Counters are zero unless statistics were enabled. */ + [[nodiscard]] auto statistics() const noexcept -> queue_statistics; + + private: + struct impl; + std::shared_ptr impl_; + + friend struct detail::search_context_access; + friend class search_session; +}; + +/** + * Keep a search context's pollers running for the lifetime of this session. + * + * Use a session for repeated low-latency search calls. Moving a context preserves active + * sessions. The host cross graph must outlive both the context and its sessions. + */ +class CUVS_EXPORT search_session { + public: + explicit search_session(search_context& context); + ~search_session(); + + search_session(search_session const&) = delete; + auto operator=(search_session const&) -> search_session& = delete; + search_session(search_session&& other) noexcept; + auto operator=(search_session&& other) noexcept -> search_session&; + + private: + void reset() noexcept; + + std::shared_ptr context_; +}; + +/** + * @brief FlowANN tiered-graph index. + * + * The index stores a compressed inner graph in device memory and the cross-graph rows in host + * memory. The dataset is a non-owning view and its backing storage must outlive the index. + */ +template > +class index : public cuvs::neighbors::index { + public: + using value_type = T; + using index_type = IdxT; + using graph_index_type = std::uint32_t; + using inner_index_type = std::uint8_t; + using dataset_view_type = DatasetViewT; + using search_params_type = flowann::search_params; + + static_assert(!raft::is_narrowing_v, + "IdxT must be able to represent all uint32_t graph indices"); + + explicit index(raft::resources const& res, + cuvs::distance::DistanceType metric = cuvs::distance::DistanceType::L2Expanded, + std::uint16_t node_per_cacheline = 4, + std::uint16_t n_bits = 16) + : metric_(metric), + node_per_cacheline_(node_per_cacheline), + n_bits_(n_bits), + inner_graph_(raft::make_device_matrix(res, 0, 0)), + cross_graph_(raft::make_host_matrix(0, 0)), + dataset_norms_(std::nullopt) + { + validate_encoding_(); + } + + index(raft::resources const& res, + cuvs::distance::DistanceType metric, + DatasetViewT dataset, + std::uint16_t node_per_cacheline = 4, + std::uint16_t n_bits = 16) + : metric_(metric), + node_per_cacheline_(node_per_cacheline), + n_bits_(n_bits), + inner_graph_(raft::make_device_matrix(res, 0, 0)), + cross_graph_(raft::make_host_matrix(0, 0)), + dataset_(dataset), + dataset_norms_(std::nullopt) + { + validate_encoding_(); + if constexpr (cuvs::neighbors::is_device_padded_dataset_view_v) { + if (metric_ == cuvs::distance::DistanceType::CosineExpanded && dataset_.n_rows() > 0) { + compute_dataset_norms_(res); + raft::resource::sync_stream(res); + } + } + } + + index(index const&) = delete; + auto operator=(index const&) -> index& = delete; + index(index&&) = default; + auto operator=(index&&) -> index& = default; + ~index() = default; + + [[nodiscard]] constexpr auto metric() const noexcept -> cuvs::distance::DistanceType + { + return metric_; + } + + [[nodiscard]] auto size() const noexcept -> IdxT + { + return static_cast(cross_graph_.extent(0)); + } + + [[nodiscard]] auto dim() const noexcept -> std::uint32_t { return dataset_.dim(); } + + [[nodiscard]] auto graph_degree() const noexcept -> std::uint32_t { return full_graph_degree_; } + + /** Uniform number of neighbors stored in the packed resident graph. Zero for legacy layouts. */ + [[nodiscard]] constexpr auto resident_degree() const noexcept -> std::uint32_t + { + return resident_degree_; + } + + /** Maximum number of host-resident neighbors stored after the count column. */ + [[nodiscard]] auto cross_degree() const noexcept -> std::uint32_t + { + return cross_graph_.extent(1) == 0 ? 0 : static_cast(cross_graph_.extent(1) - 1); + } + + [[nodiscard]] constexpr auto node_per_cacheline() const noexcept -> std::uint16_t + { + return node_per_cacheline_; + } + + [[nodiscard]] constexpr auto n_bits() const noexcept -> std::uint16_t { return n_bits_; } + + [[nodiscard]] auto dataset() const noexcept -> DatasetViewT const& { return dataset_; } + + /** Dataset norms used by cosine distance. */ + [[nodiscard]] auto dataset_norms() const noexcept + -> std::optional> + { + if (!dataset_norms_.has_value()) { return std::nullopt; } + return raft::make_const_mdspan(dataset_norms_->view()); + } + + void update_dataset(raft::resources const& res, DatasetViewT dataset) + { + if (cross_graph_.extent(0) > 0) { + RAFT_EXPECTS(dataset.n_rows() == cross_graph_.extent(0), + "Dataset and FlowANN graph must have the same number of rows"); + } + dataset_ = dataset; + dataset_norms_.reset(); + if constexpr (cuvs::neighbors::is_device_padded_dataset_view_v) { + if (metric_ == cuvs::distance::DistanceType::CosineExpanded && dataset_.n_rows() > 0) { + compute_dataset_norms_(res); + raft::resource::sync_stream(res); + } + } + } + + [[nodiscard]] auto inner_graph() const noexcept + -> raft::device_matrix_view + { + return raft::make_const_mdspan(inner_graph_.view()); + } + + /** + * Cross-graph rows have shape `[size, cross_degree + 1]`. Element zero stores the number of valid + * cross-graph neighbors in the rest of the row. `cross_degree` may be smaller than the complete + * graph degree when every node has a uniform resident degree. + */ + [[nodiscard]] auto cross_graph() const noexcept + -> raft::host_matrix_view + { + return raft::make_const_mdspan(cross_graph_.view()); + } + + [[nodiscard]] auto source_indices() const noexcept + -> std::optional> + { + if (!source_indices_.has_value()) { return std::nullopt; } + return raft::make_const_mdspan(source_indices_->view()); + } + + [[nodiscard]] auto subgraph_offsets() const noexcept + -> std::optional> + { + if (!subgraph_offsets_.has_value()) { return std::nullopt; } + return raft::make_const_mdspan(subgraph_offsets_->view()); + } + + /** Subgraph identifier for every internal node. */ + [[nodiscard]] auto subgraph_ids() const noexcept + -> std::optional> + { + if (!subgraph_ids_.has_value()) { return std::nullopt; } + return raft::make_const_mdspan(subgraph_ids_->view()); + } + + /** Optional medoid seed nodes shared by all queries. */ + [[nodiscard]] auto seeds() const noexcept + -> std::optional> + { + if (!seeds_.has_value()) { return std::nullopt; } + return raft::make_const_mdspan(seeds_->view()); + } + + /** Cached cross-graph rows for the fixed medoid seeds. */ + [[nodiscard]] auto seed_cross_graph() const noexcept + -> std::optional> + { + if (!seed_cross_graph_.has_value()) { return std::nullopt; } + return raft::make_const_mdspan(seed_cross_graph_->view()); + } + + [[nodiscard]] auto seed_lookup_keys() const noexcept + -> std::optional> + { + if (!seed_lookup_keys_.has_value()) { return std::nullopt; } + return raft::make_const_mdspan(seed_lookup_keys_->view()); + } + + [[nodiscard]] auto seed_lookup_values() const noexcept + -> std::optional> + { + if (!seed_lookup_values_.has_value()) { return std::nullopt; } + return raft::make_const_mdspan(seed_lookup_values_->view()); + } + + void update_graph( + raft::resources const& res, + raft::host_matrix_view inner_graph, + raft::host_matrix_view cross_graph) + { + auto const inferred_degree = + cross_graph.extent(1) == 0 ? 0u : static_cast(cross_graph.extent(1) - 1); + update_graph(res, inner_graph, cross_graph, inferred_degree, 0); + } + + /** Replace the graph and explicitly describe a compact uniform-residency layout. */ + void update_graph( + raft::resources const& res, + raft::host_matrix_view inner_graph, + raft::host_matrix_view cross_graph, + std::uint32_t full_graph_degree, + std::uint32_t resident_degree) + { + validate_graph_(inner_graph, cross_graph, full_graph_degree, resident_degree); + + auto new_inner = raft::make_device_matrix( + res, inner_graph.extent(0), inner_graph.extent(1)); + raft::copy(new_inner.data_handle(), + inner_graph.data_handle(), + inner_graph.size(), + raft::resource::get_cuda_stream(res)); + // A non-owning host view may refer to pinned memory released by the caller on return. + raft::resource::sync_stream(res); + + auto new_cross = raft::make_host_matrix(cross_graph.extent(0), + cross_graph.extent(1)); + std::copy_n(cross_graph.data_handle(), cross_graph.size(), new_cross.data_handle()); + + auto new_seed_cross_graph = + prepare_graph_update_(res, raft::make_const_mdspan(new_cross.view())); + inner_graph_ = std::move(new_inner); + cross_graph_ = std::move(new_cross); + full_graph_degree_ = full_graph_degree; + resident_degree_ = resident_degree; + seed_cross_graph_ = std::move(new_seed_cross_graph); + } + + /** Replace the graph while taking ownership of the host cross-graph allocation. */ + void update_graph(raft::resources const& res, + raft::host_matrix&& inner_graph, + raft::host_matrix&& cross_graph) + { + auto const inferred_degree = + cross_graph.extent(1) == 0 ? 0u : static_cast(cross_graph.extent(1) - 1); + update_graph(res, std::move(inner_graph), std::move(cross_graph), inferred_degree, 0); + } + + /** Replace the graph while taking ownership of a compact uniform-residency layout. */ + void update_graph(raft::resources const& res, + raft::host_matrix&& inner_graph, + raft::host_matrix&& cross_graph, + std::uint32_t full_graph_degree, + std::uint32_t resident_degree) + { + validate_graph_(raft::make_const_mdspan(inner_graph.view()), + raft::make_const_mdspan(cross_graph.view()), + full_graph_degree, + resident_degree); + + auto new_inner = raft::make_device_matrix( + res, inner_graph.extent(0), inner_graph.extent(1)); + raft::copy(new_inner.data_handle(), + inner_graph.data_handle(), + inner_graph.size(), + raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + + auto new_seed_cross_graph = + prepare_graph_update_(res, raft::make_const_mdspan(cross_graph.view())); + inner_graph_ = std::move(new_inner); + cross_graph_ = std::move(cross_graph); + full_graph_degree_ = full_graph_degree; + resident_degree_ = resident_degree; + seed_cross_graph_ = std::move(new_seed_cross_graph); + } + + /** Replace a variable-residency grouped graph while retaining the packed rows on the device. */ + void update_graph(raft::resources const& res, + raft::device_matrix&& inner_graph, + raft::host_matrix&& cross_graph, + bool validate_edges) + { + RAFT_EXPECTS(cross_graph.extent(0) > 0 && cross_graph.extent(1) > 1, + "FlowANN grouped cross graph must contain rows and neighbors"); + auto const full_degree = static_cast(cross_graph.extent(1) - 1); + auto const expected_rows = + raft::div_rounding_up_safe(cross_graph.extent(0), node_per_cacheline_); + auto const has_inner_graph = inner_graph.extent(0) != 0 || inner_graph.extent(1) != 0; + if (has_inner_graph) { + RAFT_EXPECTS(inner_graph.extent(0) == expected_rows, + "FlowANN grouped inner graph has an invalid row count"); + RAFT_EXPECTS(inner_graph.extent(1) >= node_per_cacheline_, + "FlowANN grouped inner graph rows cannot hold their count headers"); + } + if (dataset_.n_rows() > 0) { + RAFT_EXPECTS(dataset_.n_rows() == cross_graph.extent(0), + "Dataset and FlowANN grouped graph must have the same number of rows"); + } + for (int64_t node = 0; node < cross_graph.extent(0); ++node) { + auto const cross_count = cross_graph(node, 0); + RAFT_EXPECTS(cross_count <= full_degree, + "FlowANN grouped cross-neighbor count exceeds graph degree at row %ld", + static_cast(node)); + if (validate_edges) { + for (std::uint32_t edge = 0; edge < cross_count; ++edge) { + RAFT_EXPECTS(cross_graph(node, edge + 1) < cross_graph.extent(0), + "FlowANN grouped cross graph contains an out-of-range node at row %ld", + static_cast(node)); + } + } + } + + auto new_seed_cross_graph = + prepare_graph_update_(res, raft::make_const_mdspan(cross_graph.view())); + inner_graph_ = std::move(inner_graph); + cross_graph_ = std::move(cross_graph); + full_graph_degree_ = full_degree; + resident_degree_ = 0; + seed_cross_graph_ = std::move(new_seed_cross_graph); + } + + /** Replace a compact graph while taking ownership of both device and host allocations. */ + void update_graph(raft::resources const& res, + raft::device_matrix&& inner_graph, + raft::host_matrix&& cross_graph, + std::uint32_t full_graph_degree, + std::uint32_t resident_degree) + { + RAFT_EXPECTS(cross_graph.extent(0) > 0 && cross_graph.extent(1) > 0, + "FlowANN cross graph must contain rows and a neighbor-count column"); + RAFT_EXPECTS(full_graph_degree > 0 && resident_degree <= full_graph_degree, + "FlowANN compact graph degree metadata is invalid"); + auto const cross_degree = static_cast(cross_graph.extent(1) - 1); + RAFT_EXPECTS(resident_degree + cross_degree == full_graph_degree, + "FlowANN compact graph degrees do not add up to the complete graph degree"); + if (resident_degree == 0) { + RAFT_EXPECTS(inner_graph.extent(0) == 0 && inner_graph.extent(1) == 0, + "FlowANN zero-residency graph must not allocate packed rows"); + } else { + auto const expected_rows = + raft::div_rounding_up_safe(cross_graph.extent(0), node_per_cacheline_); + auto const required_bits = + static_cast(node_per_cacheline_) * 8 + + static_cast(node_per_cacheline_) * resident_degree * n_bits_; + RAFT_EXPECTS(inner_graph.extent(0) == expected_rows, + "FlowANN compact inner graph has an invalid row count"); + RAFT_EXPECTS(static_cast(inner_graph.extent(1)) * 8 >= required_bits, + "FlowANN compact inner graph rows are too short"); + } + if (dataset_.n_rows() > 0) { + RAFT_EXPECTS(dataset_.n_rows() == cross_graph.extent(0), + "Dataset and FlowANN compact graph must have the same number of rows"); + } + for (int64_t node = 0; node < cross_graph.extent(0); ++node) { + RAFT_EXPECTS(cross_graph(node, 0) == cross_degree, + "FlowANN uniform compact graph has an invalid cross count at row %ld", + static_cast(node)); + for (std::uint32_t edge = 0; edge < cross_degree; ++edge) { + RAFT_EXPECTS(cross_graph(node, edge + 1) < cross_graph.extent(0), + "FlowANN compact cross graph contains an out-of-range node at row %ld", + static_cast(node)); + } + } + + auto new_seed_cross_graph = + prepare_graph_update_(res, raft::make_const_mdspan(cross_graph.view())); + inner_graph_ = std::move(inner_graph); + cross_graph_ = std::move(cross_graph); + full_graph_degree_ = full_graph_degree; + resident_degree_ = resident_degree; + seed_cross_graph_ = std::move(new_seed_cross_graph); + } + + void update_source_indices( + raft::resources const& res, + raft::host_vector_view source_indices) + { + RAFT_EXPECTS(source_indices.extent(0) == static_cast(size()), + "FlowANN source-index mapping must have one entry per graph row"); + source_indices_.emplace(raft::make_device_vector(res, source_indices.extent(0))); + raft::copy(source_indices_->data_handle(), + source_indices.data_handle(), + source_indices.extent(0), + raft::resource::get_cuda_stream(res)); + } + + void update_seeds(raft::resources const& res, + raft::host_vector_view seeds) + { + RAFT_EXPECTS(seeds.extent(0) > 0, "FlowANN seeds must not be empty"); + RAFT_EXPECTS(seeds.extent(0) <= static_cast(size()), + "FlowANN seed count must not exceed the index size"); + for (int64_t i = 0; i < seeds.extent(0); ++i) { + RAFT_EXPECTS( + seeds(i) < size(), "FlowANN seed at position %ld is out of range", static_cast(i)); + } + seeds_.emplace(raft::make_device_vector(res, seeds.extent(0))); + raft::copy(seeds_->data_handle(), + seeds.data_handle(), + seeds.extent(0), + raft::resource::get_cuda_stream(res)); + seed_ids_host_.assign(seeds.data_handle(), seeds.data_handle() + seeds.extent(0)); + rebuild_seed_lookup_(res); + rebuild_seed_cross_graph_(res); + } + + void update_subgraph_layout( + raft::resources const& res, + raft::host_vector_view offsets, + raft::host_vector_view subgraph_ids) + { + RAFT_EXPECTS(offsets.extent(0) > 0, "FlowANN subgraph offsets must not be empty"); + RAFT_EXPECTS(offsets.extent(0) <= static_cast(size()), + "FlowANN subgraph count must not exceed the index size"); + RAFT_EXPECTS(subgraph_ids.extent(0) == static_cast(size()), + "FlowANN must store one subgraph identifier per graph row"); + RAFT_EXPECTS(offsets(0) == 0, "FlowANN first subgraph offset must be zero"); + for (int64_t i = 1; i < offsets.extent(0); ++i) { + RAFT_EXPECTS(offsets(i - 1) < offsets(i), + "FlowANN subgraph offsets must be strictly increasing"); + } + RAFT_EXPECTS(offsets(offsets.extent(0) - 1) < static_cast(size()), + "FlowANN subgraph offsets exceed the index size"); + for (int64_t i = 0; i < subgraph_ids.extent(0); ++i) { + RAFT_EXPECTS(subgraph_ids(i) < offsets.extent(0), + "FlowANN subgraph identifier is out of range at row %ld", + static_cast(i)); + } + + subgraph_offsets_.emplace( + raft::make_device_vector(res, offsets.extent(0))); + raft::copy(subgraph_offsets_->data_handle(), + offsets.data_handle(), + offsets.extent(0), + raft::resource::get_cuda_stream(res)); + subgraph_ids_.emplace( + raft::make_device_vector(res, subgraph_ids.extent(0))); + raft::copy(subgraph_ids_->data_handle(), + subgraph_ids.data_handle(), + subgraph_ids.extent(0), + raft::resource::get_cuda_stream(res)); + } + + private: + void validate_graph_( + raft::host_matrix_view inner_graph, + raft::host_matrix_view cross_graph, + std::uint32_t full_graph_degree, + std::uint32_t resident_degree) const + { + RAFT_EXPECTS(cross_graph.extent(1) > 0, + "FlowANN cross graph must include a neighbor-count column"); + RAFT_EXPECTS(static_cast(cross_graph.extent(0)) <= + static_cast(std::numeric_limits::max()), + "FlowANN graph row count exceeds the index type"); + auto const expected_inner_rows = + raft::div_rounding_up_safe(cross_graph.extent(0), node_per_cacheline_); + auto const has_inner_graph = inner_graph.extent(0) != 0 || inner_graph.extent(1) != 0; + if (has_inner_graph) { + RAFT_EXPECTS(inner_graph.extent(0) == expected_inner_rows, + "FlowANN inner graph must contain one packed row per node cacheline"); + RAFT_EXPECTS(inner_graph.extent(1) >= node_per_cacheline_, + "FlowANN inner graph rows are too short to contain neighbor counts"); + } else { + RAFT_EXPECTS(resident_degree == 0, + "FlowANN may omit the packed graph only when resident degree is zero"); + } + RAFT_EXPECTS(full_graph_degree > 0, "FlowANN complete graph degree must be positive"); + RAFT_EXPECTS(resident_degree <= full_graph_degree, + "FlowANN resident graph degree exceeds the complete graph degree"); + auto const cross_degree = static_cast(cross_graph.extent(1) - 1); + RAFT_EXPECTS(cross_degree <= full_graph_degree, + "FlowANN cross graph degree exceeds the complete graph degree"); + if (resident_degree > 0 || cross_degree < full_graph_degree) { + RAFT_EXPECTS(resident_degree + cross_degree == full_graph_degree, + "FlowANN compact graph degrees do not add up to the complete graph degree"); + } + if (dataset_.n_rows() > 0) { + RAFT_EXPECTS(dataset_.n_rows() == cross_graph.extent(0), + "Dataset and FlowANN cross graph must have the same number of rows"); + } + + for (int64_t node = 0; node < cross_graph.extent(0); ++node) { + auto const cross_count = cross_graph(node, 0); + RAFT_EXPECTS(cross_count <= cross_degree, + "FlowANN cross-neighbor count exceeds graph degree at row %ld", + static_cast(node)); + for (std::uint32_t edge = 0; edge < cross_count; ++edge) { + RAFT_EXPECTS(cross_graph(node, edge + 1) < cross_graph.extent(0), + "FlowANN cross graph contains an out-of-range node at row %ld", + static_cast(node)); + } + + auto const packed_row = node / node_per_cacheline_; + auto const packed_col = node % node_per_cacheline_; + auto const inner_count = + has_inner_graph ? static_cast(inner_graph(packed_row, packed_col)) : 0u; + RAFT_EXPECTS(inner_count + cross_count <= full_graph_degree, + "FlowANN inner and cross neighbors exceed graph degree at row %ld", + static_cast(node)); + } + for (int64_t packed_row = 0; packed_row < inner_graph.extent(0); ++packed_row) { + std::uint64_t payload_bits = static_cast(node_per_cacheline_) * 8; + for (std::uint16_t packed_col = 0; packed_col < node_per_cacheline_; ++packed_col) { + auto const node = packed_row * node_per_cacheline_ + packed_col; + if (node < cross_graph.extent(0)) { + payload_bits += static_cast(inner_graph(packed_row, packed_col)) * n_bits_; + } + } + RAFT_EXPECTS(payload_bits <= static_cast(inner_graph.extent(1)) * 8, + "FlowANN inner graph payload exceeds its packed row at row %ld", + static_cast(packed_row)); + } + } + + static auto seed_hash_(IdxT key, std::uint32_t capacity) -> std::uint32_t + { + std::uint64_t value = static_cast(key); + value ^= value >> 16; + value *= 0x85ebca77u; + value ^= value >> 13; + value *= 0xc2b2ae3du; + value ^= value >> 16; + return static_cast(value % capacity); + } + + void rebuild_seed_lookup_(raft::resources const& res) + { + if (seed_ids_host_.empty()) { + seed_lookup_keys_.reset(); + seed_lookup_values_.reset(); + return; + } + RAFT_EXPECTS(seed_ids_host_.size() <= + (static_cast(std::numeric_limits::max()) - 1) / 2, + "FlowANN seed lookup exceeds uint32 capacity"); + + auto const capacity = static_cast(seed_ids_host_.size() * 2 + 1); + auto const empty_key = std::numeric_limits::max(); + auto const empty_value = std::numeric_limits::max(); + std::vector host_keys(capacity, empty_key); + std::vector host_values(capacity, empty_value); + + for (std::uint32_t seed = 0; seed < seed_ids_host_.size(); ++seed) { + auto const key = seed_ids_host_[seed]; + auto slot = seed_hash_(key, capacity); + while (host_keys[slot] != empty_key) { + slot = (slot + 1) % capacity; + } + host_keys[slot] = key; + host_values[slot] = seed; + } + + seed_lookup_keys_.emplace(raft::make_device_vector(res, capacity)); + seed_lookup_values_.emplace(raft::make_device_vector(res, capacity)); + auto const stream = raft::resource::get_cuda_stream(res); + raft::copy(seed_lookup_keys_->data_handle(), host_keys.data(), capacity, stream); + raft::copy(seed_lookup_values_->data_handle(), host_values.data(), capacity, stream); + } + + // Prepare all potentially failing work before committing a graph replacement. + auto prepare_graph_update_( + raft::resources const& res, + raft::host_matrix_view incoming) + -> std::optional> + { + for (std::size_t seed = 0; seed < seed_ids_host_.size(); ++seed) { + RAFT_EXPECTS(seed_ids_host_[seed] < incoming.extent(0), + "FlowANN seed at position %zu is out of range for the new graph", + seed); + } + if (incoming.extent(0) != cross_graph_.extent(0)) { + source_indices_.reset(); + subgraph_offsets_.reset(); + subgraph_ids_.reset(); + seeds_.reset(); + seed_ids_host_.clear(); + seed_lookup_keys_.reset(); + seed_lookup_values_.reset(); + return std::nullopt; + } + if (seed_ids_host_.empty()) { return std::nullopt; } + auto host_cache = raft::make_host_matrix( + static_cast(seed_ids_host_.size()), incoming.extent(1)); + for (std::size_t seed = 0; seed < seed_ids_host_.size(); ++seed) { + std::copy_n(incoming.data_handle() + + static_cast(seed_ids_host_[seed]) * incoming.extent(1), + incoming.extent(1), + host_cache.data_handle() + seed * incoming.extent(1)); + } + auto device_cache = raft::make_device_matrix( + res, host_cache.extent(0), host_cache.extent(1)); + raft::copy(device_cache.data_handle(), + host_cache.data_handle(), + host_cache.size(), + raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + return device_cache; + } + + void rebuild_seed_cross_graph_(raft::resources const& res) + { + if (seed_ids_host_.empty()) { + seed_cross_graph_.reset(); + return; + } + RAFT_EXPECTS(cross_graph_.extent(0) > 0 && cross_graph_.extent(1) > 0, + "FlowANN seed caching requires a non-empty cross graph"); + + auto host_cache = raft::make_host_matrix( + static_cast(seed_ids_host_.size()), cross_graph_.extent(1)); + for (std::size_t seed = 0; seed < seed_ids_host_.size(); ++seed) { + auto const node = seed_ids_host_[seed]; + RAFT_EXPECTS(node < size(), "FlowANN seed at position %zu is out of range", seed); + std::copy_n( + cross_graph_.data_handle() + static_cast(node) * cross_graph_.extent(1), + cross_graph_.extent(1), + host_cache.data_handle() + seed * cross_graph_.extent(1)); + } + + seed_cross_graph_.emplace(raft::make_device_matrix( + res, host_cache.extent(0), host_cache.extent(1))); + raft::copy(seed_cross_graph_->data_handle(), + host_cache.data_handle(), + host_cache.size(), + raft::resource::get_cuda_stream(res)); + } + + void validate_encoding_() const + { + RAFT_EXPECTS(node_per_cacheline_ > 0, "FlowANN node_per_cacheline must be positive"); + RAFT_EXPECTS(n_bits_ > 0 && n_bits_ <= 32, "FlowANN n_bits must be in [1, 32]"); + } + + CUVS_EXPORT void compute_dataset_norms_(raft::resources const& res); + + cuvs::distance::DistanceType metric_; + std::uint16_t node_per_cacheline_; + std::uint16_t n_bits_; + std::uint32_t full_graph_degree_ = 0; + std::uint32_t resident_degree_ = 0; + raft::device_matrix inner_graph_; + raft::host_matrix cross_graph_; + DatasetViewT dataset_{cuvs::neighbors::detail::make_empty_dataset_view(0)}; + std::optional> dataset_norms_; + std::optional> source_indices_; + std::optional> subgraph_offsets_; + std::optional> subgraph_ids_; + std::optional> seeds_; + std::vector seed_ids_host_; + std::optional> seed_lookup_keys_; + std::optional> seed_lookup_values_; + std::optional> seed_cross_graph_; +}; + +template +using device_padded_index = index>; + +template +using vpq_f16_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>; + +template +using vpq_f16_index_bundle = + index_bundle, flowann::vpq_f16_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(T, vpq_f16_index, std::uint32_t); \ + CUVS_FLOWANN_DECLARE_SEARCH(T, vpq_f16_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..da73855226 --- /dev/null +++ b/cpp/src/neighbors/detail/flowann/flowann_build.hpp @@ -0,0 +1,204 @@ +/* + * 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 host + * overload ranks with the full-precision input and produces an owning VPQ-F16 bundle using + * `vpq_params`. Both overloads retain 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 host + * overload additionally accepts `vpq_params` for the output dataset. Each 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_EXPORT auto build(raft::resources const& res, \ + index_params const& params, \ + cuvs::neighbors::vpq_params const& vpq_params, \ + cuvs::neighbors::host_standard_dataset_view const& dataset) \ + -> vpq_f16_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. The rvalue VPQ overload releases the full CAGRA graph before + * allocating the FlowANN graph to bound peak device memory use. + */ +#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_EXPORT auto build_from_cagra( \ + raft::resources const& res, \ + build_params const& params, \ + cuvs::neighbors::cagra::device_pq_index const& cagra_index, \ + raft::host_vector_view partition_labels, \ + std::optional> seeds = std::nullopt) \ + -> vpq_f16_index_bundle; \ + CUVS_EXPORT auto build_from_cagra( \ + raft::resources const& res, \ + build_params const& params, \ + cuvs::neighbors::cagra::device_pq_index&& cagra_index, \ + std::unique_ptr>&& vpq_dataset, \ + raft::host_vector_view partition_labels, \ + std::optional> seeds = std::nullopt) \ + -> vpq_f16_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/flowann_serialize.hpp b/cpp/src/neighbors/detail/flowann/flowann_serialize.hpp new file mode 100644 index 0000000000..9a9d7adc5a --- /dev/null +++ b/cpp/src/neighbors/detail/flowann/flowann_serialize.hpp @@ -0,0 +1,939 @@ +/* + * 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 + +namespace cuvs::neighbors::cagra::experimental::flowann { + +namespace detail { + +inline constexpr int legacy_serialization_version = 5; +inline constexpr std::uint32_t legacy_vpq_dataset_tag = 3; +inline constexpr std::uint32_t legacy_vpq_codebook_bits = 8; +inline constexpr int previous_serialization_version = 1; +inline constexpr int serialization_version = 2; + +enum class serialized_dataset_kind : std::uint32_t { dense = 1, vpq_f16 = 2 }; + +enum class legacy_graph_scalar_layout { typed_enums, historical_unsigned_enums }; + +template +inline auto is_supported_metric(serialized_dataset_kind dataset_kind, + cuvs::distance::DistanceType metric) -> bool +{ + using cuvs::distance::DistanceType; + if (dataset_kind == serialized_dataset_kind::vpq_f16) { + return metric == DistanceType::L2Expanded; + } + switch (metric) { + case DistanceType::L2Expanded: + case DistanceType::InnerProduct: + case DistanceType::CosineExpanded: + case DistanceType::L1: return true; + case DistanceType::BitwiseHamming: return std::is_same_v; + default: return false; + } +} + +template +void validate_dtype_prefix(char const* serialized_dtype, char const* component) +{ + auto expected_dtype = raft::numpy_serializer::get_numpy_dtype().to_string(); + expected_dtype.resize(4); + RAFT_EXPECTS(std::memcmp(serialized_dtype, expected_dtype.data(), 4) == 0, + "FlowANN legacy %s dtype does not match the requested type", + component); +} + +template +void read_and_validate_dtype(std::istream& is, char const* component) +{ + char serialized_dtype[4]{}; + RAFT_EXPECTS(static_cast(is.read(serialized_dtype, sizeof(serialized_dtype))), + "FlowANN legacy %s is truncated before its dtype prefix", + component); + validate_dtype_prefix(serialized_dtype, component); +} + +inline auto is_legacy_combined_dtype_prefix(char const* serialized_dtype) -> bool +{ + constexpr char legacy_dtype[4] = {'|', 'u', '1', '\0'}; + return std::memcmp(serialized_dtype, legacy_dtype, sizeof(legacy_dtype)) == 0; +} + +inline void validate_matrix_size(std::uint64_t rows, + std::uint64_t cols, + std::size_t element_size, + char const* component) +{ + RAFT_EXPECTS(cols == 0 || rows <= std::numeric_limits::max() / cols, + "FlowANN legacy %s shape overflows", + component); + auto const elements = rows * cols; + RAFT_EXPECTS( + element_size == 0 || elements <= std::numeric_limits::max() / element_size, + "FlowANN legacy %s byte size overflows", + component); +} + +// File and string streams can reject impossible payloads before allocating. A non-seekable +// stream retains its original state and falls back to the ordinary checked read path. +inline void validate_remaining_payload(std::istream& is, + std::uint64_t count, + std::size_t element_size, + char const* component) +{ + auto const state = is.rdstate(); + auto const position = is.tellg(); + if (position == std::istream::pos_type(-1)) { + is.clear(state); + return; + } + is.seekg(0, std::ios::end); + auto const end = is.tellg(); + is.clear(state); + is.seekg(position); + RAFT_EXPECTS(static_cast(is), "FlowANN cannot restore the %s stream position", component); + if (end == std::istream::pos_type(-1)) { return; } + RAFT_EXPECTS(end >= position && element_size > 0 && + count <= static_cast(end - position) / element_size, + "FlowANN %s payload is truncated", + component); +} + +inline auto is_legacy_f16_dtype(raft::numpy_serializer::dtype_t dtype) -> bool +{ + // RAFT historically serialized CUDA half as kind 'e'. NumPy uses kind 'f' for float16. + return dtype.itemsize == sizeof(std::uint16_t) && (dtype.kind == 'e' || dtype.kind == 'f') && + (dtype.byteorder == '<' || dtype.byteorder == '=' || dtype.byteorder == '|'); +} + +template +auto read_host_matrix(std::istream& is, + std::uint64_t expected_rows, + std::uint64_t expected_cols, + char const* component) -> raft::host_matrix +{ + auto const header = raft::numpy_serializer::read_header(is); + RAFT_EXPECTS( + header.shape.size() == 2, "FlowANN legacy %s must be a two-dimensional array", component); + RAFT_EXPECTS(!header.fortran_order, "FlowANN legacy %s must use row-major order", component); + RAFT_EXPECTS(header.dtype == raft::numpy_serializer::get_numpy_dtype(), + "FlowANN legacy %s has an unexpected dtype", + component); + RAFT_EXPECTS(header.shape[0] == expected_rows && header.shape[1] == expected_cols, + "FlowANN legacy %s has an unexpected shape", + component); + validate_matrix_size(expected_rows, expected_cols, sizeof(T), component); + RAFT_EXPECTS(expected_rows <= static_cast(std::numeric_limits::max()) && + expected_cols <= static_cast(std::numeric_limits::max()), + "FlowANN legacy %s shape exceeds int64", + component); + + validate_remaining_payload(is, expected_rows * expected_cols, sizeof(T), component); + auto result = raft::make_host_matrix(static_cast(expected_rows), + static_cast(expected_cols)); + auto const bytes = result.size() * sizeof(T); + RAFT_EXPECTS(static_cast(is.read(reinterpret_cast(result.data_handle()), bytes)), + "FlowANN legacy %s payload is truncated", + component); + return result; +} + +template +auto read_host_vector(std::istream& is, std::uint64_t expected_size, char const* component) + -> raft::host_vector +{ + auto const header = raft::numpy_serializer::read_header(is); + RAFT_EXPECTS(header.shape.size() == 1, "FlowANN %s must be a one-dimensional array", component); + RAFT_EXPECTS(!header.fortran_order, "FlowANN %s must use row-major order", component); + RAFT_EXPECTS(header.dtype == raft::numpy_serializer::get_numpy_dtype(), + "FlowANN %s has an unexpected dtype", + component); + RAFT_EXPECTS(header.shape[0] == expected_size, "FlowANN %s has an unexpected shape", component); + validate_matrix_size(expected_size, 1, sizeof(T), component); + RAFT_EXPECTS(expected_size <= static_cast(std::numeric_limits::max()), + "FlowANN %s shape exceeds int64", + component); + + validate_remaining_payload(is, expected_size, sizeof(T), component); + auto result = raft::make_host_vector(static_cast(expected_size)); + RAFT_EXPECTS(static_cast( + is.read(reinterpret_cast(result.data_handle()), result.size() * sizeof(T))), + "FlowANN %s payload is truncated", + component); + return result; +} + +template +void serialize_common(raft::resources const& res, + std::ostream& os, + FlowannIndexT const& index, + serialized_dataset_kind dataset_kind) +{ + static_assert(std::is_same_v, + "The current FlowANN format stores uint32 indices"); + using value_type = typename FlowannIndexT::value_type; + RAFT_EXPECTS(is_supported_metric(dataset_kind, index.metric()), + "FlowANN index uses an unsupported dataset and distance-metric combination"); + auto dtype = raft::numpy_serializer::get_numpy_dtype().to_string(); + dtype.resize(4); + os.write(dtype.data(), dtype.size()); + RAFT_EXPECTS(os.good(), "FlowANN failed to write its dtype prefix"); + + auto const source_indices = index.source_indices(); + auto const subgraph_offsets = index.subgraph_offsets(); + auto const subgraph_ids = index.subgraph_ids(); + auto const seeds = index.seeds(); + RAFT_EXPECTS(index.size() > 0 && index.dim() > 0 && index.graph_degree() > 0, + "FlowANN cannot serialize an empty index"); + RAFT_EXPECTS(subgraph_offsets.has_value() && subgraph_ids.has_value(), + "FlowANN serialization requires subgraph metadata"); + RAFT_EXPECTS(index.inner_graph().extent(1) <= std::numeric_limits::max(), + "FlowANN inner graph row length exceeds the serialization format"); + RAFT_EXPECTS(subgraph_offsets->extent(0) <= std::numeric_limits::max(), + "FlowANN subgraph count exceeds the serialization format"); + RAFT_EXPECTS(!seeds.has_value() || seeds->extent(0) <= std::numeric_limits::max(), + "FlowANN seed count exceeds the serialization format"); + + raft::serialize_scalar(res, os, serialization_version); + raft::serialize_scalar(res, os, static_cast(dataset_kind)); + raft::serialize_scalar(res, os, static_cast(index.size())); + raft::serialize_scalar(res, os, index.dim()); + raft::serialize_scalar(res, os, index.graph_degree()); + raft::serialize_scalar(res, os, index.resident_degree()); + raft::serialize_scalar(res, os, index.cross_degree()); + raft::serialize_scalar(res, os, index.metric()); + raft::serialize_scalar(res, os, index.node_per_cacheline()); + raft::serialize_scalar(res, os, index.n_bits()); + raft::serialize_scalar(res, os, static_cast(index.inner_graph().extent(1))); + raft::serialize_scalar(res, os, static_cast(subgraph_offsets->extent(0))); + raft::serialize_scalar( + res, + os, + source_indices.has_value() ? static_cast(source_indices->extent(0)) : 0); + raft::serialize_scalar( + res, os, seeds.has_value() ? static_cast(seeds->extent(0)) : 0); + + raft::serialize_mdspan(res, os, index.inner_graph()); + raft::serialize_mdspan(res, os, index.cross_graph()); + if (source_indices.has_value()) { raft::serialize_mdspan(res, os, *source_indices); } + raft::serialize_mdspan(res, os, *subgraph_offsets); + raft::serialize_mdspan(res, os, *subgraph_ids); + if (seeds.has_value()) { raft::serialize_mdspan(res, os, *seeds); } +} + +struct serialized_graph_data { + std::uint32_t n_rows; + std::uint32_t dim; + std::uint32_t graph_degree; + std::uint32_t resident_degree; + std::uint32_t cross_degree; + cuvs::distance::DistanceType metric; + std::uint16_t node_per_cacheline; + std::uint16_t n_bits; + serialized_dataset_kind dataset_kind; + raft::host_matrix inner_graph; + raft::host_matrix cross_graph; + std::optional> source_indices; + raft::host_vector subgraph_offsets; + raft::host_vector subgraph_ids; + std::optional> seeds; +}; + +template +auto deserialize_common_after_version(raft::resources const& res, + std::istream& is, + serialized_dataset_kind expected_kind, + int version) -> serialized_graph_data +{ + RAFT_EXPECTS(version == previous_serialization_version || version == serialization_version, + "FlowANN serialization version mismatch: expected %d or %d, got %d", + previous_serialization_version, + serialization_version, + version); + auto const dataset_kind = + static_cast(raft::deserialize_scalar(res, is)); + RAFT_EXPECTS(dataset_kind == expected_kind, + "FlowANN serialized dataset kind does not match the requested index type"); + auto const n_rows = raft::deserialize_scalar(res, is); + auto const dim = raft::deserialize_scalar(res, is); + auto const graph_degree = raft::deserialize_scalar(res, is); + auto const resident_degree = + version == serialization_version ? raft::deserialize_scalar(res, is) : 0u; + auto const cross_degree = version == serialization_version + ? raft::deserialize_scalar(res, is) + : graph_degree; + auto const metric = raft::deserialize_scalar(res, is); + auto const node_per_cacheline = raft::deserialize_scalar(res, is); + auto const n_bits = raft::deserialize_scalar(res, is); + auto const inner_row_bytes = raft::deserialize_scalar(res, is); + auto const n_subgraphs = raft::deserialize_scalar(res, is); + auto const source_count = raft::deserialize_scalar(res, is); + auto const seed_count = raft::deserialize_scalar(res, is); + + RAFT_EXPECTS(n_rows > 0 && dim > 0 && graph_degree > 0, + "FlowANN serialized index contains a zero dimension"); + RAFT_EXPECTS(graph_degree <= std::numeric_limits::max(), + "FlowANN serialized graph degree exceeds its packed count type"); + RAFT_EXPECTS(resident_degree <= graph_degree && cross_degree <= graph_degree, + "FlowANN serialized resident or cross degree exceeds the complete graph degree"); + RAFT_EXPECTS(resident_degree + cross_degree == graph_degree, + "FlowANN serialized graph degrees are inconsistent"); + RAFT_EXPECTS(is_supported_metric(dataset_kind, metric), + "FlowANN serialized index uses an unsupported distance metric"); + RAFT_EXPECTS(node_per_cacheline > 0 && n_bits > 0 && n_bits <= 32, + "FlowANN serialized encoding parameters are invalid"); + RAFT_EXPECTS( + (resident_degree == 0 && inner_row_bytes == 0) || inner_row_bytes >= node_per_cacheline, + "FlowANN serialized inner row is too short"); + RAFT_EXPECTS(n_subgraphs > 0 && n_subgraphs <= n_rows, + "FlowANN serialized subgraph count is invalid"); + RAFT_EXPECTS(source_count == 0 || source_count == n_rows, + "FlowANN serialized source mapping has an invalid size"); + RAFT_EXPECTS(seed_count <= n_rows, "FlowANN serialized seed list is too large"); + + auto const inner_rows = inner_row_bytes == 0 + ? 0 + : raft::div_rounding_up_safe(n_rows, node_per_cacheline); + auto inner_graph = read_host_matrix(is, inner_rows, inner_row_bytes, "inner graph"); + auto cross_graph = read_host_matrix( + is, n_rows, static_cast(cross_degree) + 1, "cross graph"); + + std::optional> source_indices; + if (source_count > 0) { + source_indices.emplace( + read_host_vector(is, source_count, "source-index mapping")); + for (std::uint32_t i = 0; i < source_count; ++i) { + RAFT_EXPECTS((*source_indices)(i) < n_rows, + "FlowANN source-index mapping contains an out-of-range ID at row %u", + i); + } + } + auto subgraph_offsets = read_host_vector(is, n_subgraphs, "subgraph offsets"); + auto subgraph_ids = read_host_vector(is, n_rows, "subgraph identifiers"); + std::optional> seeds; + if (seed_count > 0) { seeds.emplace(read_host_vector(is, seed_count, "seeds")); } + + return {n_rows, + dim, + graph_degree, + resident_degree, + cross_degree, + metric, + node_per_cacheline, + n_bits, + dataset_kind, + std::move(inner_graph), + std::move(cross_graph), + std::move(source_indices), + std::move(subgraph_offsets), + std::move(subgraph_ids), + std::move(seeds)}; +} + +template +auto deserialize_common(raft::resources const& res, + std::istream& is, + serialized_dataset_kind expected_kind) -> serialized_graph_data +{ + read_and_validate_dtype(is, "index"); + auto const version = raft::deserialize_scalar(res, is); + return deserialize_common_after_version(res, is, expected_kind, version); +} + +template +void populate_deserialized_index(raft::resources const& res, + FlowannIndexT& index, + serialized_graph_data&& data) +{ + index.update_graph(res, + std::move(data.inner_graph), + std::move(data.cross_graph), + data.graph_degree, + data.resident_degree); + if (data.source_indices.has_value()) { + index.update_source_indices(res, raft::make_const_mdspan(data.source_indices->view())); + } + index.update_subgraph_layout(res, + raft::make_const_mdspan(data.subgraph_offsets.view()), + raft::make_const_mdspan(data.subgraph_ids.view())); + if (data.seeds.has_value()) { + index.update_seeds(res, raft::make_const_mdspan(data.seeds->view())); + } + raft::resource::sync_stream(res); +} + +inline auto read_host_f16_bits(std::istream& is, + std::uint64_t expected_rows, + std::uint64_t expected_cols, + char const* component) -> raft::host_matrix +{ + auto const header = raft::numpy_serializer::read_header(is); + RAFT_EXPECTS( + header.shape.size() == 2, "FlowANN legacy %s must be a two-dimensional array", component); + RAFT_EXPECTS(!header.fortran_order, "FlowANN legacy %s must use row-major order", component); + RAFT_EXPECTS(is_legacy_f16_dtype(header.dtype), + "FlowANN legacy %s must contain little-endian float16 values", + component); + RAFT_EXPECTS(header.shape[0] == expected_rows && header.shape[1] == expected_cols, + "FlowANN legacy %s has an unexpected shape", + component); + validate_matrix_size(expected_rows, expected_cols, sizeof(std::uint16_t), component); + RAFT_EXPECTS(expected_rows <= static_cast(std::numeric_limits::max()) && + expected_cols <= static_cast(std::numeric_limits::max()), + "FlowANN legacy %s shape exceeds int64", + component); + + validate_remaining_payload(is, expected_rows * expected_cols, sizeof(std::uint16_t), component); + auto result = raft::make_host_matrix(static_cast(expected_rows), + static_cast(expected_cols)); + auto const bytes = result.size() * sizeof(std::uint16_t); + RAFT_EXPECTS(static_cast(is.read(reinterpret_cast(result.data_handle()), bytes)), + "FlowANN legacy %s payload is truncated", + component); + return result; +} + +template +auto deserialize_legacy_graph_body(raft::resources const& res, + std::istream& is, + flowann::index* index, + legacy_graph_scalar_layout scalar_layout) + -> std::pair +{ + static_assert(std::is_same_v, + "The FlowANN legacy v5 graph format stores uint32 indices"); + RAFT_EXPECTS(index != nullptr, "FlowANN deserialize requires a non-null index output"); + + auto const n_rows = raft::deserialize_scalar(res, is); + auto const dim = raft::deserialize_scalar(res, is); + auto const graph_degree = raft::deserialize_scalar(res, is); + auto const metric = + scalar_layout == legacy_graph_scalar_layout::historical_unsigned_enums + ? static_cast(raft::deserialize_scalar(res, is)) + : raft::deserialize_scalar(res, is); + auto const node_per_cacheline = raft::deserialize_scalar(res, is); + auto const n_bits = raft::deserialize_scalar(res, is); + + constexpr auto dataset_kind = cuvs::neighbors::is_vpq_dataset_view_v + ? serialized_dataset_kind::vpq_f16 + : serialized_dataset_kind::dense; + RAFT_EXPECTS(is_supported_metric(dataset_kind, metric), + "FlowANN legacy graph uses an unsupported distance metric"); + RAFT_EXPECTS(n_rows > 0, "FlowANN legacy graph must not be empty"); + RAFT_EXPECTS(dim > 0, "FlowANN legacy graph dimension must be positive"); + RAFT_EXPECTS(graph_degree > 0, "FlowANN legacy graph degree must be positive"); + RAFT_EXPECTS(node_per_cacheline > 0, "FlowANN legacy node_per_cacheline must be positive"); + RAFT_EXPECTS(n_bits > 0 && n_bits <= 32, "FlowANN legacy n_bits must be in [1, 32]"); + RAFT_EXPECTS(index->dataset().n_rows() == 0 || index->dataset().n_rows() == n_rows, + "FlowANN graph and dataset row counts differ"); + RAFT_EXPECTS(index->dataset().dim() == 0 || index->dataset().dim() == dim, + "FlowANN graph and dataset dimensions differ"); + + auto const inner_header = raft::numpy_serializer::read_header(is); + RAFT_EXPECTS(inner_header.shape.size() == 2, + "FlowANN legacy inner graph must be a two-dimensional array"); + RAFT_EXPECTS(!inner_header.fortran_order, "FlowANN legacy inner graph must use row-major order"); + RAFT_EXPECTS(inner_header.dtype == raft::numpy_serializer::get_numpy_dtype(), + "FlowANN legacy inner graph must contain uint8 values"); + auto const inner_rows = inner_header.shape[0]; + auto const inner_cols = inner_header.shape[1]; + auto const expected_inner_rows = raft::div_rounding_up_safe( + static_cast(n_rows), node_per_cacheline); + RAFT_EXPECTS(inner_rows == expected_inner_rows, + "FlowANN legacy inner graph row count does not match node_per_cacheline"); + validate_matrix_size(inner_rows, inner_cols, sizeof(std::uint8_t), "inner graph"); + RAFT_EXPECTS(inner_rows <= static_cast(std::numeric_limits::max()) && + inner_cols <= static_cast(std::numeric_limits::max()), + "FlowANN legacy inner graph shape exceeds int64"); + validate_remaining_payload(is, inner_rows * inner_cols, sizeof(std::uint8_t), "inner graph"); + auto inner_graph = raft::make_host_matrix( + static_cast(inner_rows), static_cast(inner_cols)); + RAFT_EXPECTS(static_cast(is.read(reinterpret_cast(inner_graph.data_handle()), + inner_graph.size() * sizeof(std::uint8_t))), + "FlowANN legacy inner graph payload is truncated"); + + auto cross_graph = read_host_matrix( + is, n_rows, static_cast(graph_degree) + 1, "cross graph"); + auto source_indices = read_host_matrix(is, 1, n_rows, "source-index mapping"); + for (int64_t i = 0; i < source_indices.extent(1); ++i) { + RAFT_EXPECTS(source_indices(0, i) < n_rows, + "FlowANN source-index mapping contains an out-of-range ID at row %ld", + static_cast(i)); + } + + DatasetViewT dataset = index->dataset(); + *index = flowann::index(res, metric, dataset, node_per_cacheline, n_bits); + index->update_graph( + res, raft::make_const_mdspan(inner_graph.view()), raft::make_const_mdspan(cross_graph.view())); + index->update_source_indices( + res, raft::make_host_vector_view(source_indices.data_handle(), n_rows)); + raft::resource::sync_stream(res); + return {n_rows, dim}; +} + +template +void deserialize_legacy_graph(raft::resources const& res, + std::istream& is, + flowann::index* index) +{ + read_and_validate_dtype(is, "graph"); + auto const version = raft::deserialize_scalar(res, is); + RAFT_EXPECTS(version == legacy_serialization_version, + "FlowANN graph serialization version mismatch: expected %d, got %d", + legacy_serialization_version, + version); + (void)deserialize_legacy_graph_body(res, is, index, legacy_graph_scalar_layout::typed_enums); +} + +template +void deserialize_legacy_graph(raft::resources const& res, + std::string const& filename, + flowann::index* index) +{ + std::ifstream is(filename, std::ios::in | std::ios::binary); + RAFT_EXPECTS(is.is_open(), "Cannot open FlowANN graph file %s", filename.c_str()); + deserialize_legacy_graph(res, is, index); +} + +struct legacy_vpq_f16_data { + std::unique_ptr> dataset; + raft::host_vector subgraph_offsets; + raft::host_vector subgraph_ids; +}; + +template +auto deserialize_legacy_vpq_f16_body(raft::resources const& res, std::istream& is) + -> legacy_vpq_f16_data +{ + auto const dataset_tag = raft::deserialize_scalar(res, is); + RAFT_EXPECTS(dataset_tag == legacy_vpq_dataset_tag, + "FlowANN legacy dataset must contain VPQ data"); + auto const dtype = raft::deserialize_scalar(res, is); + RAFT_EXPECTS(dtype == CUDA_R_16F, "FlowANN legacy VPQ codebooks must use float16"); + + auto const n_rows = raft::deserialize_scalar(res, is); + auto const dim = raft::deserialize_scalar(res, is); + auto const vq_n_centers = raft::deserialize_scalar(res, is); + auto const pq_n_centers = raft::deserialize_scalar(res, is); + auto const pq_len = raft::deserialize_scalar(res, is); + auto const serialized_row_length = raft::deserialize_scalar(res, is); + + RAFT_EXPECTS(n_rows > 0 && dim > 0 && vq_n_centers > 0 && pq_len > 0, + "FlowANN legacy VPQ metadata contains a zero dimension"); + RAFT_EXPECTS(vq_n_centers <= (1u << 16), "FlowANN legacy VPQ stores VQ labels in 16 bits"); + RAFT_EXPECTS(pq_n_centers == (1u << legacy_vpq_codebook_bits), + "FlowANN legacy VPQ search requires 8-bit PQ codes"); + RAFT_EXPECTS(pq_len == 2 || pq_len == 4 || pq_len == 8, + "FlowANN legacy VPQ pq_len must be 2, 4, or 8"); + auto const pq_dim = raft::div_rounding_up_safe(dim, pq_len); + RAFT_EXPECTS(serialized_row_length >= sizeof(std::uint32_t) + pq_dim, + "FlowANN legacy VPQ encoded rows are too short"); + RAFT_EXPECTS(serialized_row_length <= std::numeric_limits::max() - 3, + "FlowANN legacy VPQ encoded row length cannot be aligned safely"); + // Discard legacy padding when normalizing rows for the current VPQ descriptor. + auto const payload_row_length = sizeof(std::uint32_t) + pq_dim; + auto const encoded_row_length = (payload_row_length + 3u) & ~3u; + validate_matrix_size(static_cast(n_rows), + encoded_row_length, + sizeof(std::uint8_t), + "aligned VPQ codes"); + + auto vq_bits = read_host_f16_bits(is, vq_n_centers, dim, "VQ codebook"); + auto pq_bits = read_host_f16_bits(is, pq_n_centers, pq_len, "PQ codebook"); + auto serialized_data = read_host_matrix( + is, static_cast(n_rows), serialized_row_length, "VPQ codes"); + auto host_data = raft::make_host_matrix(n_rows, encoded_row_length); + std::fill_n(host_data.data_handle(), host_data.size(), std::uint8_t{0}); + for (int64_t row = 0; row < n_rows; ++row) { + std::copy_n( + serialized_data.data_handle() + static_cast(row) * serialized_row_length, + payload_row_length, + host_data.data_handle() + static_cast(row) * encoded_row_length); + } + + auto const n_subgraphs = raft::deserialize_scalar(res, is); + RAFT_EXPECTS(n_subgraphs > 0, "FlowANN legacy VPQ dataset must contain subgraph offsets"); + auto offsets_matrix = read_host_matrix(is, 1, n_subgraphs, "subgraph offsets"); + auto offsets = raft::make_host_vector(n_subgraphs); + std::copy_n(offsets_matrix.data_handle(), n_subgraphs, offsets.data_handle()); + RAFT_EXPECTS(offsets(0) == 0, "FlowANN first subgraph offset must be zero"); + for (std::uint32_t i = 1; i < n_subgraphs; ++i) { + RAFT_EXPECTS(offsets(i - 1) <= offsets(i), "FlowANN subgraph offsets must be monotonic"); + } + RAFT_EXPECTS(offsets(n_subgraphs - 1) < static_cast(n_rows), + "FlowANN subgraph offsets exceed the dataset size"); + + auto subgraph_ids = raft::make_host_vector(n_rows); + for (int64_t row = 0; row < n_rows; ++row) { + auto* encoded = host_data.data_handle() + static_cast(row) * encoded_row_length; + auto const vq_label = + static_cast(encoded[0]) | (static_cast(encoded[1]) << 8); + RAFT_EXPECTS(vq_label < vq_n_centers, + "FlowANN VPQ row %ld has an out-of-range VQ label", + static_cast(row)); + auto const subgraph_id = + static_cast(encoded[2]) | (static_cast(encoded[3]) << 8); + RAFT_EXPECTS(subgraph_id < n_subgraphs, + "FlowANN VPQ row %ld has an out-of-range subgraph identifier", + static_cast(row)); + subgraph_ids(row) = subgraph_id; + // Current CAGRA VPQ expects a 32-bit VQ label. Legacy FlowANN used the upper 16 bits for the + // subgraph ID; preserve that metadata separately and normalize the row for the current + // distance descriptor. + encoded[2] = 0; + encoded[3] = 0; + } + + auto vq_code_book = raft::make_device_matrix(res, vq_n_centers, dim); + auto pq_code_book = raft::make_device_matrix(res, pq_n_centers, pq_len); + auto data = raft::make_device_matrix(res, n_rows, encoded_row_length); + auto const stream = raft::resource::get_cuda_stream(res).get(); + RAFT_CUDA_TRY(cudaMemcpyAsync(vq_code_book.data_handle(), + vq_bits.data_handle(), + vq_bits.size() * sizeof(std::uint16_t), + cudaMemcpyHostToDevice, + stream)); + RAFT_CUDA_TRY(cudaMemcpyAsync(pq_code_book.data_handle(), + pq_bits.data_handle(), + pq_bits.size() * sizeof(std::uint16_t), + cudaMemcpyHostToDevice, + stream)); + RAFT_CUDA_TRY(cudaMemcpyAsync(data.data_handle(), + host_data.data_handle(), + host_data.size() * sizeof(std::uint8_t), + cudaMemcpyHostToDevice, + stream)); + + return {std::make_unique>( + std::move(data), + cuvs::neighbors::device_vpq_dataset::manifest_type{ + std::move(vq_code_book), std::move(pq_code_book)}), + std::move(offsets), + std::move(subgraph_ids)}; +} + +template +auto deserialize_legacy_vpq_f16(raft::resources const& res, std::istream& is) -> legacy_vpq_f16_data +{ + read_and_validate_dtype(is, "dataset"); + auto const version = raft::deserialize_scalar(res, is); + RAFT_EXPECTS(version == legacy_serialization_version, + "FlowANN dataset serialization version mismatch: expected %d, got %d", + legacy_serialization_version, + version); + return deserialize_legacy_vpq_f16_body(res, is); +} + +template +auto deserialize_legacy_vpq_f16(raft::resources const& res, std::string const& filename) + -> legacy_vpq_f16_data +{ + std::ifstream is(filename, std::ios::in | std::ios::binary); + RAFT_EXPECTS(is.is_open(), "Cannot open FlowANN dataset file %s", filename.c_str()); + return deserialize_legacy_vpq_f16(res, is); +} + +template +auto deserialize_legacy_combined_vpq_f16_after_version(raft::resources const& res, + std::istream& is, + int version) -> vpq_f16_index_bundle +{ + static_assert(std::is_same_v, + "The FlowANN legacy v5 combined format stores uint32 indices"); + RAFT_EXPECTS(version == legacy_serialization_version, + "FlowANN combined serialization version mismatch: expected %d, got %d", + legacy_serialization_version, + version); + + flowann::vpq_f16_index index(res); + auto const [graph_rows, graph_dim] = deserialize_legacy_graph_body( + res, is, &index, legacy_graph_scalar_layout::historical_unsigned_enums); + auto const has_dataset = raft::deserialize_scalar(res, is); + RAFT_EXPECTS(has_dataset, "FlowANN legacy combined index does not contain a dataset"); + auto dataset_data = deserialize_legacy_vpq_f16_body(res, is); + RAFT_EXPECTS(graph_rows == static_cast(dataset_data.dataset->n_rows()), + "FlowANN graph and VPQ dataset row counts differ"); + RAFT_EXPECTS(graph_dim == dataset_data.dataset->dim(), + "FlowANN graph and VPQ dataset dimensions differ"); + index.update_dataset(res, dataset_data.dataset->view()); + index.update_subgraph_layout( + res, + raft::make_host_vector_view( + dataset_data.subgraph_offsets.data_handle(), dataset_data.subgraph_offsets.extent(0)), + raft::make_host_vector_view( + dataset_data.subgraph_ids.data_handle(), dataset_data.subgraph_ids.extent(0))); + raft::resource::sync_stream(res); + return {std::move(dataset_data.dataset), std::move(index)}; +} + +} // namespace detail + +/** Load the split graph and VPQ dataset emitted by the FlowANN legacy v5 tools. */ +template +auto deserialize_vpq_f16(raft::resources const& res, + std::string const& graph_filename, + std::string const& dataset_filename) -> vpq_f16_index_bundle +{ + auto dataset_data = detail::deserialize_legacy_vpq_f16(res, dataset_filename); + flowann::vpq_f16_index index(res); + index.update_dataset(res, dataset_data.dataset->view()); + detail::deserialize_legacy_graph(res, graph_filename, &index); + RAFT_EXPECTS(index.size() == static_cast(dataset_data.dataset->n_rows()), + "FlowANN graph and VPQ dataset row counts differ"); + index.update_subgraph_layout( + res, + raft::make_host_vector_view( + dataset_data.subgraph_offsets.data_handle(), dataset_data.subgraph_offsets.extent(0)), + raft::make_host_vector_view( + dataset_data.subgraph_ids.data_handle(), dataset_data.subgraph_ids.extent(0))); + raft::resource::sync_stream(res); + return {std::move(dataset_data.dataset), std::move(index)}; +} + +/** Serialize a dense FlowANN index using the current self-contained format. */ +template +void serialize(raft::resources const& res, + std::ostream& os, + device_padded_index const& index) +{ + detail::serialize_common(res, os, index, detail::serialized_dataset_kind::dense); + raft::serialize_scalar( + res, os, static_cast(index.dataset().data_view().stride(0))); + auto rows = index.dataset().data_view(); + raft::serialize_mdspan(res, + os, + raft::make_device_matrix_view( + rows.data_handle(), rows.extent(0), rows.stride(0))); + RAFT_EXPECTS(os.good(), "FlowANN failed while serializing its dense dataset"); +} + +/** Serialize a float16-VPQ FlowANN index using the current self-contained format. */ +template +void serialize(raft::resources const& res, std::ostream& os, vpq_f16_index const& index) +{ + detail::serialize_common(res, os, index, detail::serialized_dataset_kind::vpq_f16); + auto const dataset = index.dataset(); + auto const manifest = dataset.manifest_view(); + raft::serialize_scalar(res, os, manifest.vq_n_centers()); + raft::serialize_scalar(res, os, manifest.pq_n_centers()); + raft::serialize_scalar(res, os, manifest.pq_len()); + raft::serialize_scalar(res, os, static_cast(dataset.data_view().extent(1))); + raft::serialize_mdspan(res, os, manifest.vq_code_book); + raft::serialize_mdspan(res, os, manifest.pq_code_book); + raft::serialize_mdspan(res, os, dataset.data_view()); + RAFT_EXPECTS(os.good(), "FlowANN failed while serializing its VPQ dataset"); +} + +/** Serialize a FlowANN index to a file using the current self-contained format. */ +template +void serialize(raft::resources const& res, std::string const& filename, FlowannIndexT const& index) +{ + std::ofstream os(filename, std::ios::out | std::ios::binary); + RAFT_EXPECTS(os.is_open(), "Cannot open FlowANN output file %s", filename.c_str()); + serialize(res, os, index); + os.close(); + RAFT_EXPECTS(os.good(), "FlowANN failed to close output file %s", filename.c_str()); +} + +/** Deserialize a dense FlowANN index from the current self-contained format. */ +template +auto deserialize_device_padded(raft::resources const& res, std::istream& is) + -> device_padded_index_bundle +{ + static_assert(std::is_same_v, + "The current FlowANN format stores uint32 indices"); + auto graph = detail::deserialize_common(res, is, detail::serialized_dataset_kind::dense); + auto const stride = raft::deserialize_scalar(res, is); + RAFT_EXPECTS(stride >= graph.dim, "FlowANN dense dataset stride is smaller than its dimension"); + detail::validate_matrix_size(graph.n_rows, stride, sizeof(T), "dense dataset"); + detail::validate_remaining_payload( + is, static_cast(graph.n_rows) * stride, sizeof(T), "dense dataset"); + auto owned = + cuvs::neighbors::detail::make_dense_array>( + res, graph.n_rows, graph.dim); + auto const owned_stride = owned.mapping().stride(0); + if (stride == owned_stride) { + // Read canonical rows directly into their owner, including initialized padding. + // This keeps the common load path to one device allocation and one upload. + auto physical_rows = + raft::make_device_matrix_view(owned.data_handle(), graph.n_rows, stride); + raft::deserialize_mdspan(res, is, physical_rows); + } else { + // Older files can use a different physical row pitch. Preserve the logical + // values and initialize padding before cosine reduction reads complete rows. + auto storage = raft::make_device_matrix(res, graph.n_rows, stride); + raft::deserialize_mdspan(res, is, storage.view()); + auto stream = raft::resource::get_cuda_stream(res); + RAFT_CUDA_TRY(cudaMemsetAsync( + owned.data_handle(), 0, owned.mapping().required_span_size() * sizeof(T), stream.get())); + raft::copy_matrix(owned.data_handle(), + owned_stride, + storage.data_handle(), + stride, + graph.dim, + graph.n_rows, + stream); + } + auto dataset = + std::make_unique>(std::move(owned)); + device_padded_index index( + res, graph.metric, dataset->view(), graph.node_per_cacheline, graph.n_bits); + detail::populate_deserialized_index(res, index, std::move(graph)); + return {std::move(dataset), std::move(index)}; +} + +/** Deserialize a dense FlowANN index file in the current self-contained format. */ +template +auto deserialize_device_padded(raft::resources const& res, std::string const& filename) + -> device_padded_index_bundle +{ + std::ifstream is(filename, std::ios::in | std::ios::binary); + RAFT_EXPECTS(is.is_open(), "Cannot open FlowANN index file %s", filename.c_str()); + return deserialize_device_padded(res, is); +} + +/** + * @brief Deserialize a float16-VPQ FlowANN index from a stream. + * + * The loader identifies the format from the four-byte dtype prefix and serialization version. + * Current self-contained versions 1 and 2 use their typed dtype prefix. Historical combined v5 + * files use the `|u1` prefix emitted by the original FlowANN split serializer. The stream must be + * positioned at the beginning of the index. The returned bundle owns the device VPQ dataset that + * backs the index. + * + * @tparam T Requested query value type. Current formats validate `T` against their typed prefix. + * Historical combined v5 uses an untyped `|u1` format marker, so the caller selects `T` from the + * query types supported by FlowANN search. + * @tparam IdxT Index type. Only `std::uint32_t` is supported by these formats. + * @param res RAFT resources used for allocations and transfers. + * @param is Input stream containing one complete index. + * @return An owning VPQ dataset and the loaded FlowANN index. + * @throws raft::logic_error if the format, dtype, dimensions, graph, or VPQ metadata is invalid. + */ +template +auto deserialize_vpq_f16(raft::resources const& res, std::istream& is) + -> vpq_f16_index_bundle +{ + static_assert(std::is_same_v, + "FlowANN VPQ serialization formats store uint32 indices"); + char serialized_dtype[4]{}; + RAFT_EXPECTS(static_cast(is.read(serialized_dtype, sizeof(serialized_dtype))), + "FlowANN index is truncated before its dtype prefix"); + auto const version = raft::deserialize_scalar(res, is); + if (detail::is_legacy_combined_dtype_prefix(serialized_dtype) && + version == detail::legacy_serialization_version) { + return detail::deserialize_legacy_combined_vpq_f16_after_version(res, is, version); + } + detail::validate_dtype_prefix(serialized_dtype, "index"); + auto graph = detail::deserialize_common_after_version( + res, is, detail::serialized_dataset_kind::vpq_f16, version); + auto const vq_n_centers = raft::deserialize_scalar(res, is); + auto const pq_n_centers = raft::deserialize_scalar(res, is); + auto const pq_len = raft::deserialize_scalar(res, is); + auto const row_length = raft::deserialize_scalar(res, is); + RAFT_EXPECTS(vq_n_centers > 0 && pq_n_centers > 0 && pq_len > 0 && row_length > 0, + "FlowANN VPQ metadata contains a zero dimension"); + RAFT_EXPECTS(pq_n_centers == 256, "FlowANN search currently requires 8-bit VPQ codebooks"); + RAFT_EXPECTS(pq_len == 2 || pq_len == 4 || pq_len == 8, "FlowANN VPQ pq_len must be 2, 4, or 8"); + auto const pq_dim = raft::div_rounding_up_safe(graph.dim, pq_len); + auto const expected_row_length = + sizeof(std::uint32_t) * (1 + raft::div_rounding_up_safe(pq_dim, 4)); + RAFT_EXPECTS(row_length == expected_row_length, + "FlowANN VPQ encoded row length does not match its dimension and pq_len"); + detail::validate_matrix_size(vq_n_centers, graph.dim, sizeof(half), "VQ codebook"); + detail::validate_matrix_size(pq_n_centers, pq_len, sizeof(half), "PQ codebook"); + detail::validate_matrix_size(graph.n_rows, row_length, sizeof(std::uint8_t), "VPQ codes"); + + detail::validate_remaining_payload( + is, static_cast(vq_n_centers) * graph.dim, sizeof(half), "VQ codebook"); + auto vq_code_book = raft::make_device_matrix(res, vq_n_centers, graph.dim); + raft::deserialize_mdspan(res, is, vq_code_book.view()); + detail::validate_remaining_payload( + is, static_cast(pq_n_centers) * pq_len, sizeof(half), "PQ codebook"); + auto pq_code_book = raft::make_device_matrix(res, pq_n_centers, pq_len); + raft::deserialize_mdspan(res, is, pq_code_book.view()); + auto host_encoded = + detail::read_host_matrix(is, graph.n_rows, row_length, "VPQ codes"); + for (std::uint32_t row = 0; row < graph.n_rows; ++row) { + std::uint32_t vq_label; + std::memcpy(&vq_label, + host_encoded.data_handle() + static_cast(row) * row_length, + sizeof(vq_label)); + RAFT_EXPECTS(vq_label < vq_n_centers, "FlowANN VPQ row %u has an out-of-range VQ label", row); + } + auto encoded = raft::make_device_matrix(res, graph.n_rows, row_length); + raft::copy(encoded.data_handle(), + host_encoded.data_handle(), + host_encoded.size(), + raft::resource::get_cuda_stream(res)); + // Host rows must outlive the asynchronous upload. + raft::resource::sync_stream(res); + auto dataset = std::make_unique>( + std::move(encoded), + cuvs::neighbors::device_vpq_dataset::manifest_type{std::move(vq_code_book), + std::move(pq_code_book)}); + RAFT_EXPECTS(dataset->dim() == graph.dim && dataset->n_rows() == graph.n_rows, + "FlowANN VPQ dataset metadata does not match the graph"); + + vpq_f16_index index( + res, graph.metric, dataset->view(), graph.node_per_cacheline, graph.n_bits); + detail::populate_deserialized_index(res, index, std::move(graph)); + return {std::move(dataset), std::move(index)}; +} + +/** + * @brief Deserialize a float16-VPQ FlowANN index file. + * + * This overload opens `filename` and applies the same current-versus-legacy auto-detection as the + * stream overload. The returned bundle owns the device VPQ dataset that backs the index. + * + * @tparam T Requested query value type. Current formats validate `T` against their typed prefix. + * Historical combined v5 uses an untyped `|u1` format marker, so the caller selects `T` from the + * query types supported by FlowANN search. + * @tparam IdxT Index type. Only `std::uint32_t` is supported by these formats. + * @param res RAFT resources used for allocations and transfers. + * @param filename Path to a current self-contained or historical combined v5 index. + * @return An owning VPQ dataset and the loaded FlowANN index. + * @throws raft::logic_error if the file cannot be opened or its serialized contents are invalid. + */ +template +auto deserialize_vpq_f16(raft::resources const& res, std::string const& filename) + -> vpq_f16_index_bundle +{ + std::ifstream is(filename, std::ios::in | std::ios::binary); + RAFT_EXPECTS(is.is_open(), "Cannot open FlowANN index file %s", filename.c_str()); + return deserialize_vpq_f16(res, is); +} + +} // namespace cuvs::neighbors::cagra::experimental::flowann + +#endif // CUVS_ENABLE_FLOWANN_SEARCH diff --git a/cpp/src/neighbors/detail/flowann/graph_state.cuh b/cpp/src/neighbors/detail/flowann/graph_state.cuh new file mode 100644 index 0000000000..d1ffa4e3e7 --- /dev/null +++ b/cpp/src/neighbors/detail/flowann/graph_state.cuh @@ -0,0 +1,138 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * FlowANN is described in "Disentangling Graph Dependencies for Efficient + * Billion-Scale GPU Vector Search" (OSDI 2026). + */ + +#pragma once + +#include "kernel.cuh" + +#include + +namespace cuvs::neighbors::cagra::experimental::flowann::detail { + +inline constexpr std::uint32_t pending_capacity = 512; +inline constexpr std::uint32_t pending_mask = pending_capacity - 1; + +static_assert((pending_capacity & pending_mask) == 0, + "FlowANN pending capacity must be a power of two"); + +struct pending_entry { + std::uint32_t parent; + queue_reservation reservation; + std::uint16_t sync_window; +}; + +struct pending_queue { + std::uint16_t head; + std::uint16_t tail; + std::uint32_t consumed; + pending_entry entries[pending_capacity]; + + RAFT_DEVICE_INLINE_FUNCTION void initialize() + { + head = 0; + tail = 0; + consumed = 0; + } + + [[nodiscard]] RAFT_DEVICE_INLINE_FUNCTION auto size() const -> std::uint16_t + { + return static_cast((head - tail) & pending_mask); + } + + [[nodiscard]] RAFT_DEVICE_INLINE_FUNCTION auto available() const -> std::uint16_t + { + return static_cast(pending_mask - size()); + } + + [[nodiscard]] RAFT_DEVICE_INLINE_FUNCTION auto empty() const -> bool { return head == tail; } + + RAFT_DEVICE_INLINE_FUNCTION void push(pending_entry entry) + { + entries[head] = entry; + head = static_cast((head + 1) & pending_mask); + } + + [[nodiscard]] RAFT_DEVICE_INLINE_FUNCTION auto front() -> pending_entry& { return entries[tail]; } + + RAFT_DEVICE_INLINE_FUNCTION void pop() + { + tail = static_cast((tail + 1) & pending_mask); + } +}; + +template +struct graph_state { + pending_queue pending; + IndexT parent_ids[max_search_width]; + bool parent_valid[max_search_width]; + std::uint32_t parent_seed_positions[max_search_width]; + std::uint32_t parent_inner_counts[max_search_width]; + std::uint64_t parent_bit_offsets[max_search_width]; + std::uint32_t parent_subgraph_offsets[max_search_width]; + const std::uint32_t* ready_cross_data[max_search_width]; + std::uint32_t ready_cross_edge_counts[max_search_width]; + queue_reservation ready_reservations[max_search_width]; + std::uint32_t ready_cross_count; + std::uint32_t total_inner_children; + std::uint32_t total_children; + std::uint32_t previous_parent_position; + std::uint32_t iteration; + bool use_seed_cache; +}; + +#ifdef __CUDACC__ +/** Submit a complete parent batch or restore its traversal state for retry. */ +template +RAFT_DEVICE_INLINE_FUNCTION void submit_parent_requests(queue_view const& queue, + graph_state* state, + const IndexT* parent_indices, + const IndexT* internal_topk, + std::uint32_t search_width, + std::uint16_t sync_window, + IndexT* traversed_hashmap, + std::uint32_t traversed_hash_bitlen) +{ + constexpr IndexT index_msb = + cuvs::neighbors::cagra::detail::utils::gen_index_msb_1_mask::value; + std::uint32_t commands[max_search_width]; + std::uint32_t command_count = 0; + for (std::uint32_t i = 0; i < search_width; ++i) { + if (!state->parent_valid[i]) { continue; } + commands[command_count++] = static_cast(state->parent_ids[i]); + } + + queue_reservation first{}; + bool reserved = command_count == 0; + // One slot is reserved to distinguish a full pending ring from an empty one. + // Do not publish commands unless every reservation can be tracked until release. + if (command_count <= state->pending.available()) { + for (std::uint32_t attempt = 0; attempt < 200 && !reserved; ++attempt) { + reserved = try_reserve_deferred(queue, command_count, &first); + } + } + if (command_count > 0 && reserved) { + submit_batch(queue, first, commands, command_count); + for (std::uint32_t i = 0; i < command_count; ++i) { + state->pending.push({commands[i], reservation_at(first, i), sync_window}); + } + } else if (!reserved) { + // Parent selection marked these nodes as traversed. Retry their cross rows + // after making progress on older requests rather than silently losing edges. + for (std::uint32_t i = 0; i < search_width; ++i) { + if (!state->parent_valid[i]) { continue; } + const_cast(internal_topk)[parent_indices[i]] &= ~index_msb; + if (traversed_hashmap != nullptr) { + cuvs::neighbors::cagra::detail::hashmap::remove( + traversed_hashmap, traversed_hash_bitlen, state->parent_ids[i]); + } + } + } +} +#endif + +} // namespace cuvs::neighbors::cagra::experimental::flowann::detail diff --git a/cpp/src/neighbors/detail/flowann/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..1db4e9cce5 --- /dev/null +++ b/cpp/src/neighbors/detail/flowann/kernel.cuh @@ -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). + */ + +#pragma once + +#include "queue.cuh" + +#include +#include + +#include + +namespace cuvs::neighbors::cagra::experimental::flowann::detail { + +inline constexpr std::uint32_t max_search_width = 8; +inline constexpr std::uint32_t additional_buffer_width = 2; + +template +using search_single_cta_kernel_func_t = + void(uintptr_t, + DistanceT*, + std::uint32_t, + const DataT*, + std::uint32_t, + const std::uint8_t*, + std::uint32_t, + std::uint32_t, + std::uint32_t, + const std::uint32_t*, + const std::uint32_t*, + queue_view*, + std::uint32_t, + std::uint32_t, + std::uint32_t, + const SourceIndexT*, + unsigned, + std::uint64_t, + const IndexT*, + std::uint32_t, + const std::uint32_t*, + const IndexT*, + const std::uint32_t*, + std::uint32_t, + IndexT*, + std::uint32_t, + std::uint32_t, + std::uint32_t, + std::uint32_t, + std::uint32_t, + std::uint32_t, + std::uint32_t*, + std::uint32_t, + std::uint32_t, + std::uint32_t, + std::uint32_t, + const cuvs::neighbors::cagra::detail::dataset_descriptor_base_t*, + IndexT, + float, + std::uint32_t, + cuvs::neighbors::cagra::detail::cagra_sample_filter); + +template +using search_multi_cta_kernel_func_t = + void(IndexT*, + DistanceT*, + const cuvs::neighbors::cagra::detail::dataset_descriptor_base_t*, + const DataT*, + std::uint32_t, + const std::uint8_t*, + std::uint32_t, + std::uint32_t, + std::uint32_t, + const std::uint32_t*, + const std::uint32_t*, + queue_view*, + std::uint32_t, + std::uint32_t, + std::uint32_t, + 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..7fd6f7ab0f --- /dev/null +++ b/cpp/src/neighbors/detail/flowann/launcher.hpp @@ -0,0 +1,253 @@ +/* + * 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; + + if constexpr (std::is_same_v) { + planner.add_setup_workspace_device_function(dataset_desc.team_size, + dataset_desc.dataset_block_dim, + dataset_desc.pq_len, + dataset_desc.smem_dtype); + planner.add_compute_distance_device_function(dataset_desc.team_size, + dataset_desc.dataset_block_dim, + dataset_desc.pq_len, + dataset_desc.smem_dtype); + } else { + 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()); + + if (dataset_desc.is_vpq) { + using query_tag = query_type_tag_vpq_t; + using codebook_tag = codebook_tag_vpq_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 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; + if constexpr (std::is_same_v) { + planner.add_setup_workspace_device_function(dataset_desc.team_size, + dataset_desc.dataset_block_dim, + dataset_desc.pq_len, + dataset_desc.smem_dtype); + planner.add_compute_distance_device_function(dataset_desc.team_size, + dataset_desc.dataset_block_dim, + dataset_desc.pq_len, + dataset_desc.smem_dtype); + } else { + 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()); + if (dataset_desc.is_vpq) { + using query_tag = query_type_tag_vpq_t; + using codebook_tag = codebook_tag_vpq_t; + return build_multi_cta_launcher(dataset_desc, + std::move(sample_filter_udf_fragment)); + } + 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..6e32f49a88 --- /dev/null +++ b/cpp/src/neighbors/detail/flowann/queue.cuh @@ -0,0 +1,286 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * FlowANN is described in "Disentangling Graph Dependencies for Efficient + * Billion-Scale GPU Vector Search" (OSDI 2026). + */ + +#pragma once + +#include + +#include + +#include + +#include +#include +#include + +namespace cuvs::neighbors::cagra::experimental::flowann::detail { + +inline constexpr std::uint32_t queue_capacity = 1u << 16; +inline constexpr std::uint32_t queue_mask = queue_capacity - 1; +inline constexpr unsigned long long empty_command = 0ull; + +static_assert((queue_capacity & queue_mask) == 0, "FlowANN queue capacity must be a power of two"); + +// Each word carries the row ID and its complement. Empty or inconsistent +// words fail validation; every uint32 row ID remains representable. +RAFT_INLINE_FUNCTION constexpr auto make_command(std::uint32_t row_id) -> unsigned long long +{ + return (static_cast(~row_id) << 32) | row_id; +} + +RAFT_INLINE_FUNCTION constexpr auto command_ready(unsigned long long word) -> bool +{ + return static_cast(word >> 32) == ~static_cast(word); +} + +struct alignas(128) queue_state { + alignas(128) unsigned long long reservation_head; + alignas(128) unsigned long long cpu_tail; + alignas(128) unsigned long long gpu_tail; + std::uint32_t capacity; + std::uint32_t row_length; +}; + +/** Device pointers for one host-serviced cross-graph queue. */ +struct queue_view { + queue_state* state; + std::uint32_t* rows; + unsigned long long* command_words; + unsigned long long* result_sequence; + unsigned long long* released_sequence; + std::uint32_t const* abort_flag; + // Private to each device-side queue view. Avoid a PCIe read on every readiness poll. + mutable unsigned long long next_abort_check{}; + mutable bool aborted{}; +}; + +struct queue_reservation { + unsigned long long sequence; + std::uint32_t slot; +}; + +#ifdef __CUDACC__ + +RAFT_DEVICE_INLINE_FUNCTION auto load_counter(unsigned long long const* value) -> unsigned long long +{ + cuda::atomic_ref counter( + *const_cast(value)); + return counter.load(cuda::memory_order_relaxed); +} + +[[nodiscard]] RAFT_DEVICE_INLINE_FUNCTION auto reservation_at(queue_reservation first, + std::uint32_t offset) + -> queue_reservation +{ + auto const sequence = first.sequence + offset; + return {sequence, static_cast(sequence) & queue_mask}; +} + +RAFT_DEVICE_INLINE_FUNCTION auto load_published_counter(unsigned long long const* value) + -> unsigned long long +{ + cuda::atomic_ref counter( + *const_cast(value)); + return counter.load(cuda::memory_order_acquire); +} + +/** + * Reclaim only the contiguous released prefix. A failed CAS means another reclaimer advanced + * the tail; restart from that tail so no slot is returned to producers twice. + * This scalar fallback is used when a deferred producer runs out of queue capacity. + */ +RAFT_DEVICE_INLINE_FUNCTION void reclaim_released(queue_view queue) +{ + while (true) { + auto const tail = load_counter(&queue.state->gpu_tail); + auto const head = load_counter(&queue.state->reservation_head); + auto end = tail; + while (end < head) { + auto const slot = static_cast(end) & queue_mask; + cuda::atomic_ref sequence( + queue.released_sequence[slot]); + if (sequence.load(cuda::memory_order_relaxed) != end + 1) { break; } + ++end; + } + if (end == tail) { return; } + if (atomicCAS(&queue.state->gpu_tail, tail, end) == tail) { return; } + } +} + +/** + * Called by all 32 lanes of a warp after the CTA has consumed its pending rows. The first + * unreleased sequence stops reclamation, including when the scan crosses the ring boundary. + * The tail CAS also permits concurrent reclaimers and users of the ordinary release path. + */ +RAFT_DEVICE_INLINE_FUNCTION void reclaim_released_warp(queue_view queue) +{ + auto const lane = threadIdx.x & 31u; + while (true) { + auto const tail = + __shfl_sync(0xffffffffu, lane == 0 ? load_counter(&queue.state->gpu_tail) : 0ull, 0); + auto const head = + __shfl_sync(0xffffffffu, lane == 0 ? load_counter(&queue.state->reservation_head) : 0ull, 0); + auto end = tail; + while (end < head) { + auto const sequence = end + lane; + auto const slot = static_cast(sequence) & queue_mask; + cuda::atomic_ref released( + queue.released_sequence[slot]); + auto const ready = + sequence < head && released.load(cuda::memory_order_relaxed) == sequence + 1; + auto const mask = __ballot_sync(0xffffffffu, ready); + auto const count = mask == 0xffffffffu ? 32u : static_cast(__ffs(~mask) - 1); + end += count; + if (count < 32) { break; } + } + if (end == tail) { return; } + int success = lane == 0 && atomicCAS(&queue.state->gpu_tail, tail, end) == tail; + if (__shfl_sync(0xffffffffu, success, 0)) { return; } + } +} + +/** Mark a consumed row; the caller must reclaim at capacity pressure and on completion. */ +RAFT_DEVICE_INLINE_FUNCTION void release_deferred(queue_view queue, queue_reservation reservation) +{ + atomicExch(&queue.released_sequence[reservation.slot], reservation.sequence + 1); +} + +/** Try to reserve a contiguous batch of queue slots atomically. */ +RAFT_DEVICE_INLINE_FUNCTION auto try_reserve_batch(queue_view queue, + std::uint32_t count, + queue_reservation* first) -> bool +{ + auto* state = queue.state; + if (count == 0 || count > state->capacity) { return false; } + + auto const head = load_counter(&state->reservation_head); + auto const tail = load_counter(&state->gpu_tail); + if (head - tail > state->capacity - count) { return false; } + if (atomicCAS(&state->reservation_head, head, head + count) != head) { return false; } + + first->sequence = head; + first->slot = static_cast(head) & queue_mask; + return true; +} + +RAFT_DEVICE_INLINE_FUNCTION auto try_reserve_deferred(queue_view queue, + std::uint32_t count, + queue_reservation* first) -> bool +{ + auto* state = queue.state; + if (count == 0 || count > state->capacity) { return false; } + + auto const head = load_counter(&state->reservation_head); + auto const tail = load_counter(&state->gpu_tail); + if (head - tail > state->capacity - count) { + reclaim_released(queue); + return false; + } + if (atomicCAS(&state->reservation_head, head, head + count) != head) { return false; } + + first->sequence = head; + first->slot = static_cast(head) & queue_mask; + return true; +} + +/** Try to reserve one queue slot without blocking the CTA indefinitely. */ +RAFT_DEVICE_INLINE_FUNCTION auto try_reserve(queue_view queue, queue_reservation* reservation) + -> bool +{ + return try_reserve_batch(queue, 1, reservation); +} + +/** Atomically publish each complete request in a reserved batch. */ +RAFT_DEVICE_INLINE_FUNCTION void submit_batch(queue_view queue, + queue_reservation first, + std::uint32_t const* row_ids, + std::uint32_t count) +{ + // The atomically published word contains the complete request. No separate + // payload needs a system fence. The host clears this word before publishing + // the result; result acquisition and deferred reclamation gate slot reuse. + for (std::uint32_t i = 0; i < count; ++i) { + auto const reservation = reservation_at(first, i); + cuda::atomic_ref command( + queue.command_words[reservation.slot]); + static_cast(command.exchange(make_command(row_ids[i]), cuda::memory_order_relaxed)); + } +} + +/** Publish one complete cross-graph row request atomically. */ +RAFT_DEVICE_INLINE_FUNCTION void submit(queue_view queue, + queue_reservation reservation, + std::uint32_t row_id) +{ + submit_batch(queue, reservation, &row_id, 1); +} + +RAFT_DEVICE_INLINE_FUNCTION auto result_ready(queue_view const& queue, + queue_reservation reservation) -> bool +{ + if (load_published_counter(&queue.result_sequence[reservation.slot]) == + reservation.sequence + 1) { + return true; + } + if (queue.aborted) { return true; } + if (queue.abort_flag == nullptr) { return false; } + auto const now = clock64(); + if (queue.next_abort_check == 0) { + queue.next_abort_check = now + (1ull << 24); + return false; + } + if (now < queue.next_abort_check) { return false; } + // Read the host-written emergency signal at a bounded interval. Ordinary completion + // stays in device memory; continuously loading mapped host memory stalls queue traffic. + queue.next_abort_check = now + (1ull << 24); + cuda::atomic_ref aborted( + *const_cast(queue.abort_flag)); + queue.aborted = aborted.load(cuda::memory_order_acquire) != 0; + return queue.aborted; +} + +RAFT_DEVICE_INLINE_FUNCTION auto result_row(queue_view queue, queue_reservation reservation) + -> std::uint32_t const* +{ + return queue.rows + static_cast(reservation.slot) * queue.state->row_length; +} + +/** + * Mark a consumed result as released and advance the global GPU tail over consecutive releases. + */ +RAFT_DEVICE_INLINE_FUNCTION void release(queue_view queue, queue_reservation reservation) +{ + atomicExch(&queue.released_sequence[reservation.slot], reservation.sequence + 1); + + while (true) { + auto const tail = load_counter(&queue.state->gpu_tail); + auto const slot = static_cast(tail) & queue_mask; + if (load_counter(&queue.released_sequence[slot]) != tail + 1) { break; } + if (atomicCAS(&queue.state->gpu_tail, tail, tail + 1) != tail) { continue; } + } +} + +#endif // __CUDACC__ + +#ifdef CUVS_ENABLE_FLOWANN_SEARCH +struct CUVS_EXPORT search_context_access { + static auto device_queues(search_context const& context) noexcept -> queue_view*; + static auto device_id(search_context const& context) noexcept -> int; + static auto num_queues(search_context const& context) noexcept -> std::uint32_t; + static auto row_count(search_context const& context) noexcept -> std::uint64_t; + static auto row_length(search_context const& context) noexcept -> std::uint32_t; + static void acquire_search(search_context& context); + static void release_search(search_context& context); + // Internal failure injection uses the same path as an unexpected poller exception. + static void report_failure(search_context& context, std::exception_ptr error); + static auto is_paused(search_context const& context) noexcept -> bool; + static auto pending_pause_requests(search_context const& context) noexcept -> std::size_t; +}; +#endif + +} // namespace cuvs::neighbors::cagra::experimental::flowann::detail diff --git a/cpp/src/neighbors/detail/flowann/search_multi_cta.cuh b/cpp/src/neighbors/detail/flowann/search_multi_cta.cuh new file mode 100644 index 0000000000..b192915f80 --- /dev/null +++ b/cpp/src/neighbors/detail/flowann/search_multi_cta.cuh @@ -0,0 +1,416 @@ +/* + * 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); + + auto const total_smem_size = + raft::round_up_safe(sizeof(graph_state), std::size_t{128}) + + this->smem_size; + auto const properties = raft::resource::get_device_properties(res); + RAFT_EXPECTS(total_smem_size <= properties.sharedMemPerBlockOptin, + "FlowANN search needs %zu bytes of shared memory, device limit is %zu", + total_smem_size, + properties.sharedMemPerBlockOptin); + + constexpr std::uint32_t min_block_size = 64; + constexpr std::uint32_t 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_size < max_block_size && + total_smem_size > smem_per_warp_limit / 32 * block_size) { + block_size *= 2; + } + auto const properties = raft::resource::get_device_properties(res); + while (block_size < max_block_size && + 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)); + if (queries.extent(0) == 0) { return; } + + 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"); + + std::unique_ptr> padded_queries; + data_t const* query_data{}; + std::uint32_t query_stride{}; + if (cuvs::neighbors::matrix_row_width_matches_cagra_required(queries)) { + auto view = cuvs::neighbors::make_device_padded_dataset_view(res, queries); + query_data = view.data_view().data_handle(); + query_stride = view.data_view().stride(0); + } else { + padded_queries = cuvs::neighbors::make_device_padded_dataset(res, queries); + auto view = padded_queries->view(); + query_data = view.data_view().data_handle(); + query_stride = view.data_view().stride(0); + } + auto const max_batch_queries = plan.max_queries; + auto const stream = raft::resource::get_cuda_stream(res).get(); + auto const* descriptor_ptr = descriptor.dev_ptr(stream); + auto const inner_graph = index.inner_graph(); + auto const subgraph_offsets = index.subgraph_offsets().value(); + auto const subgraph_ids = index.subgraph_ids().value(); + auto const source_indices = index.source_indices(); + auto const* source_ptr = source_indices.has_value() ? source_indices->data_handle() : nullptr; + 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, stream); + for (std::uint32_t offset = 0; offset < num_queries; + offset += static_cast(max_batch_queries)) { + auto const count = + static_cast(std::min(max_batch_queries, num_queries - offset)); + auto const* 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, + query_stride, + inner_graph.data_handle(), + static_cast(inner_graph.extent(1)), + static_cast(index.node_per_cacheline()), + static_cast(index.n_bits()), + subgraph_offsets.data_handle(), + subgraph_ids.data_handle(), + search_context_access::device_queues(context), + search_context_access::num_queues(context), + 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; + auto const multi_cta_fits = + std::max(32ull + 3ull * index.graph_degree(), params.num_seeds) <= 256; + algo = params.itopk_size <= 512 && + (!multi_cta_fits || 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..eac63b4444 --- /dev/null +++ b/cpp/src/neighbors/detail/flowann/search_multi_cta_jit.cuh @@ -0,0 +1,108 @@ +/* + * 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, + std::uint32_t query_stride, + const std::uint8_t* inner_graph, + std::uint32_t inner_row_length, + std::uint32_t node_per_cacheline, + std::uint32_t n_bits, + const std::uint32_t* subgraph_offsets, + const std::uint32_t* subgraph_ids, + queue_view* queues, + std::uint32_t num_queues, + std::uint32_t 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, + query_stride); +} + +} // 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..cb8119285f --- /dev/null +++ b/cpp/src/neighbors/detail/flowann/search_multi_cta_kernel.cu.in @@ -0,0 +1,109 @@ +/* + * 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, + std::uint32_t query_stride, + const std::uint8_t* inner_graph, + std::uint32_t inner_row_length, + std::uint32_t node_per_cacheline, + std::uint32_t n_bits, + const std::uint32_t* subgraph_offsets, + const std::uint32_t* subgraph_ids, + queue_view* queues, + std::uint32_t num_queues, + std::uint32_t 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, + query_stride, + 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..6cccab4520 --- /dev/null +++ b/cpp/src/neighbors/detail/flowann/search_single_cta.cuh @@ -0,0 +1,509 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * FlowANN is described in "Disentangling Graph Dependencies for Efficient + * Billion-Scale GPU Vector Search" (OSDI 2026). + */ + +#pragma once + +#include "graph_state.cuh" +#include "launcher.hpp" +#include "queue.cuh" + +#include +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace cuvs::neighbors::cagra::experimental::flowann::detail { + +inline auto automatic_max_iterations(std::uint32_t itopk_size, + std::uint32_t search_width, + std::uint64_t dataset_size, + std::uint32_t graph_degree, + std::uint32_t min_iterations) -> std::uint32_t +{ + auto iterations = itopk_size / search_width; + std::uint64_t reachable = 1; + while (reachable < dataset_size) { + auto const expansion = std::max(2, graph_degree / 2); + if (reachable > std::numeric_limits::max() / expansion) { break; } + reachable *= expansion; + ++iterations; + } + return std::max(iterations, min_iterations); +} + +template +struct single_cta_plan + : cuvs::neighbors::cagra::detail:: + search_plan_impl { + using base_type = cuvs::neighbors::cagra::detail:: + search_plan_impl; + + std::uint32_t num_itopk_candidates{}; + + static auto planning_params(flowann::search_params const& input) + -> cuvs::neighbors::cagra::search_params + { + cuvs::neighbors::cagra::search_params result = input; + result.search_width += additional_buffer_width; + result.algo = cuvs::neighbors::cagra::search_algo::SINGLE_CTA; + result.persistent = false; + return result; + } + + single_cta_plan( + raft::resources const& res, + flowann::search_params const& params, + const cuvs::neighbors::cagra::detail::dataset_descriptor_host& + dataset_desc, + int64_t dim, + int64_t dataset_size, + int64_t graph_degree, + std::uint32_t topk) + : base_type(res, planning_params(params), dataset_desc, dim, dataset_size, graph_degree, topk) + { + // The inherited CAGRA plan sizes hashmaps and candidate storage for two additional graph rows, + // while the actual traversal width remains the user-requested value. + this->search_width = params.search_width; + if (params.max_iterations == 0) { + this->max_iterations = automatic_max_iterations(static_cast(this->itopk_size), + params.search_width, + dataset_size, + graph_degree, + params.min_iterations); + } else { + this->max_iterations = std::max(params.max_iterations, params.min_iterations); + } + set_launch_params(res, params.num_seeds); + } + + void set_launch_params(raft::resources const& res, std::uint32_t num_seeds) + { + auto const expansion_candidates = + (static_cast(this->search_width) + additional_buffer_width) * + static_cast(this->graph_degree); + auto const base_result_buffer_size = + static_cast(this->itopk_size) + expansion_candidates; + this->result_buffer_size = std::max(base_result_buffer_size, num_seeds); + num_itopk_candidates = this->result_buffer_size - static_cast(this->itopk_size); + auto const result_buffer_size_32 = + raft::round_up_safe(this->result_buffer_size, 32); + constexpr std::uint32_t max_itopk = 512; + RAFT_EXPECTS( + this->itopk_size <= max_itopk, "FlowANN itopk_size cannot be larger than %u", max_itopk); + + auto const base_smem_size = + this->dataset_desc.smem_ws_size_in_bytes + + (sizeof(IndexT) + sizeof(DistanceT)) * result_buffer_size_32 + + sizeof(IndexT) * cuvs::neighbors::cagra::detail::hashmap::get_size(this->small_hash_bitlen) + + sizeof(IndexT) * this->search_width + 4 * sizeof(std::uint32_t); + auto additional_smem_size = std::uint32_t{0}; + if (num_itopk_candidates > 256) { + auto const radix_itopk = this->itopk_size <= 256 ? 256u : 512u; + additional_smem_size = + cuvs::neighbors::cagra::detail::single_cta_search::topk_by_radix_sort::smem_size( + radix_itopk) * + sizeof(std::uint32_t); + } + if constexpr (!std::is_same_v) { + using scan_op_t = cub::WarpScan; + additional_smem_size = + std::max(additional_smem_size, sizeof(typename scan_op_t::TempStorage)); + } + this->smem_size = base_smem_size + additional_smem_size; + // graph_state is declared with __shared__ in the JIT kernel, in addition to + // the dynamic workspace passed to the launch API. Include alignment padding. + auto const total_smem_size = + raft::round_up_safe(sizeof(graph_state), std::size_t{128}) + + this->smem_size; + auto const properties = raft::resource::get_device_properties(res); + RAFT_EXPECTS(total_smem_size <= properties.sharedMemPerBlockOptin, + "FlowANN search needs %zu bytes of shared memory, device limit is %zu", + total_smem_size, + properties.sharedMemPerBlockOptin); + + constexpr std::uint32_t min_block_size = 64; + constexpr std::uint32_t min_radix_block_size = 256; + constexpr std::uint32_t max_block_size = 1024; + auto block_size = static_cast(this->thread_block_size); + if (block_size == 0) { + block_size = num_itopk_candidates > 256 ? min_radix_block_size : min_block_size; + while (num_itopk_candidates > 256 && block_size < max_block_size && + max_itopk / block_size > 4) { + block_size *= 2; + } + constexpr std::uint32_t smem_per_warp_limit = 4096; + while (block_size < max_block_size && + total_smem_size > smem_per_warp_limit / 32 * block_size) { + block_size *= 2; + } + auto const properties = raft::resource::get_device_properties(res); + while (block_size < max_block_size && + this->graph_degree * this->search_width * this->team_size >= block_size * 2 && + this->max_queries <= (1024 / (block_size * 2)) * properties.multiProcessorCount) { + block_size *= 2; + } + } + RAFT_EXPECTS(block_size >= min_block_size && block_size <= max_block_size, + "FlowANN thread_block_size must be in [64, 1024]"); + this->thread_block_size = block_size; + + this->hashmap_size = 0; + if (this->small_hash_bitlen == 0) { + this->hashmap_size = + this->max_queries * cuvs::neighbors::cagra::detail::hashmap::get_size(this->hash_bitlen); + this->hashmap.resize(this->hashmap_size, raft::resource::get_cuda_stream(res)); + } + } +}; + +template +auto make_dataset_descriptor(raft::resources const& res, + flowann::search_params& params, + flowann::index const& index) +{ + if constexpr (cuvs::neighbors::is_device_vpq_f16_dataset_view_v) { + if (params.smem_dtype == cuvs::neighbors::cagra::internal_dtype::E5M2 && + raft::getComputeCapability().first < 9) { + RAFT_LOG_WARN( + "FlowANN VPQ E5M2 smem_dtype requires native FP8 support on SM90+. Falling back to F16."); + params.smem_dtype = cuvs::neighbors::cagra::internal_dtype::F16; + } + return cuvs::neighbors::cagra::detail:: + dataset_descriptor_init_with_cache( + res, params, index.dataset(), index.metric(), nullptr); + } else { + static_assert(cuvs::neighbors::is_device_padded_dataset_view_v, + "FlowANN search supports device-padded or float16 VPQ datasets"); + if (params.smem_dtype != cuvs::neighbors::cagra::internal_dtype::F16) { + RAFT_LOG_WARN("FlowANN dense search supports only F16 smem_dtype. Falling back to F16."); + params.smem_dtype = cuvs::neighbors::cagra::internal_dtype::F16; + } + auto const dataset_norms = index.dataset_norms(); + RAFT_EXPECTS( + index.metric() != cuvs::distance::DistanceType::CosineExpanded || dataset_norms.has_value(), + "FlowANN dataset norms must be available for CosineExpanded metric"); + auto const* dataset_norms_ptr = + dataset_norms.has_value() ? dataset_norms->data_handle() : nullptr; + return cuvs::neighbors::cagra::detail:: + dataset_descriptor_init_with_cache( + res, params, index.dataset(), index.metric(), dataset_norms_ptr); + } +} + +template +void postprocess_distances(raft::resources const& res, + cuvs::distance::DistanceType metric, + raft::device_matrix_view queries, + raft::device_matrix_view distances) +{ + static_assert(std::is_same_v, + "FlowANN currently supports only float distances"); + constexpr float scale = cuvs::spatial::knn::detail::utils::config::kDivisor / + cuvs::spatial::knn::detail::utils::config::kDivisor; + + if (metric == cuvs::distance::DistanceType::CosineExpanded) { + auto query_norms = raft::make_device_vector(res, queries.extent(0)); + auto scaled_sq_op = raft::compose_op( + raft::sq_op{}, raft::div_const_op{DistanceT(scale)}, raft::cast_op()); + raft::linalg::reduce(res, + queries, + query_norms.view(), + DistanceT{0}, + false, + scaled_sq_op, + raft::add_op(), + raft::sqrt_op{}); + raft::linalg::matrix_vector_op( + res, + raft::make_const_mdspan(distances), + raft::make_const_mdspan(query_norms.view()), + distances, + raft::compose_op(raft::add_const_op{DistanceT{1}}, raft::div_checkzero_op{})); + } else { + cuvs::neighbors::ivf::detail::postprocess_distances(res, + distances.data_handle(), + distances.data_handle(), + metric, + distances.extent(0), + distances.extent(1), + scale, + true); + } +} + +class context_run_guard { + public: + explicit context_run_guard(search_context& context, cudaStream_t stream) + : context_(context), stream_(stream) + { + search_context_access::acquire_search(context_); + } + context_run_guard(context_run_guard const&) = delete; + auto operator=(context_run_guard const&) -> context_run_guard& = delete; + + ~context_run_guard() + { + if (active_) { + // A failed launch or an exception during a later batch must not stop pollers + // while an earlier kernel still waits for its host-serviced requests. + RAFT_CUDA_TRY_NO_THROW(cudaStreamSynchronize(stream_)); + try { + active_ = false; + search_context_access::release_search(context_); + } catch (...) { + } + } + } + + void pause() + { + active_ = false; + search_context_access::release_search(context_); + } + + private: + search_context& context_; + cudaStream_t stream_; + bool active_{true}; +}; + +template +void search_single_cta(raft::resources const& res, + flowann::search_params params, + flowann::index const& index, + search_context& context, + raft::device_matrix_view queries, + raft::device_matrix_view neighbors, + raft::device_matrix_view distances, + SampleFilterT sample_filter) +{ + using data_t = DataT; + using index_t = std::uint32_t; + using distance_t = float; + using source_index_t = std::uint32_t; + using filter_t = + typename cuvs::neighbors::cagra::detail::CagraSampleFilterT_Selector::type; + + static_assert( + std::is_same_v || std::is_same_v, + "FlowANN output indices must be uint32_t or int64_t"); + RAFT_EXPECTS(index.size() > 0 && index.graph_degree() > 0, + "FlowANN search requires a non-empty tiered graph"); + RAFT_EXPECTS(index.dataset().n_rows() == static_cast(index.size()), + "FlowANN search requires a dataset row for every graph node"); + RAFT_EXPECTS(index.subgraph_offsets().has_value() && index.subgraph_ids().has_value(), + "FlowANN search requires subgraph offsets and per-node subgraph IDs"); + RAFT_EXPECTS(queries.extent(1) == index.dim(), + "FlowANN query dimension must match the index dimension"); + RAFT_EXPECTS(neighbors.extent(0) == queries.extent(0) && distances.extent(0) == queries.extent(0), + "FlowANN output row count must match the number of queries"); + RAFT_EXPECTS(neighbors.extent(1) == distances.extent(1) && neighbors.extent(1) > 0, + "FlowANN neighbor and distance outputs must have the same positive k"); + RAFT_EXPECTS(queries.extent(0) <= std::numeric_limits::max(), + "FlowANN single-CTA query count exceeds uint32"); + RAFT_EXPECTS(params.search_width > 0 && params.search_width <= max_search_width, + "FlowANN search_width must be in [1, %u]", + max_search_width); + RAFT_EXPECTS(params.sync_window_scale >= 0.0f, "FlowANN sync_window_scale must not be negative"); + RAFT_EXPECTS(params.algo == cuvs::neighbors::cagra::search_algo::AUTO || + params.algo == cuvs::neighbors::cagra::search_algo::SINGLE_CTA, + "This FlowANN path supports SINGLE_CTA or AUTO"); + RAFT_EXPECTS(search_context_access::row_count(context) == index.size() && + search_context_access::row_length(context) == index.cross_degree() + 1, + "FlowANN search context does not match the index cross graph"); + + auto const requested_seeds = params.num_seeds; + auto const index_seeds = index.seeds(); + auto const seed_cross_graph = index.seed_cross_graph(); + auto const seed_lookup_keys = index.seed_lookup_keys(); + auto const seed_lookup_values = index.seed_lookup_values(); + RAFT_EXPECTS(requested_seeds == 0 || index_seeds.has_value(), + "FlowANN num_seeds requires preloaded index seeds"); + RAFT_EXPECTS(!index_seeds.has_value() || requested_seeds <= index_seeds->extent(0), + "FlowANN num_seeds exceeds the number of preloaded seeds"); + RAFT_EXPECTS(requested_seeds == 0 || seed_cross_graph.has_value(), + "FlowANN num_seeds requires cached seed cross-graph rows"); + RAFT_EXPECTS( + requested_seeds == 0 || (seed_lookup_keys.has_value() && seed_lookup_values.has_value()), + "FlowANN num_seeds requires a prebuilt seed lookup"); + auto const* seeds = requested_seeds == 0 ? nullptr : index_seeds->data_handle(); + auto const* seed_cross_ptr = requested_seeds == 0 ? nullptr : seed_cross_graph->data_handle(); + auto const* seed_lookup_keys_ptr = + requested_seeds == 0 ? nullptr : seed_lookup_keys->data_handle(); + auto const* seed_lookup_values_ptr = + requested_seeds == 0 ? nullptr : seed_lookup_values->data_handle(); + auto const seed_lookup_capacity = + requested_seeds == 0 ? 0u : static_cast(seed_lookup_keys->extent(0)); + + if (queries.extent(0) == 0) { return; } + + auto device_properties = raft::resource::get_device_properties(res); + if (params.max_queries == 0) { + params.max_queries = std::min(queries.extent(0), device_properties.maxGridSize[1]); + } + RAFT_EXPECTS(params.max_queries > 0, "FlowANN max_queries must be positive"); + + auto descriptor = make_dataset_descriptor(res, params, index); + auto const topk = static_cast(neighbors.extent(1)); + single_cta_plan plan( + res, params, descriptor, index.dim(), index.size(), index.graph_degree(), topk); + plan.check(topk); + plan.num_seeds = requested_seeds; + + auto launch_config = cuvs::neighbors::cagra::detail::single_cta_search::compute_launch_config( + plan.num_itopk_candidates, + static_cast(plan.itopk_size), + static_cast(plan.thread_block_size)); + auto launcher_filter = cuvs::neighbors::cagra::detail::set_offset(sample_filter, 0); + auto launcher = make_single_cta_launcher< + data_t, + index_t, + distance_t, + source_index_t, + cuvs::neighbors::cagra::detail::sample_filter_jit_tag_t>( + descriptor, + launch_config.topk_by_bitonic_sort, + launch_config.bitonic_sort_and_merge_multi_warps, + cuvs::neighbors::cagra::detail::make_cagra_sample_filter_udf_fragment( + launcher_filter)); + RAFT_EXPECTS(launcher != nullptr, "Failed to create the FlowANN single-CTA JIT launcher"); + + std::unique_ptr> padded_queries; + data_t const* query_data{}; + std::uint32_t query_stride{}; + if (cuvs::neighbors::matrix_row_width_matches_cagra_required(queries)) { + auto view = cuvs::neighbors::make_device_padded_dataset_view(res, queries); + query_data = view.data_view().data_handle(); + query_stride = view.data_view().stride(0); + } else { + padded_queries = cuvs::neighbors::make_device_padded_dataset(res, queries); + auto view = padded_queries->view(); + query_data = view.data_view().data_handle(); + query_stride = view.data_view().stride(0); + } + auto const max_batch_queries = plan.max_queries; + auto const stream = raft::resource::get_cuda_stream(res).get(); + auto const* descriptor_ptr = descriptor.dev_ptr(stream); + auto const inner_graph = index.inner_graph(); + auto const subgraph_offsets = index.subgraph_offsets().value(); + auto const subgraph_ids = index.subgraph_ids().value(); + auto const source_indices = index.source_indices(); + auto const* source_ptr = source_indices.has_value() ? source_indices->data_handle() : nullptr; + constexpr uintptr_t output_tag = raft::Pow2::Log2; + static_assert(output_tag <= 3, "FlowANN output index type cannot exceed 8 bytes"); + auto const num_queries = static_cast(queries.extent(0)); + using kernel_type = search_single_cta_kernel_func_t; + + context_run_guard running(context, stream); + for (std::uint32_t offset = 0; offset < num_queries; + offset += static_cast(max_batch_queries)) { + auto const count = + static_cast(std::min(max_batch_queries, num_queries - offset)); + auto const result_ptr = reinterpret_cast(neighbors.data_handle() + + static_cast(offset) * topk); + if constexpr (output_tag <= 1) { + RAFT_EXPECTS((result_ptr & 0x3) == 0, + "FlowANN result indices must be at least 4-byte aligned"); + } + auto const* query_ptr = query_data + static_cast(offset) * query_stride; + auto batch_filter = cuvs::neighbors::cagra::detail::set_offset(sample_filter, offset); + auto filter_payload = + cuvs::neighbors::cagra::detail::extract_cagra_sample_filter(batch_filter, + stream); + dim3 grid(1, count, 1); + dim3 block(static_cast(plan.thread_block_size), 1, 1); + + auto kernel_launcher = [&]() { + launcher->template dispatch( + stream, + grid, + block, + static_cast(plan.smem_size), + result_ptr | output_tag, + distances.data_handle() + static_cast(offset) * topk, + topk, + query_ptr, + query_stride, + inner_graph.data_handle(), + static_cast(inner_graph.extent(1)), + static_cast(index.node_per_cacheline()), + static_cast(index.n_bits()), + subgraph_offsets.data_handle(), + subgraph_ids.data_handle(), + search_context_access::device_queues(context), + search_context_access::num_queues(context), + index.graph_degree(), + index.cross_degree(), + source_ptr, + static_cast(plan.num_random_samplings), + plan.rand_xor_mask, + seeds, + requested_seeds, + seed_cross_ptr, + seed_lookup_keys_ptr, + seed_lookup_values_ptr, + seed_lookup_capacity, + plan.hashmap.data(), + launch_config.max_candidates, + launch_config.max_itopk, + static_cast(plan.itopk_size), + static_cast(plan.search_width), + static_cast(plan.min_iterations), + static_cast(plan.max_iterations), + static_cast(nullptr), + static_cast(plan.hash_bitlen), + static_cast(plan.small_hash_bitlen), + static_cast(plan.small_hash_reset_interval), + filter_payload.query_id_offset, + descriptor_ptr, + index.size(), + params.sync_window_scale, + params.sync_drop_threshold, + filter_payload); + }; + cuvs::neighbors::detail::safely_launch_kernel_with_smem_size( + plan.smem_size, kernel_launcher, launcher->get_kernel()); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + } + raft::resource::sync_stream(res); + running.pause(); + + postprocess_distances(res, index.metric(), queries, distances); +} + +} // namespace cuvs::neighbors::cagra::experimental::flowann::detail diff --git a/cpp/src/neighbors/detail/flowann/search_single_cta_jit.cuh b/cpp/src/neighbors/detail/flowann/search_single_cta_jit.cuh new file mode 100644 index 0000000000..e4022a25f2 --- /dev/null +++ b/cpp/src/neighbors/detail/flowann/search_single_cta_jit.cuh @@ -0,0 +1,636 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * FlowANN is described in "Disentangling Graph Dependencies for Efficient + * Billion-Scale GPU Vector Search" (OSDI 2026). + */ + +#pragma once + +#include "graph_state.cuh" + +#include +#include + +#include +#include +#include + +#include +#include +#include + +namespace cuvs::neighbors::cagra::experimental::flowann::detail { + +RAFT_DEVICE_INLINE_FUNCTION auto load_packed_neighbor(const std::uint8_t* inner_graph, + std::uint64_t bit_position, + std::uint32_t neighbor, + std::uint32_t n_bits) -> std::uint32_t +{ + bit_position += static_cast(neighbor) * n_bits; + auto const byte_position = bit_position / 8; + auto const bit_offset = static_cast(bit_position % 8); + auto const byte_count = (bit_offset + n_bits + 7) / 8; + + // The legacy 24-bit layout is byte aligned. Issue the three independent byte + // reads without a runtime loop; never read beyond the encoded neighbor. + if (n_bits == 24 && bit_offset == 0) { + return static_cast(inner_graph[byte_position]) | + (static_cast(inner_graph[byte_position + 1]) << 8) | + (static_cast(inner_graph[byte_position + 2]) << 16); + } + + std::uint64_t packed = 0; + for (std::uint32_t i = 0; i < byte_count; ++i) { + packed |= static_cast(inner_graph[byte_position + i]) << (8 * i); + } + auto const mask = n_bits == 32 ? std::uint64_t{0xffffffffu} : ((std::uint64_t{1} << n_bits) - 1); + return static_cast((packed >> bit_offset) & mask); +} + +template +struct tiered_graph_provider { + const std::uint8_t* inner_graph; + std::uint32_t inner_row_length; + std::uint32_t node_per_cacheline; + std::uint32_t n_bits; + const std::uint32_t* subgraph_offsets; + const std::uint32_t* subgraph_ids; + const IndexT* seeds; + const std::uint32_t* seed_cross_graph; + const IndexT* seed_lookup_keys; + const std::uint32_t* seed_lookup_values; + std::uint32_t seed_lookup_capacity; + std::uint32_t num_seeds; + queue_view queue; + std::uint32_t graph_degree; + std::uint32_t cross_graph_degree; + IndexT graph_size; + float sync_window_scale; + std::uint32_t sync_drop_threshold; + graph_state* state; + + RAFT_DEVICE_INLINE_FUNCTION void initialize() const + { + if (threadIdx.x == 0) { + state->pending.initialize(); + state->previous_parent_position = 0; + state->iteration = 0; + state->use_seed_cache = false; + } + } + + [[nodiscard]] RAFT_DEVICE_INLINE_FUNCTION auto candidate_capacity( + std::uint32_t search_width) const -> std::uint32_t + { + return (search_width + additional_buffer_width) * graph_degree; + } + + [[nodiscard]] RAFT_DEVICE_INLINE_FUNCTION auto result_buffer_size( + std::uint32_t base_size, std::uint32_t requested_seeds) const -> std::uint32_t + { + return seeds == nullptr ? base_size : std::max(base_size, requested_seeds); + } + + [[nodiscard]] RAFT_DEVICE_INLINE_FUNCTION auto seed_pointer(const IndexT* seeds, + std::uint32_t, + std::uint32_t) const -> const IndexT* + { + return seeds; + } + + RAFT_DEVICE_INLINE_FUNCTION auto seed_block_id(std::uint32_t cta_id, + std::uint32_t num_cta_per_query, + std::uint32_t row) const -> std::uint32_t + { + return seeds != nullptr && num_seeds > 0 ? 0 : cta_id + num_cta_per_query * row; + } + + RAFT_DEVICE_INLINE_FUNCTION auto seed_num_blocks(std::uint32_t num_cta_per_query, + std::uint32_t num_queries, + std::uint32_t num_partitions) const + -> std::uint32_t + { + return seeds != nullptr && num_seeds > 0 ? 1 : num_cta_per_query * num_queries * num_partitions; + } + + [[nodiscard]] RAFT_DEVICE_INLINE_FUNCTION auto seed_position(IndexT key) const -> std::uint32_t + { + if (seed_lookup_keys == nullptr || seed_lookup_values == nullptr || seed_lookup_capacity == 0) { + return num_seeds; + } + + std::uint64_t hash = static_cast(key); + hash ^= hash >> 16; + hash *= 0x85ebca77u; + hash ^= hash >> 13; + hash *= 0xc2b2ae3du; + hash ^= hash >> 16; + auto slot = static_cast(hash % seed_lookup_capacity); + + constexpr auto empty_key = ~IndexT{0}; + for (std::uint32_t probe = 0; probe < seed_lookup_capacity; ++probe) { + auto const stored = seed_lookup_keys[slot]; + if (stored == empty_key) { return num_seeds; } + if (stored == key) { return seed_lookup_values[slot]; } + slot = (slot + 1) % seed_lookup_capacity; + } + return num_seeds; + } + + RAFT_DEVICE_INLINE_FUNCTION auto initial_candidate_count(std::uint32_t candidate_count, + std::uint32_t requested_seeds, + std::uint32_t block_id, + std::uint32_t num_blocks) const + -> std::uint32_t + { + if (seeds == nullptr || requested_seeds == 0) { return candidate_count; } + if (block_id >= requested_seeds) { return 0; } + auto const seeds_for_block = (requested_seeds - block_id + num_blocks - 1) / num_blocks; + // The result buffer includes space for the explicit seeds. The graph degree only limits + // random initialization; applying it here silently discards later medoids. + return seeds_for_block; + } + + RAFT_DEVICE_INLINE_FUNCTION auto initial_distillation_count(std::uint32_t count, + std::uint32_t requested_seeds) const + -> std::uint32_t + { + return seeds != nullptr && requested_seeds > 0 ? 1u : count; + } + + RAFT_DEVICE_INLINE_FUNCTION auto expand( + IndexT* result_indices, + DistanceT* result_distances, + const cuvs::neighbors::cagra::detail::dataset_descriptor_base_t* + dataset_desc, + IndexT* visited_hashmap, + std::uint32_t visited_hash_bitlen, + const IndexT* parent_indices, + const IndexT* internal_topk, + std::uint32_t search_width, + IndexT* traversed_hashmap = nullptr, + std::uint32_t traversed_hash_bitlen = 0, + int* result_position = nullptr, + int max_result_position = 0) const -> std::uint32_t + { + using cuvs::neighbors::cagra::detail::device::team_sum; + using cuvs::neighbors::cagra::detail::utils::gen_index_msb_1_mask; + constexpr IndexT invalid_index = ~static_cast(0); + constexpr IndexT index_msb = gen_index_msb_1_mask::value; + + auto store_child = [&](IndexT child, std::uint32_t i) { + if (child != invalid_index && cuvs::neighbors::cagra::detail::hashmap::insert( + visited_hashmap, visited_hash_bitlen, child) == 0) { + child = invalid_index; + } else if (child != invalid_index && traversed_hashmap != nullptr && + cuvs::neighbors::cagra::detail::hashmap::search( + traversed_hashmap, traversed_hash_bitlen, child)) { + child = invalid_index; + } + if (result_position == nullptr) { + result_indices[i] = child; + } else if (child != invalid_index) { + auto const position = atomicSub(result_position, 1) - 1; + result_indices[position] = child; + } + }; + + // A one-parent compacted expansion fits in one warp. Decode and deduplicate + // its inner neighbors while warp 0 publishes requests. Compaction writes + // only the candidate tail, beyond the retained parent pool. + bool const overlap_inner = + result_position != nullptr && search_width == 1 && graph_degree <= 32; + + // Read packed-row metadata independently while warp 0 publishes cross-graph requests. + // These lanes derive parents directly; they do not depend on warp 0's shared writes. + if (threadIdx.x >= 32 && threadIdx.x - 32 < search_width) { + auto const p = threadIdx.x - 32; + auto const position = parent_indices[p]; + auto const valid = + position != invalid_index && inner_graph != nullptr && inner_row_length != 0; + std::uint32_t count = 0; + std::uint64_t bits = 0; + std::uint32_t base = 0; + if (valid) { + auto const parent = static_cast(internal_topk[position] & ~index_msb); + auto const row = static_cast(parent) / node_per_cacheline; + auto const col = static_cast(parent % node_per_cacheline); + auto const start = row * inner_row_length; + count = inner_graph[start + col]; + bits = (start + node_per_cacheline) * 8; + for (std::uint32_t j = 0; j < col; ++j) { + bits += static_cast(inner_graph[start + j]) * n_bits; + } + if (count > 0) { base = subgraph_offsets[subgraph_ids[parent]]; } + } + state->parent_inner_counts[p] = count; + state->parent_bit_offsets[p] = bits; + state->parent_subgraph_offsets[p] = base; + } + + if (overlap_inner && threadIdx.x >= 32 && threadIdx.x < 64) { + __syncwarp(); + auto const lane = threadIdx.x - 32; + if (lane < state->parent_inner_counts[0]) { + auto const local = + load_packed_neighbor(inner_graph, state->parent_bit_offsets[0], lane, n_bits); + auto const decoded = static_cast(state->parent_subgraph_offsets[0]) + local; + if (decoded < graph_size) { store_child(static_cast(decoded), lane); } + } + } + + if (threadIdx.x == 0) { + auto const current_parent_position = + parent_indices[0] == invalid_index ? 0u : static_cast(parent_indices[0]); + auto const sharp_drop = + state->previous_parent_position > current_parent_position && + state->previous_parent_position - current_parent_position >= sync_drop_threshold; + auto const window_value = + sharp_drop + ? 0u + : std::max(search_width, + static_cast(sync_window_scale * current_parent_position)); + auto const sync_window = static_cast(std::min(window_value, 0xffffu)); + + for (std::uint32_t i = 0; i < search_width; ++i) { + auto const parent_position = parent_indices[i]; + auto const valid = parent_position != invalid_index; + auto const parent = + valid ? static_cast(internal_topk[parent_position] & ~index_msb) : IndexT{}; + state->parent_ids[i] = parent; + state->parent_valid[i] = valid; + } + + state->use_seed_cache = state->iteration == 0 && seeds != nullptr && + seed_cross_graph != nullptr && seed_lookup_keys != nullptr && + seed_lookup_values != nullptr && seed_lookup_capacity > 0 && + num_seeds > 0; + if (state->use_seed_cache) { + for (std::uint32_t i = 0; i < search_width; ++i) { + if (!state->parent_valid[i]) { continue; } + auto const seed_position = this->seed_position(state->parent_ids[i]); + state->parent_seed_positions[i] = seed_position; + if (seed_position == num_seeds) { + state->use_seed_cache = false; + break; + } + } + } + + if (!state->use_seed_cache) { + submit_parent_requests(queue, + state, + parent_indices, + internal_topk, + search_width, + sync_window, + traversed_hashmap, + traversed_hash_bitlen); + } + } + // The one-parent path can consume ready rows while warp 1 prepares inner + // neighbors. Warp 0 reads the raw inner count independently below; the + // following join still protects all candidate writes and row descriptors. + if (!overlap_inner) { __syncthreads(); } + + auto decode_inner = [&](std::uint32_t i) { + IndexT child = invalid_index; + auto local = i; + std::uint32_t parent_index = 0; + for (; parent_index < search_width; ++parent_index) { + auto const count = state->parent_inner_counts[parent_index]; + if (local < count) { break; } + local -= count; + } + if (parent_index < search_width && state->parent_valid[parent_index]) { + auto const local_neighbor = + load_packed_neighbor(inner_graph, state->parent_bit_offsets[parent_index], local, n_bits); + auto const decoded = + static_cast(state->parent_subgraph_offsets[parent_index]) + local_neighbor; + if (decoded < graph_size) { child = static_cast(decoded); } + } + return child; + }; + + // The caller has acquired this row's ready sequence. A partial row remains + // pending until the rest of its neighbors fit in a later expansion. + auto consume_cross_row = [&](std::uint32_t& remaining) { + auto const& entry = state->pending.front(); + auto const* row = result_row(queue, entry.reservation); + auto edge_count = std::min(row[0], cross_graph_degree); + auto const consumed = state->pending.consumed; + edge_count = edge_count > consumed ? edge_count - consumed : 0; + auto const take = std::min(edge_count, remaining); + auto const ready = state->ready_cross_count; + state->ready_cross_data[ready] = row + 1 + consumed; + state->ready_cross_edge_counts[ready] = take; + ++state->ready_cross_count; + remaining -= take; + + if (take < edge_count) { + state->pending.consumed += take; + return false; + } + state->pending.consumed = 0; + state->ready_reservations[ready] = entry.reservation; + state->pending.pop(); + return true; + }; + + if (threadIdx.x == 0) { + std::uint32_t inner_count = 0; + if (overlap_inner) { + if (state->parent_valid[0] && inner_graph != nullptr && inner_row_length != 0) { + auto const parent = state->parent_ids[0]; + auto const row = static_cast(parent) / node_per_cacheline; + auto const col = static_cast(parent % node_per_cacheline); + inner_count = inner_graph[row * inner_row_length + col]; + } + } else { + for (std::uint32_t i = 0; i < search_width; ++i) { + inner_count += state->parent_inner_counts[i]; + } + } + state->total_inner_children = inner_count; + + auto remaining = candidate_capacity(search_width) - inner_count; + state->ready_cross_count = 0; + + if (state->use_seed_cache) { + auto const row_length = cross_graph_degree + 1; + for (std::uint32_t i = 0; + i < search_width && state->ready_cross_count < max_search_width && remaining > 0; + ++i) { + if (!state->parent_valid[i]) { continue; } + auto const* row = seed_cross_graph + + static_cast(state->parent_seed_positions[i]) * row_length; + auto const take = std::min(std::min(row[0], cross_graph_degree), remaining); + if (take == 0) { continue; } + auto const ready = state->ready_cross_count; + state->ready_cross_data[ready] = row + 1; + state->ready_cross_edge_counts[ready] = take; + ++state->ready_cross_count; + remaining -= take; + } + } else { + auto available = state->pending.available(); + auto deferred = state->pending.size(); + while (!state->pending.empty() && state->ready_cross_count < max_search_width && + remaining > 0) { + auto& entry = state->pending.front(); + auto need_wait = available < search_width || deferred > entry.sync_window; + if (need_wait) { + while (!result_ready(queue, entry.reservation)) {} + ++available; + --deferred; + } else { + if (!result_ready(queue, entry.reservation)) { break; } + } + + if (!consume_cross_row(remaining)) { break; } + } + } + state->total_children = candidate_capacity(search_width) - remaining; + state->previous_parent_position = + parent_indices[0] == invalid_index ? 0u : static_cast(parent_indices[0]); + } else if (!overlap_inner && threadIdx.x >= 32) { + std::uint32_t inner_count = 0; + for (std::uint32_t p = 0; p < search_width; ++p) { + inner_count += state->parent_inner_counts[p]; + } + for (std::uint32_t i = threadIdx.x - 32; i < inner_count; i += blockDim.x - 32) { + store_child(decode_inner(i), i); + } + } + __syncthreads(); + + // In the compacted multi-CTA path, include rows that became ready during + // inner-neighbor work. This pass never waits or expands the candidate budget. + // The single-CTA path retains its single readiness pass. + if (result_position != nullptr) { + if (threadIdx.x == 0 && !state->use_seed_cache) { + auto remaining = candidate_capacity(search_width) - state->total_children; + while (!state->pending.empty() && state->ready_cross_count < max_search_width && + remaining > 0) { + if (!result_ready(queue, state->pending.front().reservation)) { break; } + if (!consume_cross_row(remaining)) { break; } + } + state->total_children = candidate_capacity(search_width) - remaining; + } + __syncthreads(); + } + + for (std::uint32_t i = threadIdx.x + state->total_inner_children; i < state->total_children; + i += blockDim.x) { + IndexT child = invalid_index; + auto local = i - state->total_inner_children; + for (std::uint32_t ready = 0; ready < state->ready_cross_count; ++ready) { + auto const count = state->ready_cross_edge_counts[ready]; + if (local < count) { + auto const decoded = state->ready_cross_data[ready][local]; + if (decoded < graph_size) { child = static_cast(decoded); } + break; + } + local -= count; + } + store_child(child, i); + } + __syncthreads(); + + if (threadIdx.x == 0) { + if (!state->use_seed_cache) { + auto release_count = state->ready_cross_count; + if (state->pending.consumed != 0 && release_count > 0) { --release_count; } + for (std::uint32_t i = 0; i < release_count; ++i) { + release_deferred(queue, state->ready_reservations[i]); + } + } + ++state->iteration; + } + // PQ uses only the copied child IDs. The caller joins the CTA after expansion, + // before any thread can reuse iteration state or consume another queue row. + + auto const team_size_bits = dataset_desc->team_size_bitshift_from_smem(); + auto const offset = + result_position == nullptr ? 0u : static_cast(result_position[0]); + auto const max_i = raft::round_up_safe( + state->total_children, cuvs::neighbors::cagra::detail::device::warp_size >> team_size_bits); + auto const args = dataset_desc->args.load(); + auto const lead_lane = (threadIdx.x & ((1u << team_size_bits) - 1u)) == 0; + for (std::uint32_t i = threadIdx.x >> team_size_bits; i < max_i; + i += blockDim.x >> team_size_bits) { + auto const position = i + offset; + auto const valid = i < state->total_children && + (result_position == nullptr || position < max_result_position) && + result_indices[position] != invalid_index; + auto const distance = team_sum( + valid + ? cuvs::neighbors::cagra::detail::compute_distance_per_thread( + args, result_indices[position]) + : (lead_lane ? raft::upper_bound() : DistanceT{}), + team_size_bits); + __syncwarp(); + if (i < state->total_children && + (result_position == nullptr || position < max_result_position) && lead_lane) { + result_distances[position] = distance; + } + } + return state->total_children; + } + + RAFT_DEVICE_INLINE_FUNCTION auto expand( + IndexT* result_indices, + DistanceT* result_distances, + const cuvs::neighbors::cagra::detail::dataset_descriptor_base_t* + dataset_desc, + IndexT* visited_hashmap, + std::uint32_t visited_hash_bitlen, + IndexT* traversed_hashmap, + std::uint32_t traversed_hash_bitlen, + const IndexT* parent_indices, + const IndexT* internal_topk, + int* result_position, + int max_result_position) const -> std::uint32_t + { + return expand(result_indices, + result_distances, + dataset_desc, + visited_hashmap, + visited_hash_bitlen, + parent_indices, + internal_topk, + 1, + traversed_hashmap, + traversed_hash_bitlen, + result_position, + max_result_position); + } + + RAFT_DEVICE_INLINE_FUNCTION void finalize() const + { + if (threadIdx.x == 0) { + while (!state->pending.empty()) { + auto const reservation = state->pending.front().reservation; + while (!result_ready(queue, reservation)) {} + release_deferred(queue, reservation); + state->pending.pop(); + } + state->pending.consumed = 0; + } + __syncthreads(); + if (threadIdx.x < 32) { reclaim_released_warp(queue); } + __syncthreads(); + } +}; + +template +__device__ void search_kernel_jit( + uintptr_t result_indices_ptr, + DistanceT* result_distances_ptr, + std::uint32_t top_k, + const DataT* queries_ptr, + std::uint32_t query_stride, + const std::uint8_t* inner_graph, + std::uint32_t inner_row_length, + std::uint32_t node_per_cacheline, + std::uint32_t n_bits, + const std::uint32_t* subgraph_offsets, + const std::uint32_t* subgraph_ids, + queue_view* queues, + std::uint32_t num_queues, + std::uint32_t graph_degree, + std::uint32_t cross_graph_degree, + const SourceIndexT* source_indices, + unsigned num_distillation, + std::uint64_t rand_xor_mask, + const IndexT* seeds, + std::uint32_t num_seeds, + const std::uint32_t* seed_cross_graph, + const IndexT* seed_lookup_keys, + const std::uint32_t* seed_lookup_values, + std::uint32_t seed_lookup_capacity, + IndexT* visited_hashmap, + std::uint32_t max_candidates, + std::uint32_t max_itopk, + std::uint32_t internal_topk, + std::uint32_t search_width, + std::uint32_t min_iterations, + std::uint32_t max_iterations, + std::uint32_t* num_executed_iterations, + std::uint32_t hash_bitlen, + std::uint32_t small_hash_bitlen, + std::uint32_t small_hash_reset_interval, + std::uint32_t query_id_offset, + const cuvs::neighbors::cagra::detail::dataset_descriptor_base_t* + dataset_desc, + IndexT graph_size, + float sync_window_scale, + std::uint32_t sync_drop_threshold, + cuvs::neighbors::cagra::detail::cagra_sample_filter filter_payload) +{ + __shared__ graph_state state; + auto const query_id = blockIdx.y; + auto const queue_id = query_id % num_queues; + tiered_graph_provider provider{inner_graph, + inner_row_length, + node_per_cacheline, + n_bits, + subgraph_offsets, + subgraph_ids, + seeds, + seed_cross_graph, + seed_lookup_keys, + seed_lookup_values, + seed_lookup_capacity, + num_seeds, + queues[queue_id], + graph_degree, + cross_graph_degree, + graph_size, + sync_window_scale, + sync_drop_threshold, + &state}; + cuvs::neighbors::cagra::detail::single_cta_search::search_core( + result_indices_ptr, + result_distances_ptr, + top_k, + queries_ptr, + provider, + graph_degree, + source_indices, + num_distillation, + rand_xor_mask, + seeds, + num_seeds, + visited_hashmap, + max_candidates, + max_itopk, + internal_topk, + search_width, + min_iterations, + max_iterations, + num_executed_iterations, + hash_bitlen, + small_hash_bitlen, + small_hash_reset_interval, + query_id, + query_id_offset, + dataset_desc, + filter_payload, + graph_size, + query_stride); +} + +} // namespace cuvs::neighbors::cagra::experimental::flowann::detail diff --git a/cpp/src/neighbors/detail/flowann/search_single_cta_kernel.cu.in b/cpp/src/neighbors/detail/flowann/search_single_cta_kernel.cu.in new file mode 100644 index 0000000000..7e0fea9789 --- /dev/null +++ b/cpp/src/neighbors/detail/flowann/search_single_cta_kernel.cu.in @@ -0,0 +1,127 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * FlowANN is described in "Disentangling Graph Dependencies for Efficient + * Billion-Scale GPU Vector Search" (OSDI 2026). + */ + +#include +#include + +#include + +namespace { + +constexpr bool k_topk_by_bitonic_sort = @topk_by_bitonic_sort@; +constexpr bool k_bitonic_sort_and_merge_multi_warps = @bitonic_sort_and_merge_multi_warps@; + +using data_t = @data_type@; +using index_t = @index_type@; +using distance_t = @distance_type@; +using source_index_t = @source_index_type@; +using dataset_desc_base = + cuvs::neighbors::cagra::detail::dataset_descriptor_base_t; +using filter_payload_t = cuvs::neighbors::cagra::detail::cagra_sample_filter; +using queue_view = cuvs::neighbors::cagra::experimental::flowann::detail::queue_view; + +} // namespace + +namespace cuvs::neighbors::cagra::experimental::flowann::detail { + +extern "C" __global__ __launch_bounds__(1024, 1) void flowann_search_single_cta( + uintptr_t result_indices, + distance_t* distances, + std::uint32_t top_k, + const data_t* queries, + std::uint32_t query_stride, + const std::uint8_t* inner_graph, + std::uint32_t inner_row_length, + std::uint32_t node_per_cacheline, + std::uint32_t n_bits, + const std::uint32_t* subgraph_offsets, + const std::uint32_t* subgraph_ids, + queue_view* queues, + std::uint32_t num_queues, + std::uint32_t graph_degree, + std::uint32_t cross_graph_degree, + const source_index_t* source_indices, + unsigned num_distillation, + std::uint64_t rand_xor_mask, + const index_t* seeds, + std::uint32_t num_seeds, + const std::uint32_t* seed_cross_graph, + const index_t* seed_lookup_keys, + const std::uint32_t* seed_lookup_values, + std::uint32_t seed_lookup_capacity, + index_t* visited_hashmap, + std::uint32_t max_candidates, + std::uint32_t max_itopk, + std::uint32_t internal_topk, + std::uint32_t search_width, + std::uint32_t min_iterations, + std::uint32_t max_iterations, + std::uint32_t* num_executed_iterations, + std::uint32_t hash_bitlen, + std::uint32_t small_hash_bitlen, + std::uint32_t small_hash_reset_interval, + std::uint32_t query_id_offset, + const dataset_desc_base* dataset_desc, + index_t graph_size, + float sync_window_scale, + std::uint32_t sync_drop_threshold, + filter_payload_t filter_payload) +{ + search_kernel_jit(result_indices, + distances, + top_k, + queries, + query_stride, + inner_graph, + inner_row_length, + node_per_cacheline, + n_bits, + subgraph_offsets, + subgraph_ids, + queues, + num_queues, + graph_degree, + cross_graph_degree, + source_indices, + num_distillation, + rand_xor_mask, + seeds, + num_seeds, + seed_cross_graph, + seed_lookup_keys, + seed_lookup_values, + seed_lookup_capacity, + visited_hashmap, + max_candidates, + max_itopk, + internal_topk, + search_width, + min_iterations, + max_iterations, + num_executed_iterations, + hash_bitlen, + small_hash_bitlen, + small_hash_reset_interval, + query_id_offset, + dataset_desc, + graph_size, + sync_window_scale, + sync_drop_threshold, + filter_payload); +} + +static_assert( + std::is_same_v>); + +} // namespace cuvs::neighbors::cagra::experimental::flowann::detail diff --git a/cpp/src/neighbors/detail/flowann/search_single_cta_matrix.json b/cpp/src/neighbors/detail/flowann/search_single_cta_matrix.json new file mode 100644 index 0000000000..ba1334f11d --- /dev/null +++ b/cpp/src/neighbors/detail/flowann/search_single_cta_matrix.json @@ -0,0 +1,21 @@ +[ + { + "_data": [ + {"data_type": "float", "data_abbrev": "f"}, + {"data_type": "__half", "data_abbrev": "h"}, + {"data_type": "uint8_t", "data_abbrev": "u8"}, + {"data_type": "int8_t", "data_abbrev": "i8"} + ], + "_source_index": [{"source_index_type": "uint32_t", "source_index_abbrev": "u32"}], + "_index": [{"index_type": "uint32_t", "index_abbrev": "u32"}], + "_distance": [{"distance_type": "float", "distance_abbrev": "f"}], + "_topk_by_bitonic": [ + {"topk_by_bitonic_sort": "true", "topk_by_bitonic_sort_str": "topk_by_bitonic_sort"}, + {"topk_by_bitonic_sort": "false", "topk_by_bitonic_sort_str": "no_topk_by_bitonic_sort"} + ], + "_bitonic_sort_and_merge_multi_warps": [ + {"bitonic_sort_and_merge_multi_warps": "true", "bitonic_sort_and_merge_multi_warps_str": "bitonic_sort_and_merge_multi_warps"}, + {"bitonic_sort_and_merge_multi_warps": "false", "bitonic_sort_and_merge_multi_warps_str": "no_bitonic_sort_and_merge_multi_warps"} + ] + } +] diff --git a/cpp/src/neighbors/flowann.cu b/cpp/src/neighbors/flowann.cu new file mode 100644 index 0000000000..150a062da0 --- /dev/null +++ b/cpp/src/neighbors/flowann.cu @@ -0,0 +1,166 @@ +/* + * 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 = cuvs::neighbors::dense_row_major_view(dataset_); + if (!dataset_norms_.has_value() || dataset_norms_->extent(0) != dataset.extent(0)) { + dataset_norms_.reset(); + dataset_norms_.emplace(raft::make_device_vector(res, dataset.extent(0))); + } + + constexpr float scale = cuvs::spatial::knn::detail::utils::config::kDivisor / + cuvs::spatial::knn::detail::utils::config::kDivisor; + auto scaled_sq_op = + raft::compose_op(raft::sq_op{}, raft::div_const_op{scale}, raft::cast_op()); + raft::linalg::reduce(res, + dataset, + dataset_norms_->view(), + 0.0f, + false, + scaled_sq_op, + raft::add_op(), + raft::sqrt_op{}); + } +} + +namespace detail { + +template +void search_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) \ + CUVS_FLOWANN_DEFINE_SEARCH(T, vpq_f16_index, std::uint32_t) \ + CUVS_FLOWANN_DEFINE_SEARCH(T, vpq_f16_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..2960afdcf8 --- /dev/null +++ b/cpp/src/neighbors/flowann_build.cu @@ -0,0 +1,2673 @@ +/* + * 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< + typename decltype(std::declval().data_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.data_view(); + auto const dataset_stride = dataset.data_view().stride(0); + 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< + typename decltype(std::declval().data_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.data_view(); + auto const dataset_stride = dataset.data_view().stride(0); + 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.data_view().data_handle()[ids[i] * dataset.data_view().stride(0) + 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.data_view().data_handle(), + device_ids.data_handle(), + 0, + ids.size(), + dim, + dataset.data_view().stride(0)); + 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.data_view() + .data_handle()[std::uint64_t{source} * dataset.data_view().stride(0) + 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.data_view().data_handle(), + device_ids->data_handle(), + 0, + size, + dim, + dataset.data_view().stride(0)); + 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< + typename decltype(std::declval().data_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.data_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.data_view().stride(0); + 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.data_view().stride(0)); + 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.data_view().stride(0); + 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.data_view().stride(0)); + 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.data_view(); + auto reordered = + cuvs::neighbors::detail::make_dense_array>( + res, source.n_rows(), source.dim()); + gather_device_rows(res, + reordered.data_handle(), + source_data.data_handle(), + source_data.extent(0), + static_cast(source_data.stride(0)), + new_to_old); + return std::make_unique>(std::move(reordered)); +} + +auto reorder_vpq_dataset(raft::resources const& res, + cuvs::neighbors::device_vpq_dataset_view const& source, + std::vector const& new_to_old) + -> std::unique_ptr> +{ + auto const source_data = source; + auto vq_code_book = raft::make_device_matrix( + res, + source_data.manifest_view().vq_code_book.extent(0), + source_data.manifest_view().vq_code_book.extent(1)); + auto pq_code_book = raft::make_device_matrix( + res, + source_data.manifest_view().pq_code_book.extent(0), + source_data.manifest_view().pq_code_book.extent(1)); + auto encoded = raft::make_device_matrix( + res, source_data.data_view().extent(0), source_data.data_view().extent(1)); + auto const stream = raft::resource::get_cuda_stream(res); + raft::copy(vq_code_book.data_handle(), + source_data.manifest_view().vq_code_book.data_handle(), + source_data.manifest_view().vq_code_book.size(), + stream); + raft::copy(pq_code_book.data_handle(), + source_data.manifest_view().pq_code_book.data_handle(), + source_data.manifest_view().pq_code_book.size(), + stream); + gather_device_rows(res, + encoded.data_handle(), + source_data.data_view().data_handle(), + source_data.data_view().extent(0), + static_cast(source_data.data_view().extent(1)), + new_to_old); + return std::make_unique>( + std::move(encoded), + cuvs::neighbors::device_vpq_dataset::manifest_type{std::move(vq_code_book), + std::move(pq_code_book)}); +} + +auto reorder_vpq_dataset_consuming( + raft::resources const& res, + std::unique_ptr> source, + std::vector const& new_to_old) + -> std::unique_ptr> +{ + RAFT_EXPECTS(source != nullptr, "FlowANN VPQ conversion requires a dataset owner"); + auto const n_rows = source->data_view().extent(0); + auto const row_length = source->data_view().extent(1); + RAFT_EXPECTS(n_rows == static_cast(new_to_old.size()), + "FlowANN VPQ dataset and graph row counts differ"); + + auto host_source = raft::make_host_matrix(n_rows, row_length); + auto const stream = raft::resource::get_cuda_stream(res); + raft::copy(host_source.data_handle(), + source->data_view().data_handle(), + source->data_view().size(), + stream); + raft::resource::sync_stream(res); + + // The dataset API exposes codebook views rather than detachable owners. Copy only + // the bounded codebooks, then release the old encoded rows before allocating new rows. + auto manifest = source->const_manifest_view(); + auto vq_code_book = raft::make_device_matrix( + res, manifest.vq_code_book.extent(0), manifest.vq_code_book.extent(1)); + auto pq_code_book = raft::make_device_matrix( + res, manifest.pq_code_book.extent(0), manifest.pq_code_book.extent(1)); + raft::copy(vq_code_book.data_handle(), + manifest.vq_code_book.data_handle(), + manifest.vq_code_book.size(), + stream); + raft::copy(pq_code_book.data_handle(), + manifest.pq_code_book.data_handle(), + manifest.pq_code_book.size(), + stream); + raft::resource::sync_stream(res); + source.reset(); + + auto host_reordered = raft::make_host_matrix(n_rows, row_length); +#ifdef _OPENMP +#pragma omp parallel for schedule(static) +#endif + for (int64_t new_row = 0; new_row < n_rows; ++new_row) { + std::copy_n( + host_source.data_handle() + + static_cast(new_to_old[static_cast(new_row)]) * row_length, + row_length, + host_reordered.data_handle() + static_cast(new_row) * row_length); + } + host_source = raft::make_host_matrix(0, 0); + + auto encoded = raft::make_device_matrix(res, n_rows, row_length); + raft::copy(encoded.data_handle(), host_reordered.data_handle(), host_reordered.size(), stream); + raft::resource::sync_stream(res); + return std::make_unique>( + std::move(encoded), + cuvs::neighbors::device_vpq_dataset::manifest_type{std::move(vq_code_book), + std::move(pq_code_book)}); +} + +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->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.data_view(); + auto owned_rows = + cuvs::neighbors::detail::make_dense_array>( + res, dataset.n_rows(), dataset.dim()); + raft::copy(owned_rows.data_handle(), + source_rows.data_handle(), + owned_rows.mapping().required_span_size(), + raft::resource::get_cuda_stream(res)); + auto dataset_owner = + std::make_unique>(std::move(owned_rows)); + flowann::device_padded_index index(res, + params.cagra_params.metric, + dataset_owner->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)}; +} + +template +auto build_vpq_integrated(raft::resources const& res, + index_params const& params, + cuvs::neighbors::vpq_params const& vpq_params, + cuvs::neighbors::host_standard_dataset_view const& dataset) + -> vpq_f16_index_bundle +{ + detail::reject_ace(params); + RAFT_EXPECTS(params.cagra_params.metric == cuvs::distance::DistanceType::L2Expanded, + "FlowANN VPQ-F16 integrated build requires L2Expanded"); + + 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 unordered_dataset = std::make_unique>( + cuvs::preprocessing::quantize::pq::make_vpq_dataset(res, vpq_params, dataset)); + auto source_indices = conversion.new_to_old; + auto dataset_owner = + detail::reorder_vpq_dataset(res, unordered_dataset->view(), conversion.new_to_old); + unordered_dataset.reset(); + flowann::vpq_f16_index index(res, + params.cagra_params.metric, + dataset_owner->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, vpq=%.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 dataset_owner = std::make_unique>( + cuvs::preprocessing::quantize::pq::make_vpq_dataset(res, vpq_params, dataset)); + flowann::vpq_f16_index index(res, + params.cagra_params.metric, + dataset_owner->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, vpq=%.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); \ + } \ + auto build(raft::resources const& res, \ + index_params const& params, \ + cuvs::neighbors::vpq_params const& vpq_params, \ + cuvs::neighbors::host_standard_dataset_view const& dataset) \ + -> vpq_f16_index_bundle \ + { \ + return build_vpq_integrated(res, params, vpq_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->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)}; +} + +template +auto build_vpq_from_cagra( + raft::resources const& res, + build_params const& params, + cuvs::neighbors::cagra::device_pq_index const& cagra_index, + raft::host_vector_view partition_labels, + std::optional> seeds) + -> vpq_f16_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_vpq_dataset(res, cagra_index.dataset(), conversion.new_to_old); + flowann::vpq_f16_index index( + res, cagra_index.metric(), 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)}; +} + +template +auto build_vpq_from_cagra_consuming( + raft::resources const& res, + build_params const& params, + cuvs::neighbors::cagra::device_pq_index&& cagra_index, + std::unique_ptr>&& vpq_dataset, + raft::host_vector_view partition_labels, + std::optional> seeds) + -> vpq_f16_index_bundle +{ + RAFT_EXPECTS(vpq_dataset != nullptr, "FlowANN VPQ conversion requires a dataset owner"); + RAFT_EXPECTS(vpq_dataset->n_rows() > 0, "FlowANN VPQ dataset must not be empty"); + RAFT_EXPECTS(cagra_index.dataset().n_rows() == vpq_dataset->n_rows(), + "FlowANN CAGRA index does not view the supplied VPQ dataset"); + RAFT_EXPECTS( + cagra_index.dataset().data_view().data_handle() == vpq_dataset->data_view().data_handle() && + cagra_index.dataset().manifest_view().vq_code_book.data_handle() == + vpq_dataset->const_manifest_view().vq_code_book.data_handle() && + cagra_index.dataset().manifest_view().pq_code_book.data_handle() == + vpq_dataset->const_manifest_view().pq_code_book.data_handle(), + "FlowANN CAGRA index does not view the supplied VPQ dataset owner"); + + auto const metric = cagra_index.metric(); + 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); + cagra_index = cuvs::neighbors::cagra::device_pq_index(res, metric); + + auto dataset = + detail::reorder_vpq_dataset_consuming(res, std::move(vpq_dataset), conversion.new_to_old); + flowann::vpq_f16_index index( + res, metric, 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); \ + } \ + auto build_from_cagra( \ + raft::resources const& res, \ + build_params const& params, \ + cuvs::neighbors::cagra::device_pq_index const& cagra_index, \ + raft::host_vector_view partition_labels, \ + std::optional> seeds) \ + -> vpq_f16_index_bundle \ + { \ + return build_vpq_from_cagra(res, params, cagra_index, partition_labels, seeds); \ + } \ + auto build_from_cagra( \ + raft::resources const& res, \ + build_params const& params, \ + cuvs::neighbors::cagra::device_pq_index&& cagra_index, \ + std::unique_ptr>&& vpq_dataset, \ + raft::host_vector_view partition_labels, \ + std::optional> seeds) \ + -> vpq_f16_index_bundle \ + { \ + return build_vpq_from_cagra_consuming( \ + res, params, std::move(cagra_index), std::move(vpq_dataset), 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..f9fc171601 --- /dev/null +++ b/cpp/src/neighbors/flowann_queue.cpp @@ -0,0 +1,999 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * FlowANN is described in "Disentangling Graph Dependencies for Efficient + * Billion-Scale GPU Vector Search" (OSDI 2026). + */ + +#include + +#include "detail/flowann/queue.cuh" + +#include +#include +#include + +#ifdef CUVS_FLOWANN_USE_GDRCOPY +#include +#include +#endif + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#if defined(__linux__) +#include +#include +#include +#endif + +namespace cuvs::neighbors::cagra::experimental::flowann { +namespace { + +constexpr auto align_up(std::size_t value, std::size_t alignment) -> std::size_t +{ + return (value + alignment - 1) & ~(alignment - 1); +} + +struct queue_layout { + std::size_t state; + std::size_t rows; + std::size_t command_words; + std::size_t result_sequence; + std::size_t released_sequence; + std::size_t total_size; + + explicit queue_layout(std::uint32_t row_length) + { + constexpr std::size_t alignment = 128; + state = 0; + rows = align_up(sizeof(detail::queue_state), alignment); + + RAFT_EXPECTS(row_length <= std::numeric_limits::max() / detail::queue_capacity / + sizeof(std::uint32_t), + "FlowANN queue row storage size overflows"); + auto const row_bytes = + static_cast(detail::queue_capacity) * row_length * sizeof(std::uint32_t); + command_words = align_up(rows + row_bytes, alignment); + result_sequence = + align_up(command_words + detail::queue_capacity * sizeof(unsigned long long), alignment); + released_sequence = + align_up(result_sequence + detail::queue_capacity * sizeof(unsigned long long), alignment); + total_size = + align_up(released_sequence + detail::queue_capacity * sizeof(unsigned long long), alignment); + } +}; + +template +auto add_bytes(void* base, std::size_t offset) -> T* +{ + return reinterpret_cast(static_cast(base) + offset); +} + +class cuda_allocation { + public: + cuda_allocation() = default; + cuda_allocation(cuda_allocation const&) = delete; + auto operator=(cuda_allocation const&) -> cuda_allocation& = delete; + + ~cuda_allocation() + { + if (allocation_ != nullptr) { RAFT_CUDA_TRY_NO_THROW(cudaFree(allocation_)); } + } + + void allocate(std::size_t size, std::size_t alignment = 1) + { + RAFT_EXPECTS(allocation_ == nullptr, "FlowANN queue allocation is already initialized"); + RAFT_EXPECTS(alignment > 0 && (alignment & (alignment - 1)) == 0, + "FlowANN queue allocation alignment must be a power of two"); + RAFT_EXPECTS(size <= std::numeric_limits::max() - (alignment - 1), + "FlowANN queue allocation size overflows"); + RAFT_CUDA_TRY(cudaMalloc(&allocation_, size + alignment - 1)); + auto const address = reinterpret_cast(allocation_); + ptr_ = reinterpret_cast(align_up(address, alignment)); + } + + [[nodiscard]] auto get() const noexcept -> void* { return ptr_; } + + private: + void* allocation_{}; + void* ptr_{}; +}; + +class cuda_stream { + public: + cuda_stream() = default; + cuda_stream(cuda_stream const&) = delete; + auto operator=(cuda_stream const&) -> cuda_stream& = delete; + + ~cuda_stream() + { + if (stream_ != nullptr) { RAFT_CUDA_TRY_NO_THROW(cudaStreamDestroy(stream_)); } + } + + void create() + { + RAFT_EXPECTS(stream_ == nullptr, "FlowANN queue stream is already initialized"); + RAFT_CUDA_TRY(cudaStreamCreateWithFlags(&stream_, cudaStreamNonBlocking)); + } + + [[nodiscard]] auto get() const noexcept -> cudaStream_t { return stream_; } + + private: + cudaStream_t stream_{}; +}; + +// A host-written flag remains publishable even if CUDA copies fail in a poller. +class mapped_abort_flag { + public: + mapped_abort_flag() + { + RAFT_CUDA_TRY(cudaHostAlloc(&host_, sizeof(std::uint32_t), cudaHostAllocMapped)); + try { + *host_ = 0; + RAFT_CUDA_TRY(cudaHostGetDevicePointer(&device_, host_, 0)); + } catch (...) { + RAFT_CUDA_TRY_NO_THROW(cudaFreeHost(host_)); + throw; + } + } + mapped_abort_flag(mapped_abort_flag const&) = delete; + auto operator=(mapped_abort_flag const&) -> mapped_abort_flag& = delete; + ~mapped_abort_flag() { RAFT_CUDA_TRY_NO_THROW(cudaFreeHost(host_)); } + void abort() noexcept + { + std::atomic_ref(*host_).store(1, std::memory_order_release); + } + auto device_ptr() const noexcept -> std::uint32_t const* { return device_; } + + private: + std::uint32_t* host_{}; + std::uint32_t* device_{}; +}; + +#ifndef CUVS_FLOWANN_USE_GDRCOPY +class cuda_host_buffer { + public: + cuda_host_buffer() = default; + cuda_host_buffer(cuda_host_buffer const&) = delete; + auto operator=(cuda_host_buffer const&) -> cuda_host_buffer& = delete; + + ~cuda_host_buffer() + { + if (ptr_ != nullptr) { RAFT_CUDA_TRY_NO_THROW(cudaFreeHost(ptr_)); } + } + + void allocate(std::size_t size) + { + RAFT_EXPECTS(ptr_ == nullptr, "FlowANN host copy buffer is already initialized"); + RAFT_CUDA_TRY(cudaMallocHost(&ptr_, size)); + size_ = size; + } + + [[nodiscard]] auto get(std::size_t size) const -> void* + { + RAFT_EXPECTS(size <= size_, "FlowANN transfer exceeds the host copy buffer"); + return ptr_; + } + + private: + void* ptr_{}; + std::size_t size_{}; +}; +#endif + +#ifdef CUVS_FLOWANN_USE_GDRCOPY +class gdrcopy_mapping { + public: + gdrcopy_mapping(void* device_ptr, std::size_t size) : size_(size) + { + try { + unsigned int sync_memops = 1; + RAFT_EXPECTS(cuPointerSetAttribute(&sync_memops, + CU_POINTER_ATTRIBUTE_SYNC_MEMOPS, + reinterpret_cast(device_ptr)) == CUDA_SUCCESS, + "Failed to enable synchronous memory operations for a FlowANN GDRCopy queue"); + + handle_ = gdr_open(); + RAFT_EXPECTS(handle_ != nullptr, + "Failed to open GDRCopy; ensure the gdrdrv kernel module is loaded"); + RAFT_EXPECTS( + gdr_pin_buffer( + handle_, reinterpret_cast(device_ptr), size_, 0, 0, &memory_handle_) == 0, + "Failed to pin a FlowANN queue with GDRCopy"); + pinned_ = true; + RAFT_EXPECTS(gdr_map(handle_, memory_handle_, &mapped_base_, size_) == 0, + "Failed to map a FlowANN queue with GDRCopy"); + mapped_ = true; + + gdr_info_t info{}; + RAFT_EXPECTS(gdr_get_info(handle_, memory_handle_, &info) == 0, + "Failed to query a FlowANN GDRCopy mapping"); + auto const device_address = reinterpret_cast(device_ptr); + RAFT_EXPECTS(device_address >= info.va, + "GDRCopy returned an invalid FlowANN mapping address"); + auto const offset = device_address - info.va; + RAFT_EXPECTS(offset <= size_, "GDRCopy returned an invalid FlowANN mapping offset"); + mapped_device_ptr_ = static_cast(mapped_base_) + offset; + } catch (...) { + reset(); + throw; + } + } + + gdrcopy_mapping(gdrcopy_mapping const&) = delete; + auto operator=(gdrcopy_mapping const&) -> gdrcopy_mapping& = delete; + + ~gdrcopy_mapping() { reset(); } + + void read(void* dst, std::size_t offset, std::size_t size) + { + RAFT_EXPECTS(offset <= size_ && size <= size_ - offset, + "FlowANN GDRCopy read exceeds its queue mapping"); + RAFT_EXPECTS(gdr_copy_from_mapping(memory_handle_, dst, mapped_device_ptr_ + offset, size) == 0, + "Failed to read a FlowANN GDRCopy queue"); + } + + void write(std::size_t offset, void const* src, std::size_t size) + { + RAFT_EXPECTS(offset <= size_ && size <= size_ - offset, + "FlowANN GDRCopy write exceeds its queue mapping"); + RAFT_EXPECTS(gdr_copy_to_mapping(memory_handle_, mapped_device_ptr_ + offset, src, size) == 0, + "Failed to write a FlowANN GDRCopy queue"); + } + + private: + void reset() noexcept + { + if (mapped_) { + static_cast(gdr_unmap(handle_, memory_handle_, mapped_base_, size_)); + mapped_ = false; + } + if (pinned_) { + static_cast(gdr_unpin_buffer(handle_, memory_handle_)); + pinned_ = false; + } + if (handle_ != nullptr) { + static_cast(gdr_close(handle_)); + handle_ = nullptr; + } + } + + gdr_t handle_{}; + gdr_mh_t memory_handle_{}; + void* mapped_base_{}; + std::byte* mapped_device_ptr_{}; + std::size_t size_{}; + bool pinned_{}; + bool mapped_{}; +}; +#endif + +#if defined(__linux__) +auto read_topology_value(int cpu, char const* name) -> int +{ + std::ifstream input("/sys/devices/system/cpu/cpu" + std::to_string(cpu) + "/topology/" + name); + int value = -1; + input >> value; + return input ? value : -1; +} + +auto physical_cpus() -> std::vector +{ + cpu_set_t allowed; + CPU_ZERO(&allowed); + if (sched_getaffinity(0, sizeof(allowed), &allowed) != 0) { return {}; } + + std::vector> cores; + std::vector cpus; + for (int cpu = 0; cpu < CPU_SETSIZE; ++cpu) { + if (!CPU_ISSET(cpu, &allowed)) { continue; } + auto package = read_topology_value(cpu, "physical_package_id"); + auto core = read_topology_value(cpu, "core_id"); + if (package < 0 || core < 0) { + package = 0; + core = cpu; + } + auto const identity = std::tuple{package, core}; + if (std::find(cores.begin(), cores.end(), identity) == cores.end()) { + cores.push_back(identity); + cpus.push_back(cpu); + } + } + return cpus; +} + +void bind_to_physical_cpu(std::uint32_t queue_id) +{ + static auto const cpus = physical_cpus(); + if (cpus.empty()) { return; } + auto const cpu = cpus[queue_id % cpus.size()]; + cpu_set_t selected; + CPU_ZERO(&selected); + CPU_SET(cpu, &selected); + auto const error = pthread_setaffinity_np(pthread_self(), sizeof(selected), &selected); + if (error != 0) { + RAFT_LOG_WARN( + "FlowANN queue %u could not bind its poller to CPU %d (error=%d)", queue_id, cpu, error); + } + if (queue_id >= cpus.size()) { + RAFT_LOG_WARN("FlowANN queue %u reuses physical CPU %d because only %zu are available", + queue_id, + cpu, + cpus.size()); + } +} +#else +void bind_to_physical_cpu(std::uint32_t) {} +#endif + +class host_queue { + public: + static constexpr std::uint32_t max_batch_size = 128; + + host_queue(std::uint32_t const* cross_graph, + std::uint64_t row_count, + std::uint32_t row_length, + bool collect_statistics) + : cross_graph_(cross_graph), + row_count_(row_count), + row_length_(row_length), + collect_statistics_(collect_statistics), + layout_(row_length), + empty_commands_(max_batch_size, detail::empty_command), + command_staging_(max_batch_size), + result_sequence_staging_(max_batch_size), + row_staging_(static_cast(max_batch_size) * row_length) + { +#ifdef CUVS_FLOWANN_USE_GDRCOPY + // gdrdrv rejects mappings whose device address is not aligned to its 64 KiB GPU page. + // cudaMalloc can return a suballocation after the graph build, so align within an owning + // allocation instead of relying on cudaMalloc's current address. + constexpr std::size_t gpu_page_size = 64u * 1024u; + RAFT_EXPECTS( + layout_.total_size <= std::numeric_limits::max() - (gpu_page_size - 1), + "FlowANN GDRCopy mapping size overflows"); + mapped_size_ = align_up(layout_.total_size, gpu_page_size); + storage_.allocate(mapped_size_, gpu_page_size); +#else + storage_.allocate(layout_.total_size); +#endif + RAFT_CUDA_TRY(cudaMemset(storage_.get(), 0, layout_.total_size)); + + detail::queue_state state{}; + state.capacity = detail::queue_capacity; + state.row_length = row_length_; + RAFT_CUDA_TRY(cudaMemcpy(add_bytes(storage_.get(), layout_.state), + &state, + sizeof(state), + cudaMemcpyHostToDevice)); + + view_.state = add_bytes(storage_.get(), layout_.state); + view_.rows = add_bytes(storage_.get(), layout_.rows); + view_.command_words = add_bytes(storage_.get(), layout_.command_words); + view_.result_sequence = add_bytes(storage_.get(), layout_.result_sequence); + view_.released_sequence = + add_bytes(storage_.get(), layout_.released_sequence); + view_.abort_flag = abort_flag_.device_ptr(); + +#ifdef CUVS_FLOWANN_USE_GDRCOPY + gdrcopy_.emplace(storage_.get(), mapped_size_); +#else + copy_stream_.create(); + // Pageable copies may synchronize inside the driver and contend with lazy + // kernel initialization. Allocate once, before starting any pollers, and + // use pinned memory for every transfer on the nonblocking copy stream. + copy_staging_.allocate(std::max(row_staging_.size() * sizeof(std::uint32_t), + command_staging_.size() * sizeof(unsigned long long))); +#endif + } + + host_queue(host_queue const&) = delete; + auto operator=(host_queue const&) -> host_queue& = delete; + + [[nodiscard]] auto view() const noexcept -> detail::queue_view { return view_; } + void abort() noexcept { abort_flag_.abort(); } + auto consume_invalid_requests() noexcept -> std::uint64_t + { + return invalid_requests_.exchange(0, std::memory_order_relaxed); + } + + [[nodiscard]] auto poll_batch() -> bool + { + if (collect_statistics_) { polls_.fetch_add(1, std::memory_order_relaxed); } + auto const first_slot = static_cast(cpu_tail_) & detail::queue_mask; + // Sparse traffic needs only a short prefix. Read the rest only when the + // whole prefix is ready, retaining full service batches under load. + constexpr std::uint32_t prefix_size = 8; + read_ring(command_staging_.data(), layout_.command_words, first_slot, prefix_size); + if (!detail::command_ready(command_staging_[0])) { + if (collect_statistics_) { empty_polls_.fetch_add(1, std::memory_order_relaxed); } + return false; + } + + std::uint32_t count = 1; + while (count < prefix_size && detail::command_ready(command_staging_[count])) { + ++count; + } + if (count == prefix_size) { + auto const next_slot = (first_slot + prefix_size) & detail::queue_mask; + read_ring(command_staging_.data() + prefix_size, + layout_.command_words, + next_slot, + max_batch_size - prefix_size); + while (count < max_batch_size && detail::command_ready(command_staging_[count])) { + ++count; + } + } + + for (std::uint32_t i = 0; i < count; ++i) { + auto const row_id = static_cast(command_staging_[i]); + auto* destination = row_staging_.data() + static_cast(i) * row_length_; + if (row_id < row_count_) { + std::copy_n( + cross_graph_ + static_cast(row_id) * row_length_, row_length_, destination); + } else { + // Complete malformed requests with an empty row so GPU waiters can finish. + std::fill_n(destination, row_length_, 0u); + invalid_requests_.fetch_add(1, std::memory_order_relaxed); + } + result_sequence_staging_[i] = cpu_tail_ + i + 1; + } + write_ring(row_staging_.data(), layout_.rows, first_slot, count, row_length_); + // Clearing the request is ordered before result publication. A producer + // cannot reuse this slot until it acquires and releases that result. + write_ring(empty_commands_.data(), layout_.command_words, first_slot, count); + std::atomic_thread_fence(std::memory_order_release); + write_ring(result_sequence_staging_.data(), layout_.result_sequence, first_slot, count); + + cpu_tail_ += count; + if (collect_statistics_) { processed_commands_.fetch_add(count, std::memory_order_relaxed); } + return true; + } + + void publish_progress() + { + // cpu_tail is an observational snapshot, not a GPU synchronization input: + // consumers use result_sequence and producers reclaim capacity via gpu_tail. + // Publish before acknowledging a pause instead of writing on every batch. + write(layout_.state + offsetof(detail::queue_state, cpu_tail), &cpu_tail_, sizeof(cpu_tail_)); + } + + [[nodiscard]] auto statistics() const noexcept -> queue_statistics + { + return {polls_.load(std::memory_order_relaxed), + empty_polls_.load(std::memory_order_relaxed), + processed_commands_.load(std::memory_order_relaxed)}; + } + + private: + template + void read_ring(T* destination, + std::size_t base_offset, + std::uint32_t first_slot, + std::uint32_t count) + { + auto const first_count = std::min(count, detail::queue_capacity - first_slot); + read(destination, + base_offset + static_cast(first_slot) * sizeof(T), + static_cast(first_count) * sizeof(T)); + if (first_count < count) { + read(destination + first_count, + base_offset, + static_cast(count - first_count) * sizeof(T)); + } + } + + template + void write_ring(T const* source, + std::size_t base_offset, + std::uint32_t first_slot, + std::uint32_t count, + std::uint32_t elements_per_slot = 1) + { + auto const first_count = std::min(count, detail::queue_capacity - first_slot); + auto const slot_bytes = static_cast(elements_per_slot) * sizeof(T); + write(base_offset + static_cast(first_slot) * slot_bytes, + source, + static_cast(first_count) * slot_bytes); + if (first_count < count) { + write(base_offset, + source + static_cast(first_count) * elements_per_slot, + static_cast(count - first_count) * slot_bytes); + } + } + + void read(void* dst, std::size_t offset, std::size_t size) + { +#ifdef CUVS_FLOWANN_USE_GDRCOPY + gdrcopy_->read(dst, offset, size); +#else + auto* staging = copy_staging_.get(size); + RAFT_CUDA_TRY(cudaMemcpyAsync(staging, + add_bytes(storage_.get(), offset), + size, + cudaMemcpyDeviceToHost, + copy_stream_.get())); + RAFT_CUDA_TRY(cudaStreamSynchronize(copy_stream_.get())); + std::memcpy(dst, staging, size); +#endif + } + + void write(std::size_t offset, void const* src, std::size_t size) + { +#ifdef CUVS_FLOWANN_USE_GDRCOPY + gdrcopy_->write(offset, src, size); +#else + auto* staging = copy_staging_.get(size); + std::memcpy(staging, src, size); + RAFT_CUDA_TRY(cudaMemcpyAsync(add_bytes(storage_.get(), offset), + staging, + size, + cudaMemcpyHostToDevice, + copy_stream_.get())); + RAFT_CUDA_TRY(cudaStreamSynchronize(copy_stream_.get())); +#endif + } + + std::uint32_t const* cross_graph_; + std::uint64_t row_count_; + std::uint32_t row_length_; + bool collect_statistics_; + queue_layout layout_; + mapped_abort_flag abort_flag_; + cuda_allocation storage_; + detail::queue_view view_{}; + unsigned long long cpu_tail_{}; + std::atomic polls_{}; + std::atomic empty_polls_{}; + std::atomic processed_commands_{}; + std::atomic invalid_requests_{}; + std::vector empty_commands_; + std::vector command_staging_; + std::vector result_sequence_staging_; + std::vector row_staging_; +#ifdef CUVS_FLOWANN_USE_GDRCOPY + std::size_t mapped_size_{}; + std::optional gdrcopy_; +#else + cuda_stream copy_stream_; + cuda_host_buffer copy_staging_; +#endif +}; + +} // namespace + +struct search_context::impl { + explicit impl(raft::host_matrix_view cross_graph, + queue_params params) + : params_(params), + row_count_(static_cast(cross_graph.extent(0))), + row_length_(static_cast(cross_graph.extent(1))) + { + RAFT_CUDA_TRY(cudaGetDevice(&device_id_)); + RAFT_EXPECTS(cross_graph.extent(0) > 0, "FlowANN cross graph must not be empty"); + RAFT_EXPECTS(cross_graph.extent(1) > 0, + "FlowANN cross graph must contain a neighbor-count column"); + RAFT_EXPECTS(params.num_queues > 0, "FlowANN must use at least one command queue"); + RAFT_EXPECTS(cross_graph.extent(1) <= std::numeric_limits::max(), + "FlowANN cross-graph row length exceeds uint32"); + + queues_.reserve(params.num_queues); + std::vector views; + views.reserve(params.num_queues); + for (std::uint32_t i = 0; i < params.num_queues; ++i) { + queues_.push_back( + std::make_unique(cross_graph.data_handle(), + static_cast(cross_graph.extent(0)), + static_cast(cross_graph.extent(1)), + params.collect_statistics)); + views.push_back(queues_.back()->view()); + } + + RAFT_CUDA_TRY(cudaMalloc(&device_views_, views.size() * sizeof(detail::queue_view))); + try { + RAFT_CUDA_TRY(cudaMemcpy(device_views_, + views.data(), + views.size() * sizeof(detail::queue_view), + cudaMemcpyHostToDevice)); + + pollers_.reserve(queues_.size()); + for (std::uint32_t i = 0; i < queues_.size(); ++i) { + pollers_.emplace_back([this, queue = queues_[i].get(), i](std::stop_token stop_token) { + poll(stop_token, *queue, i); + }); + } + + std::unique_lock lock(pause_mutex_); + pause_cv_.wait(lock, [this] { return paused_pollers_ + failed_pollers_ == pollers_.size(); }); + auto error = poller_error_; + lock.unlock(); + if (error != nullptr) { std::rethrow_exception(error); } + } catch (...) { + stop(); + throw; + } + } + + ~impl() { stop(); } + + void report_failure(std::exception_ptr error) noexcept + { + std::lock_guard lock(pause_mutex_); + for (auto const& queue : queues_) { + queue->abort(); + } + if (poller_error_ == nullptr) { poller_error_ = error; } + pause_cv_.notify_all(); + } + + void throw_if_failed() + { + if (poller_error_ != nullptr) { std::rethrow_exception(poller_error_); } + } + + void throw_if_invalid_requests() const + { + bool invalid = false; + for (auto const& queue : queues_) { + invalid |= queue->consume_invalid_requests() != 0; + } + RAFT_EXPECTS(!invalid, "FlowANN queue received an out-of-range cross-graph row"); + } + + [[nodiscard]] auto should_run() const noexcept -> bool + { + return manually_resumed_ || active_sessions_ > 0 || active_searches_ > 0; + } + + void update_pollers(std::unique_lock& lock) + { + while (true) { + throw_if_failed(); + if (should_run()) { + if (paused_.load(std::memory_order_acquire)) { + paused_.store(false, std::memory_order_release); + pause_cv_.notify_all(); + } + pause_cv_.wait(lock, [this] { + return !should_run() || paused_pollers_ == 0 || poller_error_ != nullptr; + }); + throw_if_failed(); + if (should_run() && paused_pollers_ == 0) { return; } + } else { + if (!paused_.load(std::memory_order_acquire)) { + paused_.store(true, std::memory_order_release); + pause_cv_.notify_all(); + } + pause_cv_.wait(lock, [this] { + return should_run() || paused_pollers_ + failed_pollers_ == pollers_.size() || + poller_error_ != nullptr; + }); + throw_if_failed(); + if (!should_run() && paused_pollers_ + failed_pollers_ == pollers_.size()) { return; } + } + } + } + + void stop() noexcept + { + for (auto& poller : pollers_) { + poller.request_stop(); + } + { + std::lock_guard lock(pause_mutex_); + paused_.store(false, std::memory_order_release); + } + pause_cv_.notify_all(); + pollers_.clear(); + if (device_views_ != nullptr) { + RAFT_CUDA_TRY_NO_THROW(cudaFree(device_views_)); + device_views_ = nullptr; + } + } + + void pause() + { + std::unique_lock lock(pause_mutex_); + RAFT_EXPECTS(active_sessions_ == 0, "Cannot pause while a FlowANN search session is active"); + manually_resumed_ = false; + ++pending_pause_requests_; + try { + pause_cv_.wait(lock, [this] { return active_searches_ == 0 || poller_error_ != nullptr; }); + throw_if_failed(); + update_pollers(lock); + } catch (...) { + --pending_pause_requests_; + pause_cv_.notify_all(); + throw; + } + --pending_pause_requests_; + pause_cv_.notify_all(); + throw_if_invalid_requests(); + } + + void resume() + { + std::unique_lock lock(pause_mutex_); + pause_cv_.wait(lock, + [this] { return pending_pause_requests_ == 0 || poller_error_ != nullptr; }); + throw_if_failed(); + throw_if_invalid_requests(); + manually_resumed_ = true; + update_pollers(lock); + } + + void acquire_search() + { + std::unique_lock lock(pause_mutex_); + pause_cv_.wait(lock, + [this] { return pending_pause_requests_ == 0 || poller_error_ != nullptr; }); + throw_if_failed(); + throw_if_invalid_requests(); + manually_resumed_ = false; + ++active_searches_; + try { + update_pollers(lock); + } catch (...) { + --active_searches_; + pause_cv_.notify_all(); + throw; + } + } + + void release_search() + { + std::unique_lock lock(pause_mutex_); + RAFT_EXPECTS(active_searches_ > 0, "Cannot release an inactive FlowANN search lease"); + --active_searches_; + if (active_searches_ == 0) { + manually_resumed_ = false; + pause_cv_.notify_all(); + update_pollers(lock); + } + throw_if_failed(); + throw_if_invalid_requests(); + } + + void begin_session() + { + std::unique_lock lock(pause_mutex_); + pause_cv_.wait(lock, + [this] { return pending_pause_requests_ == 0 || poller_error_ != nullptr; }); + throw_if_failed(); + throw_if_invalid_requests(); + ++active_sessions_; + try { + update_pollers(lock); + } catch (...) { + --active_sessions_; + throw; + } + } + + void end_session() + { + std::unique_lock lock(pause_mutex_); + RAFT_EXPECTS(active_sessions_ > 0, "Cannot end an inactive FlowANN search session"); + --active_sessions_; + update_pollers(lock); + } + + void poll(std::stop_token stop_token, host_queue& queue, std::uint32_t queue_id) noexcept + { + try { + RAFT_CUDA_TRY(cudaSetDevice(device_id_)); + if (params_.pin_pollers) { + static std::atomic next_poller{0}; + bind_to_physical_cpu(next_poller.fetch_add(1, std::memory_order_relaxed)); + } + while (!stop_token.stop_requested()) { + if (paused_.load(std::memory_order_acquire)) { + queue.publish_progress(); + std::unique_lock lock(pause_mutex_); + ++paused_pollers_; + pause_cv_.notify_all(); + pause_cv_.wait(lock, [this, &stop_token] { + return stop_token.stop_requested() || !paused_.load(std::memory_order_acquire); + }); + --paused_pollers_; + pause_cv_.notify_all(); + if (stop_token.stop_requested()) { break; } + } + + auto const processed = queue.poll_batch(); + if (!processed) { +#if defined(__x86_64__) || defined(__i386__) + for (std::uint32_t i = 0; i < params_.empty_pause; ++i) { + __asm__ __volatile__("pause"); + } +#else + std::this_thread::yield(); +#endif + } + } + } catch (...) { + report_failure(std::current_exception()); + std::lock_guard lock(pause_mutex_); + ++failed_pollers_; + pause_cv_.notify_all(); + } + } + + [[nodiscard]] auto statistics() const noexcept -> queue_statistics + { + queue_statistics total{}; + for (auto const& queue : queues_) { + auto const value = queue->statistics(); + total.polls += value.polls; + total.empty_polls += value.empty_polls; + total.processed_commands += value.processed_commands; + } + return total; + } + + queue_params params_; + int device_id_{}; + std::uint64_t row_count_; + std::uint32_t row_length_; + std::vector> queues_; + std::vector pollers_; + detail::queue_view* device_views_{}; + std::atomic paused_{true}; + mutable std::mutex pause_mutex_; + std::condition_variable pause_cv_; + std::size_t paused_pollers_{}; + std::size_t failed_pollers_{}; + std::size_t active_searches_{}; + std::size_t active_sessions_{}; + bool manually_resumed_{}; + std::size_t pending_pause_requests_{}; + std::exception_ptr poller_error_{}; +}; + +search_context::search_context( + raft::host_matrix_view cross_graph, + queue_params params) + : impl_(std::make_shared(cross_graph, params)) +{ +} + +search_context::~search_context() = default; + +search_context::search_context(search_context&&) noexcept = default; + +auto search_context::operator=(search_context&&) noexcept -> search_context& = default; + +search_session::search_session(search_context& context) : context_(context.impl_) +{ + RAFT_EXPECTS(context.impl_ != nullptr, "Cannot start a session on a moved-from search context"); + context_->begin_session(); +} + +search_session::~search_session() { reset(); } + +search_session::search_session(search_session&& other) noexcept + : context_(std::move(other.context_)) +{ +} + +auto search_session::operator=(search_session&& other) noexcept -> search_session& +{ + if (this != &other) { + reset(); + context_ = std::move(other.context_); + } + return *this; +} + +void search_session::reset() noexcept +{ + if (context_ == nullptr) { return; } + auto context = std::move(context_); + try { + context->end_session(); + } catch (...) { + } +} + +void search_context::pause() +{ + RAFT_EXPECTS(impl_ != nullptr, "Cannot pause a moved-from FlowANN search context"); + impl_->pause(); +} + +void search_context::resume() +{ + RAFT_EXPECTS(impl_ != nullptr, "Cannot resume a moved-from FlowANN search context"); + impl_->resume(); +} + +auto search_context::num_queues() const noexcept -> std::uint32_t +{ + return impl_ == nullptr ? 0 : static_cast(impl_->queues_.size()); +} + +auto search_context::statistics() const noexcept -> queue_statistics +{ + return impl_ == nullptr ? queue_statistics{} : impl_->statistics(); +} + +auto detail::search_context_access::device_queues(search_context const& context) noexcept + -> queue_view* +{ + return context.impl_ == nullptr ? nullptr : context.impl_->device_views_; +} + +auto detail::search_context_access::device_id(search_context const& context) noexcept -> int +{ + return context.impl_ == nullptr ? -1 : context.impl_->device_id_; +} + +auto detail::search_context_access::num_queues(search_context const& context) noexcept + -> std::uint32_t +{ + return context.num_queues(); +} + +auto detail::search_context_access::row_count(search_context const& context) noexcept + -> std::uint64_t +{ + return context.impl_ == nullptr ? 0 : context.impl_->row_count_; +} + +auto detail::search_context_access::row_length(search_context const& context) noexcept + -> std::uint32_t +{ + return context.impl_ == nullptr ? 0 : context.impl_->row_length_; +} + +void detail::search_context_access::acquire_search(search_context& context) +{ + RAFT_EXPECTS(context.impl_ != nullptr, "Cannot search with a moved-from FlowANN search context"); + context.impl_->acquire_search(); +} + +void detail::search_context_access::release_search(search_context& context) +{ + RAFT_EXPECTS(context.impl_ != nullptr, "Cannot search with a moved-from FlowANN search context"); + context.impl_->release_search(); +} + +void detail::search_context_access::report_failure(search_context& context, + std::exception_ptr error) +{ + RAFT_EXPECTS(context.impl_ != nullptr, "Cannot fail a moved-from FlowANN search context"); + context.impl_->report_failure(std::move(error)); +} + +auto detail::search_context_access::is_paused(search_context const& context) noexcept -> bool +{ + return context.impl_ == nullptr || context.impl_->paused_.load(std::memory_order_acquire); +} + +auto detail::search_context_access::pending_pause_requests(search_context const& context) noexcept + -> std::size_t +{ + if (context.impl_ == nullptr) { return 0; } + std::lock_guard lock(context.impl_->pause_mutex_); + return context.impl_->pending_pause_requests_; +} + +} // namespace cuvs::neighbors::cagra::experimental::flowann diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index f132bb21a3..0672424662 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -267,6 +267,41 @@ 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 + neighbors/ann_cagra/flowann_serialize.cu + ) + if(CUVS_ENABLE_FLOWANN_BUILD) + list(APPEND CUVS_FLOWANN_TEST_SOURCES neighbors/ann_cagra/flowann_build.cu + neighbors/ann_cagra/cagra_tiered.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/cagra_tiered.cu b/cpp/tests/neighbors/ann_cagra/cagra_tiered.cu new file mode 100644 index 0000000000..6c1d72a4af --- /dev/null +++ b/cpp/tests/neighbors/ann_cagra/cagra_tiered.cu @@ -0,0 +1,158 @@ +/* + * 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 + +namespace cagra = cuvs::neighbors::cagra; + +TEST(CagraTiered, OwnsDataAcrossMoveAndSerialization) +{ + constexpr int64_t rows = 128, dim = 8, nq = 4, k = 4; + raft::resources res; + auto host = raft::make_host_matrix(rows, dim); + for (int64_t i = 0; i < rows; ++i) { + for (int64_t j = 0; j < dim; ++j) { + host(i, j) = static_cast((i * 31 + j * 5) % 149) * .01f; + } + } + auto queries = raft::make_device_matrix(res, nq, dim); + raft::copy( + queries.data_handle(), host.data_handle(), nq * dim, raft::resource::get_cuda_stream(res)); + auto neighbors = raft::make_device_matrix(res, nq, k); + auto distances = raft::make_device_matrix(res, nq, k); + cagra::search_params sp; + sp.algo = cagra::search_algo::SINGLE_CTA; + sp.itopk_size = 32; + sp.max_iterations = 32; + sp.tiered.emplace(); + // Include CPU-only, mixed and entirely resident representations. + for (size_t budget : {size_t{0}, size_t{512}, size_t{4096}}) { + cagra::index_params bp; + bp.graph_degree = 8; + bp.intermediate_graph_degree = 16; + bp.graph_build_params = + cagra::graph_build_params::nn_descent_params(16, cuvs::distance::DistanceType::L2Expanded); + bp.graph_storage = cagra::graph_storage_kind::tiered; + bp.tiered.emplace(); + bp.tiered->device_graph_budget_bytes = budget; + auto input = cuvs::neighbors::make_device_padded_dataset(res, host.view()); + auto original = cagra::build(res, bp, input->view()); + input.reset(); + auto idx = std::move(original); + EXPECT_EQ(idx.graph_storage(), cagra::graph_storage_kind::tiered); + EXPECT_EQ(idx.size(), rows); + EXPECT_EQ(idx.dim(), dim); + EXPECT_EQ(idx.graph_degree(), 8); + EXPECT_THROW((void)idx.graph(), raft::exception); + auto run = [&](auto const& index) { + cagra::search(res, + sp, + index, + raft::make_const_mdspan(queries.view()), + neighbors.view(), + distances.view()); + std::vector ids(nq * k); + raft::copy( + ids.data(), neighbors.data_handle(), ids.size(), raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + for (auto id : ids) + EXPECT_LT(id, rows); + for (int64_t i = 0; i < nq; ++i) + EXPECT_EQ(ids[i * k], i); + return ids; + }; + auto expected = run(idx); + sp.tiered->keep_pollers_running = true; + EXPECT_EQ(run(idx), expected); + sp.tiered->keep_pollers_running = false; + EXPECT_EQ(run(idx), expected); + std::stringstream saved(std::ios::in | std::ios::out | std::ios::binary); + cagra::serialize(res, saved, idx, true); + // Destroy both the original index and its polling context before loading. + idx = cagra::device_padded_index(res); + cagra::deserialize(res, saved, &idx); + EXPECT_EQ(idx.graph_storage(), cagra::graph_storage_kind::tiered); + EXPECT_EQ(run(idx), expected); + auto replacement = cuvs::neighbors::make_device_padded_dataset(res, host.view()); + EXPECT_THROW(cagra::update_dataset(res, std::move(idx), replacement->view()), raft::exception); + EXPECT_EQ(run(idx), expected); + } +} + +TEST(CagraTiered, CompressedBuildRoundTrip) +{ + constexpr int64_t rows = 512, dim = 16, nq = 4, k = 8; + raft::resources res; + auto host = raft::make_host_matrix(rows, dim); + for (int64_t i = 0; i < rows; ++i) { + for (int64_t j = 0; j < dim; ++j) { + host(i, j) = static_cast((i * 23 + j * 11) % 509) * .01f; + } + } + cagra::index_params bp; + bp.graph_degree = 8; + bp.intermediate_graph_degree = 16; + bp.graph_build_params = + cagra::graph_build_params::nn_descent_params(16, cuvs::distance::DistanceType::L2Expanded); + bp.graph_storage = cagra::graph_storage_kind::tiered; + bp.tiered.emplace(); + bp.tiered->device_graph_budget_bytes = 4096; + cuvs::neighbors::vpq_params pq; + pq.pq_dim = 4; + pq.vq_n_centers = 8; + pq.kmeans_n_iters = 2; + pq.max_train_points_per_pq_code = 16; + pq.max_train_points_per_vq_cluster = 16; + auto idx = + cagra::build(res, bp, pq, cuvs::neighbors::make_host_standard_dataset_view(host.view())); + EXPECT_EQ(idx.graph_storage(), cagra::graph_storage_kind::tiered); + auto queries = raft::make_device_matrix(res, nq, dim); + raft::copy( + queries.data_handle(), host.data_handle(), nq * dim, raft::resource::get_cuda_stream(res)); + auto neighbors = raft::make_device_matrix(res, nq, k); + auto distances = raft::make_device_matrix(res, nq, k); + cagra::search_params sp; + sp.algo = cagra::search_algo::MULTI_CTA; + sp.itopk_size = 64; + sp.max_iterations = 32; + auto run = [&] { + cagra::search( + res, sp, idx, raft::make_const_mdspan(queries.view()), neighbors.view(), distances.view()); + std::vector ids(nq * k); + raft::copy( + ids.data(), neighbors.data_handle(), ids.size(), raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + for (auto id : ids) { + EXPECT_GE(id, 0); + EXPECT_LT(id, rows); + } + return ids; + }; + auto expected = run(); + auto repeated = run(); + std::stringstream saved(std::ios::in | std::ios::out | std::ios::binary); + cagra::serialize(res, saved, idx); + idx = cagra::device_pq_index(res); + cagra::deserialize(res, saved, &idx); + std::stringstream reserialized(std::ios::in | std::ios::out | std::ios::binary); + cagra::serialize(res, reserialized, idx); + EXPECT_EQ(saved.str(), reserialized.str()); + auto restored = run(); + for (auto const* result : {&expected, &repeated, &restored}) { + for (int64_t q = 0; q < nq; ++q) { + EXPECT_NE(std::find(result->begin() + q * k, result->begin() + (q + 1) * k, q), + result->begin() + (q + 1) * k); + } + } +} 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..5650dcad93 --- /dev/null +++ b/cpp/tests/neighbors/ann_cagra/flowann_build.cu @@ -0,0 +1,927 @@ +/* + * 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->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->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->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->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, IntegratedHostBuildCreatesSearchableVpqIndex) +{ + 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 * 23 + column * 11) % 127) * 0.01f; + } + } + auto dataset = cuvs::neighbors::make_host_standard_dataset_view(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); + auto const bits = flowann::detail::required_id_bits(rows); + auto const row_bytes = flowann::detail::packed_row_bytes(2, bits, 4); + params.device_graph_budget_bytes = (rows / 2) * row_bytes; + + cuvs::neighbors::vpq_params vpq_params; + vpq_params.pq_dim = 4; + vpq_params.vq_n_centers = 8; + vpq_params.kmeans_n_iters = 2; + vpq_params.max_train_points_per_pq_code = 16; + vpq_params.max_train_points_per_vq_cluster = 16; + auto built = flowann::build(res, params, vpq_params, dataset); + ASSERT_EQ(built.dataset->n_rows(), rows); + ASSERT_EQ(built.dataset->dim(), dim); + ASSERT_EQ(built.index.graph_degree(), 8); + ASSERT_EQ(built.index.resident_degree(), 4); + ASSERT_EQ(built.index.cross_degree(), 4); + + 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; + search_params.smem_dtype = cuvs::neighbors::cagra::internal_dtype::F16; + 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, IntegratedHostBuildGroupsAndReordersVpqIndex) +{ + constexpr int64_t rows = 260; + constexpr int64_t dim = 16; + raft::resources res; + auto host = raft::make_host_matrix(rows, dim); + for (int64_t row = 0; row < rows; ++row) { + auto const group = row / 65; + for (int64_t column = 0; column < dim; ++column) { + host(row, column) = + static_cast(group * 100 + column) + static_cast(row % 65) * 0.001f; + } + } + auto dataset = cuvs::neighbors::make_host_standard_dataset_view(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 / 2) * 64; + // Small groups force cross edges, so the average-local-degree heuristic would underfill rows. + params.grouping.n_groups = 32; + params.grouping.n_bits = 24; + params.grouping.training_rows = 128; + params.grouping.assignment_batch_rows = 64; + params.grouping.kmeans_n_iters = 4; + params.grouping.validate = true; + + cuvs::neighbors::vpq_params vpq_params; + vpq_params.pq_dim = 4; + vpq_params.vq_n_centers = 8; + vpq_params.kmeans_n_iters = 2; + vpq_params.max_train_points_per_pq_code = 16; + vpq_params.max_train_points_per_vq_cluster = 16; + auto built = flowann::build(res, params, vpq_params, dataset); + + EXPECT_EQ(built.dataset->n_rows(), rows); + EXPECT_EQ(built.dataset->dim(), dim); + 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), 50); // Two headers plus 16 24-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)); + raft::resource::sync_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); +} + +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->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); + + if (iterative) { + // Upstream iterative CAGRA now returns a device graph and requires padded device input. + // FlowANN's host builder must stage that input and bring the graph back for tiering. + cuvs::neighbors::vpq_params vpq_params; + vpq_params.pq_dim = 4; + vpq_params.vq_n_centers = 8; + vpq_params.kmeans_n_iters = 2; + vpq_params.max_train_points_per_pq_code = 16; + vpq_params.max_train_points_per_vq_cluster = 16; + auto host_built = flowann::build( + res, params, vpq_params, cuvs::neighbors::make_host_standard_dataset_view(host.view())); + EXPECT_EQ(host_built.index.size(), rows); + EXPECT_EQ(host_built.index.graph_degree(), 8); + EXPECT_EQ(host_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->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->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); + } + } +} + +TEST(FlowannBuild, IterativeConstructionClampsOversizedDegrees) +{ + raft::resources res; + constexpr int64_t rows = 32, dim = 16; + auto host = raft::make_host_matrix(rows, dim); + for (int64_t row = 0; row < rows; ++row) { + for (int64_t col = 0; col < dim; ++col) { + host(row, col) = static_cast((row * 37 + col * 13) % 211) * 0.01f; + } + } + auto dataset = cuvs::neighbors::make_device_padded_dataset(res, host.view()); + flowann::index_params params; + params.grouping.enabled = false; + params.cagra_params.graph_degree = 64; + params.cagra_params.intermediate_graph_degree = 128; + params.cagra_params.graph_build_params = + cuvs::neighbors::cagra::graph_build_params::iterative_search_params{}; + auto built = flowann::build(res, params, dataset->view()); + ASSERT_EQ(built.index.size(), rows); + ASSERT_EQ(built.index.graph_degree(), rows - 1); + auto resident = cuvs::neighbors::cagra::build(res, params.cagra_params, dataset->view()); + EXPECT_EQ(resident.graph_degree(), rows - 1); + auto cross = built.index.cross_graph(); + ASSERT_EQ(cross.extent(1), rows); + for (int64_t row = 0; row < rows; ++row) { + ASSERT_EQ(cross(row, 0), rows - 1); + for (int64_t col = 1; col < rows; ++col) { + EXPECT_LT(cross(row, col), 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..e8991c79c0 --- /dev/null +++ b/cpp/tests/neighbors/ann_cagra/flowann_queue.cu @@ -0,0 +1,839 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * FlowANN is described in "Disentangling Graph Dependencies for Efficient + * Billion-Scale GPU Vector Search" (OSDI 2026). + */ + +#include + +#include + +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include + +namespace flowann = cuvs::neighbors::cagra::experimental::flowann; + +namespace { + +RAFT_KERNEL queue_round_trip(flowann::detail::queue_view* queues, + std::uint32_t num_queues, + std::uint32_t row_count, + std::uint32_t row_length, + std::uint32_t* output) +{ + if (threadIdx.x != 0) { return; } + + auto const command = blockIdx.x % row_count; + auto queue = queues[blockIdx.x % num_queues]; + flowann::detail::queue_reservation reservation{}; + while (!flowann::detail::try_reserve(queue, &reservation)) {} + flowann::detail::submit(queue, reservation, command); + while (!flowann::detail::result_ready(queue, reservation)) {} + auto const* row = flowann::detail::result_row(queue, reservation); + for (std::uint32_t i = 0; i < row_length; ++i) { + output[static_cast(blockIdx.x) * row_length + i] = row[i]; + } + flowann::detail::release(queue, reservation); +} + +RAFT_KERNEL mixed_release_queue_round_trip(flowann::detail::queue_view* queues, + std::uint32_t num_queues, + std::uint32_t row_count, + std::uint32_t row_length, + std::uint32_t round, + std::uint32_t batch, + std::uint32_t* output) +{ + auto queue = queues[blockIdx.x % num_queues]; + if (threadIdx.x == 0) { + std::uint32_t ids[7]; + for (std::uint32_t i = 0; i < batch; ++i) { + ids[i] = (blockIdx.x * batch + i + round) % row_count; + } + flowann::detail::queue_reservation first{}; + while (!flowann::detail::try_reserve_deferred(queue, batch, &first)) {} + flowann::detail::submit_batch(queue, first, ids, batch); + for (std::uint32_t i = 0; i < batch; ++i) { + auto const reservation = flowann::detail::reservation_at(first, i); + while (!flowann::detail::result_ready(queue, reservation)) {} + auto const* row = flowann::detail::result_row(queue, reservation); + for (std::uint32_t col = 0; col < row_length; ++col) { + output[(static_cast(blockIdx.x) * batch + i) * row_length + col] = row[col]; + } + if (((blockIdx.x + i) & 1u) == 0) { + flowann::detail::release_deferred(queue, reservation); + } else { + flowann::detail::release(queue, reservation); + } + } + } + __syncthreads(); + flowann::detail::reclaim_released_warp(queue); +} + +TEST(FlowannQueue, CommandEncodingPreservesRowIdsAndRejectsPartialPublication) +{ + EXPECT_FALSE(flowann::detail::command_ready(flowann::detail::empty_command)); + for (std::uint32_t row : {0u, 1u, 65535u, 65536u, 0xabcdefu, 0x7fffffffu, 0xffffffffu}) { + auto const word = flowann::detail::make_command(row); + ASSERT_TRUE(flowann::detail::command_ready(word)); + EXPECT_EQ(static_cast(word), row); + // Reading either 32-bit half before the other is published must not + // produce a ready word carrying a different request. + for (auto partial : {word & 0xffffffff00000000ull, word & 0xffffffffull}) { + if (flowann::detail::command_ready(partial)) { EXPECT_EQ(partial, word); } + } + for (unsigned bit = 0; bit < 64; ++bit) { + EXPECT_FALSE(flowann::detail::command_ready(word ^ (1ull << bit))); + } + } +} + +void check_round_trip_pause_resume(std::uint32_t num_queues) +{ + raft::resources res; + constexpr std::uint32_t row_count = 17; + constexpr std::uint32_t row_length = 9; + constexpr std::uint32_t commands = 257; + + auto cross_graph = raft::make_host_matrix(row_count, row_length); + for (std::uint32_t row = 0; row < row_count; ++row) { + for (std::uint32_t col = 0; col < row_length; ++col) { + cross_graph(row, col) = row * 1000 + col; + } + } + + flowann::queue_params params{}; + params.num_queues = num_queues; + params.empty_pause = 8; + params.collect_statistics = true; + flowann::search_context context(raft::make_const_mdspan(cross_graph.view()), params); + auto output = raft::make_device_matrix(res, commands, row_length); + + context.resume(); + queue_round_trip<<>>( + flowann::detail::search_context_access::device_queues(context), + context.num_queues(), + row_count, + row_length, + output.data_handle()); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + raft::resource::sync_stream(res); + context.pause(); + + auto host_output = raft::make_host_matrix(commands, row_length); + raft::copy(host_output.data_handle(), + output.data_handle(), + output.size(), + raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + for (std::uint32_t command = 0; command < commands; ++command) { + auto const row = command % row_count; + for (std::uint32_t col = 0; col < row_length; ++col) { + EXPECT_EQ(host_output(command, col), cross_graph(row, col)); + } + } + EXPECT_EQ(context.statistics().processed_commands, commands); + + context.resume(); + context.pause(); +} + +TEST(FlowannQueue, RoundTripPauseResume) { check_round_trip_pause_resume(2); } + +// Also run this test alone in a fresh process to exercise lazy kernel loading +// while many idle CUDA-copy pollers are already issuing transfers. +TEST(FlowannQueue, ColdLaunchWithManyPollers) { check_round_trip_pause_resume(44); } + +TEST(FlowannQueue, PollersUseContextCudaDevice) +{ + if (raft::device_setter::get_device_count() < 2) { + GTEST_SKIP() << "Requires at least two CUDA devices"; + } + + raft::device_setter device_scope(1); + raft::resources res; + constexpr std::uint32_t row_count = 4; + constexpr std::uint32_t row_length = 3; + + auto cross_graph = raft::make_host_matrix(row_count, row_length); + for (std::uint32_t row = 0; row < row_count; ++row) { + for (std::uint32_t col = 0; col < row_length; ++col) { + cross_graph(row, col) = row * 10 + col; + } + } + + flowann::search_context context(raft::make_const_mdspan(cross_graph.view())); + auto output = raft::make_device_matrix(res, row_count, row_length); + context.resume(); + queue_round_trip<<>>( + flowann::detail::search_context_access::device_queues(context), + context.num_queues(), + row_count, + row_length, + output.data_handle()); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + raft::resource::sync_stream(res); + context.pause(); + + auto host_output = raft::make_host_matrix(row_count, row_length); + raft::copy(host_output.data_handle(), + output.data_handle(), + output.size(), + raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + for (std::uint32_t row = 0; row < row_count; ++row) { + for (std::uint32_t col = 0; col < row_length; ++col) { + EXPECT_EQ(host_output(row, col), cross_graph(row, col)); + } + } +} + +TEST(FlowannQueue, RealPollersPreservePayloadsAcrossWrapWithMixedReleases) +{ + raft::resources res; + auto const stream = raft::resource::get_cuda_stream(res); + constexpr std::uint32_t row_count = 17; + constexpr std::uint32_t row_length = 9; + constexpr std::uint32_t blocks = 4099; + constexpr std::uint32_t max_batch = 7; + constexpr std::uint32_t rounds = 65; + static_assert(static_cast(blocks) * rounds > + static_cast(flowann::detail::queue_capacity) * 4); + + auto cross_graph = raft::make_host_matrix(row_count, row_length); + for (std::uint32_t row = 0; row < row_count; ++row) { + for (std::uint32_t col = 0; col < row_length; ++col) { + cross_graph(row, col) = row * 1000 + col; + } + } + auto output = + raft::make_device_matrix(res, blocks * max_batch, row_length); + auto host_output = raft::make_host_matrix(blocks * max_batch, row_length); + + for (std::uint32_t num_queues : {1u, 4u}) { + for (std::uint32_t batch : {1u, max_batch}) { + auto const commands = blocks * batch; + flowann::queue_params params{}; + params.num_queues = num_queues; + params.empty_pause = 0; + params.collect_statistics = true; + flowann::search_context context(raft::make_const_mdspan(cross_graph.view()), params); + { + flowann::search_session session(context); + for (std::uint32_t round = 0; round < rounds; ++round) { + mixed_release_queue_round_trip<<>>( + flowann::detail::search_context_access::device_queues(context), + num_queues, + row_count, + row_length, + round, + batch, + output.data_handle()); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + // Finish the host-serviced kernel before copying into pageable host memory. + // A pageable D2H copy can block inside the driver while waiting for this + // kernel, preventing the CUDA-copy poller from submitting its transfers. + raft::resource::sync_stream(res); + raft::copy(host_output.data_handle(), + output.data_handle(), + static_cast(commands) * row_length, + stream); + raft::resource::sync_stream(res); + for (std::uint32_t command = 0; command < commands; ++command) { + auto const row = (command + round) % row_count; + for (std::uint32_t col = 0; col < row_length; ++col) { + if (host_output(command, col) != cross_graph(row, col)) { + FAIL() << "payload mismatch with " << num_queues << " queues at round " << round + << ", command " << command << ", column " << col; + } + } + } + } + } + + EXPECT_TRUE(flowann::detail::search_context_access::is_paused(context)); + EXPECT_EQ(context.statistics().processed_commands, + static_cast(commands) * rounds); + std::vector queues(num_queues); + RAFT_CUDA_TRY(cudaMemcpy(queues.data(), + flowann::detail::search_context_access::device_queues(context), + queues.size() * sizeof(flowann::detail::queue_view), + cudaMemcpyDeviceToHost)); + for (auto const& queue : queues) { + flowann::detail::queue_state state{}; + RAFT_CUDA_TRY(cudaMemcpy(&state, queue.state, sizeof(state), cudaMemcpyDeviceToHost)); + EXPECT_EQ(state.gpu_tail, state.reservation_head); + EXPECT_EQ(state.cpu_tail, state.reservation_head); + } + } + } +} + +TEST(FlowannQueue, PollerStopsAtUnpublishedCommandsAndResumesFullBursts) +{ + constexpr std::uint32_t row_count = 17; + constexpr std::uint32_t row_length = 33; + constexpr std::uint32_t commands = 140; + auto graph = raft::make_host_matrix(row_count, row_length); + for (std::uint32_t row = 0; row < row_count; ++row) { + for (std::uint32_t col = 0; col < row_length; ++col) { + graph(row, col) = row * 1000 + col; + } + } + + // Exercise holes within the short prefix, at its boundary, and in a full batch. + for (std::uint32_t hole : {5u, 8u, 127u}) { + SCOPED_TRACE(hole); + flowann::queue_params params{}; + params.num_queues = 1; + params.empty_pause = 0; + params.collect_statistics = true; + flowann::search_context context(raft::make_const_mdspan(graph.view()), params); + flowann::detail::queue_view queue{}; + RAFT_CUDA_TRY(cudaMemcpy(&queue, + flowann::detail::search_context_access::device_queues(context), + sizeof(queue), + cudaMemcpyDeviceToHost)); + unsigned long long const reserved = commands; + RAFT_CUDA_TRY(cudaMemcpy( + &queue.state->reservation_head, &reserved, sizeof(reserved), cudaMemcpyHostToDevice)); + std::vector ids(commands); + std::vector sequences(commands); + for (std::uint32_t i = 0; i < commands; ++i) { + ids[i] = i % row_count; + sequences[i] = + i == hole ? flowann::detail::empty_command : flowann::detail::make_command(ids[i]); + } + RAFT_CUDA_TRY(cudaMemcpy(queue.command_words, + sequences.data(), + sequences.size() * sizeof(sequences[0]), + cudaMemcpyHostToDevice)); + auto wait_for = [&](std::uint64_t expected) { + auto const deadline = std::chrono::steady_clock::now() + std::chrono::seconds(5); + while (context.statistics().processed_commands < expected && + std::chrono::steady_clock::now() < deadline) { + std::this_thread::yield(); + } + return context.statistics().processed_commands; + }; + + context.resume(); + auto const first_count = wait_for(hole); + context.pause(); + ASSERT_EQ(first_count, hole); + ASSERT_EQ(context.statistics().processed_commands, hole); + flowann::detail::queue_state state{}; + RAFT_CUDA_TRY(cudaMemcpy(&state, queue.state, sizeof(state), cudaMemcpyDeviceToHost)); + EXPECT_EQ(state.cpu_tail, hole); + std::vector published(commands); + RAFT_CUDA_TRY(cudaMemcpy(published.data(), + queue.result_sequence, + published.size() * sizeof(published[0]), + cudaMemcpyDeviceToHost)); + for (std::uint32_t i = 0; i < commands; ++i) { + EXPECT_EQ(published[i], i < hole ? i + 1ull : 0ull); + } + + auto const ready = flowann::detail::make_command(ids[hole]); + RAFT_CUDA_TRY( + cudaMemcpy(queue.command_words + hole, &ready, sizeof(ready), cudaMemcpyHostToDevice)); + context.resume(); + auto const final_count = wait_for(commands); + context.pause(); + ASSERT_EQ(final_count, commands); + RAFT_CUDA_TRY(cudaMemcpy(&state, queue.state, sizeof(state), cudaMemcpyDeviceToHost)); + EXPECT_EQ(state.cpu_tail, commands); + std::vector cleared(commands); + RAFT_CUDA_TRY(cudaMemcpy(cleared.data(), + queue.command_words, + cleared.size() * sizeof(cleared[0]), + cudaMemcpyDeviceToHost)); + for (auto word : cleared) { + EXPECT_EQ(word, flowann::detail::empty_command); + } + std::vector returned(commands * row_length); + RAFT_CUDA_TRY(cudaMemcpy( + returned.data(), queue.rows, returned.size() * sizeof(returned[0]), cudaMemcpyDeviceToHost)); + for (std::uint32_t i = 0; i < commands; ++i) { + for (std::uint32_t col = 0; col < row_length; ++col) { + EXPECT_EQ(returned[i * row_length + col], graph(ids[i], col)); + } + } + } +} + +RAFT_KERNEL deferred_capacity_reclaim(flowann::detail::queue_view q, unsigned* errors) +{ + if (threadIdx.x != 0 || blockIdx.x != 0) return; + for (unsigned i = 0; i < flowann::detail::queue_capacity * 3 + 17; i++) { + flowann::detail::queue_reservation r{}; + unsigned tries = 0; + while (!flowann::detail::try_reserve_deferred(q, 1, &r)) { + if (++tries > 4) { + *errors = 1; + return; + } + } + if (r.sequence != i) { + *errors = 2; + return; + } + flowann::detail::release_deferred(q, r); + } + flowann::detail::reclaim_released(q); + if (q.state->gpu_tail != q.state->reservation_head) *errors = 3; +} +RAFT_KERNEL deferred_release_hole(flowann::detail::queue_view q, unsigned* errors) +{ + if (threadIdx.x == 0) { + q.state->gpu_tail = 0; + q.state->reservation_head = 3; + q.released_sequence[0] = 0; + flowann::detail::release_deferred(q, {1, 1}); + flowann::detail::release_deferred(q, {2, 2}); + } + __syncthreads(); + flowann::detail::reclaim_released_warp(q); + if (threadIdx.x == 0 && q.state->gpu_tail != 0) *errors = 4; + if (threadIdx.x == 0) flowann::detail::release_deferred(q, {0, 0}); + __syncthreads(); + flowann::detail::reclaim_released_warp(q); + if (threadIdx.x == 0 && q.state->gpu_tail != 3) *errors = 5; +} + +RAFT_KERNEL reserve_deferred_concurrently(flowann::detail::queue_view queue, + std::uint64_t expected_base, + std::uint32_t reservations_per_block, + flowann::detail::queue_reservation* reservations, + std::uint32_t* valid, + std::uint32_t* sequence_owners, + std::uint32_t* errors) +{ + auto const reservation_index = blockIdx.x; + if (threadIdx.x == 0) { + flowann::detail::queue_reservation first{}; + bool reserved = false; + for (std::uint32_t attempt = 0; attempt < 4096 && !reserved; ++attempt) { + reserved = flowann::detail::try_reserve_deferred(queue, reservations_per_block, &first); + } + reservations[reservation_index] = first; + valid[reservation_index] = reserved; + if (!reserved) { atomicOr(errors, 1u); } + } + __syncthreads(); + + if (valid[reservation_index] == 0) { return; } + auto const first = reservations[reservation_index]; + for (std::uint32_t i = threadIdx.x; i < reservations_per_block; i += blockDim.x) { + auto const reservation = flowann::detail::reservation_at(first, i); + if (reservation.sequence < expected_base || + reservation.sequence >= + expected_base + static_cast(gridDim.x) * reservations_per_block) { + atomicOr(errors, 2u); + continue; + } + auto const offset = static_cast(reservation.sequence - expected_base); + if (atomicCAS(&sequence_owners[offset], ~std::uint32_t{0}, reservation_index) != + ~std::uint32_t{0}) { + atomicOr(errors, 4u); + } + } +} + +RAFT_KERNEL release_deferred_except_first(flowann::detail::queue_view queue, + std::uint64_t first_sequence, + std::uint32_t reservations_per_block, + flowann::detail::queue_reservation const* reservations, + std::uint32_t const* valid) +{ + if (valid[blockIdx.x] != 0) { + auto const first = reservations[blockIdx.x]; + for (std::uint32_t i = threadIdx.x; i < reservations_per_block; i += blockDim.x) { + auto const reservation = flowann::detail::reservation_at(first, i); + if (reservation.sequence != first_sequence) { + flowann::detail::release_deferred(queue, reservation); + } + } + } + __syncthreads(); + if (threadIdx.x < 32) { flowann::detail::reclaim_released_warp(queue); } +} + +RAFT_KERNEL release_first_and_reclaim(flowann::detail::queue_view queue, + std::uint64_t first_sequence, + std::uint32_t reservations_per_block, + flowann::detail::queue_reservation const* reservations, + std::uint32_t const* valid) +{ + if (valid[blockIdx.x] != 0) { + auto const first = reservations[blockIdx.x]; + if (first.sequence <= first_sequence && + first_sequence < first.sequence + reservations_per_block && threadIdx.x == 0) { + flowann::detail::release_deferred( + queue, + flowann::detail::reservation_at( + first, static_cast(first_sequence - first.sequence))); + } + } + __syncthreads(); + if (threadIdx.x < 32) { flowann::detail::reclaim_released_warp(queue); } +} + +RAFT_KERNEL check_deferred_queue_counters(flowann::detail::queue_view queue, + std::uint64_t expected_tail, + std::uint32_t error_bit, + std::uint32_t* errors) +{ + if (threadIdx.x == 0 && blockIdx.x == 0 && + (queue.state->gpu_tail != expected_tail || + (error_bit != 0 && queue.state->reservation_head != expected_tail))) { + atomicOr(errors, error_bit == 0 ? 8u : error_bit); + } +} + +TEST(FlowannQueue, DeferredReclamationHandlesCapacityWrapAndUnreleasedHoles) +{ + raft::resources res; + auto const stream = raft::resource::get_cuda_stream(res); + auto cross = raft::make_host_matrix(1, 2); + cross(0, 0) = 1; + cross(0, 1) = 0; + flowann::search_context context(raft::make_const_mdspan(cross.view())); + flowann::detail::queue_view queue; + RAFT_CUDA_TRY(cudaMemcpy(&queue, + flowann::detail::search_context_access::device_queues(context), + sizeof(queue), + cudaMemcpyDeviceToHost)); + auto errors = raft::make_device_matrix(res, 1, 1); + RAFT_CUDA_TRY(cudaMemsetAsync(errors.data_handle(), 0, sizeof(std::uint32_t), stream.get())); + deferred_capacity_reclaim<<<1, 1, 0, stream.get()>>>(queue, errors.data_handle()); + deferred_release_hole<<<1, 32, 0, stream.get()>>>(queue, errors.data_handle()); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + std::uint32_t host_error{}; + raft::copy(&host_error, errors.data_handle(), 1, stream); + raft::resource::sync_stream(res); + EXPECT_EQ(host_error, 0u); +} + +TEST(FlowannQueue, ConcurrentDeferredReclamationPreservesOwnershipAcrossWrap) +{ + raft::resources res; + auto const stream = raft::resource::get_cuda_stream(res); + auto cross = raft::make_host_matrix(1, 2); + cross(0, 0) = 1; + cross(0, 1) = 0; + flowann::search_context context(raft::make_const_mdspan(cross.view())); + flowann::detail::queue_view queue; + RAFT_CUDA_TRY(cudaMemcpy(&queue, + flowann::detail::search_context_access::device_queues(context), + sizeof(queue), + cudaMemcpyDeviceToHost)); + + constexpr std::uint32_t blocks = 255; + constexpr std::uint32_t reservations_per_block = 257; + constexpr std::uint32_t reservations_per_wave = blocks * reservations_per_block; + static_assert(reservations_per_wave == flowann::detail::queue_capacity - 1); + + auto reservations = + raft::make_device_matrix(res, blocks, 1); + auto valid = raft::make_device_matrix(res, blocks, 1); + auto owners = raft::make_device_matrix(res, reservations_per_wave, 1); + auto errors = raft::make_device_matrix(res, 1, 1); + RAFT_CUDA_TRY(cudaMemsetAsync(errors.data_handle(), 0, sizeof(std::uint32_t), stream.get())); + + std::uint64_t expected_base = 0; + for (std::uint32_t wave = 0; wave < 3; ++wave) { + RAFT_CUDA_TRY( + cudaMemsetAsync(valid.data_handle(), 0, valid.size() * sizeof(std::uint32_t), stream.get())); + RAFT_CUDA_TRY(cudaMemsetAsync( + owners.data_handle(), 0xff, owners.size() * sizeof(std::uint32_t), stream.get())); + reserve_deferred_concurrently<<>>(queue, + expected_base, + reservations_per_block, + reservations.data_handle(), + valid.data_handle(), + owners.data_handle(), + errors.data_handle()); + release_deferred_except_first<<>>(queue, + expected_base, + reservations_per_block, + reservations.data_handle(), + valid.data_handle()); + check_deferred_queue_counters<<<1, 1, 0, stream.get()>>>( + queue, expected_base, 0, errors.data_handle()); + release_first_and_reclaim<<>>(queue, + expected_base, + reservations_per_block, + reservations.data_handle(), + valid.data_handle()); + expected_base += reservations_per_wave; + check_deferred_queue_counters<<<1, 1, 0, stream.get()>>>( + queue, expected_base, 16u, errors.data_handle()); + } + RAFT_CUDA_TRY(cudaPeekAtLastError()); + + std::uint32_t host_error{}; + raft::copy(&host_error, errors.data_handle(), 1, stream); + raft::resource::sync_stream(res); + EXPECT_EQ(host_error, 0u); +} + +TEST(FlowannQueue, SearchLeaseConsumesRawResume) +{ + constexpr std::uint32_t row_count = 4; + constexpr std::uint32_t row_length = 3; + auto cross_graph = raft::make_host_matrix(row_count, row_length); + std::fill_n(cross_graph.data_handle(), cross_graph.size(), 0u); + + flowann::search_context context(raft::make_const_mdspan(cross_graph.view())); + EXPECT_TRUE(flowann::detail::search_context_access::is_paused(context)); + + context.resume(); + EXPECT_FALSE(flowann::detail::search_context_access::is_paused(context)); + flowann::detail::search_context_access::acquire_search(context); + flowann::detail::search_context_access::release_search(context); + EXPECT_TRUE(flowann::detail::search_context_access::is_paused(context)); + + context.resume(); + context.pause(); + EXPECT_TRUE(flowann::detail::search_context_access::is_paused(context)); + + flowann::detail::search_context_access::acquire_search(context); + context.resume(); + flowann::detail::search_context_access::release_search(context); + EXPECT_TRUE(flowann::detail::search_context_access::is_paused(context)); +} + +TEST(FlowannQueue, SearchSessionKeepsPollersRunningAcrossSearchLease) +{ + constexpr std::uint32_t row_count = 4; + constexpr std::uint32_t row_length = 3; + auto cross_graph = raft::make_host_matrix(row_count, row_length); + std::fill_n(cross_graph.data_handle(), cross_graph.size(), 0u); + + flowann::search_context context(raft::make_const_mdspan(cross_graph.view())); + { + flowann::search_session session(context); + EXPECT_FALSE(flowann::detail::search_context_access::is_paused(context)); + flowann::detail::search_context_access::acquire_search(context); + flowann::detail::search_context_access::release_search(context); + EXPECT_FALSE(flowann::detail::search_context_access::is_paused(context)); + } + EXPECT_TRUE(flowann::detail::search_context_access::is_paused(context)); +} + +TEST(FlowannQueue, PauseBlocksNewSearchesUntilPollersReachQuiescence) +{ + using namespace std::chrono_literals; + constexpr std::uint32_t row_count = 4; + constexpr std::uint32_t row_length = 3; + auto cross_graph = raft::make_host_matrix(row_count, row_length); + std::fill_n(cross_graph.data_handle(), cross_graph.size(), 0u); + + flowann::search_context context(raft::make_const_mdspan(cross_graph.view())); + flowann::detail::search_context_access::acquire_search(context); + std::atomic pauses_returned{}; + std::jthread first_pauser([&] { + context.pause(); + pauses_returned.fetch_add(1, std::memory_order_release); + }); + + auto const deadline = std::chrono::steady_clock::now() + 2s; + while (flowann::detail::search_context_access::pending_pause_requests(context) < 1 && + std::chrono::steady_clock::now() < deadline) { + std::this_thread::yield(); + } + std::jthread second_pauser([&] { + context.pause(); + pauses_returned.fetch_add(1, std::memory_order_release); + }); + auto const second_deadline = std::chrono::steady_clock::now() + 2s; + while (flowann::detail::search_context_access::pending_pause_requests(context) < 2 && + std::chrono::steady_clock::now() < second_deadline) { + std::this_thread::yield(); + } + auto const pending_pauses = + flowann::detail::search_context_access::pending_pause_requests(context); + + std::atomic later_search_started{}; + std::atomic later_search_acquired{}; + std::jthread later_search([&] { + later_search_started.store(true, std::memory_order_release); + flowann::detail::search_context_access::acquire_search(context); + later_search_acquired.store(true, std::memory_order_release); + flowann::detail::search_context_access::release_search(context); + }); + while (!later_search_started.load(std::memory_order_acquire)) { + std::this_thread::yield(); + } + std::this_thread::sleep_for(20ms); + if (pending_pauses == 2) { EXPECT_FALSE(later_search_acquired.load(std::memory_order_acquire)); } + + flowann::detail::search_context_access::release_search(context); + first_pauser.join(); + second_pauser.join(); + later_search.join(); + EXPECT_EQ(pending_pauses, 2); + EXPECT_EQ(pauses_returned.load(std::memory_order_acquire), 2); + EXPECT_TRUE(later_search_acquired.load(std::memory_order_acquire)); + EXPECT_TRUE(flowann::detail::search_context_access::is_paused(context)); +} + +TEST(FlowannQueue, SearchLeasePausesAfterLastSearch) +{ + constexpr std::uint32_t row_count = 4; + constexpr std::uint32_t row_length = 3; + auto cross_graph = raft::make_host_matrix(row_count, row_length); + std::fill_n(cross_graph.data_handle(), cross_graph.size(), 0u); + + flowann::search_context context(raft::make_const_mdspan(cross_graph.view())); + flowann::detail::search_context_access::acquire_search(context); + flowann::detail::search_context_access::acquire_search(context); + EXPECT_FALSE(flowann::detail::search_context_access::is_paused(context)); + + flowann::detail::search_context_access::release_search(context); + EXPECT_FALSE(flowann::detail::search_context_access::is_paused(context)); + flowann::detail::search_context_access::release_search(context); + EXPECT_TRUE(flowann::detail::search_context_access::is_paused(context)); +} + +} // namespace + +TEST(FlowannQueue, SessionSurvivesContextMove) +{ + auto graph = raft::make_host_matrix(2, 2); + graph(0, 0) = graph(1, 0) = 1; + graph(0, 1) = 1; + graph(1, 1) = 0; + flowann::search_context context(raft::make_const_mdspan(graph.view())); + auto session = std::make_unique(context); + flowann::search_context moved(std::move(context)); + session.reset(); + EXPECT_NO_THROW(moved.pause()); + EXPECT_TRUE(flowann::detail::search_context_access::is_paused(moved)); + + flowann::search_context destination(raft::make_const_mdspan(graph.view())); + session = std::make_unique(moved); + destination = std::move(moved); + session.reset(); + EXPECT_NO_THROW(destination.pause()); + EXPECT_TRUE(flowann::detail::search_context_access::is_paused(destination)); +} + +RAFT_KERNEL invalid_queue_request(flowann::detail::queue_view* queues, std::uint32_t* output) +{ + if (threadIdx.x != 0) { return; } + auto queue = queues[0]; + flowann::detail::queue_reservation reservation{}; + if (!flowann::detail::try_reserve(queue, &reservation)) { + output[0] = 0; + return; + } + flowann::detail::submit(queue, reservation, 0xffffffffu); + auto const start = clock64(); + while (!flowann::detail::result_ready(queue, reservation) && clock64() - start < 100000000ull) {} + auto const ready = flowann::detail::result_ready(queue, reservation); + output[0] = ready; + output[1] = ready ? flowann::detail::result_row(queue, reservation)[0] : 0xffffffffu; + if (ready) { flowann::detail::release(queue, reservation); } +} + +TEST(FlowannQueue, InvalidRequestsCompleteAndReportAnError) +{ + raft::resources res; + auto graph = raft::make_host_matrix(2, 2); + graph(0, 0) = graph(1, 0) = 1; + graph(0, 1) = 1; + graph(1, 1) = 0; + flowann::search_context context(raft::make_const_mdspan(graph.view())); + auto output = raft::make_device_vector(res, 2); + context.resume(); + invalid_queue_request<<<1, 1, 0, raft::resource::get_cuda_stream(res).get()>>>( + flowann::detail::search_context_access::device_queues(context), output.data_handle()); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + raft::resource::sync_stream(res); + std::uint32_t observed[2]{}; + raft::copy(observed, output.data_handle(), 2, raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + EXPECT_EQ(observed[0], 1); + EXPECT_EQ(observed[1], 0); + EXPECT_ANY_THROW(context.pause()); + // Reporting consumes the malformed-request count. Reuse this same queue/context. + EXPECT_NO_THROW(context.resume()); + EXPECT_NO_THROW(context.pause()); + EXPECT_NO_THROW(flowann::detail::search_context_access::acquire_search(context)); + EXPECT_NO_THROW(flowann::detail::search_context_access::release_search(context)); +} + +TEST(FlowannQueue, EveryLeaseReportsPollerFailure) +{ + auto graph = raft::make_host_matrix(4, 1); + std::fill_n(graph.data_handle(), graph.size(), 0u); + flowann::search_context context(raft::make_const_mdspan(graph.view())); + flowann::detail::search_context_access::acquire_search(context); + flowann::detail::search_context_access::acquire_search(context); + flowann::detail::search_context_access::report_failure( + context, std::make_exception_ptr(std::runtime_error("injected poller failure"))); + EXPECT_THROW(flowann::detail::search_context_access::release_search(context), std::runtime_error); + EXPECT_THROW(flowann::detail::search_context_access::release_search(context), std::runtime_error); +} + +namespace { +RAFT_KERNEL wait_for_host_abort(flowann::detail::queue_view queue, std::uint32_t* completed) +{ + if (threadIdx.x != 0) { return; } + flowann::detail::queue_reservation reservation{0, 0}; + auto const start = clock64(); + bool ready = false; + while (!(ready = flowann::detail::result_ready(queue, reservation)) && + clock64() - start < 100000000ull) {} + // One queue view must drain every outstanding reservation after the first abort read. + for (std::uint32_t i = 1; i < 8; ++i) { + ready = flowann::detail::result_ready(queue, {i, i}) && ready; + } + *completed = ready; +} +} // namespace + +TEST(FlowannQueue, HostAbortCompletesAWaitWithoutCudaCopies) +{ + raft::resources res; + auto graph = raft::make_host_matrix(4, 1); + std::fill_n(graph.data_handle(), graph.size(), 0u); + flowann::search_context context(raft::make_const_mdspan(graph.view())); + flowann::detail::queue_view queue{}; + RAFT_CUDA_TRY(cudaMemcpy(&queue, + flowann::detail::search_context_access::device_queues(context), + sizeof(queue), + cudaMemcpyDeviceToHost)); + cudaPointerAttributes attributes{}; + RAFT_CUDA_TRY(cudaPointerGetAttributes(&attributes, queue.abort_flag)); + auto* host_flag = static_cast(attributes.hostPointer); + ASSERT_NE(host_flag, nullptr); + auto completed = raft::make_device_vector(res, 1); + std::atomic_ref(*host_flag).store(1, std::memory_order_release); + wait_for_host_abort<<<1, 1, 0, raft::resource::get_cuda_stream(res).get()>>>( + queue, completed.data_handle()); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + raft::resource::sync_stream(res); + std::uint32_t actual{}; + raft::copy(&actual, completed.data_handle(), 1, raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + EXPECT_EQ(actual, 1u); +} diff --git a/cpp/tests/neighbors/ann_cagra/flowann_search.cu b/cpp/tests/neighbors/ann_cagra/flowann_search.cu new file mode 100644 index 0000000000..023937ce3d --- /dev/null +++ b/cpp/tests/neighbors/ann_cagra/flowann_search.cu @@ -0,0 +1,655 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * FlowANN is described in "Disentangling Graph Dependencies for Efficient + * Billion-Scale GPU Vector Search" (OSDI 2026). + */ + +#include + +#include +#include + +#include +#include +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace flowann = cuvs::neighbors::cagra::experimental::flowann; + +namespace { + +void pack_two_nodes(raft::host_matrix& inner, + int64_t row, + std::uint8_t first_neighbor, + std::uint8_t second_neighbor, + std::uint32_t packed_bits) +{ + inner(row, 0) = 1; + inner(row, 1) = 1; + for (int64_t col = 2; col < inner.extent(1); ++col) { + inner(row, col) = 0; + } + for (std::uint32_t node = 0; node < 2; ++node) { + auto const value = node == 0 ? first_neighbor : second_neighbor; + for (std::uint32_t bit = 0; bit < packed_bits; ++bit) { + auto const position = node * packed_bits + bit; + inner(row, 2 + position / 8) |= ((std::uint32_t{value} >> bit) & 1u) << (position % 8); + } + } +} + +struct search_result { + raft::host_matrix neighbors; + raft::host_matrix distances; + flowann::queue_statistics queue_statistics; + bool pollers_paused; +}; + +auto run_dense_search(cuvs::neighbors::cagra::search_algo algo, + std::vector const& query_values, + std::vector const& seeds = {}, + std::uint32_t max_iterations = 16, + bool keep_pollers_running = false, + std::uint32_t max_batch_queries = 0, + std::uint32_t itopk_size = 32, + std::uint32_t search_width = 1, + std::uint32_t num_queues = 2, + bool alternate_modes = false, + std::uint32_t packed_bits = 2, + std::uint32_t row_padding = 0, + std::uint32_t query_dim = 4, + std::uint32_t complete_degree = 2) -> search_result +{ + raft::resources res; + constexpr int64_t rows = 8; + auto const dim = static_cast(query_dim); + + auto host_dataset = raft::make_host_matrix(rows, dim); + for (int64_t row = 0; row < rows; ++row) { + auto const group = row < 4 ? 0.0f : 10.0f; + host_dataset(row, 0) = group + static_cast(row % 4); + for (int64_t col = 1; col < dim; ++col) { + host_dataset(row, col) = 0.0f; + } + } + auto device_dataset = raft::make_device_matrix(res, rows, dim); + raft::copy(device_dataset.data_handle(), + host_dataset.data_handle(), + host_dataset.size(), + raft::resource::get_cuda_stream(res)); + std::unique_ptr> dataset_owner; + auto dataset = [&]() -> cuvs::neighbors::device_padded_dataset_view { + if (cuvs::neighbors::matrix_row_width_matches_cagra_required(device_dataset.view())) { + return cuvs::neighbors::make_device_padded_dataset_view(res, device_dataset.view()); + } + dataset_owner = cuvs::neighbors::make_device_padded_dataset(res, device_dataset.view()); + return dataset_owner->view(); + }(); + + flowann::device_padded_index index( + res, cuvs::distance::DistanceType::L2Expanded, dataset, 2, packed_bits); + auto inner = + raft::make_host_matrix(4, 2 + (2 * packed_bits + 7) / 8 + row_padding); + pack_two_nodes(inner, 0, 1, 2, packed_bits); + pack_two_nodes(inner, 1, 3, 0, packed_bits); + pack_two_nodes(inner, 2, 1, 2, packed_bits); + pack_two_nodes(inner, 3, 3, 0, packed_bits); + + auto cross = raft::make_host_matrix(rows, complete_degree); + for (std::uint32_t row = 0; row < rows; ++row) { + cross(row, 0) = complete_degree - 1; + for (std::uint32_t column = 1; column < complete_degree; ++column) { + cross(row, column) = row < 4 ? row + 4 : row - 4; + } + } + index.update_graph(res, + raft::make_const_mdspan(inner.view()), + raft::make_const_mdspan(cross.view()), + complete_degree, + 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 + +namespace { +RAFT_KERNEL check_pending_backpressure(flowann::detail::queue_view queue, + std::uint32_t mode, + std::uint32_t* output) +{ + if (threadIdx.x != 0) { return; } + __shared__ flowann::detail::graph_state state; + __shared__ std::uint32_t traversed[32]; + std::uint32_t topk[8]; + std::uint32_t positions[8]; + for (std::uint32_t i = 0; i < 32; ++i) { + traversed[i] = ~std::uint32_t{0}; + } + state.pending.initialize(); + // Mode 0 has one free pending slot; mode 1 has a full device queue; + // mode 2 has capacity in both and must publish the complete batch. + state.pending.head = mode == 0 ? 510 : 0; + queue.state->reservation_head = mode == 1 ? flowann::detail::queue_capacity : 0; + queue.state->gpu_tail = 0; + queue.state->capacity = flowann::detail::queue_capacity; + queue.released_sequence[0] = 0; + for (std::uint32_t i = 0; i < 8; ++i) { + queue.command_words[i] = flowann::detail::empty_command; + topk[i] = (i + 1) | 0x80000000u; + positions[i] = i; + state.parent_ids[i] = i + 1; + state.parent_valid[i] = true; + cuvs::neighbors::cagra::detail::hashmap::insert(traversed, 5, i + 1); + } + flowann::detail::submit_parent_requests(queue, &state, positions, topk, 8, 1, traversed, 5); + output[0] = state.pending.size(); + output[1] = queue.state->reservation_head; + output[2] = 0; + output[3] = 0; + for (std::uint32_t i = 0; i < 8; ++i) { + output[2] += (topk[i] & 0x80000000u) != 0; + output[3] += + cuvs::neighbors::cagra::detail::hashmap::search(traversed, 5, i + 1); + } + output[4] = 0; + for (std::uint32_t i = 0; i < 8; ++i) { + output[4] += queue.command_words[i] == flowann::detail::make_command(i + 1); + } +} +} // namespace + +TEST(FlowannSearch, RetriesParentsWithoutOverwritingPendingRequests) +{ + raft::resources res; + auto queue_state = raft::make_device_vector(res, 1); + auto commands = raft::make_device_vector(res, 8); + auto released = raft::make_device_vector(res, 1); + auto output = raft::make_device_vector(res, 5); + flowann::detail::queue_view queue{}; + queue.state = queue_state.data_handle(); + queue.command_words = commands.data_handle(); + queue.released_sequence = released.data_handle(); + for (std::uint32_t mode = 0; mode < 3; ++mode) { + check_pending_backpressure<<<1, 1, 0, raft::resource::get_cuda_stream(res).get()>>>( + queue, mode, output.data_handle()); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + std::uint32_t actual[5]; + raft::copy(actual, output.data_handle(), 5, raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + EXPECT_EQ(actual[0], mode == 0 ? 510 : mode == 1 ? 0 : 8); + EXPECT_EQ(actual[1], mode == 1 ? flowann::detail::queue_capacity : mode == 2 ? 8 : 0); + EXPECT_EQ(actual[2], mode == 2 ? 8 : 0); + EXPECT_EQ(actual[3], mode == 2 ? 8 : 0); + EXPECT_EQ(actual[4], mode == 2 ? 8 : 0); + } +} + +TEST(FlowannSearch, EmptyQueriesAreANoOp) +{ + EXPECT_NO_THROW(run_dense_search(cuvs::neighbors::cagra::search_algo::SINGLE_CTA, {})); +} + +TEST(FlowannSearch, PaddedQueriesMatchAcrossBatchBoundaries) +{ + for (auto algo : {cuvs::neighbors::cagra::search_algo::SINGLE_CTA, + cuvs::neighbors::cagra::search_algo::MULTI_CTA}) { + for (std::uint32_t dim : {5u, 33u}) { + auto one = run_dense_search( + algo, {0.1f, 12.2f, 1.1f, 10.3f, 2.2f}, {}, 16, false, 1, 32, 1, 2, false, 2, 0, dim); + auto batched = run_dense_search( + algo, {0.1f, 12.2f, 1.1f, 10.3f, 2.2f}, {}, 16, false, 3, 32, 1, 2, false, 2, 0, dim); + for (int64_t row = 0; row < 5; ++row) { + for (int64_t rank = 0; rank < 2; ++rank) { + EXPECT_EQ(batched.neighbors(row, rank), one.neighbors(row, rank)); + EXPECT_NEAR(batched.distances(row, rank), one.distances(row, rank), 1e-4f); + } + } + } + EXPECT_NO_THROW(run_dense_search(algo, {})); + } +} + +TEST(FlowannSearch, AutoFallsBackForGraphsTooWideForMultiCta) +{ + auto result = run_dense_search(cuvs::neighbors::cagra::search_algo::AUTO, + {0.1f}, + {}, + 16, + false, + 1, + 32, + 1, + 2, + false, + 2, + 0, + 4, + 75); + EXPECT_EQ(result.neighbors(0, 0), 7); + EXPECT_NEAR(result.distances(0, 0), 0.01f, 1e-4f); +} + +TEST(FlowannSearch, PaddedSingleCtaQueriesMatchAcrossBatches) +{ + for (std::uint32_t dim : {5u, 33u}) { + auto one = run_dense_search(cuvs::neighbors::cagra::search_algo::SINGLE_CTA, + {0.1f, 12.2f, 1.1f, 10.3f, 2.2f}, + {}, + 16, + false, + 1, + 32, + 1, + 2, + false, + 2, + 0, + dim); + auto batched = run_dense_search(cuvs::neighbors::cagra::search_algo::SINGLE_CTA, + {0.1f, 12.2f, 1.1f, 10.3f, 2.2f}, + {}, + 16, + false, + 3, + 32, + 1, + 2, + false, + 2, + 0, + dim); + for (int64_t row = 0; row < 5; ++row) { + for (int64_t rank = 0; rank < 2; ++rank) { + EXPECT_EQ(batched.neighbors(row, rank), one.neighbors(row, rank)); + EXPECT_NEAR(batched.distances(row, rank), one.distances(row, rank), 1e-4f); + } + } + } +} + +TEST(FlowannSearch, PaddedMultiCtaQueriesMatchAcrossBatches) +{ + for (std::uint32_t dim : {5u, 33u}) { + auto one = run_dense_search(cuvs::neighbors::cagra::search_algo::MULTI_CTA, + {0.1f, 12.2f, 1.1f, 10.3f, 2.2f}, + {}, + 16, + false, + 1, + 32, + 1, + 2, + false, + 2, + 0, + dim); + auto batched = run_dense_search(cuvs::neighbors::cagra::search_algo::MULTI_CTA, + {0.1f, 12.2f, 1.1f, 10.3f, 2.2f}, + {}, + 16, + false, + 3, + 32, + 1, + 2, + false, + 2, + 0, + dim); + for (int64_t row = 0; row < 5; ++row) { + for (int64_t rank = 0; rank < 2; ++rank) { + EXPECT_EQ(batched.neighbors(row, rank), one.neighbors(row, rank)); + EXPECT_NEAR(batched.distances(row, rank), one.distances(row, rank), 1e-4f); + } + } + } +} + +TEST_P(FlowannSearchModesTest, MappedInvalidIdsMatchTheOutputWidth) +{ + raft::resources res; + auto host = raft::make_host_matrix(4, 4); + std::fill_n(host.data_handle(), host.size(), 0.0f); + auto dataset = cuvs::neighbors::make_device_padded_dataset(res, host.view()); + flowann::device_padded_index index( + res, cuvs::distance::DistanceType::L2Expanded, dataset->view()); + auto inner = raft::make_host_matrix(0, 0); + auto cross = raft::make_host_matrix(4, 2); + auto ids = raft::make_host_vector(4); + for (std::uint32_t row = 0; row < 4; ++row) { + cross(row, 0) = 1; + cross(row, 1) = (row + 1) % 4; + ids(row) = 3 - row; + } + index.update_graph( + res, raft::make_const_mdspan(inner.view()), raft::make_const_mdspan(cross.view()), 1, 0); + index.update_source_indices(res, raft::make_const_mdspan(ids.view())); + auto offsets = raft::make_host_vector(1); + auto groups = raft::make_host_vector(4); + offsets(0) = 0; + std::fill_n(groups.data_handle(), groups.size(), 0u); + index.update_subgraph_layout( + res, raft::make_const_mdspan(offsets.view()), raft::make_const_mdspan(groups.view())); + auto queries = raft::make_device_matrix(res, 2, 4); + auto neighbors = raft::make_device_matrix(res, 2, 2); + auto distances = raft::make_device_matrix(res, 2, 2); + auto bits = raft::make_device_vector(res, 1); + RAFT_CUDA_TRY(cudaMemsetAsync(queries.data_handle(), + 0, + queries.size() * sizeof(float), + raft::resource::get_cuda_stream(res).get())); + RAFT_CUDA_TRY(cudaMemsetAsync( + bits.data_handle(), 0, sizeof(std::uint32_t), raft::resource::get_cuda_stream(res).get())); + cuvs::neighbors::filtering::bitset_filter filter( + cuvs::core::bitset_view(bits.data_handle(), 4)); + flowann::search_context context(index.cross_graph()); + flowann::search_params params; + params.algo = GetParam(); + params.itopk_size = 32; + params.max_iterations = 16; + flowann::search(res, + params, + index, + context, + raft::make_const_mdspan(queries.view()), + neighbors.view(), + distances.view(), + filter); + int64_t actual[4]; + raft::copy(actual, neighbors.data_handle(), 4, raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + for (auto id : actual) { + EXPECT_EQ(id, int64_t{-1}); + } +} diff --git a/cpp/tests/neighbors/ann_cagra/flowann_serialize.cu b/cpp/tests/neighbors/ann_cagra/flowann_serialize.cu new file mode 100644 index 0000000000..bc28e97bae --- /dev/null +++ b/cpp/tests/neighbors/ann_cagra/flowann_serialize.cu @@ -0,0 +1,833 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * FlowANN is described in "Disentangling Graph Dependencies for Efficient + * Billion-Scale GPU Vector Search" (OSDI 2026). + */ + +#include + +#include +#include +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace flowann = cuvs::neighbors::cagra::experimental::flowann; + +namespace { + +template +void write_dtype_prefix(std::ostream& os) +{ + auto dtype = raft::numpy_serializer::get_numpy_dtype().to_string(); + dtype.resize(4); + os.write(dtype.data(), dtype.size()); +} + +template +void write_host_matrix(std::ostream& os, raft::host_matrix const& matrix) +{ + raft::numpy_serializer::write_header( + os, + {raft::numpy_serializer::get_numpy_dtype(), + false, + {static_cast(matrix.extent(0)), static_cast(matrix.extent(1))}}); + os.write(reinterpret_cast(matrix.data_handle()), matrix.size() * sizeof(T)); +} + +void write_host_f16_bits(std::ostream& os, raft::host_matrix const& matrix) +{ + raft::numpy_serializer::write_header( + os, + {{'<', 'e', sizeof(std::uint16_t)}, + false, + {static_cast(matrix.extent(0)), static_cast(matrix.extent(1))}}); + os.write(reinterpret_cast(matrix.data_handle()), + matrix.size() * sizeof(std::uint16_t)); +} + +auto f16_bits(float value) -> std::uint16_t +{ + auto encoded = __float2half(value); + std::uint16_t result{}; + static_assert(sizeof(result) == sizeof(encoded)); + std::memcpy(&result, &encoded, sizeof(result)); + return result; +} + +void write_legacy_dataset_body(raft::resources const& res, + std::ostream& stream, + bool invalid_vq_label = false, + std::uint32_t encoded_length = 6) +{ + constexpr int64_t rows = 4; + constexpr std::uint32_t dim = 4; + constexpr std::uint32_t vq_centers = 2; + constexpr std::uint32_t pq_centers = 256; + constexpr std::uint32_t pq_len = 2; + + raft::serialize_scalar(res, stream, flowann::detail::legacy_vpq_dataset_tag); + raft::serialize_scalar(res, stream, CUDA_R_16F); + raft::serialize_scalar(res, stream, rows); + raft::serialize_scalar(res, stream, dim); + raft::serialize_scalar(res, stream, vq_centers); + raft::serialize_scalar(res, stream, pq_centers); + raft::serialize_scalar(res, stream, pq_len); + raft::serialize_scalar(res, stream, encoded_length); + + auto vq = raft::make_host_matrix(vq_centers, dim); + auto pq = raft::make_host_matrix(pq_centers, pq_len); + std::fill_n(vq.data_handle(), vq.size(), f16_bits(0.0f)); + std::fill_n(pq.data_handle(), pq.size(), f16_bits(0.0f)); + vq(1, 0) = f16_bits(10.0f); + pq(1, 0) = f16_bits(1.0f); + write_host_f16_bits(stream, vq); + write_host_f16_bits(stream, pq); + + auto codes = raft::make_host_matrix(rows, encoded_length); + for (int64_t row = 0; row < rows; ++row) { + codes(row, 0) = static_cast(invalid_vq_label && row == 0 ? vq_centers : row / 2); + codes(row, 1) = 0; + codes(row, 2) = static_cast(row / 2); + codes(row, 3) = 0; + codes(row, 4) = static_cast(row % 2); + codes(row, 5) = 0; + } + write_host_matrix(stream, codes); + + constexpr std::uint32_t n_subgraphs = 2; + raft::serialize_scalar(res, stream, n_subgraphs); + auto offsets = raft::make_host_matrix(1, n_subgraphs); + offsets(0, 0) = 0; + offsets(0, 1) = 2; + write_host_matrix(stream, offsets); +} + +auto make_legacy_dataset(raft::resources const& res, + bool truncate = false, + bool invalid_vq_label = false) -> std::stringstream +{ + std::stringstream stream(std::ios::in | std::ios::out | std::ios::binary); + write_dtype_prefix(stream); + raft::serialize_scalar(res, stream, flowann::detail::legacy_serialization_version); + write_legacy_dataset_body(res, stream, invalid_vq_label); + auto bytes = stream.str(); + if (truncate) { bytes.resize(bytes.size() - 1); } + return std::stringstream(bytes, std::ios::in | std::ios::out | std::ios::binary); +} + +enum class legacy_metric_header { typed_enum, historical_unsigned_enum }; + +void write_legacy_graph_body(raft::resources const& res, + std::ostream& stream, + legacy_metric_header metric_header = legacy_metric_header::typed_enum) +{ + constexpr std::uint32_t rows = 4; + constexpr std::uint32_t dim = 4; + constexpr std::uint32_t degree = 2; + constexpr std::uint16_t nodes_per = 2; + constexpr std::uint16_t index_bits = 2; + + raft::serialize_scalar(res, stream, rows); + raft::serialize_scalar(res, stream, dim); + raft::serialize_scalar(res, stream, degree); + if (metric_header == legacy_metric_header::historical_unsigned_enum) { + raft::serialize_scalar( + res, stream, static_cast(cuvs::distance::DistanceType::L2Expanded)); + } else { + raft::serialize_scalar(res, stream, cuvs::distance::DistanceType::L2Expanded); + } + raft::serialize_scalar(res, stream, nodes_per); + raft::serialize_scalar(res, stream, index_bits); + + auto inner = raft::make_host_matrix(2, 4); + inner(0, 0) = 1; + inner(0, 1) = 1; + inner(0, 2) = 1; + inner(0, 3) = 0; + inner(1, 0) = 1; + inner(1, 1) = 1; + inner(1, 2) = 1; + inner(1, 3) = 0; + write_host_matrix(stream, inner); + + auto cross = raft::make_host_matrix(rows, degree + 1); + for (std::uint32_t row = 0; row < rows; ++row) { + cross(row, 0) = 1; + cross(row, 1) = (row + 2) % rows; + cross(row, 2) = 0; + } + write_host_matrix(stream, cross); + + auto mapping = raft::make_host_matrix(1, rows); + mapping(0, 0) = 3; + mapping(0, 1) = 2; + mapping(0, 2) = 1; + mapping(0, 3) = 0; + write_host_matrix(stream, mapping); +} + +auto make_legacy_graph(raft::resources const& res) -> std::stringstream +{ + std::stringstream stream(std::ios::in | std::ios::out | std::ios::binary); + write_dtype_prefix(stream); + raft::serialize_scalar(res, stream, flowann::detail::legacy_serialization_version); + write_legacy_graph_body(res, stream); + stream.seekg(0); + return stream; +} + +auto make_legacy_combined_vpq(raft::resources const& res, std::uint32_t row_length = 6) + -> std::stringstream +{ + std::stringstream stream(std::ios::in | std::ios::out | std::ios::binary); + char const historical_dtype[4] = {'|', 'u', '1', '\0'}; + stream.write(historical_dtype, sizeof(historical_dtype)); + raft::serialize_scalar(res, stream, flowann::detail::legacy_serialization_version); + write_legacy_graph_body(res, stream, legacy_metric_header::historical_unsigned_enum); + raft::serialize_scalar(res, stream, true); + write_legacy_dataset_body(res, stream, false, row_length); + stream.seekg(0); + return stream; +} + +struct compact_test_graph { + raft::host_matrix inner; + raft::host_matrix cross; +}; + +template +auto install_compact_test_graph(raft::resources const& res, IndexT& index) -> compact_test_graph +{ + auto inner = raft::make_host_matrix(2, 3); + std::fill_n(inner.data_handle(), inner.size(), std::uint8_t{0}); + inner(0, 0) = 1; + inner(0, 1) = 1; + inner(1, 0) = 1; + inner(1, 1) = 1; + auto cross = raft::make_host_matrix(4, 2); + for (std::uint32_t row = 0; row < 4; ++row) { + cross(row, 0) = 1; + cross(row, 1) = (row + 1) % 4; + } + index.update_graph( + res, raft::make_const_mdspan(inner.view()), raft::make_const_mdspan(cross.view()), 2, 1); + return {std::move(inner), std::move(cross)}; +} + +TEST(FlowannSerialize, LoadsAndNormalizesLegacyVpq) +{ + raft::resources res; + auto stream = make_legacy_dataset(res); + auto data = flowann::detail::deserialize_legacy_vpq_f16(res, stream); + raft::resource::sync_stream(res); + + ASSERT_EQ(data.dataset->n_rows(), 4); + ASSERT_EQ(data.dataset->dim(), 4); + ASSERT_EQ(data.subgraph_offsets.extent(0), 2); + EXPECT_EQ(data.subgraph_offsets(0), 0); + EXPECT_EQ(data.subgraph_offsets(1), 2); + EXPECT_EQ(data.subgraph_ids(0), 0); + EXPECT_EQ(data.subgraph_ids(1), 0); + EXPECT_EQ(data.subgraph_ids(2), 1); + EXPECT_EQ(data.subgraph_ids(3), 1); + + auto host_codes = + raft::make_host_matrix(4, data.dataset->data_view().extent(1)); + raft::copy(host_codes.data_handle(), + data.dataset->data_view().data_handle(), + host_codes.size(), + raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + for (int64_t row = 0; row < host_codes.extent(0); ++row) { + EXPECT_EQ(host_codes(row, 2), 0); + EXPECT_EQ(host_codes(row, 3), 0); + } +} + +TEST(FlowannSerialize, RejectsTruncatedLegacyVpq) +{ + raft::resources res; + auto stream = make_legacy_dataset(res, true); + EXPECT_ANY_THROW(flowann::detail::deserialize_legacy_vpq_f16(res, stream)); +} + +TEST(FlowannSerialize, RejectsOutOfRangeLegacyVqLabel) +{ + raft::resources res; + auto stream = make_legacy_dataset(res, false, true); + EXPECT_ANY_THROW(flowann::detail::deserialize_legacy_vpq_f16(res, stream)); +} + +TEST(FlowannSerialize, LoadsLegacyTieredGraph) +{ + raft::resources res; + auto stream = make_legacy_graph(res); + flowann::device_padded_index index(res); + flowann::detail::deserialize_legacy_graph(res, stream, &index); + EXPECT_EQ(index.size(), 4); + EXPECT_EQ(index.graph_degree(), 2); + EXPECT_EQ(index.node_per_cacheline(), 2); + EXPECT_EQ(index.n_bits(), 2); + ASSERT_TRUE(index.source_indices().has_value()); +} + +TEST(FlowannSerialize, LoadsLegacyCombinedV5Index) +{ + raft::resources res; + auto stream = make_legacy_combined_vpq(res); + auto bundle = flowann::deserialize_vpq_f16(res, stream); + raft::resource::sync_stream(res); + + EXPECT_EQ(bundle.index.size(), 4); + EXPECT_EQ(bundle.index.dim(), 4); + EXPECT_EQ(bundle.index.graph_degree(), 2); + EXPECT_EQ(bundle.index.node_per_cacheline(), 2); + EXPECT_EQ(bundle.index.n_bits(), 2); + ASSERT_TRUE(bundle.index.source_indices().has_value()); + ASSERT_TRUE(bundle.index.subgraph_offsets().has_value()); + ASSERT_TRUE(bundle.index.subgraph_ids().has_value()); + EXPECT_EQ(bundle.index.subgraph_offsets()->extent(0), 2); + EXPECT_EQ(bundle.index.subgraph_ids()->extent(0), 4); +} + +TEST(FlowannSerialize, RoundTripsCurrentVpqFormat) +{ + raft::resources res; + constexpr std::uint32_t rows = 4; + constexpr std::uint32_t dim = 4; + constexpr std::uint32_t vq_centers = 2; + constexpr std::uint32_t pq_centers = 256; + constexpr std::uint32_t pq_len = 2; + constexpr std::uint32_t row_length = 8; + + auto host_vq = raft::make_host_matrix(vq_centers, dim); + auto host_pq = raft::make_host_matrix(pq_centers, pq_len); + auto host_encoded = raft::make_host_matrix(rows, row_length); + std::fill_n(host_vq.data_handle(), host_vq.size(), __float2half(0.0f)); + std::fill_n(host_pq.data_handle(), host_pq.size(), __float2half(0.0f)); + std::fill_n(host_encoded.data_handle(), host_encoded.size(), std::uint8_t{0}); + host_vq(1, 0) = __float2half(10.0f); + host_pq(1, 0) = __float2half(1.0f); + for (std::uint32_t row = 0; row < rows; ++row) { + host_encoded(row, 0) = static_cast(row / 2); + host_encoded(row, 4) = static_cast(row % 2); + } + + auto device_vq = raft::make_device_matrix(res, vq_centers, dim); + auto device_pq = raft::make_device_matrix(res, pq_centers, pq_len); + auto device_encoded = raft::make_device_matrix(res, rows, row_length); + auto const stream = raft::resource::get_cuda_stream(res); + raft::copy(device_vq.data_handle(), host_vq.data_handle(), host_vq.size(), stream); + raft::copy(device_pq.data_handle(), host_pq.data_handle(), host_pq.size(), stream); + raft::copy(device_encoded.data_handle(), host_encoded.data_handle(), host_encoded.size(), stream); + auto dataset = std::make_unique>( + std::move(device_encoded), + cuvs::neighbors::device_vpq_dataset::manifest_type{std::move(device_vq), + std::move(device_pq)}); + flowann::vpq_f16_index index( + res, cuvs::distance::DistanceType::L2Expanded, dataset->view(), 2, 2); + + auto graph = install_compact_test_graph(res, index); + + auto source_indices = raft::make_host_vector(rows); + auto subgraph_ids = raft::make_host_vector(rows); + for (std::uint32_t row = 0; row < rows; ++row) { + source_indices(row) = rows - row - 1; + subgraph_ids(row) = row / 2; + } + index.update_source_indices(res, raft::make_const_mdspan(source_indices.view())); + auto offsets = raft::make_host_vector(2); + offsets(0) = 0; + offsets(1) = 2; + index.update_subgraph_layout( + res, raft::make_const_mdspan(offsets.view()), raft::make_const_mdspan(subgraph_ids.view())); + auto seeds = raft::make_host_vector(2); + seeds(0) = 1; + seeds(1) = 2; + index.update_seeds(res, raft::make_const_mdspan(seeds.view())); + + std::stringstream serialized(std::ios::in | std::ios::out | std::ios::binary); + flowann::serialize(res, serialized, index); + serialized.seekg(0); + auto restored = flowann::deserialize_vpq_f16(res, serialized); + raft::resource::sync_stream(res); + + EXPECT_EQ(restored.index.size(), rows); + EXPECT_EQ(restored.index.dim(), dim); + EXPECT_EQ(restored.index.graph_degree(), 2); + EXPECT_EQ(restored.index.resident_degree(), 1); + EXPECT_EQ(restored.index.cross_degree(), 1); + EXPECT_EQ(restored.index.metric(), cuvs::distance::DistanceType::L2Expanded); + EXPECT_EQ(restored.index.node_per_cacheline(), 2); + EXPECT_EQ(restored.index.n_bits(), 2); + ASSERT_TRUE(restored.index.source_indices().has_value()); + ASSERT_TRUE(restored.index.subgraph_offsets().has_value()); + ASSERT_TRUE(restored.index.subgraph_ids().has_value()); + ASSERT_TRUE(restored.index.seeds().has_value()); + EXPECT_EQ(restored.index.source_indices()->extent(0), rows); + EXPECT_EQ(restored.index.subgraph_offsets()->extent(0), 2); + EXPECT_EQ(restored.index.subgraph_ids()->extent(0), rows); + EXPECT_EQ(restored.index.seeds()->extent(0), 2); + + auto restored_inner = raft::make_host_matrix(2, 3); + auto restored_source_indices = raft::make_host_vector(rows); + auto restored_offsets = raft::make_host_vector(2); + auto restored_subgraph_ids = raft::make_host_vector(rows); + auto restored_seeds = raft::make_host_vector(2); + auto restored_vq = raft::make_host_matrix(vq_centers, dim); + auto restored_pq = raft::make_host_matrix(pq_centers, pq_len); + auto restored_encoded = raft::make_host_matrix(rows, row_length); + raft::copy(restored_inner.data_handle(), + restored.index.inner_graph().data_handle(), + restored_inner.size(), + stream); + raft::copy(restored_source_indices.data_handle(), + restored.index.source_indices()->data_handle(), + restored_source_indices.size(), + stream); + raft::copy(restored_offsets.data_handle(), + restored.index.subgraph_offsets()->data_handle(), + restored_offsets.size(), + stream); + raft::copy(restored_subgraph_ids.data_handle(), + restored.index.subgraph_ids()->data_handle(), + restored_subgraph_ids.size(), + stream); + raft::copy(restored_seeds.data_handle(), + restored.index.seeds()->data_handle(), + restored_seeds.size(), + stream); + raft::copy(restored_vq.data_handle(), + restored.dataset->manifest_view().vq_code_book.data_handle(), + restored_vq.size(), + stream); + raft::copy(restored_pq.data_handle(), + restored.dataset->manifest_view().pq_code_book.data_handle(), + restored_pq.size(), + stream); + raft::copy(restored_encoded.data_handle(), + restored.dataset->data_view().data_handle(), + restored_encoded.size(), + stream); + raft::resource::sync_stream(res); + EXPECT_EQ(std::memcmp(graph.inner.data_handle(), + restored_inner.data_handle(), + graph.inner.size() * sizeof(std::uint8_t)), + 0); + EXPECT_EQ(std::memcmp(graph.cross.data_handle(), + restored.index.cross_graph().data_handle(), + graph.cross.size() * sizeof(std::uint32_t)), + 0); + EXPECT_EQ(std::memcmp(source_indices.data_handle(), + restored_source_indices.data_handle(), + source_indices.size() * sizeof(std::uint32_t)), + 0); + EXPECT_EQ(std::memcmp(offsets.data_handle(), + restored_offsets.data_handle(), + offsets.size() * sizeof(std::uint32_t)), + 0); + EXPECT_EQ(std::memcmp(subgraph_ids.data_handle(), + restored_subgraph_ids.data_handle(), + subgraph_ids.size() * sizeof(std::uint32_t)), + 0); + EXPECT_EQ( + std::memcmp( + seeds.data_handle(), restored_seeds.data_handle(), seeds.size() * sizeof(std::uint32_t)), + 0); + EXPECT_EQ( + std::memcmp(host_vq.data_handle(), restored_vq.data_handle(), host_vq.size() * sizeof(half)), + 0); + EXPECT_EQ( + std::memcmp(host_pq.data_handle(), restored_pq.data_handle(), host_pq.size() * sizeof(half)), + 0); + EXPECT_EQ(std::memcmp(host_encoded.data_handle(), + restored_encoded.data_handle(), + host_encoded.size() * sizeof(std::uint8_t)), + 0); +} + +TEST(FlowannSerialize, RoundTripsCompactDegreeMetadata) +{ + raft::resources res; + auto host_dataset = raft::make_host_matrix(4, 4); + std::fill_n(host_dataset.data_handle(), host_dataset.size(), 0.0f); + auto dataset = cuvs::neighbors::make_device_padded_dataset(res, host_dataset.view()); + flowann::device_padded_index index( + res, cuvs::distance::DistanceType::L2Expanded, dataset->view(), 2, 2); + + install_compact_test_graph(res, index); + auto offsets = raft::make_host_vector(1); + auto ids = raft::make_host_vector(4); + offsets(0) = 0; + std::fill_n(ids.data_handle(), ids.size(), std::uint32_t{0}); + index.update_subgraph_layout( + res, raft::make_const_mdspan(offsets.view()), raft::make_const_mdspan(ids.view())); + + std::stringstream stream(std::ios::in | std::ios::out | std::ios::binary); + flowann::serialize(res, stream, index); + stream.seekg(0); + auto restored = flowann::deserialize_device_padded(res, stream); + + EXPECT_EQ(restored.index.graph_degree(), 2); + EXPECT_EQ(restored.index.resident_degree(), 1); + EXPECT_EQ(restored.index.cross_degree(), 1); + EXPECT_EQ(restored.index.cross_graph().extent(1), 2); +} + +TEST(FlowannSerialize, LoadsPreviousSelfContainedFormat) +{ + raft::resources res; + constexpr std::uint32_t rows = 4; + constexpr std::uint32_t dim = 4; + constexpr std::uint32_t degree = 2; + constexpr std::uint16_t nodes_per = 2; + constexpr std::uint16_t index_bits = 2; + std::stringstream stream(std::ios::in | std::ios::out | std::ios::binary); + write_dtype_prefix(stream); + raft::serialize_scalar(res, stream, flowann::detail::previous_serialization_version); + raft::serialize_scalar( + res, stream, static_cast(flowann::detail::serialized_dataset_kind::dense)); + raft::serialize_scalar(res, stream, rows); + raft::serialize_scalar(res, stream, dim); + raft::serialize_scalar(res, stream, degree); + raft::serialize_scalar(res, stream, cuvs::distance::DistanceType::L2Expanded); + raft::serialize_scalar(res, stream, nodes_per); + raft::serialize_scalar(res, stream, index_bits); + raft::serialize_scalar(res, stream, std::uint32_t{3}); + raft::serialize_scalar(res, stream, std::uint32_t{1}); + raft::serialize_scalar(res, stream, std::uint32_t{0}); + raft::serialize_scalar(res, stream, std::uint32_t{0}); + + auto inner = raft::make_host_matrix(2, 3); + std::fill_n(inner.data_handle(), inner.size(), std::uint8_t{0}); + inner(0, 0) = 1; + inner(0, 1) = 1; + inner(1, 0) = 1; + inner(1, 1) = 1; + auto cross = raft::make_host_matrix(rows, degree + 1); + for (std::uint32_t row = 0; row < rows; ++row) { + cross(row, 0) = 1; + cross(row, 1) = (row + 1) % rows; + cross(row, 2) = 0; + } + auto offsets = raft::make_host_vector(1); + auto ids = raft::make_host_vector(rows); + offsets(0) = 0; + std::fill_n(ids.data_handle(), ids.size(), std::uint32_t{0}); + raft::serialize_mdspan(res, stream, inner.view()); + raft::serialize_mdspan(res, stream, cross.view()); + raft::serialize_mdspan(res, stream, offsets.view()); + raft::serialize_mdspan(res, stream, ids.view()); + + auto dataset = raft::make_host_matrix(rows, dim); + std::fill_n(dataset.data_handle(), dataset.size(), 0.0f); + raft::serialize_scalar(res, stream, dim); + raft::serialize_mdspan(res, stream, dataset.view()); + stream.seekg(0); + + auto restored = flowann::deserialize_device_padded(res, stream); + EXPECT_EQ(restored.index.graph_degree(), degree); + EXPECT_EQ(restored.index.resident_degree(), 0); + EXPECT_EQ(restored.index.cross_degree(), degree); +} + +TEST(FlowannSerialize, SearchesLegacyCombinedVpqWithAndWithoutFixedSeeds) +{ + raft::resources res; + auto combined_stream = make_legacy_combined_vpq(res); + auto bundle = flowann::deserialize_vpq_f16(res, combined_stream); + auto& index = bundle.index; + + auto host_queries = raft::make_host_matrix(2, 4); + std::fill_n(host_queries.data_handle(), host_queries.size(), 0.0f); + host_queries(0, 0) = 0.1f; + host_queries(1, 0) = 10.9f; + auto queries = raft::make_device_matrix(res, 2, 4); + 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, 2, 1); + auto distances = raft::make_device_matrix(res, 2, 1); + flowann::search_context context(index.cross_graph()); + + flowann::search_params params{}; + params.algo = cuvs::neighbors::cagra::search_algo::SINGLE_CTA; + params.itopk_size = 32; + params.search_width = 1; + params.max_iterations = 8; + params.max_queries = 2; + params.smem_dtype = cuvs::neighbors::cagra::internal_dtype::F16; + + for (bool use_seeds : {false, true}) { + if (use_seeds) { + auto seeds = raft::make_host_vector(2); + seeds(0) = 1; + seeds(1) = 2; + index.update_seeds( + res, raft::make_host_vector_view(seeds.data_handle(), 2)); + params.num_seeds = 2; + } + + flowann::search(res, + params, + index, + context, + raft::make_const_mdspan(queries.view()), + neighbors.view(), + distances.view()); + auto host_neighbors = raft::make_host_matrix(2, 1); + auto host_distances = raft::make_host_matrix(2, 1); + 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); + + EXPECT_EQ(host_neighbors(0, 0), 3); + EXPECT_EQ(host_neighbors(1, 0), 0); + EXPECT_NEAR(host_distances(0, 0), 0.01f, 1e-3f); + EXPECT_NEAR(host_distances(1, 0), 0.01f, 1e-3f); + } +} + +} // namespace + +TEST(FlowannSerialize, RejectsCurrentOutOfRangeSourceIds) +{ + raft::resources res; + auto legacy = make_legacy_combined_vpq(res); + auto bundle = flowann::deserialize_vpq_f16(res, legacy); + std::uint32_t ids[] = {4, 1, 2, 3}; + bundle.index.update_source_indices( + res, raft::make_host_vector_view(ids, 4)); + std::stringstream current(std::ios::in | std::ios::out | std::ios::binary); + flowann::serialize(res, current, bundle.index); + current.seekg(0); + EXPECT_ANY_THROW(flowann::deserialize_vpq_f16(res, current)); +} + +TEST(FlowannSerialize, RejectsCurrentOutOfRangeVqLabels) +{ + raft::resources res; + auto legacy = make_legacy_combined_vpq(res); + auto bundle = flowann::deserialize_vpq_f16(res, legacy); + std::uint32_t invalid = bundle.dataset->const_manifest_view().vq_n_centers(); + raft::copy(bundle.dataset->data_view().data_handle(), + reinterpret_cast(&invalid), + sizeof(invalid), + raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + std::stringstream current(std::ios::in | std::ios::out | std::ios::binary); + flowann::serialize(res, current, bundle.index); + current.seekg(0); + EXPECT_ANY_THROW(flowann::deserialize_vpq_f16(res, current)); +} + +TEST(FlowannSerialize, NormalizesLegacyRowPaddingForCurrentRoundTrip) +{ + raft::resources res; + auto legacy = make_legacy_combined_vpq(res, 12); + auto bundle = flowann::deserialize_vpq_f16(res, legacy); + EXPECT_EQ(bundle.dataset->data_view().extent(1), 8); + std::stringstream current(std::ios::in | std::ios::out | std::ios::binary); + flowann::serialize(res, current, bundle.index); + current.seekg(0); + auto restored = flowann::deserialize_vpq_f16(res, current); + EXPECT_EQ(restored.dataset->data_view().extent(1), 8); + EXPECT_EQ(restored.index.size(), bundle.index.size()); +} + +TEST(FlowannSerialize, RejectsUnsupportedLegacyMetrics) +{ + raft::resources res; + for (auto metric : {cuvs::distance::DistanceType::InnerProduct, + static_cast(999)}) { + auto stream = make_legacy_graph(res); + // Every scalar has its own NumPy header. Walk the header to replace only the metric. + stream.seekg(4); + (void)raft::deserialize_scalar(res, stream); + for (int field = 0; field < 3; ++field) { + (void)raft::deserialize_scalar(res, stream); + } + stream.seekp(stream.tellg()); + raft::serialize_scalar(res, stream, metric); + stream.seekg(0); + flowann::vpq_f16_index index(res); + EXPECT_ANY_THROW(flowann::detail::deserialize_legacy_graph(res, stream, &index)); + } +} + +TEST(FlowannSerialize, RoundTripsPaddedDenseRows) +{ + raft::resources res; + auto host = raft::make_host_matrix(4, 5); + for (int64_t row = 0; row < 4; ++row) { + for (int64_t col = 0; col < 5; ++col) { + host(row, col) = 10.0f * row + col; + } + } + auto dataset = cuvs::neighbors::make_device_padded_dataset(res, host.view()); + flowann::device_padded_index index( + res, cuvs::distance::DistanceType::CosineExpanded, dataset->view(), 2, 2); + install_compact_test_graph(res, index); + auto offsets = raft::make_host_vector(1); + auto ids = raft::make_host_vector(4); + offsets(0) = 0; + std::fill_n(ids.data_handle(), ids.size(), std::uint32_t{0}); + index.update_subgraph_layout( + res, raft::make_const_mdspan(offsets.view()), raft::make_const_mdspan(ids.view())); + std::stringstream stream(std::ios::in | std::ios::out | std::ios::binary); + flowann::serialize(res, stream, index); + stream.seekg(0); + auto restored = flowann::deserialize_device_padded(res, stream); + EXPECT_EQ(restored.index.dim(), 5); + auto rows = restored.dataset->data_view(); + EXPECT_EQ(rows.extent(1), 5); + EXPECT_EQ(rows.stride(0), 8); + auto actual = raft::make_host_matrix(4, 5); + raft::copy_matrix(actual.data_handle(), + int64_t{5}, + rows.data_handle(), + rows.stride(0), + int64_t{5}, + int64_t{4}, + raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + for (int64_t row = 0; row < 4; ++row) { + for (int64_t col = 0; col < 5; ++col) { + EXPECT_EQ(actual(row, col), host(row, col)); + } + } + auto physical = raft::make_host_matrix(4, 8); + raft::copy(physical.data_handle(), + rows.data_handle(), + physical.size(), + raft::resource::get_cuda_stream(res)); + float expected_norms[4], actual_norms[4]; + raft::copy( + expected_norms, index.dataset_norms()->data_handle(), 4, raft::resource::get_cuda_stream(res)); + raft::copy(actual_norms, + restored.index.dataset_norms()->data_handle(), + 4, + raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + for (int64_t row = 0; row < 4; ++row) { + EXPECT_FLOAT_EQ(actual_norms[row], expected_norms[row]); + for (int64_t col = 5; col < 8; ++col) { + EXPECT_EQ(physical(row, col), 0.0f); + } + } + std::stringstream reserialized(std::ios::in | std::ios::out | std::ios::binary); + flowann::serialize(res, reserialized, restored.index); + EXPECT_EQ(reserialized.str(), stream.str()); +} + +TEST(FlowannSerialize, LoadsDenseRowsWithOlderPhysicalPitch) +{ + raft::resources res; + auto host = raft::make_host_matrix(4, 5); + auto physical = raft::make_host_matrix(4, 12); + std::fill_n(physical.data_handle(), physical.size(), 9.0f); + for (int64_t row = 0; row < 4; ++row) { + for (int64_t col = 0; col < 5; ++col) { + host(row, col) = physical(row, col) = 1.0f + row + col; + } + } + auto dataset = cuvs::neighbors::make_device_padded_dataset(res, host.view()); + flowann::device_padded_index index( + res, cuvs::distance::DistanceType::CosineExpanded, dataset->view(), 2, 2); + install_compact_test_graph(res, index); + auto offsets = raft::make_host_vector(1); + auto ids = raft::make_host_vector(4); + offsets(0) = 0; + std::fill_n(ids.data_handle(), ids.size(), std::uint32_t{0}); + index.update_subgraph_layout( + res, raft::make_const_mdspan(offsets.view()), raft::make_const_mdspan(ids.view())); + std::stringstream stream(std::ios::in | std::ios::out | std::ios::binary); + flowann::detail::serialize_common( + res, stream, index, flowann::detail::serialized_dataset_kind::dense); + raft::serialize_scalar(res, stream, std::uint32_t{12}); + raft::serialize_mdspan(res, stream, physical.view()); + stream.seekg(0); + auto restored = flowann::deserialize_device_padded(res, stream); + auto rows = restored.dataset->data_view(); + EXPECT_EQ(rows.extent(1), 5); + EXPECT_EQ(rows.stride(0), 8); + auto actual = raft::make_host_matrix(4, 8); + raft::copy( + actual.data_handle(), rows.data_handle(), actual.size(), raft::resource::get_cuda_stream(res)); + float expected_norms[4], actual_norms[4]; + raft::copy( + expected_norms, index.dataset_norms()->data_handle(), 4, raft::resource::get_cuda_stream(res)); + raft::copy(actual_norms, + restored.index.dataset_norms()->data_handle(), + 4, + raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + for (int64_t row = 0; row < 4; ++row) { + EXPECT_FLOAT_EQ(actual_norms[row], expected_norms[row]); + for (int64_t col = 0; col < 5; ++col) { + EXPECT_EQ(actual(row, col), host(row, col)); + } + for (int64_t col = 5; col < 8; ++col) { + EXPECT_EQ(actual(row, col), 0.0f); + } + } +} + +TEST(FlowannSerialize, RejectsTruncatedDensePayloadBeforeDeviceAllocation) +{ + raft::resources res; + auto host = raft::make_host_matrix(4, 4); + std::fill_n(host.data_handle(), host.size(), 0.0f); + auto dataset = cuvs::neighbors::make_device_padded_dataset(res, host.view()); + flowann::device_padded_index index( + res, cuvs::distance::DistanceType::L2Expanded, dataset->view(), 2, 2); + install_compact_test_graph(res, index); + auto offsets = raft::make_host_vector(1); + auto ids = raft::make_host_vector(4); + offsets(0) = 0; + std::fill_n(ids.data_handle(), ids.size(), 0u); + index.update_subgraph_layout( + res, raft::make_const_mdspan(offsets.view()), raft::make_const_mdspan(ids.view())); + std::stringstream stream(std::ios::in | std::ios::out | std::ios::binary); + flowann::detail::serialize_common( + res, stream, index, flowann::detail::serialized_dataset_kind::dense); + constexpr std::uint32_t large_dim = 1u << 28; + raft::serialize_scalar(res, stream, large_dim); + // Replace only the logical dimension in the common header; omit the dataset payload. + stream.seekg(4); + (void)raft::deserialize_scalar(res, stream); + (void)raft::deserialize_scalar(res, stream); // kind + (void)raft::deserialize_scalar(res, stream); // rows + stream.seekp(stream.tellg()); + raft::serialize_scalar(res, stream, large_dim); + stream.seekg(0); + rmm::mr::limiting_resource_adaptor limited(rmm::mr::cuda_memory_resource{}, std::size_t{1} << 20); + struct restore_resource { + cuda::mr::any_resource previous; + ~restore_resource() { rmm::mr::set_current_device_resource(previous); } + } restore{rmm::mr::set_current_device_resource(limited)}; + // A bad header must produce a format error, rather than requesting a multi-GB allocation. + EXPECT_THROW((flowann::deserialize_device_padded(res, stream)), raft::logic_error); +} diff --git a/cpp/tests/neighbors/ann_cagra/flowann_storage.cu b/cpp/tests/neighbors/ann_cagra/flowann_storage.cu new file mode 100644 index 0000000000..03f418c29a --- /dev/null +++ b/cpp/tests/neighbors/ann_cagra/flowann_storage.cu @@ -0,0 +1,161 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * FlowANN is described in "Disentangling Graph Dependencies for Efficient + * Billion-Scale GPU Vector Search" (OSDI 2026). + */ + +#include + +#include +#include + +#include + +namespace flowann = cuvs::neighbors::cagra::experimental::flowann; + +TEST(FlowannStorage, OwnsCrossGraphAndPreservesItAcrossMove) +{ + raft::resources res; + flowann::device_padded_index index(res); + auto inner = raft::make_host_matrix(0, 0); + auto cross = raft::make_host_matrix(2, 2); + cross(0, 0) = 1; + cross(0, 1) = 1; + cross(1, 0) = 1; + cross(1, 1) = 0; + index.update_graph( + res, raft::make_const_mdspan(inner.view()), raft::make_const_mdspan(cross.view()), 1, 0); + cross(0, 1) = 0; + auto moved = std::move(index); + EXPECT_EQ(moved.size(), 2); + EXPECT_EQ(moved.graph_degree(), 1); + EXPECT_EQ(moved.resident_degree(), 0); + EXPECT_EQ(moved.cross_graph()(0, 1), 1); + raft::resource::sync_stream(res); +} + +TEST(FlowannStorage, RejectsInvalidCrossEdgesBeforeReplacingGraph) +{ + raft::resources res; + flowann::device_padded_index index(res); + auto inner = raft::make_host_matrix(0, 0); + auto cross = raft::make_host_matrix(2, 2); + cross(0, 0) = 1; + cross(0, 1) = 1; + cross(1, 0) = 1; + cross(1, 1) = 0; + index.update_graph( + res, raft::make_const_mdspan(inner.view()), raft::make_const_mdspan(cross.view()), 1, 0); + cross(0, 1) = 2; + EXPECT_ANY_THROW(index.update_graph( + res, raft::make_const_mdspan(inner.view()), raft::make_const_mdspan(cross.view()), 1, 0)); + EXPECT_EQ(index.cross_graph()(0, 1), 1); + EXPECT_EQ(index.size(), 2); + raft::resource::sync_stream(res); +} + +TEST(FlowannStorage, ComputesCosineNormsWithoutSearchImplementation) +{ + raft::resources res; + auto host = raft::make_host_matrix(2, 2); + host(0, 0) = 3.0f; + host(0, 1) = 4.0f; + host(1, 0) = 0.0f; + host(1, 1) = 2.0f; + auto dataset = cuvs::neighbors::make_device_padded_dataset(res, host.view()); + flowann::device_padded_index index( + res, cuvs::distance::DistanceType::CosineExpanded, dataset->view()); + ASSERT_TRUE(index.dataset_norms().has_value()); + float norms[2]{}; + raft::copy(norms, index.dataset_norms()->data_handle(), 2, raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + EXPECT_FLOAT_EQ(norms[0], 5.0f); + EXPECT_FLOAT_EQ(norms[1], 2.0f); +} + +TEST(FlowannStorage, RejectsStoredSeedsBeforeReplacingGraph) +{ + for (int overload = 0; overload < 4; ++overload) { + raft::resources res; + flowann::device_padded_index index(res); + auto inner = raft::make_host_matrix(0, 0); + auto original = raft::make_host_matrix(4, 2); + for (std::uint32_t row = 0; row < 4; ++row) { + original(row, 0) = 1; + original(row, 1) = (row + 1) % 4; + } + index.update_graph( + res, raft::make_const_mdspan(inner.view()), raft::make_const_mdspan(original.view()), 1, 0); + std::uint32_t seed = 3; + index.update_seeds(res, raft::make_host_vector_view(&seed, 1)); + auto replacement = raft::make_host_matrix(2, 2); + replacement(0, 0) = replacement(1, 0) = 1; + replacement(0, 1) = 1; + replacement(1, 1) = 0; + if (overload == 0) { + EXPECT_ANY_THROW(index.update_graph(res, + raft::make_const_mdspan(inner.view()), + raft::make_const_mdspan(replacement.view()), + 1, + 0)); + } else if (overload == 1) { + EXPECT_ANY_THROW(index.update_graph(res, std::move(inner), std::move(replacement), 1, 0)); + } else { + auto device_inner = raft::make_device_matrix(res, 0, 0); + if (overload == 2) { + EXPECT_ANY_THROW( + index.update_graph(res, std::move(device_inner), std::move(replacement), true)); + } else { + EXPECT_ANY_THROW( + index.update_graph(res, std::move(device_inner), std::move(replacement), 1, 0)); + } + } + ASSERT_EQ(index.size(), 4) << "overload " << overload; + EXPECT_EQ(index.cross_graph()(3, 1), 0); + ASSERT_TRUE(index.seeds().has_value()); + raft::resource::sync_stream(res); + } +} + +TEST(FlowannStorage, RowCountChangeClearsDependentMappings) +{ + raft::resources res; + flowann::device_padded_index index(res); + auto inner = raft::make_host_matrix(0, 0); + auto graph = raft::make_host_matrix(4, 2); + for (std::uint32_t row = 0; row < 4; ++row) { + graph(row, 0) = 1; + graph(row, 1) = (row + 1) % 4; + } + index.update_graph( + res, raft::make_const_mdspan(inner.view()), raft::make_const_mdspan(graph.view()), 1, 0); + std::uint32_t ids[4] = {3, 2, 1, 0}; + std::uint32_t groups[4] = {0, 0, 0, 0}; + std::uint32_t offset = 0; + std::uint32_t seed = 1; + index.update_source_indices(res, + raft::make_host_vector_view(ids, 4)); + index.update_subgraph_layout( + res, + raft::make_host_vector_view(&offset, 1), + raft::make_host_vector_view(groups, 4)); + index.update_seeds(res, raft::make_host_vector_view(&seed, 1)); + auto larger = raft::make_host_matrix(6, 2); + for (std::uint32_t row = 0; row < 6; ++row) { + larger(row, 0) = 1; + larger(row, 1) = (row + 1) % 6; + } + index.update_graph( + res, raft::make_const_mdspan(inner.view()), raft::make_const_mdspan(larger.view()), 1, 0); + EXPECT_EQ(index.size(), 6); + EXPECT_FALSE(index.source_indices().has_value()); + EXPECT_FALSE(index.subgraph_offsets().has_value()); + EXPECT_FALSE(index.subgraph_ids().has_value()); + EXPECT_FALSE(index.seeds().has_value()); + EXPECT_FALSE(index.seed_cross_graph().has_value()); + EXPECT_FALSE(index.seed_lookup_keys().has_value()); + EXPECT_FALSE(index.seed_lookup_values().has_value()); + raft::resource::sync_stream(res); +} diff --git a/examples/cpp/CMakeLists.txt b/examples/cpp/CMakeLists.txt index 053099e429..5338403603 100644 --- a/examples/cpp/CMakeLists.txt +++ b/examples/cpp/CMakeLists.txt @@ -23,6 +23,12 @@ rapids_cuda_init_architectures(test_cuvs) project(test_cuvs LANGUAGES CXX CUDA) find_package(Threads) +option(CUVS_ENABLE_FLOWANN_SEARCH "Build the experimental FlowANN search implementation" OFF) +option(CUVS_ENABLE_FLOWANN_BUILD "Build the experimental FlowANN integrated builder" OFF) +if(CUVS_ENABLE_FLOWANN_BUILD AND NOT CUVS_ENABLE_FLOWANN_SEARCH) + message(FATAL_ERROR "CUVS_ENABLE_FLOWANN_BUILD requires CUVS_ENABLE_FLOWANN_SEARCH") +endif() + set(CMAKE_CXX_STANDARD 20) set(CMAKE_CXX_STANDARD_REQUIRED ON) set(CMAKE_CUDA_STANDARD 20) @@ -55,6 +61,35 @@ add_executable(IVF_PQ_EXAMPLE src/ivf_pq_example.cu) add_executable(VAMANA_EXAMPLE src/vamana_example.cu) add_executable(SCANN_EXAMPLE src/scann_example.cu) add_executable(SCANN_TWO_LEVEL_EXAMPLE src/scann_two_level_example.cu) +# These tools use private implementation headers from the matching cuVS source tree. +if(CUVS_ENABLE_FLOWANN_SEARCH) + get_target_property(flowann_cuvs_imported cuvs::cuvs IMPORTED) + if(flowann_cuvs_imported) + set(flowann_private_includes "${CMAKE_CURRENT_SOURCE_DIR}/../../cpp/src") + else() + get_target_property(flowann_cuvs_source cuvs::cuvs SOURCE_DIR) + set(flowann_private_includes "${flowann_cuvs_source}/src") + endif() +endif() + +if(CUVS_ENABLE_FLOWANN_SEARCH AND NOT TARGET FLOWANN_BENCHMARK) + find_package(OpenMP REQUIRED COMPONENTS CXX CUDA) + add_executable(FLOWANN_BENCHMARK src/flowann_benchmark.cu) + target_include_directories(FLOWANN_BENCHMARK PRIVATE "${flowann_private_includes}") + target_link_libraries( + FLOWANN_BENCHMARK + PRIVATE cuvs::cuvs OpenMP::OpenMP_CXX OpenMP::OpenMP_CUDA + $ $ + ) +endif() +if(CUVS_ENABLE_FLOWANN_BUILD AND NOT TARGET FLOWANN_BUILD_BENCHMARK) + add_executable(FLOWANN_BUILD_BENCHMARK src/flowann_build_benchmark.cu) + target_include_directories(FLOWANN_BUILD_BENCHMARK PRIVATE "${flowann_private_includes}") + target_link_libraries( + FLOWANN_BUILD_BENCHMARK PRIVATE cuvs::cuvs OpenMP::OpenMP_CUDA + $ + ) +endif() # `$` is a generator expression that ensures that targets are # installed in a conda environment, if one exists diff --git a/examples/cpp/src/flowann_benchmark.cu b/examples/cpp/src/flowann_benchmark.cu new file mode 100644 index 0000000000..e7448633eb --- /dev/null +++ b/examples/cpp/src/flowann_benchmark.cu @@ -0,0 +1,856 @@ +/* + * 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 +#ifdef _OPENMP +#include +#endif + +#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 flowann = cuvs::neighbors::cagra::experimental::flowann; +using clock_type = std::chrono::steady_clock; + +namespace { + +struct benchmark_options { + std::string index_path; + std::string legacy_graph_path; + std::string legacy_dataset_path; + std::string base_path; + std::string queries_path; + std::string ground_truth_path; + std::string seeds_path; + + std::uint32_t gpu_id = 0; + std::uint32_t batch_size = 1; + std::uint32_t query_count = 0; + std::uint32_t top_k = 10; + std::uint32_t candidate_count = 32; + std::uint32_t team_size = 0; + std::uint32_t thread_block_size = 0; + std::uint32_t itopk_size = 64; + std::uint32_t search_width = 1; + std::uint32_t min_iterations = 0; + std::uint32_t max_iterations = 0; + std::uint32_t num_queues = 1; + std::uint32_t empty_pause = 64; + std::uint32_t warmup_rounds = 1; + std::uint32_t measured_rounds = 3; + std::uint32_t cpu_threads = 1; + std::uint32_t sync_drop_threshold = 51; + std::uint64_t pool_size_mib = 2048; + float sync_window_scale = 10.0f; + std::optional num_seeds; + + cuvs::neighbors::cagra::search_algo algorithm = cuvs::neighbors::cagra::search_algo::AUTO; + cuvs::neighbors::cagra::internal_dtype smem_dtype = cuvs::neighbors::cagra::internal_dtype::F16; + bool use_search_session = true; + bool rerank = true; + bool collect_statistics = false; + bool include_query_transfer = false; +}; + +template +struct host_matrix_file { + std::uint32_t rows{}; + std::uint32_t cols{}; + std::vector values; +}; + +template +class mapped_matrix_file { + public: + explicit mapped_matrix_file(std::string const& path) + { + fd_ = open(path.c_str(), O_RDONLY); + if (fd_ < 0) { throw std::runtime_error("cannot open " + path); } + struct stat info{}; + if (fstat(fd_, &info) != 0 || info.st_size < 8) { + throw std::runtime_error("cannot stat or invalid matrix file " + path); + } + bytes_ = static_cast(info.st_size); + mapping_ = mmap(nullptr, bytes_, PROT_READ, MAP_PRIVATE, fd_, 0); + if (mapping_ == MAP_FAILED) { + mapping_ = nullptr; + throw std::runtime_error("cannot mmap " + path); + } + std::memcpy(&rows_, mapping_, sizeof(rows_)); + std::memcpy(&cols_, static_cast(mapping_) + sizeof(rows_), sizeof(cols_)); + auto const elements = static_cast(rows_) * cols_; + auto const expected = std::uint64_t{8} + elements * sizeof(T); + if (rows_ == 0 || cols_ == 0 || expected != bytes_) { + throw std::runtime_error("invalid tightly packed matrix file " + path); + } + } + + mapped_matrix_file(mapped_matrix_file const&) = delete; + auto operator=(mapped_matrix_file const&) -> mapped_matrix_file& = delete; + + ~mapped_matrix_file() + { + if (mapping_ != nullptr) { munmap(mapping_, bytes_); } + if (fd_ >= 0) { close(fd_); } + } + + [[nodiscard]] auto rows() const noexcept -> std::uint32_t { return rows_; } + [[nodiscard]] auto cols() const noexcept -> std::uint32_t { return cols_; } + [[nodiscard]] auto data() const noexcept -> T const* + { + return reinterpret_cast(static_cast(mapping_) + 8); + } + + private: + int fd_{-1}; + void* mapping_{}; + std::size_t bytes_{}; + std::uint32_t rows_{}; + std::uint32_t cols_{}; +}; + +struct round_result { + double search_ms{}; + double rerank_prep_ms{}; + double rerank_ms{}; + double total_ms{}; + double recall = std::numeric_limits::quiet_NaN(); + double qps{}; + double p50_ms{}; + double p95_ms{}; + double p99_ms{}; +}; + +class scoped_device_resource { + public: + template + explicit scoped_device_resource(Resource& resource) + : previous_(rmm::mr::set_current_device_resource(resource)) + { + } + + ~scoped_device_resource() { rmm::mr::set_current_device_resource(std::move(previous_)); } + + scoped_device_resource(scoped_device_resource const&) = delete; + auto operator=(scoped_device_resource const&) -> scoped_device_resource& = delete; + + private: + cuda::mr::any_resource previous_; +}; + +void print_usage(char const* program) +{ + std::printf( + "FlowANN VPQ search benchmark\n\n" + "Usage:\n" + " %s --index INDEX --queries QUERIES [options]\n" + " %s --legacy-graph GRAPH --legacy-dataset DATASET --queries QUERIES [options]\n\n" + "Index and input paths:\n" + " --index PATH Self-contained index or legacy combined v5 split.bin\n" + " --legacy-graph PATH Legacy split graph file\n" + " --legacy-dataset PATH Legacy VPQ-F16 dataset-only file\n" + " --queries PATH Query matrix in fbin format\n" + " --ground-truth PATH Optional ground truth matrix in ibin format\n" + " --base PATH Base vectors in fbin format, required for CPU rerank\n" + " --seeds PATH Optional raw uint32 medoid seed file\n\n" + "Search parameters:\n" + " --gpu-id N CUDA device (default: 0)\n" + " --batch-size N Queries per search call (default: 1)\n" + " --query-count N Prefix to benchmark, 0 means all (default: 0)\n" + " --top-k N Recall and final result width (default: 10)\n" + " --candidate-count N FlowANN candidates before rerank (default: 32)\n" + " --algorithm NAME auto, single_cta, or multi_cta (default: auto)\n" + " --team-size N Distance team size, 0 means auto (default: 0)\n" + " --thread-block-size N CUDA threads per block, 0 means auto (default: 0)\n" + " --itopk-size N Internal top-k size (default: 64)\n" + " --search-width N Search width (default: 1)\n" + " --min-iterations N Minimum graph iterations (default: 0)\n" + " --max-iterations N Maximum graph iterations, 0 means auto (default: 0)\n" + " --num-seeds N Number of loaded seeds to use; default is all loaded seeds\n" + " --smem-dtype NAME f16 or e5m2 (default: f16)\n" + " --sync-window-scale X Deferred request window scale (default: 10)\n" + " --sync-drop-threshold N Forced synchronization threshold (default: 51)\n\n" + "Queue, session, and measurement parameters:\n" + " --num-queues N CPU-serviced queue count (default: 1)\n" + " --empty-pause N Pause instructions after an empty poll (default: 64)\n" + " --cpu-threads N OpenMP threads used by CPU rerank (default: 1)\n" + " --warmup-rounds N Complete warmup passes (default: 1)\n" + " --measured-rounds N Complete measured passes (default: 3)\n" + " --pool-size-mib N Managed-memory RMM pool, 0 disables it (default: 2048)\n" + " --queue-statistics Collect and print queue counters\n" + " --no-session Disable the default persistent search_session\n" + " --no-rerank Return FlowANN candidates without CPU exact rerank\n" + " --include-query-transfer Include query upload in search and total latency\n" + " --help Show this message\n\n" + "fbin/ibin files contain uint32 rows and columns followed by row-major values.\n" + "The selected query count must be divisible by the batch size.\n", + program, + program); +} + +template +auto parse_number(std::string_view text, std::string_view option) -> T +{ + T value{}; + auto const* begin = text.data(); + auto const* end = begin + text.size(); + auto [position, error] = std::from_chars(begin, end, value); + if (error != std::errc{} || position != end) { + throw std::invalid_argument("invalid value for " + std::string(option) + ": " + + std::string(text)); + } + return value; +} + +template <> +auto parse_number(std::string_view text, std::string_view option) -> float +{ + std::string value_text{text}; + char* end = nullptr; + auto const value = std::strtof(value_text.c_str(), &end); + if (end == value_text.c_str() || end != value_text.c_str() + value_text.size() || + !std::isfinite(value)) { + throw std::invalid_argument("invalid value for " + std::string(option) + ": " + value_text); + } + return value; +} + +auto parse_algorithm(std::string_view value) -> cuvs::neighbors::cagra::search_algo +{ + if (value == "auto") { return cuvs::neighbors::cagra::search_algo::AUTO; } + if (value == "single_cta") { return cuvs::neighbors::cagra::search_algo::SINGLE_CTA; } + if (value == "multi_cta") { return cuvs::neighbors::cagra::search_algo::MULTI_CTA; } + throw std::invalid_argument("--algorithm must be auto, single_cta, or multi_cta"); +} + +auto parse_smem_dtype(std::string_view value) -> cuvs::neighbors::cagra::internal_dtype +{ + if (value == "f16") { return cuvs::neighbors::cagra::internal_dtype::F16; } + if (value == "e5m2") { return cuvs::neighbors::cagra::internal_dtype::E5M2; } + throw std::invalid_argument("--smem-dtype must be f16 or e5m2"); +} + +auto parse_options(int argc, char** argv) -> std::optional +{ + benchmark_options options; + auto require_value = [&](int& position, std::string_view option) -> std::string_view { + if (++position >= argc) { + throw std::invalid_argument(std::string(option) + " requires a value"); + } + return argv[position]; + }; + + for (int i = 1; i < argc; ++i) { + std::string_view option = argv[i]; + if (option == "--help" || option == "-h") { + print_usage(argv[0]); + return std::nullopt; + } else if (option == "--index") { + options.index_path = require_value(i, option); + } else if (option == "--legacy-graph") { + options.legacy_graph_path = require_value(i, option); + } else if (option == "--legacy-dataset") { + options.legacy_dataset_path = require_value(i, option); + } else if (option == "--base") { + options.base_path = require_value(i, option); + } else if (option == "--queries") { + options.queries_path = require_value(i, option); + } else if (option == "--ground-truth") { + options.ground_truth_path = require_value(i, option); + } else if (option == "--seeds") { + options.seeds_path = require_value(i, option); + } else if (option == "--gpu-id") { + options.gpu_id = parse_number(require_value(i, option), option); + } else if (option == "--batch-size") { + options.batch_size = parse_number(require_value(i, option), option); + } else if (option == "--query-count") { + options.query_count = parse_number(require_value(i, option), option); + } else if (option == "--top-k") { + options.top_k = parse_number(require_value(i, option), option); + } else if (option == "--candidate-count") { + options.candidate_count = parse_number(require_value(i, option), option); + } else if (option == "--algorithm") { + options.algorithm = parse_algorithm(require_value(i, option)); + } else if (option == "--team-size") { + options.team_size = parse_number(require_value(i, option), option); + } else if (option == "--thread-block-size") { + options.thread_block_size = parse_number(require_value(i, option), option); + } else if (option == "--itopk-size") { + options.itopk_size = parse_number(require_value(i, option), option); + } else if (option == "--search-width") { + options.search_width = parse_number(require_value(i, option), option); + } else if (option == "--min-iterations") { + options.min_iterations = parse_number(require_value(i, option), option); + } else if (option == "--max-iterations") { + options.max_iterations = parse_number(require_value(i, option), option); + } else if (option == "--num-seeds") { + options.num_seeds = parse_number(require_value(i, option), option); + } else if (option == "--smem-dtype") { + options.smem_dtype = parse_smem_dtype(require_value(i, option)); + } else if (option == "--sync-window-scale") { + options.sync_window_scale = parse_number(require_value(i, option), option); + } else if (option == "--sync-drop-threshold") { + options.sync_drop_threshold = parse_number(require_value(i, option), option); + } else if (option == "--num-queues") { + options.num_queues = parse_number(require_value(i, option), option); + } else if (option == "--empty-pause") { + options.empty_pause = parse_number(require_value(i, option), option); + } else if (option == "--cpu-threads") { + options.cpu_threads = parse_number(require_value(i, option), option); + } else if (option == "--warmup-rounds") { + options.warmup_rounds = parse_number(require_value(i, option), option); + } else if (option == "--measured-rounds") { + options.measured_rounds = parse_number(require_value(i, option), option); + } else if (option == "--pool-size-mib") { + options.pool_size_mib = parse_number(require_value(i, option), option); + } else if (option == "--queue-statistics") { + options.collect_statistics = true; + } else if (option == "--no-session") { + options.use_search_session = false; + } else if (option == "--no-rerank") { + options.rerank = false; + } else if (option == "--include-query-transfer") { + options.include_query_transfer = true; + } else { + throw std::invalid_argument("unknown option: " + std::string(option)); + } + } + + auto const has_single_file_index = !options.index_path.empty(); + auto const has_legacy_index = + !options.legacy_graph_path.empty() || !options.legacy_dataset_path.empty(); + if (has_single_file_index == has_legacy_index) { + throw std::invalid_argument( + "select exactly one index mode: --index, or --legacy-graph plus --legacy-dataset"); + } + if (has_legacy_index && + (options.legacy_graph_path.empty() || options.legacy_dataset_path.empty())) { + throw std::invalid_argument("legacy index mode requires both files"); + } + if (options.queries_path.empty()) { throw std::invalid_argument("--queries is required"); } + if (options.rerank && options.base_path.empty()) { + throw std::invalid_argument("--base is required unless --no-rerank is used"); + } + if (options.batch_size == 0 || options.top_k == 0 || options.candidate_count == 0 || + options.itopk_size == 0 || options.search_width == 0 || options.num_queues == 0 || + options.cpu_threads == 0 || options.measured_rounds == 0) { + throw std::invalid_argument( + "batch, result, queue, thread, and measured-round values must be positive"); + } + if (options.candidate_count < options.top_k) { + throw std::invalid_argument("--candidate-count must be at least --top-k"); + } + if (options.candidate_count > options.itopk_size) { + throw std::invalid_argument("--candidate-count must not exceed --itopk-size"); + } + if (options.sync_window_scale <= 0.0f) { + throw std::invalid_argument("--sync-window-scale must be positive"); + } + if (options.cpu_threads > static_cast(std::numeric_limits::max())) { + throw std::invalid_argument("--cpu-threads exceeds the OpenMP limit"); + } + if (options.pool_size_mib > std::numeric_limits::max() / (1024ULL * 1024ULL)) { + throw std::invalid_argument("--pool-size-mib is too large"); + } + return options; +} + +template +auto load_matrix_file(std::string const& path, std::uint32_t row_limit = 0) -> host_matrix_file +{ + std::ifstream input(path, std::ios::binary); + if (!input) { throw std::runtime_error("cannot open " + path); } + + host_matrix_file result; + input.read(reinterpret_cast(&result.rows), sizeof(result.rows)); + input.read(reinterpret_cast(&result.cols), sizeof(result.cols)); + if (!input || result.rows == 0 || result.cols == 0) { + throw std::runtime_error("invalid matrix header in " + path); + } + if (row_limit != 0) { result.rows = std::min(result.rows, row_limit); } + + auto const count = static_cast(result.rows) * result.cols; + if (count > std::numeric_limits::max() / sizeof(T)) { + throw std::runtime_error("matrix size overflows in " + path); + } + result.values.resize(count); + input.read(reinterpret_cast(result.values.data()), + static_cast(count * sizeof(T))); + if (!input) { throw std::runtime_error("truncated matrix in " + path); } + return result; +} + +auto load_u32_vector(std::string const& path) -> std::vector +{ + std::ifstream input(path, std::ios::binary | std::ios::ate); + if (!input) { throw std::runtime_error("cannot open " + path); } + auto const bytes = input.tellg(); + if (bytes <= 0 || bytes % sizeof(std::uint32_t) != 0) { + throw std::runtime_error("invalid raw uint32 file " + path); + } + std::vector values(static_cast(bytes) / sizeof(std::uint32_t)); + input.seekg(0); + input.read(reinterpret_cast(values.data()), static_cast(bytes)); + if (!input) { throw std::runtime_error("truncated raw uint32 file " + path); } + return values; +} + +auto milliseconds(clock_type::time_point begin, clock_type::time_point end) -> double +{ + return std::chrono::duration(end - begin).count(); +} + +auto percentile(std::vector values, double fraction) -> double +{ + if (values.empty()) { return 0.0; } + std::sort(values.begin(), values.end()); + auto const rank = + static_cast(std::ceil(fraction * static_cast(values.size()))); + return values[std::min(values.size() - 1, std::max(1, rank) - 1)]; +} + +auto median(std::vector values) -> double +{ + if (values.empty()) { return 0.0; } + std::sort(values.begin(), values.end()); + auto const middle = values.size() / 2; + return values.size() % 2 == 0 ? (values[middle - 1] + values[middle]) / 2.0 : values[middle]; +} + +auto recall_at_k(std::vector const& actual, + host_matrix_file const& truth, + std::uint32_t top_k) -> double +{ + std::uint64_t matches = 0; + auto const query_count = actual.size() / top_k; + for (std::size_t query = 0; query < query_count; ++query) { + for (std::uint32_t i = 0; i < top_k; ++i) { + auto const id = actual[query * top_k + i]; + for (std::uint32_t j = 0; j < top_k; ++j) { + if (id == truth.values[query * truth.cols + j]) { + ++matches; + break; + } + } + } + } + return static_cast(matches) / static_cast(query_count * top_k); +} + +auto algorithm_name(cuvs::neighbors::cagra::search_algo algorithm) -> char const* +{ + switch (algorithm) { + case cuvs::neighbors::cagra::search_algo::SINGLE_CTA: return "single_cta"; + case cuvs::neighbors::cagra::search_algo::MULTI_CTA: return "multi_cta"; + case cuvs::neighbors::cagra::search_algo::AUTO: return "auto"; + default: return "unsupported"; + } +} + +auto run_round(raft::device_resources& resources, + cuvs::neighbors::cagra::device_pq_index& index, + benchmark_options const& options, + mapped_matrix_file const* base, + host_matrix_file const& queries, + std::optional> const& truth) -> round_result +{ + auto const batch_count = queries.rows / options.batch_size; + cuvs::neighbors::cagra::search_params params{}; + params.tiered.emplace(); + params.tiered->num_queues = options.num_queues; + params.tiered->empty_pause = options.empty_pause; + params.tiered->collect_statistics = options.collect_statistics; + params.tiered->keep_pollers_running = options.use_search_session; + params.algo = options.algorithm; + params.team_size = options.team_size; + params.thread_block_size = options.thread_block_size; + params.search_width = options.search_width; + params.itopk_size = options.itopk_size; + params.min_iterations = options.min_iterations; + params.max_iterations = options.max_iterations; + params.max_queries = options.batch_size; + params.tiered->num_seeds = options.num_seeds.value_or(0); + params.smem_dtype = options.smem_dtype; + params.tiered->sync_window_scale = options.sync_window_scale; + params.tiered->sync_drop_threshold = options.sync_drop_threshold; + + auto query_device = + raft::make_device_matrix(resources, options.batch_size, queries.cols); + auto candidates_device = raft::make_device_matrix( + resources, options.batch_size, options.candidate_count); + auto candidate_distances_device = raft::make_device_matrix( + resources, options.batch_size, options.candidate_count); + + std::vector candidates_host(static_cast(options.batch_size) * + options.candidate_count); + std::vector final_indices(static_cast(options.batch_size) * + options.top_k); + std::vector final_distances(static_cast(options.batch_size) * options.top_k); + std::vector all_results(static_cast(queries.rows) * options.top_k); + std::vector batch_latencies; + batch_latencies.reserve(batch_count); + + double search_total = 0.0; + double rerank_prep_total = 0.0; + double rerank_total = 0.0; + double total = 0.0; + auto const stream = raft::resource::get_cuda_stream(resources); + + for (std::uint32_t batch = 0; batch < batch_count; ++batch) { + auto const query_offset = static_cast(batch) * options.batch_size * queries.cols; + auto const before_upload = clock_type::now(); + raft::copy(query_device.data_handle(), + queries.values.data() + query_offset, + static_cast(options.batch_size) * queries.cols, + stream); + raft::resource::sync_stream(resources); + + auto const begin = options.include_query_transfer ? before_upload : clock_type::now(); + cuvs::neighbors::cagra::search(resources, + params, + index, + raft::make_const_mdspan(query_device.view()), + candidates_device.view(), + candidate_distances_device.view()); + raft::resource::sync_stream(resources); + auto const after_search = clock_type::now(); + + raft::copy( + candidates_host.data(), candidates_device.data_handle(), candidates_host.size(), stream); + raft::resource::sync_stream(resources); + auto const row_limit = base != nullptr ? base->rows() : index.size(); + auto const has_invalid_candidates = std::any_of( + candidates_host.begin(), candidates_host.end(), [row_limit](std::uint32_t candidate) { + return candidate >= row_limit; + }); + auto const after_prep = clock_type::now(); + + if (options.rerank) { + auto const dataset_view = raft::make_host_matrix_view( + base->data(), base->rows(), base->cols()); + auto const query_view = raft::make_host_matrix_view( + queries.values.data() + query_offset, options.batch_size, queries.cols); + auto const candidate_view = + raft::make_host_matrix_view( + candidates_host.data(), options.batch_size, options.candidate_count); + auto const index_view = raft::make_host_matrix_view( + final_indices.data(), options.batch_size, options.top_k); + auto const distance_view = raft::make_host_matrix_view( + final_distances.data(), options.batch_size, options.top_k); + if (!has_invalid_candidates) { + cuvs::neighbors::refine(resources, + dataset_view, + query_view, + candidate_view, + index_view, + distance_view, + index.metric()); + } else { + // Refine cannot dereference sentinel IDs. Compact only valid candidates for + // each affected row and preserve missing neighbors as invalid output slots. + std::fill( + final_indices.begin(), final_indices.end(), std::numeric_limits::max()); + std::fill( + final_distances.begin(), final_distances.end(), std::numeric_limits::infinity()); + std::vector valid; + valid.reserve(options.candidate_count); + for (std::uint32_t query = 0; query < options.batch_size; ++query) { + valid.clear(); + auto const offset = static_cast(query) * options.candidate_count; + for (std::uint32_t i = 0; i < options.candidate_count; ++i) { + auto const candidate = candidates_host[offset + i]; + if (candidate < row_limit) { valid.push_back(candidate); } + } + if (valid.empty()) { continue; } + auto const k = std::min(options.top_k, valid.size()); + cuvs::neighbors::refine( + resources, + dataset_view, + raft::make_host_matrix_view( + queries.values.data() + query_offset + static_cast(query) * queries.cols, + 1, + queries.cols), + raft::make_host_matrix_view( + valid.data(), 1, valid.size()), + raft::make_host_matrix_view( + final_indices.data() + static_cast(query) * options.top_k, 1, k), + raft::make_host_matrix_view( + final_distances.data() + static_cast(query) * options.top_k, 1, k), + index.metric()); + } + } + } else { + for (std::uint32_t query = 0; query < options.batch_size; ++query) { + std::copy_n( + candidates_host.data() + static_cast(query) * options.candidate_count, + options.top_k, + final_indices.data() + static_cast(query) * options.top_k); + } + } + auto const end = clock_type::now(); + + auto const result_offset = static_cast(batch) * options.batch_size * options.top_k; + std::copy(final_indices.begin(), final_indices.end(), all_results.begin() + result_offset); + + auto const search_ms = milliseconds(begin, after_search); + auto const prep_ms = milliseconds(after_search, after_prep); + auto const rerank_ms = milliseconds(after_prep, end); + auto const batch_ms = milliseconds(begin, end); + search_total += search_ms; + rerank_prep_total += prep_ms; + rerank_total += rerank_ms; + total += batch_ms; + batch_latencies.push_back(batch_ms); + } + + round_result result; + result.search_ms = search_total / batch_count; + result.rerank_prep_ms = rerank_prep_total / batch_count; + result.rerank_ms = rerank_total / batch_count; + result.total_ms = total / batch_count; + result.qps = static_cast(queries.rows) / (total / 1000.0); + result.p50_ms = percentile(batch_latencies, 0.50); + result.p95_ms = percentile(batch_latencies, 0.95); + result.p99_ms = percentile(batch_latencies, 0.99); + if (truth) { result.recall = recall_at_k(all_results, *truth, options.top_k); } + return result; +} + +void print_round(char const* kind, + std::uint32_t round, + std::uint32_t query_count, + benchmark_options const& options, + round_result const& result) +{ + std::printf( + "%s round=%u batch=%u query_count=%u search_ms=%.6f rerank_prep_ms=%.6f rerank_ms=%.6f " + "total_ms=%.6f recall=%.6f qps=%.3f p50_ms=%.6f p95_ms=%.6f p99_ms=%.6f\n", + kind, + round, + options.batch_size, + query_count, + result.search_ms, + result.rerank_prep_ms, + result.rerank_ms, + result.total_ms, + result.recall, + result.qps, + result.p50_ms, + result.p95_ms, + result.p99_ms); +} + +auto run_benchmark(benchmark_options options) -> int +{ + raft::device_resources resources; + auto const load_index_begin = clock_type::now(); + auto bundle = [&]() -> flowann::vpq_f16_index_bundle { + if (!options.index_path.empty()) { + return flowann::deserialize_vpq_f16(resources, options.index_path); + } + return flowann::deserialize_vpq_f16( + resources, options.legacy_graph_path, options.legacy_dataset_path); + }(); + + std::vector seeds; + if (!options.seeds_path.empty()) { + seeds = load_u32_vector(options.seeds_path); + bundle.index.update_seeds( + resources, + raft::make_host_vector_view(seeds.data(), seeds.size())); + raft::resource::sync_stream(resources); + } + auto const index_seeds = bundle.index.seeds(); + auto const available_seeds = + index_seeds.has_value() ? static_cast(index_seeds->extent(0)) : 0u; + if (!options.num_seeds) { options.num_seeds = available_seeds; } + if (*options.num_seeds > available_seeds) { + throw std::invalid_argument("--num-seeds exceeds the number of seeds stored in the index"); + } + + std::printf( + "LOAD_INDEX seconds=%.3f rows=%u dim=%u degree=%u resident_degree=%u cross_degree=%u " + "resident_bytes=%zu cross_bytes=%zu node_per_cacheline=%u n_bits=%u\n", + milliseconds(load_index_begin, clock_type::now()) / 1000.0, + bundle.index.size(), + bundle.index.dim(), + bundle.index.graph_degree(), + bundle.index.resident_degree(), + bundle.index.cross_degree(), + bundle.index.inner_graph().size() * sizeof(std::uint8_t), + bundle.index.cross_graph().size() * sizeof(std::uint32_t), + bundle.index.node_per_cacheline(), + bundle.index.n_bits()); + + auto const load_input_begin = clock_type::now(); + auto queries = load_matrix_file(options.queries_path, options.query_count); + std::unique_ptr> base; + if (options.rerank) { base = std::make_unique>(options.base_path); } + std::optional> truth; + if (!options.ground_truth_path.empty()) { + truth.emplace(load_matrix_file(options.ground_truth_path, queries.rows)); + } + + if (queries.cols != bundle.index.dim()) { + throw std::invalid_argument("query dimension does not match the index"); + } + if (queries.rows % options.batch_size != 0) { + throw std::invalid_argument("selected query count is not divisible by --batch-size"); + } + if (base && (base->rows() != bundle.index.size() || base->cols() != queries.cols)) { + throw std::invalid_argument("base matrix shape does not match the index"); + } + if (truth && (truth->rows != queries.rows || truth->cols < options.top_k)) { + throw std::invalid_argument("ground truth must cover every query and at least --top-k columns"); + } + + std::printf("LOAD_INPUT seconds=%.3f queries=%u dim=%u base_rows=%u gt_width=%u seeds=%u\n", + milliseconds(load_input_begin, clock_type::now()) / 1000.0, + queries.rows, + queries.cols, + base ? base->rows() : 0, + truth ? truth->cols : 0, + available_seeds); + std::printf( + "CONFIG gpu=%u batch=%u top_k=%u candidates=%u algorithm=%s team_size=%u itopk=%u " + "thread_block_size=%u width=%u min_iterations=%u max_iterations=%u " + "seeds=%u queues=%u empty_pause=%u " + "cpu_threads=%u session=%s rerank=%s queue_statistics=%s sync_window_scale=%.3f " + "sync_drop_threshold=%u\n", + options.gpu_id, + options.batch_size, + options.top_k, + options.candidate_count, + algorithm_name(options.algorithm), + options.team_size, + options.itopk_size, + options.thread_block_size, + options.search_width, + options.min_iterations, + options.max_iterations, + *options.num_seeds, + options.num_queues, + options.empty_pause, + options.cpu_threads, + options.use_search_session ? "on" : "off", + options.rerank ? "on" : "off", + options.collect_statistics ? "on" : "off", + options.sync_window_scale, + options.sync_drop_threshold); + + auto index = cuvs::neighbors::cagra::detail::adopt_tiered(resources, std::move(bundle)); + + for (std::uint32_t round = 1; round <= options.warmup_rounds; ++round) { + auto const result = run_round(resources, index, options, base.get(), queries, truth); + print_round("WARMUP", round, queries.rows, options, result); + } + + std::vector search_values; + std::vector prep_values; + std::vector rerank_values; + std::vector total_values; + std::vector recall_values; + std::vector qps_values; + for (std::uint32_t round = 1; round <= options.measured_rounds; ++round) { + auto const result = run_round(resources, index, options, base.get(), queries, truth); + print_round("MEASURE", round, queries.rows, options, result); + search_values.push_back(result.search_ms); + prep_values.push_back(result.rerank_prep_ms); + rerank_values.push_back(result.rerank_ms); + total_values.push_back(result.total_ms); + if (truth) { recall_values.push_back(result.recall); } + qps_values.push_back(result.qps); + } + + std::printf( + "SUMMARY batch=%u search_median_ms=%.6f rerank_prep_median_ms=%.6f " + "rerank_median_ms=%.6f total_median_ms=%.6f recall_median=%.6f qps_median=%.3f " + "search_min_ms=%.6f search_max_ms=%.6f\n", + options.batch_size, + median(search_values), + median(prep_values), + median(rerank_values), + median(total_values), + truth ? median(recall_values) : std::numeric_limits::quiet_NaN(), + median(qps_values), + *std::min_element(search_values.begin(), search_values.end()), + *std::max_element(search_values.begin(), search_values.end())); + + auto const statistics = cuvs::neighbors::cagra::detail::tiered_statistics(index); + std::printf("QUEUE_STATS polls=%llu empty_polls=%llu commands=%llu\n", + static_cast(statistics.polls), + static_cast(statistics.empty_polls), + static_cast(statistics.processed_commands)); + return 0; +} + +} // namespace + +int main(int argc, char** argv) +{ + try { + auto options = parse_options(argc, argv); + if (!options) { return 0; } + + RAFT_CUDA_TRY(cudaSetDevice(options->gpu_id)); +#ifdef _OPENMP + omp_set_num_threads(static_cast(options->cpu_threads)); +#else + if (options->cpu_threads != 1) { + throw std::invalid_argument("--cpu-threads requires an OpenMP-enabled build"); + } +#endif + if (options->pool_size_mib == 0) { return run_benchmark(*options); } + + auto const pool_bytes = options->pool_size_mib * 1024ULL * 1024ULL; + rmm::mr::managed_memory_resource managed_resource; + rmm::mr::pool_memory_resource pool_resource(managed_resource, pool_bytes); + scoped_device_resource resource_scope(pool_resource); + return run_benchmark(*options); + } catch (std::exception const& error) { + std::fprintf(stderr, "FLOWANN_BENCHMARK error: %s\n", error.what()); + std::fprintf(stderr, "Use --help for command-line documentation.\n"); + return 2; + } +} diff --git a/examples/cpp/src/flowann_build_benchmark.cu b/examples/cpp/src/flowann_build_benchmark.cu new file mode 100644 index 0000000000..f5b8f0c46c --- /dev/null +++ b/examples/cpp/src/flowann_build_benchmark.cu @@ -0,0 +1,331 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * FlowANN is described in "Disentangling Graph Dependencies for Efficient + * Billion-Scale GPU Vector Search" (OSDI 2026). + */ + +#include +#include + +#include + +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace flowann = cuvs::neighbors::cagra::experimental::flowann; +using clock_type = std::chrono::steady_clock; + +namespace { + +struct options { + std::string base_path; + std::string output_path; + std::string grouping_graph_path; + std::uint16_t grouping_bits = 0; + std::uint32_t grouping_count = 0; + std::uint64_t device_graph_budget_bytes{}; + std::uint32_t graph_degree = 32; + std::uint32_t intermediate_graph_degree = 96; + std::uint32_t vpq_dim = 192; + std::uint32_t vpq_bits = 8; + std::uint32_t vq_centers = 10'000; + std::uint32_t num_seeds = 0; + std::uint64_t seed_training_rows = 0; + std::uint64_t seed = 0x9e3779b97f4a7c15ULL; + bool guarantee_connectivity = false; + bool legacy_reference_build = false; +}; + +class mapped_fbin { + public: + explicit mapped_fbin(std::string const& path) + { + fd_ = open(path.c_str(), O_RDONLY); + if (fd_ < 0) { throw std::runtime_error("cannot open base file: " + path); } + struct stat info{}; + if (fstat(fd_, &info) != 0 || info.st_size < 8) { + throw std::runtime_error("cannot stat or invalid base file: " + path); + } + bytes_ = static_cast(info.st_size); + mapping_ = mmap(nullptr, bytes_, PROT_READ, MAP_PRIVATE, fd_, 0); + if (mapping_ == MAP_FAILED) { + mapping_ = nullptr; + throw std::runtime_error("cannot mmap base file: " + path); + } + std::memcpy(&rows_, mapping_, sizeof(rows_)); + std::memcpy(&dim_, static_cast(mapping_) + sizeof(rows_), sizeof(dim_)); + auto const expected = + std::uint64_t{8} + static_cast(rows_) * dim_ * sizeof(float); + if (rows_ == 0 || dim_ == 0 || expected != bytes_) { + throw std::runtime_error("base file is not a tightly packed float fbin: " + path); + } + } + + mapped_fbin(mapped_fbin const&) = delete; + auto operator=(mapped_fbin const&) -> mapped_fbin& = delete; + + ~mapped_fbin() + { + if (mapping_ != nullptr) { munmap(mapping_, bytes_); } + if (fd_ >= 0) { close(fd_); } + } + + [[nodiscard]] auto rows() const noexcept -> std::uint32_t { return rows_; } + [[nodiscard]] auto dim() const noexcept -> std::uint32_t { return dim_; } + [[nodiscard]] auto data() const noexcept -> float const* + { + return reinterpret_cast(static_cast(mapping_) + 8); + } + + private: + int fd_{-1}; + void* mapping_{}; + std::size_t bytes_{}; + std::uint32_t rows_{}; + std::uint32_t dim_{}; +}; + +template +auto parse_number(char const* text, char const* option) -> T +{ + char* end = nullptr; + if constexpr (std::is_floating_point_v) { + auto value = std::strtod(text, &end); + if (end == text || *end != '\0') { throw std::invalid_argument(option); } + return static_cast(value); + } else { + errno = 0; + auto value = std::strtoull(text, &end, 10); + if (end == text || *end != '\0' || text[0] == '-' || errno == ERANGE || + value > std::numeric_limits::max()) { + throw std::invalid_argument(option); + } + return static_cast(value); + } +} + +void print_usage(char const* program) +{ + std::printf( + "Usage: %s --base BASE.fbin --output INDEX --device-graph-budget-bytes N [options]\n" + " --grouping-graph GRAPH.ibin partition an existing graph; output raw uint32 labels\n" + " --grouping-bits N 0 = automatic 4-bit-aligned width, capped at 24 " + "(default)\n" + " --grouping-count N 0 = automatic (default)\n" + " --graph-degree N default 32\n" + " --intermediate-graph-degree N default 96\n" + " --vpq-dim N default 192\n" + " --vpq-bits N default 8\n" + " --vq-centers N default 10000\n" + " --num-seeds N default 0\n" + " --seed-training-rows N default 0 (automatic)\n" + " --seed N deterministic sample rotation\n" + " --legacy-reference-build use explicit legacy-reference IVF-PQ build parameters\n" + " --guarantee-connectivity enable CAGRA MST connectivity optimization\n", + program); +} + +auto parse_options(int argc, char** argv) -> options +{ + options result; + for (int i = 1; i < argc; ++i) { + auto const option = std::string_view(argv[i]); + auto value = [&]() -> char const* { + if (++i >= argc) { throw std::invalid_argument("missing value for " + std::string(option)); } + return argv[i]; + }; + if (option == "--base") { + result.base_path = value(); + } else if (option == "--grouping-graph") { + result.grouping_graph_path = value(); + } else if (option == "--grouping-bits") { + result.grouping_bits = parse_number(value(), option.data()); + } else if (option == "--grouping-count") { + result.grouping_count = parse_number(value(), option.data()); + } else if (option == "--output") { + result.output_path = value(); + } else if (option == "--device-graph-budget-bytes") { + result.device_graph_budget_bytes = parse_number(value(), option.data()); + } else if (option == "--graph-degree") { + result.graph_degree = parse_number(value(), option.data()); + } else if (option == "--intermediate-graph-degree") { + result.intermediate_graph_degree = parse_number(value(), option.data()); + } else if (option == "--vpq-dim") { + result.vpq_dim = parse_number(value(), option.data()); + } else if (option == "--vpq-bits") { + result.vpq_bits = parse_number(value(), option.data()); + } else if (option == "--vq-centers") { + result.vq_centers = parse_number(value(), option.data()); + } else if (option == "--num-seeds") { + result.num_seeds = parse_number(value(), option.data()); + } else if (option == "--seed-training-rows") { + result.seed_training_rows = parse_number(value(), option.data()); + } else if (option == "--seed") { + result.seed = parse_number(value(), option.data()); + } else if (option == "--legacy-reference-build") { + result.legacy_reference_build = true; + } else if (option == "--guarantee-connectivity") { + result.guarantee_connectivity = true; + } else if (option == "--help") { + print_usage(argv[0]); + std::exit(0); + } else { + throw std::invalid_argument("unknown option: " + std::string(option)); + } + } + if (result.base_path.empty() || result.output_path.empty()) { + throw std::invalid_argument("--base and --output are required"); + } + return result; +} + +auto seconds(clock_type::time_point begin, clock_type::time_point end) -> double +{ + return std::chrono::duration(end - begin).count(); +} + +} // namespace + +int main(int argc, char** argv) +{ + try { + auto const args = parse_options(argc, argv); + auto base = mapped_fbin(args.base_path); + auto matrix = raft::make_host_matrix_view( + base.data(), static_cast(base.rows()), static_cast(base.dim())); + auto dataset = cuvs::neighbors::make_host_standard_dataset_view(matrix); + + flowann::index_params params; + params.grouping.n_bits = args.grouping_bits; + params.grouping.n_groups = args.grouping_count; + params.device_graph_budget_bytes = args.device_graph_budget_bytes; + params.cagra_params.metric = cuvs::distance::DistanceType::L2Expanded; + params.cagra_params.graph_degree = args.graph_degree; + params.cagra_params.intermediate_graph_degree = args.intermediate_graph_degree; + params.cagra_params.guarantee_connectivity = args.guarantee_connectivity; + auto graph_build = cuvs::neighbors::cagra::graph_build_params::ivf_pq_params( + raft::matrix_extent(base.rows(), base.dim()), params.cagra_params.metric); + if (args.legacy_reference_build) { + // Match the old graph-building benchmark explicitly, independently of search VPQ. + graph_build = cuvs::neighbors::cagra::graph_build_params::ivf_pq_params{}; + graph_build.build_params.metric = params.cagra_params.metric; + graph_build.build_params.n_lists = 1024; + graph_build.build_params.kmeans_n_iters = 20; + graph_build.build_params.kmeans_trainset_fraction = 0.1; + graph_build.build_params.pq_dim = 192; + graph_build.build_params.pq_bits = 8; + graph_build.build_params.codebook_kind = cuvs::neighbors::ivf_pq::codebook_gen::PER_SUBSPACE; + graph_build.build_params.force_random_rotation = false; + graph_build.build_params.max_train_points_per_pq_code = 256; + graph_build.search_params.n_probes = 20; + graph_build.search_params.lut_dtype = CUDA_R_32F; + graph_build.search_params.internal_distance_dtype = CUDA_R_32F; + graph_build.search_params.coarse_search_dtype = CUDA_R_32F; + // Both profiles request at least intermediate_degree + 1 initial candidates. + graph_build.refinement_rate = 1.0f; + } + params.cagra_params.graph_build_params = graph_build; + std::printf( + "GRAPH_BUILD profile=%s n_lists=%u n_probes=%u pq_dim=%u pq_bits=%u " + "kmeans_n_iters=%u train_fraction=%.9f refinement_rate=%.3f " + "lut_dtype=%d internal_dtype=%d coarse_dtype=%d\n", + args.legacy_reference_build ? "legacy-reference" : "dataset-default", + graph_build.build_params.n_lists, + graph_build.search_params.n_probes, + graph_build.build_params.pq_dim, + graph_build.build_params.pq_bits, + graph_build.build_params.kmeans_n_iters, + graph_build.build_params.kmeans_trainset_fraction, + graph_build.refinement_rate, + static_cast(graph_build.search_params.lut_dtype), + static_cast(graph_build.search_params.internal_distance_dtype), + static_cast(graph_build.search_params.coarse_search_dtype)); + std::fflush(stdout); + params.num_seeds = args.num_seeds; + params.seed_training_rows = args.seed_training_rows; + params.seed = args.seed; + + cuvs::neighbors::vpq_params vpq; + vpq.pq_dim = args.vpq_dim; + vpq.pq_bits = args.vpq_bits; + vpq.vq_n_centers = args.vq_centers; + + std::printf( + "CONFIG rows=%u dim=%u graph_degree=%u intermediate_degree=%u graph_budget_bytes=%llu " + "vpq_dim=%u vpq_bits=%u vq_centers=%u seeds=%u seed_training_rows=%llu " + "guarantee_connectivity=%s\n", + base.rows(), + base.dim(), + args.graph_degree, + args.intermediate_graph_degree, + static_cast(args.device_graph_budget_bytes), + args.vpq_dim, + args.vpq_bits, + args.vq_centers, + args.num_seeds, + static_cast(args.seed_training_rows), + args.guarantee_connectivity ? "true" : "false"); + + raft::device_resources resources; + if (!args.grouping_graph_path.empty()) { + auto graph = mapped_fbin(args.grouping_graph_path); + if (graph.rows() != base.rows()) { throw std::invalid_argument("graph/base row mismatch"); } + auto graph_view = raft::make_host_matrix_view( + reinterpret_cast(graph.data()), graph.rows(), graph.dim()); + auto const begin = clock_type::now(); + auto grouped = flowann::group_graph(resources, params, graph_view, dataset); + auto const end = clock_type::now(); + auto* output = std::fopen(args.output_path.c_str(), "wb"); + if (!output) { throw std::runtime_error("cannot open grouping output"); } + auto const written = std::fwrite( + grouped.labels.data_handle(), sizeof(std::uint32_t), grouped.labels.size(), output); + auto const closed = std::fclose(output); + if (written != grouped.labels.size() || closed != 0) { + throw std::runtime_error("cannot write grouping output"); + } + std::printf("GROUPING_RESULT seconds=%.6f n_bits=%u n_groups=%u rows=%u\n", + seconds(begin, end), + grouped.n_bits, + grouped.n_groups, + base.rows()); + return 0; + } + auto const build_begin = clock_type::now(); + auto built = flowann::build(resources, params, vpq, dataset); + auto const build_end = clock_type::now(); + flowann::serialize(resources, args.output_path, built.index); + auto const serialize_end = clock_type::now(); + std::printf( + "BUILD_RESULT build_seconds=%.6f serialize_seconds=%.6f total_seconds=%.6f " + "resident_degree=%u cross_degree=%u n_bits=%u resident_bytes=%zu cross_bytes=%zu\n", + seconds(build_begin, build_end), + seconds(build_end, serialize_end), + seconds(build_begin, serialize_end), + built.index.resident_degree(), + built.index.cross_degree(), + built.index.n_bits(), + built.index.inner_graph().size() * sizeof(std::uint8_t), + built.index.cross_graph().size() * sizeof(std::uint32_t)); + return 0; + } catch (std::exception const& error) { + std::fprintf(stderr, "FLOWANN_BUILD_BENCHMARK error: %s\n", error.what()); + return 1; + } +} diff --git a/fern/pages/c_api/c-api-neighbors-cagra.md b/fern/pages/c_api/c-api-neighbors-cagra.md index c6ceb153d5..305fa009b6 100644 --- a/fern/pages/c_api/c-api-neighbors-cagra.md +++ b/fern/pages/c_api/c-api-neighbors-cagra.md @@ -116,10 +116,101 @@ struct cuvsAceParams { | `max_host_memory_gb` | `double` | Maximum host memory to use for ACE build in GiB.

When set to 0 (default), uses available host memory. When set to a positive value, limits host memory usage to the specified amount. Useful for testing or when running alongside other memory-intensive processes. | | `max_gpu_memory_gb` | `double` | Maximum GPU memory to use for ACE build in GiB.

When set to 0 (default), uses available GPU memory. When set to a positive value, limits GPU memory usage to the specified amount. Useful for testing or when running alongside other memory-intensive processes. | + +### cuvsCagraGraphStorage + +Experimental CAGRA graph storage mode. + +```c +enum cuvsCagraGraphStorage { + CUVS_CAGRA_GRAPH_DEVICE = 0, + CUVS_CAGRA_GRAPH_TIERED = 1 +}; +``` + +**Values** + +| Name | Value | +| --- | --- | +| `CUVS_CAGRA_GRAPH_DEVICE` | `0` | +| `CUVS_CAGRA_GRAPH_TIERED` | `1` | + + +### cuvsCagraTieredGraphParams + +Tiered graph settings. The GPU budget covers packed edges only. + +```c +struct cuvsCagraTieredGraphParams { + size_t device_graph_budget_bytes; + uint16_t node_per_cacheline; + bool grouping_enabled; + uint32_t n_groups; + uint16_t n_bits; + double balance_tolerance; + size_t training_rows; + size_t assignment_batch_rows; + uint32_t kmeans_n_iters; + bool validate; + uint32_t num_seeds; + size_t seed_training_rows; + uint64_t seed; +}; +``` + +**Fields** + +| Name | Type | Description | +| --- | --- | --- | +| `device_graph_budget_bytes` | `size_t` | Maximum packed resident graph bytes, excluding datasets and queues. Zero stores all edges on host. | +| `node_per_cacheline` | `uint16_t` | Consecutive nodes per packed row: [1, 64] with grouping, [1, 32] without grouping. | +| `grouping_enabled` | `bool` | Group and reorder nodes for locality; false retains the rank-only layout. | +| `n_groups` | `uint32_t` | Final group count, at most 65,536. Zero uses one group if IDs fit; otherwise ceil(N * (1 + balance_tolerance) / 2^b). Each k-means node has at most 16 children. | +| `n_bits` | `uint16_t` | Local-ID width. Zero selects min(24, max(4, align_up_4(ceil(log2(N))))). Explicit widths must be in [1, 32]. | +| `balance_tolerance` | `double` | Allowed relative deviation from average group size, in the open interval (0, 1). Hard ID capacity still applies. | +| `training_rows` | `size_t` | Maximum sample rows per k-means node. Zero shares a 160,000-row budget by subtree leaf count, bounded by 256 to 10,000 rows per child and actual node size. | +| `assignment_batch_rows` | `size_t` | Rows assigned per GPU batch. Zero targets a 256 MiB vector buffer; distance/label buffers and host capacity-repair storage are additional. | +| `kmeans_n_iters` | `uint32_t` | Positive number of balanced k-means training iterations. | +| `validate` | `bool` | Verify encoded edges after rearrangement; intended for tests and small builds. | +| `num_seeds` | `uint32_t` | Number of medoid seeds generated at build time; zero disables seed generation. | +| `seed_training_rows` | `size_t` | Rows sampled to train seed centroids. Zero selects an automatic limit. | +| `seed` | `uint64_t` | Deterministic seed for rotating the uniform k-means training sample. | + + +### cuvsCagraTieredSearchParams + +Tiered search settings; queue state is retained by the index. + +```c +struct cuvsCagraTieredSearchParams { + uint32_t num_seeds; + float sync_window_scale; + uint32_t sync_drop_threshold; + uint32_t num_queues; + uint32_t empty_pause; + bool collect_statistics; + bool keep_pollers_running; +}; +``` + +**Fields** + +| Name | Type | Description | +| --- | --- | --- | +| `num_seeds` | `uint32_t` | Maximum number of stored medoid seeds used per query; zero uses random initialization. | +| `sync_window_scale` | `float` | Scale for the adaptive deferred-cross-edge synchronization window. | +| `sync_drop_threshold` | `uint32_t` | Parent-position drop that forces synchronization of a submitted cross-edge request. | +| `num_queues` | `uint32_t` | Positive number of independent GPU-to-CPU command queues. | +| `empty_pause` | `uint32_t` | Number of CPU pause instructions after an empty poll; zero disables this pause. | +| `collect_statistics` | `bool` | Collect aggregate queue and polling counters. | +| `keep_pollers_running` | `bool` | Keep pollers active between calls for low latency. Consumes host CPU while idle. | + ### cuvsCagraIndexParams -Supplemental parameters to build CAGRA Index +Supplemental parameters to build a CAGRA index. + +Allocate with cuvsCagraIndexParamsCreate(); stack allocation and by-value copies are not covered by the C ABI stability guarantee. ```c struct cuvsCagraIndexParams { @@ -129,6 +220,8 @@ struct cuvsCagraIndexParams { enum cuvsCagraGraphBuildAlgo build_algo; size_t nn_descent_niter; void* graph_build_params; + enum cuvsCagraGraphStorage graph_storage; + struct cuvsCagraTieredGraphParams tiered; }; ``` @@ -142,6 +235,8 @@ struct cuvsCagraIndexParams { | `build_algo` | [`enum cuvsCagraGraphBuildAlgo`](/api-reference/c-api-neighbors-cagra#cuvscagragraphbuildalgo) | ANN algorithm to build knn graph. | | `nn_descent_niter` | `size_t` | Number of Iterations to run if building with NN_DESCENT | | `graph_build_params` | `void*` | Optional: specify graph build params based on build_algo
- IVF_PQ: cuvsIvfPqParams_t
- ACE: cuvsAceParams_t
- Others: nullptr | +| `graph_storage` | [`enum cuvsCagraGraphStorage`](/api-reference/c-api-neighbors-cagra#cuvscagragraphstorage) | Graph representation selected for construction. The default is CUVS_CAGRA_GRAPH_DEVICE. | +| `tiered` | [`struct cuvsCagraTieredGraphParams`](/api-reference/c-api-neighbors-cagra#cuvscagratieredgraphparams) | Tiered layout settings used when graph_storage is CUVS_CAGRA_GRAPH_TIERED. | ### cuvsCagraMergeAlgo @@ -199,7 +294,8 @@ struct cuvsCagraMergeParams { Allocate CAGRA Index params, and populate with default values ```c -cuvsError_t cuvsCagraIndexParamsCreate(cuvsCagraIndexParams_t* params); +cuvsError_t +cuvsCagraIndexParamsCreate(cuvsCagraIndexParams_t *params); ``` **Parameters** @@ -218,7 +314,8 @@ cuvsError_t cuvsCagraIndexParamsCreate(cuvsCagraIndexParams_t* params); De-allocate CAGRA Index params ```c -cuvsError_t cuvsCagraIndexParamsDestroy(cuvsCagraIndexParams_t params); +cuvsError_t +cuvsCagraIndexParamsDestroy(cuvsCagraIndexParams_t params); ``` **Parameters** @@ -237,7 +334,8 @@ cuvsError_t cuvsCagraIndexParamsDestroy(cuvsCagraIndexParams_t params); Allocate CAGRA merge params and populate them with AUTO defaults. ```c -cuvsError_t cuvsCagraMergeParamsCreate(cuvsCagraMergeParams_t* params); +cuvsError_t +cuvsCagraMergeParamsCreate(cuvsCagraMergeParams_t *params); ``` **Parameters** @@ -256,7 +354,8 @@ cuvsError_t cuvsCagraMergeParamsCreate(cuvsCagraMergeParams_t* params); De-allocate CAGRA merge params. ```c -cuvsError_t cuvsCagraMergeParamsDestroy(cuvsCagraMergeParams_t params); +cuvsError_t +cuvsCagraMergeParamsDestroy(cuvsCagraMergeParams_t params); ``` **Parameters** @@ -275,7 +374,8 @@ cuvsError_t cuvsCagraMergeParamsDestroy(cuvsCagraMergeParams_t params); Allocate CAGRA Compression params, and populate with default values. ```c -cuvsError_t cuvsCagraCompressionParamsCreate(cuvsCagraCompressionParams_t* params); +cuvsError_t +cuvsCagraCompressionParamsCreate(cuvsCagraCompressionParams_t *params); ``` Deprecated: Use `cuvsPqParamsCreate`. This compatibility API is planned for removal in the 27.02 ABI-breaking release. @@ -296,7 +396,8 @@ Deprecated: Use `cuvsPqParamsCreate`. This compatibility API is planned for remo De-allocate CAGRA Compression params. ```c -cuvsError_t cuvsCagraCompressionParamsDestroy(cuvsCagraCompressionParams_t params); +cuvsError_t +cuvsCagraCompressionParamsDestroy(cuvsCagraCompressionParams_t params); ``` Deprecated: Use `cuvsPqParamsDestroy`. This compatibility API is planned for removal in the 27.02 ABI-breaking release. @@ -355,12 +456,9 @@ cuvsError_t cuvsAceParamsDestroy(cuvsAceParams_t params); Create CAGRA index parameters similar to an HNSW index ```c -cuvsError_t cuvsCagraIndexParamsFromHnswParams(cuvsCagraIndexParams_t params, -int64_t n_rows, -int64_t dim, -int M, -int ef_construction, -enum cuvsCagraHnswHeuristicType heuristic, +cuvsError_t cuvsCagraIndexParamsFromHnswParams( +cuvsCagraIndexParams_t params, int64_t n_rows, int64_t dim, int M, +int ef_construction, enum cuvsCagraHnswHeuristicType heuristic, cuvsDistanceType metric); ``` @@ -388,12 +486,9 @@ This factory function creates CAGRA parameters that yield a graph compatible wit Create CAGRA index parameters heuristically tuned for a dataset ```c -cuvsError_t cuvsCagraIndexParamsFromDataset(cuvsCagraIndexParams_t params, -int64_t n_rows, -int64_t dim, -size_t graph_degree, -cuvsDistanceType metric, -size_t build_quality); +cuvsError_t cuvsCagraIndexParamsFromDataset( +cuvsCagraIndexParams_t params, int64_t n_rows, int64_t dim, +size_t graph_degree, cuvsDistanceType metric, size_t build_quality); ``` This factory function selects the graph build algorithm and its parameters based on the shape of the dataset. @@ -438,7 +533,8 @@ struct cuvsCagraExtendParams { Allocate CAGRA Extend params, and populate with default values ```c -cuvsError_t cuvsCagraExtendParamsCreate(cuvsCagraExtendParams_t* params); +cuvsError_t +cuvsCagraExtendParamsCreate(cuvsCagraExtendParams_t *params); ``` **Parameters** @@ -457,7 +553,8 @@ cuvsError_t cuvsCagraExtendParamsCreate(cuvsCagraExtendParams_t* params); De-allocate CAGRA Extend params ```c -cuvsError_t cuvsCagraExtendParamsDestroy(cuvsCagraExtendParams_t params); +cuvsError_t +cuvsCagraExtendParamsDestroy(cuvsCagraExtendParams_t params); ``` **Parameters** @@ -568,6 +665,8 @@ struct cuvsCagraSearchParams { bool persistent; float persistent_lifetime; float persistent_device_usage; + bool use_tiered_params; + struct cuvsCagraTieredSearchParams tiered; }; ``` @@ -591,6 +690,8 @@ struct cuvsCagraSearchParams { | `persistent` | `bool` | Whether to use the persistent version of the kernel (only SINGLE_CTA is supported a.t.m.) | | `persistent_lifetime` | `float` | Persistent kernel: time in seconds before the kernel stops if no requests received. | | `persistent_device_usage` | `float` | Set the fraction of maximum grid size used by persistent kernel. Value 1.0 means the kernel grid size is maximum possible for the selected device. The value must be greater than 0.0 and not greater than 1.0.

One may need to run other kernels alongside this persistent kernel. This parameter can be used to reduce the grid size of the persistent kernel to leave a few SMs idle.
Note: running any other work on GPU alongside with the persistent kernel makes the setup fragile.
- Running another kernel in another thread usually works, but no progress guaranteed
- Any CUDA allocations block the context (this issue may be obscured by using pools)
- Memory copies to not-pinned host memory may block the context

Even when we know there are no other kernels working at the same time, setting kDeviceUsage to 1.0 surprisingly sometimes hurts performance. Proceed with care. If you suspect this is an issue, you can reduce this number to ~0.9 without a significant impact on the throughput. | +| `use_tiered_params` | `bool` | Enable the tiered settings below. Allocate this struct with cuvsCagraSearchParamsCreate(). | +| `tiered` | [`struct cuvsCagraTieredSearchParams`](/api-reference/c-api-neighbors-cagra#cuvscagratieredsearchparams) | Tiered queue and search settings used only when use_tiered_params is true. | ### cuvsCagraSearchParamsCreate @@ -598,7 +699,8 @@ struct cuvsCagraSearchParams { Allocate CAGRA search params, and populate with default values ```c -cuvsError_t cuvsCagraSearchParamsCreate(cuvsCagraSearchParams_t* params); +cuvsError_t +cuvsCagraSearchParamsCreate(cuvsCagraSearchParams_t *params); ``` **Parameters** @@ -617,7 +719,8 @@ cuvsError_t cuvsCagraSearchParamsCreate(cuvsCagraSearchParams_t* params); De-allocate CAGRA search params ```c -cuvsError_t cuvsCagraSearchParamsDestroy(cuvsCagraSearchParams_t params); +cuvsError_t +cuvsCagraSearchParamsDestroy(cuvsCagraSearchParams_t params); ``` **Parameters** @@ -697,7 +800,8 @@ cuvsError_t cuvsCagraIndexDestroy(cuvsCagraIndex_t index); Get dimension of the CAGRA index ```c -cuvsError_t cuvsCagraIndexGetDims(cuvsCagraIndex_t index, int64_t* dim); +cuvsError_t cuvsCagraIndexGetDims(cuvsCagraIndex_t index, +int64_t *dim); ``` **Parameters** @@ -717,7 +821,8 @@ cuvsError_t cuvsCagraIndexGetDims(cuvsCagraIndex_t index, int64_t* dim); Get size of the CAGRA index ```c -cuvsError_t cuvsCagraIndexGetSize(cuvsCagraIndex_t index, int64_t* size); +cuvsError_t cuvsCagraIndexGetSize(cuvsCagraIndex_t index, +int64_t *size); ``` **Parameters** @@ -731,21 +836,47 @@ cuvsError_t cuvsCagraIndexGetSize(cuvsCagraIndex_t index, int64_t* size); [`cuvsError_t`](/api-reference/c-api-core-c-api#cuvserror-t) + +### cuvsCagraIndexGetGraphStorage + +Get the graph representation of a CAGRA index. + +```c +cuvsError_t cuvsCagraIndexGetGraphStorage( +cuvsCagraIndex_t index, enum cuvsCagraGraphStorage *storage); +``` + +Returns the selected storage kind without exposing the device and host edge buffers. + +**Parameters** + +| Name | Direction | Type | Description | +| --- | --- | --- | --- | +| `index` | in | [`cuvsCagraIndex_t`](/api-reference/c-api-neighbors-cagra#cuvscagraindex) | Initialized CAGRA index handle. | +| `storage` | out | [`enum cuvsCagraGraphStorage*`](/api-reference/c-api-neighbors-cagra#cuvscagragraphstorage) | Receives CUVS_CAGRA_GRAPH_DEVICE or CUVS_CAGRA_GRAPH_TIERED. | + +**Returns** + +[`cuvsError_t`](/api-reference/c-api-core-c-api#cuvserror-t) + ### cuvsCagraIndexGetGraphDegree -Get graph degree of the CAGRA index +Get the complete graph degree of a CAGRA index. ```c -cuvsError_t cuvsCagraIndexGetGraphDegree(cuvsCagraIndex_t index, int64_t* graph_degree); +cuvsError_t cuvsCagraIndexGetGraphDegree(cuvsCagraIndex_t index, +int64_t *graph_degree); ``` +For a tiered index this includes both resident and host edges. + **Parameters** | Name | Direction | Type | Description | | --- | --- | --- | --- | -| `index` | in | [`cuvsCagraIndex_t`](/api-reference/c-api-neighbors-cagra#cuvscagraindex) | CAGRA index | -| `graph_degree` | out | `int64_t*` | return graph degree | +| `index` | in | [`cuvsCagraIndex_t`](/api-reference/c-api-neighbors-cagra#cuvscagraindex) | Initialized CAGRA index handle. | +| `graph_degree` | out | `int64_t*` | Receives the complete graph degree. | **Returns** @@ -757,7 +888,8 @@ cuvsError_t cuvsCagraIndexGetGraphDegree(cuvsCagraIndex_t index, int64_t* graph_ Returns a view of the CAGRA dataset ```c -cuvsError_t cuvsCagraIndexGetDataset(cuvsCagraIndex_t index, DLManagedTensor* dataset); +cuvsError_t cuvsCagraIndexGetDataset(cuvsCagraIndex_t index, +DLManagedTensor *dataset); ``` This function returns a non-owning view of the CAGRA dataset. The output will be referencing device memory that is directly used in CAGRA, without copying the dataset at all. This means that the output is only valid as long as the CAGRA index is alive, and once cuvsCagraIndexDestroy is called on the cagra index - the returned dataset view will be invalid. @@ -781,7 +913,8 @@ Note that the DLManagedTensor dataset returned will have an associated 'deleter' Returns a view of the CAGRA graph ```c -cuvsError_t cuvsCagraIndexGetGraph(cuvsCagraIndex_t index, DLManagedTensor* graph); +cuvsError_t cuvsCagraIndexGetGraph(cuvsCagraIndex_t index, +DLManagedTensor *graph); ``` This function returns a non-owning view of the CAGRA graph. The output will be referencing device memory that is directly used in CAGRA, without copying the graph at all. This means that the output is only valid as long as the CAGRA index is alive, and once cuvsCagraIndexDestroy is called on the cagra index - the returned graph view will be invalid. @@ -862,6 +995,35 @@ A `CUVS_DATASET_LAYOUT_PQ` dataset created by `cuvsDatasetMakePQ` builds an iter [`cuvsError_t`](/api-reference/c-api-core-c-api#cuvserror-t) + +### cuvsCagraBuildCompressed + +Build an owning tiered graph and VPQ-F16 dataset from host standard vectors. + +```c +cuvsError_t cuvsCagraBuildCompressed(cuvsResources_t res, +cuvsCagraIndexParams_t params, +cuvsPqParams_t compression, +cuvsDataset_t dataset, +cuvsCagraIndex_t index); +``` + +Requires graph_storage=CUVS_CAGRA_GRAPH_TIERED and enabled FlowANN building. The caller retains the input vectors and may release them after this call. The resulting index owns its reordered compressed dataset. + +**Parameters** + +| Name | Direction | Type | Description | +| --- | --- | --- | --- | +| `res` | in | [`cuvsResources_t`](/api-reference/c-api-core-c-api#cuvsresources-t) | Resources used for construction, compression, and device transfers. | +| `params` | in | [`cuvsCagraIndexParams_t`](/api-reference/c-api-neighbors-cagra#cuvscagraindexparams) | CAGRA graph construction parameters selecting tiered storage. | +| `compression` | in | `cuvsPqParams_t` | VPQ compression parameters using 8-bit PQ codes and FP16 codebooks. | +| `dataset` | in | `cuvsDataset_t` | Host standard dataset with float32, float16, int8, or uint8 values. | +| `index` | out | [`cuvsCagraIndex_t`](/api-reference/c-api-neighbors-cagra#cuvscagraindex) | Output handle initialized with cuvsCagraIndexCreate. | + +**Returns** + +[`cuvsError_t`](/api-reference/c-api-core-c-api#cuvserror-t) + ## C API for CUDA ANN Graph-based nearest neighbor search @@ -874,13 +1036,10 @@ Search a CAGRA index with a `DLManagedTensor` which has underlying `DLDeviceType 3. `distances`: `kDLDataType.code == kDLFloat` and `kDLDataType.bits = 32` ```c -cuvsError_t cuvsCagraSearch(cuvsResources_t res, -cuvsCagraSearchParams_t params, -cuvsCagraIndex_t index, -DLManagedTensor* queries, -DLManagedTensor* neighbors, -DLManagedTensor* distances, -cuvsFilter filter); +cuvsError_t cuvsCagraSearch( +cuvsResources_t res, cuvsCagraSearchParams_t params, cuvsCagraIndex_t index, +DLManagedTensor *queries, DLManagedTensor *neighbors, +DLManagedTensor *distances, cuvsFilter filter); ``` **Parameters** @@ -905,14 +1064,11 @@ cuvsFilter filter); Search multiple CAGRA index partitions concurrently and return the global top-k per query. ```c -cuvsError_t cuvsCagraSearchMultiPartition(cuvsResources_t res, -cuvsCagraSearchParams_t params, -uint32_t num_partitions, -cuvsCagraIndex_t* indices, -DLManagedTensor* queries, -DLManagedTensor* partition_ids, -DLManagedTensor* neighbors, -DLManagedTensor* distances, +cuvsError_t cuvsCagraSearchMultiPartition( +cuvsResources_t res, cuvsCagraSearchParams_t params, +uint32_t num_partitions, cuvsCagraIndex_t *indices, +DLManagedTensor *queries, DLManagedTensor *partition_ids, +DLManagedTensor *neighbors, DLManagedTensor *distances, cuvsFilter* filters); ``` @@ -973,9 +1129,8 @@ Experimental, both the API and the serialization format are subject to change. Save the CAGRA graph and its attached dataset to file. ```c -cuvsError_t cuvsCagraSerializeGraphAndDataset(cuvsResources_t res, -const char* filename, -cuvsCagraIndex_t index); +cuvsError_t cuvsCagraSerializeGraphAndDataset( +cuvsResources_t res, const char *filename, cuvsCagraIndex_t index); ``` The index stores a non-owning dataset view. The caller must keep the memory of the dataset backing that view alive while this function runs. Returns CUVS_ERROR without modifying the destination file if the index has no attached dataset. PQ and BBQ datasets are not serialized by this function. @@ -1054,14 +1209,15 @@ Experimental, both the API and the serialization format are subject to change. Load the CAGRA graph and dataset from file. ```c -cuvsError_t cuvsCagraDeserializeGraphAndDataset(cuvsResources_t res, -const char* filename, -cuvsCagraIndex_t index, +cuvsError_t cuvsCagraDeserializeGraphAndDataset( +cuvsResources_t res, const char *filename, cuvsCagraIndex_t index, cuvsDataset_t* out_dataset); ``` The returned owning dataset preserves the serialized host/device memory type and standard/padded layout. The index stores a non-owning view into it, so the caller must keep the dataset alive while the index uses it and destroy it separately with cuvsDatasetDestroy. Only a device-padded result is immediately searchable through the C API; attach a caller-owned device-padded view with cuvsCagraUpdateDataset for any other kind. The output pointer must point to a null handle on entry; deserialization acts as a factory and transfers ownership of the allocated dataset handle on success. Returns CUVS_ERROR when the file has no dataset; the index and output handle are unchanged on failure. +For a tiered file the index owns the reordered dataset internally, and out_dataset remains null. The loaded index is immediately searchable and has no separately owned dataset handle. + Experimental, both the API and the serialization format are subject to change. **Parameters** @@ -1115,8 +1271,7 @@ Merge multiple CAGRA indices into a single CAGRA index. cuvsError_t cuvsCagraMerge(cuvsResources_t res, cuvsCagraIndexParams_t params, cuvsCagraIndex_t* indices, -size_t num_indices, -cuvsFilter filter, +size_t num_indices, cuvsFilter filter, cuvsDataset_t merged_dataset, cuvsCagraIndex_t output_index); ``` @@ -1154,13 +1309,10 @@ Example: Merge multiple CAGRA indices with explicit merge parameters. ```c -cuvsError_t cuvsCagraMergeWithParams(cuvsResources_t res, -cuvsCagraIndexParams_t params, -cuvsCagraMergeParams_t merge_params, -cuvsCagraIndex_t* indices, -size_t num_indices, -cuvsFilter filter, -cuvsDataset_t merged_dataset, +cuvsError_t cuvsCagraMergeWithParams( +cuvsResources_t res, cuvsCagraIndexParams_t params, +cuvsCagraMergeParams_t merge_params, cuvsCagraIndex_t *indices, +size_t num_indices, cuvsFilter filter, cuvsDataset_t merged_dataset, cuvsCagraIndex_t output_index); ``` diff --git a/fern/pages/cpp_api/cpp-api-neighbors-cagra.md b/fern/pages/cpp_api/cpp-api-neighbors-cagra.md index 9c3f67dec8..aff56fe384 100644 --- a/fern/pages/cpp_api/cpp-api-neighbors-cagra.md +++ b/fern/pages/cpp_api/cpp-api-neighbors-cagra.md @@ -31,6 +31,150 @@ enum class search_algo { | `MULTI_KERNEL` | `2` | | `AUTO` | `100` | + +### neighbors::cagra::graph_storage_kind + +Graph representation used by a CAGRA index. + +```cpp +enum class graph_storage_kind { + device = 0, + tiered = 1 +}; +``` + +**Values** + +| Name | Value | +| --- | --- | +| `device` | `0` | +| `tiered` | `1` | + + +### neighbors::cagra::tiered_graph_params + +Experimental GPU/CPU graph layout. The budget excludes datasets and search workspaces. + +```cpp +struct tiered_graph_params { + std::size_t device_graph_budget_bytes; + std::uint16_t node_per_cacheline; + bool grouping_enabled; + std::uint32_t n_groups; + std::uint16_t n_bits; + double balance_tolerance; + std::size_t training_rows; + std::size_t assignment_batch_rows; + std::uint32_t kmeans_n_iters; + bool validate; + std::uint32_t num_seeds; + std::size_t seed_training_rows; + std::uint64_t seed; +}; +``` + +**Fields** + +| Name | Type | Description | +| --- | --- | --- | +| `device_graph_budget_bytes` | `std::size_t` | Maximum packed resident graph bytes, excluding datasets and queues. Zero stores all edges on host. | +| `node_per_cacheline` | `std::uint16_t` | Consecutive nodes per packed row: [1, 64] with grouping, [1, 32] without grouping. | +| `grouping_enabled` | `bool` | Group and reorder nodes for locality; false retains the rank-only layout. | +| `n_groups` | `std::uint32_t` | Final group count, at most 65,536. Zero uses one group if IDs fit; otherwise ceil(N * (1 + balance_tolerance) / 2^b). Each k-means node has at most 16 children. | +| `n_bits` | `std::uint16_t` | Local-ID width. Zero selects min(24, max(4, align_up_4(ceil(log2(N))))). Explicit widths must be in [1, 32]. | +| `balance_tolerance` | `double` | Allowed relative deviation from average group size, in the open interval (0, 1). Hard ID capacity still applies. | +| `training_rows` | `std::size_t` | Maximum sample rows per k-means node. Zero shares a 160,000-row budget by subtree leaf count, bounded by 256 to 10,000 rows per child and actual node size. | +| `assignment_batch_rows` | `std::size_t` | Rows assigned per GPU batch. Zero targets a 256 MiB vector buffer; distance/label buffers and host capacity-repair storage are additional. | +| `kmeans_n_iters` | `std::uint32_t` | Positive number of balanced k-means training iterations. | +| `validate` | `bool` | Verify encoded edges after rearrangement; intended for tests and small builds. | +| `num_seeds` | `std::uint32_t` | Number of medoid seeds generated at build time; zero disables seed generation. | +| `seed_training_rows` | `std::size_t` | Rows sampled to train seed centroids. Zero selects an automatic limit. | +| `seed` | `std::uint64_t` | Deterministic seed for rotating the uniform k-means training sample. | + + +### neighbors::cagra::tiered_search_params + +Experimental tiered search and reusable host queue settings. + +```cpp +struct tiered_search_params { + std::uint32_t num_seeds; + float sync_window_scale; + std::uint32_t sync_drop_threshold; + std::uint32_t num_queues; + std::uint32_t empty_pause; + bool collect_statistics; + bool keep_pollers_running; +}; +``` + +**Fields** + +| Name | Type | Description | +| --- | --- | --- | +| `num_seeds` | `std::uint32_t` | Maximum number of stored medoid seeds used per query; zero uses random initialization. | +| `sync_window_scale` | `float` | Scale for the adaptive deferred-cross-edge synchronization window. | +| `sync_drop_threshold` | `std::uint32_t` | Parent-position drop that forces synchronization of a submitted cross-edge request. | +| `num_queues` | `std::uint32_t` | Positive number of independent GPU-to-CPU command queues. | +| `empty_pause` | `std::uint32_t` | Number of CPU pause instructions after an empty poll; zero disables this pause. | +| `collect_statistics` | `bool` | Collect aggregate queue and polling counters. | +| `keep_pollers_running` | `bool` | Keep pollers active between calls for low latency. Consumes host CPU while idle. | + + +### neighbors::cagra::search_params + +CAGRA search parameters. + +**Note:** The optional tiered settings change this struct's C++ ABI, including builds with
FlowANN disabled. Rebuild consumers with headers matching the linked cuVS library. + +```cpp +struct search_params : cuvs::neighbors::search_params { + size_t max_queries; + size_t itopk_size; + size_t max_iterations; + search_algo algo; + size_t team_size; + size_t search_width; + size_t min_iterations; + size_t thread_block_size; + hash_mode hashmap_mode; + size_t hashmap_min_bitlen; + float hashmap_max_fill_rate; + uint32_t num_random_samplings; + uint64_t rand_xor_mask; + bool persistent; + float persistent_lifetime; + float persistent_device_usage; + float filtering_rate; + internal_dtype smem_dtype; + std::optional tiered; +}; +``` + +**Fields** + +| Name | Type | Description | +| --- | --- | --- | +| `max_queries` | `size_t` | Maximum number of queries to search at the same time (batch size). Auto select when 0. | +| `itopk_size` | `size_t` | Number of intermediate search results retained during the search.

This is the main knob to adjust trade off between accuracy and search speed. Higher values improve the search accuracy. | +| `max_iterations` | `size_t` | Upper limit of search iterations. Auto select when 0. | +| `algo` | [`search_algo`](/api-reference/cpp-api-neighbors-cagra#neighbors-cagra-search-algo) | Which search implementation to use. | +| `team_size` | `size_t` | Number of threads used to calculate a single distance. 4, 8, 16, or 32. | +| `search_width` | `size_t` | Number of graph nodes to select as the starting point for the search in each iteration. aka search width? | +| `min_iterations` | `size_t` | Lower limit of search iterations. | +| `thread_block_size` | `size_t` | Thread block size. 0, 64, 128, 256, 512, 1024. Auto selection when 0. | +| `hashmap_mode` | `hash_mode` | Hashmap type. Auto selection when AUTO. | +| `hashmap_min_bitlen` | `size_t` | Lower limit of hashmap bit length. More than 8. | +| `hashmap_max_fill_rate` | `float` | Upper limit of hashmap fill rate. More than 0.1, less than 0.9. | +| `num_random_samplings` | `uint32_t` | Number of iterations of initial random seed node selection. 1 or more. | +| `rand_xor_mask` | `uint64_t` | Bit mask used for initial random seed node selection. | +| `persistent` | `bool` | Whether to use the persistent version of the kernel (only SINGLE_CTA is supported a.t.m.) | +| `persistent_lifetime` | `float` | Persistent kernel: time in seconds before the kernel stops if no requests received. | +| `persistent_device_usage` | `float` | Set the fraction of maximum grid size used by persistent kernel. Value 1.0 means the kernel grid size is maximum possible for the selected device. The value must be greater than 0.0 and not greater than 1.0.

One may need to run other kernels alongside this persistent kernel. This parameter can be used to reduce the grid size of the persistent kernel to leave a few SMs idle.
Note: running any other work on GPU alongside with the persistent kernel makes the setup fragile.
- Running another kernel in another thread usually works, but no progress guaranteed
- Any CUDA allocations block the context (this issue may be obscured by using pools)
- Memory copies to not-pinned host memory may block the context

Even when we know there are no other kernels working at the same time, setting kDeviceUsage to 1.0 surprisingly sometimes hurts performance. Proceed with care. If you suspect this is an issue, you can reduce this number to ~0.9 without a significant impact on the throughput. | +| `filtering_rate` | `float` | A parameter indicating the rate of nodes to be filtered-out, when filtering is used. The value must be equal to or greater than 0.0 and less than 1.0. Default value is negative, in which case the filtering rate is automatically calculated when possible. For `filtering::udf_filter`, CAGRA uses `udf_filter::filtering_rate` when this value is negative. If both values are negative, CAGRA assumes 0.0 because a UDF's selectivity cannot be inferred from the source string. | +| `smem_dtype` | `internal_dtype` | Data type of the query vector and codebook table on shared memory. Currently, only VPQ supports FP8. * | +| `tiered` | [`std::optional`](/api-reference/cpp-api-neighbors-cagra#neighbors-cagra-tiered-search-params) | Settings for tiered indexes. Must be unset for ordinary device-graph indexes; search throws otherwise. Tiered indexes use default settings when this is unset. | + ## CAGRA index build parameters @@ -100,6 +244,37 @@ enum class hnsw_heuristic_type : uint32_t { | `SIMILAR_SEARCH_PERFORMANCE` | `0` | | `SAME_GRAPH_FOOTPRINT` | `1` | + +### neighbors::cagra::index_params + +CAGRA index construction parameters. + +**Note:** The graph storage selector and optional tiered settings change this struct's C++ ABI,
including builds with FlowANN disabled. Rebuild consumers against matching cuVS headers. + +```cpp +struct index_params : cuvs::neighbors::index_params { + size_t intermediate_graph_degree; + size_t graph_degree; + graph_build_params_t graph_build_params; + bool guarantee_connectivity; + bool attach_dataset_on_build; + graph_storage_kind graph_storage; + std::optional tiered; +}; +``` + +**Fields** + +| Name | Type | Description | +| --- | --- | --- | +| `intermediate_graph_degree` | `size_t` | Degree of input graph for pruning. | +| `graph_degree` | `size_t` | Degree of output graph. | +| `graph_build_params` | [`graph_build_params_t`](/api-reference/cpp-api-neighbors-cagra#graph-build-params-t) | Parameters for graph building.

Set ivf_pq_params, nn_descent_params, ace_params, or iterative_search_params to select the graph build algorithm and control their parameters. The default (std::monostate) is to use a heuristic to decide the algorithm and its parameters. | +| `guarantee_connectivity` | `bool` | Whether to use MST optimization to guarantee graph connectivity. | +| `attach_dataset_on_build` | `bool` | Whether to attach the dataset to the index after graph construction, i.e.:

- `true` (default) means `build` attaches the input dataset as a **non-owning view** to the index. The caller is responsible for keeping the underlying dataset storage alive for as long as the index is used. A device-backed index is ready to search immediately; a host-backed index retains the dataset for operations such as serialization but is not searchable.
- `false` means `build` only builds the graph and the caller is expected to attach a dataset separately via `cuvs::neighbors::cagra::update_dataset` before searching.

Unlike the legacy behavior, no copy of the dataset is made: the index always stores a view. Setting `attach_dataset_on_build = false` is useful when the caller needs to apply specific memory placement or transformation (e.g. moving to managed memory) before attaching.

Host indexes are not directly searchable. Call the type-changing `update_dataset` with a user-provided device-padded dataset view to obtain a search-ready `device_padded_index`. Disk-based ACE builds manage file-backed dataset state separately and ignore this flag. | +| `graph_storage` | [`graph_storage_kind`](/api-reference/cpp-api-neighbors-cagra#neighbors-cagra-graph-storage-kind) | Ordinary device graph by default; tiered storage is experimental and opt-in. | +| `tiered` | [`std::optional`](/api-reference/cpp-api-neighbors-cagra#neighbors-cagra-tiered-graph-params) | A zero tiered graph budget stores all graph edges on the host. | + ### neighbors::cagra::index_params::graph_build_heuristic @@ -156,7 +331,7 @@ Usage example: **Returns** -`static cagra::index_params` +[`static cagra::index_params`](/api-reference/cpp-api-neighbors-cagra#neighbors-cagra-index-params) ### neighbors::cagra::index_params::from_hnsw_params @@ -190,7 +365,7 @@ Usage example: **Returns** -`static cagra::index_params` +[`static cagra::index_params`](/api-reference/cpp-api-neighbors-cagra#neighbors-cagra-index-params) ## Types @@ -332,10 +507,23 @@ Number of rows represented by the graph. `IdxT` + +### neighbors::cagra::index::graph_storage + +Return the actual graph representation. + +```cpp +[[nodiscard]] auto graph_storage() const noexcept -> graph_storage_kind; +``` + +**Returns** + +[`graph_storage_kind`](/api-reference/cpp-api-neighbors-cagra#neighbors-cagra-graph-storage-kind) + ### neighbors::cagra::index::dataset -Non-owning dataset binding stored by the index. +Dataset binding; a tiered index owns its reordered dataset internally. ```cpp [[nodiscard]] inline auto dataset() const noexcept -> DatasetViewT const&; @@ -348,13 +536,15 @@ Non-owning dataset binding stored by the index. ### neighbors::cagra::index::graph -neighborhood graph [size, graph-degree] +Device neighborhood graph with shape [size, graph_degree]. ```cpp -[[nodiscard]] inline auto graph() const noexcept +[[nodiscard]] inline auto graph() const -> raft::device_matrix_view; ``` +Throws raft::logic_error when graph_storage() is tiered. Tiered graphs have separate device and host edge storage and cannot be represented by a single device adjacency matrix. + **Returns** `raft::device_matrix_view` @@ -699,7 +889,7 @@ The returned index accepts float queries and stores a non-owning copy of `datase | Name | Direction | Type | Description | | --- | --- | --- | --- | | `res` | in | `raft::resources const&` | raft resources | -| `params` | in | `const cuvs::neighbors::cagra::index_params&` | CAGRA index build parameters | +| `params` | in | [`const cuvs::neighbors::cagra::index_params&`](/api-reference/cpp-api-neighbors-cagra#neighbors-cagra-index-params) | CAGRA index build parameters | | `dataset` | in | `cuvs::neighbors::device_vpq_dataset_view` | device VPQ dataset view | **Returns** @@ -724,7 +914,7 @@ cuvs::neighbors::device_padded_dataset_view dataset) | Name | Direction | Type | Description | | --- | --- | --- | --- | | `res` | in | `raft::resources const&` | raft resources | -| `params` | in | `const cuvs::neighbors::cagra::index_params&` | CAGRA index build parameters | +| `params` | in | [`const cuvs::neighbors::cagra::index_params&`](/api-reference/cpp-api-neighbors-cagra#neighbors-cagra-index-params) | CAGRA index build parameters | | `dataset` | in | `cuvs::neighbors::device_padded_dataset_view` | device padded dataset view [n_rows, dim] | **Returns** @@ -749,7 +939,7 @@ cuvs::neighbors::device_standard_dataset_view dataset) | Name | Direction | Type | Description | | --- | --- | --- | --- | | `res` | in | `raft::resources const&` | raft resources | -| `params` | in | `const cuvs::neighbors::cagra::index_params&` | CAGRA index build parameters | +| `params` | in | [`const cuvs::neighbors::cagra::index_params&`](/api-reference/cpp-api-neighbors-cagra#neighbors-cagra-index-params) | CAGRA index build parameters | | `dataset` | in | `cuvs::neighbors::device_standard_dataset_view` | device standard dataset view [n_rows, dim] | **Returns** @@ -774,7 +964,7 @@ cuvs::neighbors::host_padded_dataset_view dataset) | Name | Direction | Type | Description | | --- | --- | --- | --- | | `res` | in | `raft::resources const&` | raft resources | -| `params` | in | `const cuvs::neighbors::cagra::index_params&` | CAGRA index build parameters | +| `params` | in | [`const cuvs::neighbors::cagra::index_params&`](/api-reference/cpp-api-neighbors-cagra#neighbors-cagra-index-params) | CAGRA index build parameters | | `dataset` | in | `cuvs::neighbors::host_padded_dataset_view` | host padded dataset view [n_rows, dim] | **Returns** @@ -799,7 +989,7 @@ cuvs::neighbors::host_standard_dataset_view dataset) | Name | Direction | Type | Description | | --- | --- | --- | --- | | `res` | in | `raft::resources const&` | raft resources | -| `params` | in | `const cuvs::neighbors::cagra::index_params&` | CAGRA index build parameters | +| `params` | in | [`const cuvs::neighbors::cagra::index_params&`](/api-reference/cpp-api-neighbors-cagra#neighbors-cagra-index-params) | CAGRA index build parameters | | `dataset` | in | `cuvs::neighbors::host_standard_dataset_view` | host standard dataset view [n_rows, dim] | **Returns** @@ -824,7 +1014,7 @@ cuvs::neighbors::device_padded_dataset_view dataset) | Name | Direction | Type | Description | | --- | --- | --- | --- | | `res` | in | `raft::resources const&` | raft resources | -| `params` | in | `const cuvs::neighbors::cagra::index_params&` | CAGRA index build parameters | +| `params` | in | [`const cuvs::neighbors::cagra::index_params&`](/api-reference/cpp-api-neighbors-cagra#neighbors-cagra-index-params) | CAGRA index build parameters | | `dataset` | in | `cuvs::neighbors::device_padded_dataset_view` | device padded dataset view [n_rows, dim] | **Returns** @@ -849,7 +1039,7 @@ cuvs::neighbors::device_standard_dataset_view dataset) | Name | Direction | Type | Description | | --- | --- | --- | --- | | `res` | in | `raft::resources const&` | raft resources | -| `params` | in | `const cuvs::neighbors::cagra::index_params&` | CAGRA index build parameters | +| `params` | in | [`const cuvs::neighbors::cagra::index_params&`](/api-reference/cpp-api-neighbors-cagra#neighbors-cagra-index-params) | CAGRA index build parameters | | `dataset` | in | `cuvs::neighbors::device_standard_dataset_view` | device standard dataset view [n_rows, dim] | **Returns** @@ -874,7 +1064,7 @@ cuvs::neighbors::host_padded_dataset_view dataset) | Name | Direction | Type | Description | | --- | --- | --- | --- | | `res` | in | `raft::resources const&` | raft resources | -| `params` | in | `const cuvs::neighbors::cagra::index_params&` | CAGRA index build parameters | +| `params` | in | [`const cuvs::neighbors::cagra::index_params&`](/api-reference/cpp-api-neighbors-cagra#neighbors-cagra-index-params) | CAGRA index build parameters | | `dataset` | in | `cuvs::neighbors::host_padded_dataset_view` | host padded dataset view [n_rows, dim] | **Returns** @@ -899,7 +1089,7 @@ cuvs::neighbors::host_standard_dataset_view dataset) | Name | Direction | Type | Description | | --- | --- | --- | --- | | `res` | in | `raft::resources const&` | raft resources | -| `params` | in | `const cuvs::neighbors::cagra::index_params&` | CAGRA index build parameters | +| `params` | in | [`const cuvs::neighbors::cagra::index_params&`](/api-reference/cpp-api-neighbors-cagra#neighbors-cagra-index-params) | CAGRA index build parameters | | `dataset` | in | `cuvs::neighbors::host_standard_dataset_view` | host standard dataset view [n_rows, dim] | **Returns** @@ -924,7 +1114,7 @@ cuvs::neighbors::device_padded_dataset_view dataset) | Name | Direction | Type | Description | | --- | --- | --- | --- | | `res` | in | `raft::resources const&` | raft resources | -| `params` | in | `const cuvs::neighbors::cagra::index_params&` | CAGRA index build parameters | +| `params` | in | [`const cuvs::neighbors::cagra::index_params&`](/api-reference/cpp-api-neighbors-cagra#neighbors-cagra-index-params) | CAGRA index build parameters | | `dataset` | in | `cuvs::neighbors::device_padded_dataset_view` | device padded dataset view [n_rows, dim] | **Returns** @@ -949,7 +1139,7 @@ cuvs::neighbors::device_standard_dataset_view dataset) | Name | Direction | Type | Description | | --- | --- | --- | --- | | `res` | in | `raft::resources const&` | raft resources | -| `params` | in | `const cuvs::neighbors::cagra::index_params&` | CAGRA index build parameters | +| `params` | in | [`const cuvs::neighbors::cagra::index_params&`](/api-reference/cpp-api-neighbors-cagra#neighbors-cagra-index-params) | CAGRA index build parameters | | `dataset` | in | `cuvs::neighbors::device_standard_dataset_view` | device standard dataset view [n_rows, dim] | **Returns** @@ -974,7 +1164,7 @@ cuvs::neighbors::host_padded_dataset_view dataset) | Name | Direction | Type | Description | | --- | --- | --- | --- | | `res` | in | `raft::resources const&` | raft resources | -| `params` | in | `const cuvs::neighbors::cagra::index_params&` | CAGRA index build parameters | +| `params` | in | [`const cuvs::neighbors::cagra::index_params&`](/api-reference/cpp-api-neighbors-cagra#neighbors-cagra-index-params) | CAGRA index build parameters | | `dataset` | in | `cuvs::neighbors::host_padded_dataset_view` | host padded dataset view [n_rows, dim] | **Returns** @@ -999,7 +1189,7 @@ cuvs::neighbors::host_standard_dataset_view dataset) | Name | Direction | Type | Description | | --- | --- | --- | --- | | `res` | in | `raft::resources const&` | raft resources | -| `params` | in | `const cuvs::neighbors::cagra::index_params&` | CAGRA index build parameters | +| `params` | in | [`const cuvs::neighbors::cagra::index_params&`](/api-reference/cpp-api-neighbors-cagra#neighbors-cagra-index-params) | CAGRA index build parameters | | `dataset` | in | `cuvs::neighbors::host_standard_dataset_view` | host standard dataset view [n_rows, dim] | **Returns** @@ -1024,7 +1214,7 @@ cuvs::neighbors::device_padded_dataset_view dataset) | Name | Direction | Type | Description | | --- | --- | --- | --- | | `res` | in | `raft::resources const&` | raft resources | -| `params` | in | `const cuvs::neighbors::cagra::index_params&` | CAGRA index build parameters | +| `params` | in | [`const cuvs::neighbors::cagra::index_params&`](/api-reference/cpp-api-neighbors-cagra#neighbors-cagra-index-params) | CAGRA index build parameters | | `dataset` | in | `cuvs::neighbors::device_padded_dataset_view` | device padded dataset view [n_rows, dim] | **Returns** @@ -1049,7 +1239,7 @@ cuvs::neighbors::device_standard_dataset_view dataset) | Name | Direction | Type | Description | | --- | --- | --- | --- | | `res` | in | `raft::resources const&` | raft resources | -| `params` | in | `const cuvs::neighbors::cagra::index_params&` | CAGRA index build parameters | +| `params` | in | [`const cuvs::neighbors::cagra::index_params&`](/api-reference/cpp-api-neighbors-cagra#neighbors-cagra-index-params) | CAGRA index build parameters | | `dataset` | in | `cuvs::neighbors::device_standard_dataset_view` | device standard dataset view [n_rows, dim] | **Returns** @@ -1074,7 +1264,7 @@ cuvs::neighbors::host_padded_dataset_view dataset) | Name | Direction | Type | Description | | --- | --- | --- | --- | | `res` | in | `raft::resources const&` | raft resources | -| `params` | in | `const cuvs::neighbors::cagra::index_params&` | CAGRA index build parameters | +| `params` | in | [`const cuvs::neighbors::cagra::index_params&`](/api-reference/cpp-api-neighbors-cagra#neighbors-cagra-index-params) | CAGRA index build parameters | | `dataset` | in | `cuvs::neighbors::host_padded_dataset_view` | host padded dataset view [n_rows, dim] | **Returns** @@ -1099,7 +1289,7 @@ cuvs::neighbors::host_standard_dataset_view dataset) | Name | Direction | Type | Description | | --- | --- | --- | --- | | `res` | in | `raft::resources const&` | raft resources | -| `params` | in | `const cuvs::neighbors::cagra::index_params&` | CAGRA index build parameters | +| `params` | in | [`const cuvs::neighbors::cagra::index_params&`](/api-reference/cpp-api-neighbors-cagra#neighbors-cagra-index-params) | CAGRA index build parameters | | `dataset` | in | `cuvs::neighbors::host_standard_dataset_view` | host standard dataset view [n_rows, dim] | **Returns** @@ -1128,7 +1318,7 @@ The returned index cannot be searched: CAGRA has no BBQ search kernels. Call the | Name | Direction | Type | Description | | --- | --- | --- | --- | | `res` | in | `raft::resources const&` | raft resources | -| `params` | in | `const cuvs::neighbors::cagra::index_params&` | CAGRA index build parameters | +| `params` | in | [`const cuvs::neighbors::cagra::index_params&`](/api-reference/cpp-api-neighbors-cagra#neighbors-cagra-index-params) | CAGRA index build parameters | | `dataset` | in | `cuvs::neighbors::device_bbq_dataset_view` | device BBQ dataset view [n_rows, dim] | **Returns** @@ -1139,7 +1329,7 @@ built `device_bbq_index` **Additional overload:** `neighbors::cagra::build` -cuvs::neighbors::device_bbq_dataset_view<float, int64_t> dataset) +Build from a device BBQ-quantized dataset view of type `half`. ```cpp auto build(raft::resources const& res, @@ -1148,21 +1338,25 @@ cuvs::neighbors::device_bbq_dataset_view dataset) -> cuvs::neighbors::cagra::device_bbq_index; ``` +Construction uses quantized codes and requires nn-descent. The metric must match the quantizer corrections and be L2Expanded, L2SqrtExpanded, CosineExpanded, or InnerProduct. The returned index has no BBQ search support; use the type-changing `update_dataset` with an uncompressed device-padded dataset before searching the resulting graph. + **Parameters** | Name | Direction | Type | Description | | --- | --- | --- | --- | -| `res` | | `raft::resources const&` | | -| `params` | | `const cuvs::neighbors::cagra::index_params&` | | -| `dataset` | | `cuvs::neighbors::device_bbq_dataset_view` | | +| `res` | in | `raft::resources const&` | RAFT resources used for graph construction. | +| `params` | in | [`const cuvs::neighbors::cagra::index_params&`](/api-reference/cpp-api-neighbors-cagra#neighbors-cagra-index-params) | CAGRA index build parameters. | +| `dataset` | in | `cuvs::neighbors::device_bbq_dataset_view` | Device BBQ dataset view with shape [n_rows, dim]. | **Returns** `cuvs::neighbors::cagra::device_bbq_index` +Built `device_bbq_index`. + **Additional overload:** `neighbors::cagra::build` -cuvs::neighbors::device_bbq_dataset_view<float, int64_t> dataset) +Build from a device BBQ-quantized dataset view of type `int8_t`. ```cpp auto build(raft::resources const& res, @@ -1171,21 +1365,25 @@ cuvs::neighbors::device_bbq_dataset_view dataset) -> cuvs::neighbors::cagra::device_bbq_index; ``` +Construction uses quantized codes and requires nn-descent. The metric must match the quantizer corrections and be L2Expanded, L2SqrtExpanded, CosineExpanded, or InnerProduct. The returned index has no BBQ search support; use the type-changing `update_dataset` with an uncompressed device-padded dataset before searching the resulting graph. + **Parameters** | Name | Direction | Type | Description | | --- | --- | --- | --- | -| `res` | | `raft::resources const&` | | -| `params` | | `const cuvs::neighbors::cagra::index_params&` | | -| `dataset` | | `cuvs::neighbors::device_bbq_dataset_view` | | +| `res` | in | `raft::resources const&` | RAFT resources used for graph construction. | +| `params` | in | [`const cuvs::neighbors::cagra::index_params&`](/api-reference/cpp-api-neighbors-cagra#neighbors-cagra-index-params) | CAGRA index build parameters. | +| `dataset` | in | `cuvs::neighbors::device_bbq_dataset_view` | Device BBQ dataset view with shape [n_rows, dim]. | **Returns** `cuvs::neighbors::cagra::device_bbq_index` +Built `device_bbq_index`. + **Additional overload:** `neighbors::cagra::build` -cuvs::neighbors::device_bbq_dataset_view<float, int64_t> dataset) +Build from a device BBQ-quantized dataset view of type `uint8_t`. ```cpp auto build(raft::resources const& res, @@ -1194,18 +1392,22 @@ cuvs::neighbors::device_bbq_dataset_view dataset) -> cuvs::neighbors::cagra::device_bbq_index; ``` +Construction uses quantized codes and requires nn-descent. The metric must match the quantizer corrections and be L2Expanded, L2SqrtExpanded, CosineExpanded, or InnerProduct. The returned index has no BBQ search support; use the type-changing `update_dataset` with an uncompressed device-padded dataset before searching the resulting graph. + **Parameters** | Name | Direction | Type | Description | | --- | --- | --- | --- | -| `res` | | `raft::resources const&` | | -| `params` | | `const cuvs::neighbors::cagra::index_params&` | | -| `dataset` | | `cuvs::neighbors::device_bbq_dataset_view` | | +| `res` | in | `raft::resources const&` | RAFT resources used for graph construction. | +| `params` | in | [`const cuvs::neighbors::cagra::index_params&`](/api-reference/cpp-api-neighbors-cagra#neighbors-cagra-index-params) | CAGRA index build parameters. | +| `dataset` | in | `cuvs::neighbors::device_bbq_dataset_view` | Device BBQ dataset view with shape [n_rows, dim]. | **Returns** `cuvs::neighbors::cagra::device_bbq_index` +Built `device_bbq_index`. + ## CAGRA extend functions diff --git a/fern/pages/neighbors/cagra.md b/fern/pages/neighbors/cagra.md index 20dea7ce39..15aa391f7a 100644 --- a/fern/pages/neighbors/cagra.md +++ b/fern/pages/neighbors/cagra.md @@ -1461,3 +1461,75 @@ $$ - `result_size = 8,000 B = 0.0076 MB` - `workspace_size = query_size + result_size = 0.40 MB` - `total search memory ~= 3906.25 + 244.14 + 0.40 = 4150.79 MB` + +## Experimental tiered graph storage (FlowANN) + +CAGRA can store its graph across GPU and host memory using the FlowANN implementation. +Enable `CUVS_ENABLE_FLOWANN_SEARCH` and `CUVS_ENABLE_FLOWANN_BUILD` when building cuVS. +`CUVS_FLOWANN_USE_GDRCOPY` defaults to `ON`. With FlowANN search enabled, this default +requires GDRCopy to be installed and discoverable at CMake configuration time. Set +`-DCUVS_FLOWANN_USE_GDRCOPY=OFF` to select the CUDA-copy transport instead. These options are +not required for ordinary CAGRA. The public entry points remain CAGRA's `build`, `search`, +`serialize`, and `deserialize`; there is no separate public FlowANN index type. + +Set `index_params.graph_storage` to `graph_storage_kind::tiered`. The optional +`tiered_graph_params` controls the resident graph budget, grouping, and seed generation. +The budget counts packed graph edges, excluding the dataset, queues, and search workspace. +Zero places all graph edges on the host. Construction still needs memory for the complete +intermediate graph. Automatic grouping uses hierarchical k-means with capacity repair; +`n_groups=0` and `n_bits=0` select the existing size-dependent defaults. +The integrated tiered builder does not support ACE graph construction. Dense C++ builds +require a device-padded dataset; the compressed overload requires host standard vectors. + +```cpp +cuvs::neighbors::cagra::index_params params; +params.graph_storage = cuvs::neighbors::cagra::graph_storage_kind::tiered; +params.tiered.emplace(); +params.tiered->device_graph_budget_bytes = 1ULL << 30; +// dense_view is a device_padded_dataset_view. +auto index = cuvs::neighbors::cagra::build(resources, params, dense_view); + +cuvs::neighbors::cagra::search_params search_params; +search_params.itopk_size = 256; +search_params.max_iterations = 32; +search_params.tiered.emplace(); +search_params.tiered->keep_pollers_running = true; +cuvs::neighbors::cagra::search(resources, search_params, index, queries, ids, distances); +``` + +The index owns the reordered dataset and its queue runtime. Releasing the build input or moving +the index does not invalidate that data. Search returns original dataset IDs. Queue workers are +reused across calls; `keep_pollers_running=true` also keeps the workers active between calls, +which consumes CPU while the index is idle. Searches on one index are serialized and complete +before returning. Use the CUDA device on which the index was created. + +For VPQ-F16, use the overload `cagra::build(resources, params, vpq_params, host_view)`. +It builds and groups using the full-precision host vectors, then produces an owning compressed +index. + +Tiered serialization stores both the graph and reordered dataset, with a distinct format marker. +Graph-only serialization is unsupported. The C++ VPQ serialization overload stores the complete +payload for tiered storage and retains its existing graph-only behavior for ordinary CAGRA-Q. +A tiered C index owns its loaded dataset internally, so `cuvsCagraDeserializeGraphAndDataset` +leaves `out_dataset` null. +Historical split import remains available inside the native benchmark importer. + +Tiered storage does not expose a device adjacency matrix through `graph()`. Dataset replacement, +extension, merge, multi-partition search, and GPU persistent search are unsupported and report +errors. Ordinary device graph storage retains the existing defaults and behavior. Rebuild consumers +against the updated headers: the C/C++ parameter and index layouts have changed. + + +## Migrating existing C++ applications + +The FlowANN integration changes the C++ layouts of `cuvs::neighbors::cagra::index`, +`index_params`, and `search_params`. This also applies when FlowANN is disabled at build time. +Rebuild applications and any libraries that pass or instantiate these types, using headers +from the same cuVS version as the linked library. Do not combine previously compiled CAGRA +objects with the new library. Ordinary device-graph searches keep their existing defaults; +leave `search_params::tiered` unset for those indexes. + +A tiered index stores device and host edges separately. Check `index::graph_storage()` before +calling `index::graph()`, which throws for tiered indexes. Use the tiered search and persistence +APIs for that representation. The C++ layout change does not by itself require a C ABI SONAME +change; the C API follows its separate ABI policy. diff --git a/fern/scripts/generate_api_reference.py b/fern/scripts/generate_api_reference.py index 22542ca044..9cb9767e3d 100755 --- a/fern/scripts/generate_api_reference.py +++ b/fern/scripts/generate_api_reference.py @@ -3620,6 +3620,11 @@ def split_description_breaks(line: str) -> list[str]: def parse_doxygen_kind(declaration: str) -> str: untemplated = strip_leading_cpp_templates(declaration) + if re.fullmatch( + r"\s*(?:struct|class)\s+[\w:]+(?:\s+|\s*[*&]+\s*)[A-Za-z_]\w*\s*;?\s*", + untemplated, + ): + return "member" if re.search(r"^\s*(?:typedef\s+)?(?:struct|class)\b", untemplated) and ( "{" in declaration or not untemplated.lstrip().startswith("typedef") ): @@ -4063,12 +4068,13 @@ def parse_plain_struct_members( continue if declaration.startswith((",", ":")): continue + if parse_doxygen_kind(declaration) == "struct": + continue if not declaration or declaration.startswith( ( "enum ", "friend ", "static_assert", - "struct ", "template ", "typedef ", "using ", diff --git a/fern/scripts/test_generate_api_reference.py b/fern/scripts/test_generate_api_reference.py new file mode 100644 index 0000000000..61d560ca62 --- /dev/null +++ b/fern/scripts/test_generate_api_reference.py @@ -0,0 +1,64 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 + +"""Regression checks for elaborated C field types in the native API reference.""" + +import importlib.util +from pathlib import Path +import sys +import unittest + +spec = importlib.util.spec_from_file_location( + "fern_api_reference_under_test", + Path(__file__).with_name("generate_api_reference.py"), +) +api = importlib.util.module_from_spec(spec) +sys.modules[spec.name] = api +spec.loader.exec_module(api) + + +class StructFieldDocumentationTest(unittest.TestCase): + def test_elaborated_field_is_a_member(self): + self.assertEqual( + api.parse_doxygen_kind("struct Options tiered;"), "member" + ) + self.assertEqual( + api.parse_doxygen_kind("struct Options* tiered;"), "member" + ) + + def test_structure_declarations_keep_their_kind(self): + self.assertEqual(api.parse_doxygen_kind("struct Options;"), "struct") + self.assertEqual( + api.parse_doxygen_kind("struct Options { int value; };"), "struct" + ) + + def test_documented_field_keeps_its_type_and_description(self): + entry = api.parse_doxygen_entry( + "Parameters", + "struct Params { /** Tiered settings. */ struct Options tiered; };", + "test.h", + 1, + ) + members = api.parse_struct_members(entry) + self.assertEqual([m.name for m in members], ["tiered"]) + self.assertEqual(api.member_c_type(members[0]), "struct Options") + self.assertIn("Tiered settings", api.member_description(members[0])) + self.assertIn( + "struct Options tiered;", + api.render_compound_signature(entry, members, []), + ) + + def test_plain_field_keeps_its_type(self): + entry = api.parse_doxygen_entry( + "Parameters", + "struct Params { struct Options tiered; };", + "test.h", + 1, + ) + members = api.parse_struct_members(entry) + self.assertEqual([m.name for m in members], ["tiered"]) + self.assertEqual(api.member_c_type(members[0]), "struct Options") + + +if __name__ == "__main__": + unittest.main() diff --git a/python/cuvs_bench/cuvs_bench/backends/__init__.py b/python/cuvs_bench/cuvs_bench/backends/__init__.py index a09e6fd1fe..be94f1fb73 100644 --- a/python/cuvs_bench/cuvs_bench/backends/__init__.py +++ b/python/cuvs_bench/cuvs_bench/backends/__init__.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """ @@ -25,11 +25,13 @@ ) from .cpp_gbench import CppGoogleBenchmarkBackend +from .flowann import FlowannBackend from .opensearch import OpenSearchBackend # Auto-register built-in backends _registry = get_registry() _registry.register("cpp_gbench", CppGoogleBenchmarkBackend) +_registry.register("flowann", FlowannBackend) _registry.register("opensearch", OpenSearchBackend) __all__ = [ @@ -45,5 +47,6 @@ "get_backend", # Built-in backends "CppGoogleBenchmarkBackend", + "FlowannBackend", "OpenSearchBackend", ] diff --git a/python/cuvs_bench/cuvs_bench/backends/flowann.py b/python/cuvs_bench/cuvs_bench/backends/flowann.py new file mode 100644 index 0000000000..ef905a7304 --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/backends/flowann.py @@ -0,0 +1,251 @@ +# 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). +"""Search existing FlowANN indexes through FLOWANN_BENCHMARK.""" + +import json +import math +import os +from pathlib import Path +import re +import shlex +import statistics +import subprocess +import uuid + +import numpy as np + +from .base import BenchmarkBackend, SearchResult + + +def parse_measurements(text, expected_rounds): + """Read measured rounds only, rejecting incomplete or invalid runs.""" + rounds = [] + for line in text.splitlines(): + if not line.startswith("MEASURE "): + continue + fields = dict(re.findall(r"(\w+)=([^\s]+)", line)) + row = {} + for key in ( + "round", + "batch", + "query_count", + "search_ms", + "rerank_prep_ms", + "rerank_ms", + "total_ms", + "recall", + "qps", + "p50_ms", + "p95_ms", + "p99_ms", + ): + if key not in fields: + raise ValueError(f"Missing benchmark measurement field: {key}") + try: + row[key] = float(fields[key]) + except ValueError as error: + raise ValueError( + f"Benchmark measurement field {key} must be numeric" + ) from error + if ( + not math.isfinite(row["query_count"]) + or row["query_count"] <= 0 + or not row["query_count"].is_integer() + ): + raise ValueError( + "Measurement query_count must be a positive integer" + ) + if not all(math.isfinite(value) for value in row.values()): + raise ValueError("Non-finite benchmark measurement") + if ( + not 0 <= row["recall"] <= 1 + or row["qps"] <= 0 + or row["total_ms"] <= 0 + ): + raise ValueError("Invalid recall or timing") + rounds.append(row) + if len(rounds) != expected_rounds: + raise ValueError( + f"Expected {expected_rounds} measured rounds, got {len(rounds)}" + ) + # Native programs use different starting round numbers. + round_ids = sorted(row["round"] for row in rounds) + first = round_ids[0] + if ( + first < 0 + or not first.is_integer() + or round_ids != list(range(int(first), int(first) + expected_rounds)) + ): + raise ValueError( + "Measured round IDs must be consecutive nonnegative integers" + ) + if len({row["query_count"] for row in rounds}) != 1: + raise ValueError( + "Measurement query_count must be consistent across rounds" + ) + return rounds + + +class FlowannBackend(BenchmarkBackend): + """Keep native warmup/session semantics and return cuvs-bench results. + + config requires executable_path and output_dir. Optional command_prefix + (an argv list) supports taskset; env supplies process-local environment. + IndexConfig.file is an existing index, never an index to rebuild. + """ + + @property + def algo(self): + return self.config.get("algo", "flowann") + + def build(self, dataset, indexes, force=False, dry_run=False): + raise NotImplementedError( + "This backend searches existing indexes only" + ) + + def search( + self, + dataset, + indexes, + k, + batch_size=10000, + mode="latency", + force=False, + search_threads=None, + dry_run=False, + ): + if mode != "latency": + raise ValueError( + "Use latency mode; batch size controls query batching" + ) + executable = str(Path(self.config["executable_path"]).resolve()) + output = Path(self.config["output_dir"]) + results = [] + for index in indexes: + for params in index.search_params: + options = dict(params) + forbidden = { + "index", + "base", + "queries", + "ground_truth", + "batch_size", + "top_k", + "no_rerank", + "no_session", + "queue_statistics", + } + if forbidden.intersection(options): + raise ValueError( + "Search parameters override required benchmark semantics" + ) + options.setdefault("warmup_rounds", 1) + options.setdefault("measured_rounds", 3) + options.setdefault("query_count", 10000) + if ( + options["warmup_rounds"] < 1 + or options["measured_rounds"] < 1 + ): + raise ValueError( + "Warmup and measurement rounds must be positive" + ) + if search_threads is not None: + if ( + "cpu_threads" in options + and options["cpu_threads"] != search_threads + ): + raise ValueError("Conflicting CPU thread counts") + options["cpu_threads"] = search_threads + required = { + "index": index.file, + "base": dataset.base_file, + "queries": dataset.query_file, + "ground_truth": dataset.groundtruth_neighbors_file, + "top_k": k, + "batch_size": batch_size, + } + if any(value is None for value in required.values()): + raise ValueError( + "Index, base, queries and ground truth paths are required" + ) + command = list(self.config.get("command_prefix", [])) + [ + executable + ] + for key, value in (required | options).items(): + if not re.fullmatch(r"[a-z][a-z0-9_]*", key): + raise ValueError(f"Invalid option name: {key}") + flag = "--" + key.replace("_", "-") + if isinstance(value, bool): + if value: + command.append(flag) + else: + command += [flag, str(value)] + if dry_run: + print(shlex.join(command)) + continue + output.mkdir(parents=True, exist_ok=True) + stem = f"{index.algo}-bs{batch_size}-{uuid.uuid4().hex[:12]}" + log = output / (stem + ".log") + metadata = { + "command": command, + "params": options, + "index": index.file, + "batch_size": batch_size, + } + (output / (stem + ".command.json")).write_text( + json.dumps(metadata, indent=2) + "\n" + ) + env = os.environ | self.config.get("env", {}) + with log.open("w") as stream: + subprocess.run( + command, + stdout=stream, + stderr=subprocess.STDOUT, + check=True, + env=env, + ) + rounds = parse_measurements( + log.read_text(), options["measured_rounds"] + ) + if any(row["batch"] != batch_size for row in rounds): + raise ValueError( + "Output batch does not match requested batch" + ) + medians = { + key: statistics.median(row[key] for row in rounds) + for key in rounds[0] + } + query_count = int(rounds[0]["query_count"]) + result = SearchResult( + neighbors=np.empty((0, k), dtype=np.int64), + distances=np.empty((0, k), dtype=np.float32), + search_time_ms=medians["total_ms"] + * query_count + / batch_size, + queries_per_second=medians["qps"], + recall=medians["recall"], + algorithm=index.algo, + search_params=[options], + latency_percentiles={ + key: medians[key] + for key in ("p50_ms", "p95_ms", "p99_ms") + }, + metadata={ + "native_log": str(log), + "rounds": rounds, + "batch_latency_ms": medians["total_ms"], + "recall_min": min(row["recall"] for row in rounds), + "query_count": query_count, + "batch_size": batch_size, + "index": index.file, + "command": command, + }, + ) + (output / (stem + ".json")).write_text( + json.dumps(result.to_json(), indent=2) + "\n" + ) + results.append(result) + return results diff --git a/python/cuvs_bench/cuvs_bench/tests/test_flowann_backend.py b/python/cuvs_bench/cuvs_bench/tests/test_flowann_backend.py new file mode 100644 index 0000000000..c1d573b735 --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/tests/test_flowann_backend.py @@ -0,0 +1,134 @@ +# 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). +import pytest +from cuvs_bench.backends.flowann import parse_measurements + + +def measurement( + kind="MEASURE", recall="0.95", round_id=0, query_count="10000" +): + return ( + f"{kind} round={round_id} batch=1 query_count={query_count} " + "search_ms=1 rerank_prep_ms=0.1 " + f"rerank_ms=0.2 total_ms=1.3 recall={recall} qps=769.23 " + "p50_ms=1.2 p95_ms=1.5 p99_ms=2" + ) + + +def test_excludes_warmup_and_summary(): + text = "\n".join( + [ + measurement("WARMUP", "0.1"), + measurement(), + "SUMMARY recall_median=0.1", + ] + ) + assert parse_measurements(text, 1)[0]["recall"] == 0.95 + + +@pytest.mark.parametrize( + "text,count", + [ + (measurement(recall="nan"), 1), + (measurement(recall="1.1"), 1), + (measurement(), 3), + (measurement() + "\n" + measurement(), 2), + (measurement(round_id=-1), 1), + (measurement(round_id=0) + "\n" + measurement(round_id=2), 2), + ], +) +def test_rejects_invalid_or_incomplete_results(text, count): + with pytest.raises(ValueError): + parse_measurements(text, count) + + +@pytest.mark.parametrize("first", [0, 1, 2]) +def test_accepts_native_numbering_after_warmup(first): + text = "\n".join(measurement(round_id=i) for i in range(first, first + 3)) + assert len(parse_measurements(text, 3)) == 3 + + +@pytest.mark.parametrize("query_count", ["0", "1.5", "nan", "invalid"]) +def test_rejects_invalid_query_count(query_count): + with pytest.raises(ValueError, match="query_count"): + parse_measurements(measurement(query_count=query_count), 1) + + +def test_rejects_missing_query_count(): + text = measurement().replace(" query_count=10000", "") + with pytest.raises(ValueError, match="Missing.*query_count"): + parse_measurements(text, 1) + + +def test_rejects_inconsistent_query_count(): + text = "\n".join( + [ + measurement(round_id=0, query_count="10000"), + measurement(round_id=1, query_count="5000"), + ] + ) + with pytest.raises(ValueError, match="consistent"): + parse_measurements(text, 2) + + +@pytest.mark.parametrize( + "requested_query_count,actual_query_count", + [(0, 8000), (20000, 8000)], +) +def test_backend_uses_actual_query_count( + tmp_path, monkeypatch, requested_query_count, actual_query_count +): + from cuvs_bench.backends import Dataset, get_backend + from cuvs_bench.orchestrator.config_loaders import IndexConfig + import cuvs_bench.backends.flowann as module + + def fake_run(command, stdout, **kwargs): + assert "--include-query-transfer" in command + assert "True" not in command + assert command[command.index("--cpu-threads") + 1] == "128" + stdout.write(measurement("WARMUP", "0.1") + "\n") + for i in range(3): + stdout.write( + measurement(round_id=i, query_count=str(actual_query_count)) + + "\n" + ) + + monkeypatch.setattr(module.subprocess, "run", fake_run) + backend = get_backend( + "flowann", + config={ + "name": "test", + "executable_path": "/unused/test_binary", + "output_dir": str(tmp_path), + }, + ) + dataset = Dataset( + "test", + base_file="base", + query_file="queries", + groundtruth_neighbors_file="truth", + ) + index = IndexConfig( + "test", + "flowann", + {}, + [ + { + "include_query_transfer": True, + "query_count": requested_query_count, + } + ], + "index", + ) + (result,) = backend.search( + dataset, [index], k=10, batch_size=1, search_threads=128 + ) + assert result.recall == 0.95 + assert result.search_time_ms == 1.3 * actual_query_count + assert len(result.metadata["rounds"]) == 3 + assert result.metadata["batch_latency_ms"] == 1.3 + assert result.queries_per_second == 769.23 + assert result.metadata["query_count"] == actual_query_count diff --git a/rust/cuvs-sys/src/bindings.rs b/rust/cuvs-sys/src/bindings.rs index e723abaaea..68d66479ce 100644 --- a/rust/cuvs-sys/src/bindings.rs +++ b/rust/cuvs-sys/src/bindings.rs @@ -1118,6 +1118,46 @@ const _: () = { [::std::mem::offset_of!(cuvsAceParams, max_gpu_memory_gb) - 40usize]; }; pub type cuvsAceParams_t = *mut cuvsAceParams; +#[repr(u32)] +#[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)] +pub enum cuvsCagraGraphStorage { + CUVS_CAGRA_GRAPH_DEVICE = 0, + CUVS_CAGRA_GRAPH_TIERED = 1, +} +#[repr(C)] +#[derive(Debug, Copy, Clone)] +pub struct cuvsCagraTieredGraphParams { + pub device_graph_budget_bytes: usize, + pub node_per_cacheline: u16, + pub grouping_enabled: bool, + pub n_groups: u32, + pub n_bits: u16, + pub balance_tolerance: f64, + pub training_rows: usize, + pub assignment_batch_rows: usize, + pub kmeans_n_iters: u32, + pub validate: bool, + pub num_seeds: u32, + pub seed_training_rows: usize, + pub seed: u64, +} +#[repr(C)] +#[derive(Debug, Copy, Clone)] +pub struct cuvsCagraTieredSearchParams { + pub num_seeds: u32, + pub sync_window_scale: f32, + pub sync_drop_threshold: u32, + pub num_queues: u32, + pub empty_pause: u32, + pub collect_statistics: bool, + pub keep_pollers_running: bool, +} +const _: () = { + assert!(::std::mem::size_of::() == 80); + assert!(::std::mem::align_of::() == 8); + assert!(::std::mem::size_of::() == 24); + assert!(::std::mem::align_of::() == 4); +}; #[repr(C)] #[derive(Debug, Copy, Clone)] pub struct cuvsCagraIndexParams { @@ -1127,10 +1167,12 @@ pub struct cuvsCagraIndexParams { pub build_algo: cuvsCagraGraphBuildAlgo, pub nn_descent_niter: usize, pub graph_build_params: *mut ::std::os::raw::c_void, + pub graph_storage: cuvsCagraGraphStorage, + pub tiered: cuvsCagraTieredGraphParams, } #[allow(clippy::unnecessary_operation, clippy::identity_op)] const _: () = { - ["Size of cuvsCagraIndexParams"][::std::mem::size_of::() - 48usize]; + ["Size of cuvsCagraIndexParams"][::std::mem::size_of::() - 136usize]; ["Alignment of cuvsCagraIndexParams"][::std::mem::align_of::() - 8usize]; ["Offset of field: cuvsCagraIndexParams::metric"] [::std::mem::offset_of!(cuvsCagraIndexParams, metric) - 0usize]; @@ -1251,10 +1293,12 @@ pub struct cuvsCagraSearchParams { pub persistent: bool, pub persistent_lifetime: f32, pub persistent_device_usage: f32, + pub use_tiered_params: bool, + pub tiered: cuvsCagraTieredSearchParams, } #[allow(clippy::unnecessary_operation, clippy::identity_op)] const _: () = { - ["Size of cuvsCagraSearchParams"][::std::mem::size_of::() - 112usize]; + ["Size of cuvsCagraSearchParams"][::std::mem::size_of::() - 136usize]; ["Alignment of cuvsCagraSearchParams"] [::std::mem::align_of::() - 8usize]; ["Offset of field: cuvsCagraSearchParams::max_queries"] @@ -3082,3 +3126,25 @@ unsafe extern "C" { out_idx: *mut DLManagedTensor, ) -> cuvsError_t; } + +const _: () = { + assert!(::std::mem::offset_of!(cuvsCagraIndexParams, graph_storage) == 48); + assert!(::std::mem::offset_of!(cuvsCagraIndexParams, tiered) == 56); + assert!(::std::mem::offset_of!(cuvsCagraSearchParams, use_tiered_params) == 108); + assert!(::std::mem::offset_of!(cuvsCagraSearchParams, tiered) == 112); +}; +unsafe extern "C" { + #[must_use] + pub fn cuvsCagraIndexGetGraphStorage( + index: cuvsCagraIndex_t, + storage: *mut cuvsCagraGraphStorage, + ) -> cuvsError_t; + #[must_use] + pub fn cuvsCagraBuildCompressed( + res: cuvsResources_t, + params: cuvsCagraIndexParams_t, + compression: cuvsCagraCompressionParams_t, + dataset: cuvsDataset_t, + index: cuvsCagraIndex_t, + ) -> cuvsError_t; +} diff --git a/rust/cuvs/README.md b/rust/cuvs/README.md new file mode 100644 index 0000000000..bb5c5c390a --- /dev/null +++ b/rust/cuvs/README.md @@ -0,0 +1,45 @@ +# Experimental tiered CAGRA storage + +Build the native library with `CUVS_ENABLE_FLOWANN_SEARCH=ON` and +`CUVS_ENABLE_FLOWANN_BUILD=ON`. `CUVS_FLOWANN_USE_GDRCOPY=OFF` selects the +CUDA-copy transport. Ordinary CAGRA keeps its default device graph storage. + +Select tiered storage through the existing CAGRA parameter builder: + +```rust,ignore +use cuvs::neighbors::cagra::{GraphStorage, Index, IndexParams, TieredGraphParams}; + +let params = IndexParams::builder() + .nn_descent() + .graph_storage(GraphStorage::Tiered) + .tiered(TieredGraphParams { + device_graph_budget_bytes: 1 << 30, + ..Default::default() + }) + .build()?; +let index = Index::build_tiered(&resources, ¶ms, &device_padded_dataset)?; +``` + +`build_tiered` owns its reordered vectors, so the input can be released after +construction. `build_compressed` accepts a host standard dataset view and +`CompressionParams` to produce an owning VPQ-F16 index. Use +`DatasetView::standard` to request standard layout explicitly, including when +rows already meet CAGRA's alignment rule. + +`SearchParams::builder().tiered(TieredSearchParams { ..Default::default() })` +configures queues and synchronization. Keeping pollers running between calls +consumes CPU while the index is idle. Use the CUDA device that created the index. + +Serialize tiered indexes with `include_dataset=true`. After +`Index::deserialize_graph_and_dataset`, `graph_storage()` reports `Tiered` and +`has_dataset()` is false because the native index owns the vectors internally. +The loaded index is ready for search. Ordinary graph-only files still require +an attached searchable dataset. + +Run the native FlowANN lifecycle tests against a library built with both +FlowANN options enabled: + +```sh +cargo test -p cuvs tiered_dense_build_search_and_round_trip -- --ignored +cargo test -p cuvs tiered_compressed_build_search_and_round_trip -- --ignored +``` diff --git a/rust/cuvs/src/dataset.rs b/rust/cuvs/src/dataset.rs index 35f9c4faf9..e63b4cde4c 100644 --- a/rust/cuvs/src/dataset.rs +++ b/rust/cuvs/src/dataset.rs @@ -118,6 +118,24 @@ pub struct DatasetView<'a> { } impl<'a> DatasetView<'a> { + /// Borrow a tensor as a standard row-major view, including aligned rows. + /// + /// This explicit layout is required by tiered CAGRA VPQ construction from + /// host vectors. The view borrows its tensor and never copies its data. + pub fn standard(res: &Resources, dataset: &'a T) -> Result + where + T: AsDlTensor + ?Sized, + { + let tensor = dataset.as_dl_tensor()?; + let mut tensor_c = tensor.to_c(); + let handle = unsafe { + init_handle(|out| { + ffi::cuvsDatasetMakeStandardView(res.handle(), tensor_c.as_mut_ptr(), out) + })? + }; + Ok(Self { handle, _dataset: PhantomData }) + } + /// Borrow a tensor as the host/device and padded/standard view matching its /// DLPack shape/strides (CAGRA row-width rule). pub fn new(res: &Resources, dataset: &'a T) -> Result diff --git a/rust/cuvs/src/neighbors/cagra/index.rs b/rust/cuvs/src/neighbors/cagra/index.rs index a93c90ca4d..467d6cebed 100644 --- a/rust/cuvs/src/neighbors/cagra/index.rs +++ b/rust/cuvs/src/neighbors/cagra/index.rs @@ -6,7 +6,7 @@ use std::marker::PhantomData; use std::path::Path; -use super::{CagraError, IndexParams, SearchParams}; +use super::{CagraError, CompressionParams, GraphStorage, IndexParams, SearchParams}; use crate::dataset::private::Sealed as _; use crate::dataset::{CuvsDataset, Dataset, DatasetKind, DatasetView}; use crate::dlpack::{AsDlTensor, AsDlTensorMut, DLTensorView, DLTensorViewMut}; @@ -53,7 +53,8 @@ pub struct Index<'d> { /// /// A file serialized without vectors yields `dataset == None` and must have /// matching storage attached before search. Field order is significant: the -/// native index is destroyed before its dataset owner. +/// native index is destroyed before its dataset owner. Tiered indexes own +/// vectors internally and return no separate dataset owner. #[derive(Debug)] pub struct DeserializedIndex { handle: IndexHandle, @@ -89,6 +90,61 @@ impl<'d> Index<'d> { Ok(Index { handle, _dataset: PhantomData }) } + /// Build an owning tiered dense index from a device-padded dataset. + /// + /// The input may be released after this call because tiered storage owns + /// its reordered vectors. `params` must select [`GraphStorage::Tiered`]. + pub fn build_tiered( + res: &Resources, + params: &IndexParams, + dataset: &D, + ) -> Result> + where + D: CuvsDataset + ?Sized, + { + if unsafe { (*params.handle()).graph_storage } + != ffi::cuvsCagraGraphStorage::CUVS_CAGRA_GRAPH_TIERED + { + return Err(CagraError::Validation( + "build_tiered requires GraphStorage::Tiered".into(), + )); + } + let handle = Self::build_handle(res, params, dataset.raw_dataset_handle())?; + Ok(Index { handle, _dataset: PhantomData }) + } + + /// Build a tiered VPQ-F16 index from full-precision host standard vectors. + /// + /// The native index owns its reordered compressed dataset. The input may be + /// released after this call. Enable FlowANN building in the native library + /// and select [`GraphStorage::Tiered`] in `params`. + pub fn build_compressed( + res: &Resources, + params: &IndexParams, + compression: &CompressionParams, + dataset: &D, + ) -> Result> + where + D: CuvsDataset + ?Sized, + { + let handle = IndexHandle::new()?; + check_cuvs(unsafe { + ffi::cuvsCagraBuildCompressed( + res.handle(), + params.handle(), + compression.handle(), + dataset.raw_dataset_handle(), + handle.raw(), + ) + })?; + Ok(Index { handle, _dataset: PhantomData }) + } + + /// Return the index's native graph storage mode. + pub fn graph_storage(&self) -> Result { + graph_storage_impl(&self.handle) + } + fn build_handle( res: &Resources, params: &IndexParams, @@ -216,7 +272,8 @@ impl<'d> Index<'d> { serialize_to_hnswlib_impl(&self.handle, res, filename.as_ref()) } - /// Load only the graph, ignoring any dataset stored in the file. + /// Load only the graph for ordinary device storage. + /// Tiered files always load the complete owning graph and dataset. pub fn deserialize_graph>( res: &Resources, filename: P, @@ -245,17 +302,33 @@ impl<'d> Index<'d> { &mut out, ) })?; - Ok(DeserializedIndex { handle, dataset: Some(Dataset::from_raw(out)?) }) + let dataset = if out.is_null() { + if graph_storage_impl(&handle)? != GraphStorage::Tiered { + return Err(CagraError::Validation( + "native load returned no dataset for a device graph".into(), + )); + } + None + } else { + Some(Dataset::from_raw(out)?) + }; + Ok(DeserializedIndex { handle, dataset }) } } impl DeserializedIndex { + /// Return the index's native graph storage mode. + pub fn graph_storage(&self) -> Result { + graph_storage_impl(&self.handle) + } + /// Borrow the dataset owner when the serialized file included vectors. pub fn dataset(&self) -> Option<&D> { self.dataset.as_ref() } - /// Whether the serialized file included vector storage. + /// Whether a separate dataset owner was returned by deserialization. + /// Tiered indexes keep vector storage inside the native index instead. pub fn has_dataset(&self) -> bool { self.dataset.is_some() } @@ -352,6 +425,9 @@ impl DeserializedIndex { } fn require_dataset(&self) -> Result<()> { + if self.graph_storage()? == GraphStorage::Tiered { + return Ok(()); + } let Some(dataset) = self.dataset.as_ref() else { return Err(CagraError::Validation( "cannot search a graph-only index without an attached dataset".to_string(), @@ -366,6 +442,12 @@ impl DeserializedIndex { } } +fn graph_storage_impl(handle: &IndexHandle) -> Result { + let mut storage = ffi::cuvsCagraGraphStorage::CUVS_CAGRA_GRAPH_DEVICE; + check_cuvs(unsafe { ffi::cuvsCagraIndexGetGraphStorage(handle.raw(), &mut storage) })?; + Ok(storage.into()) +} + fn search_impl( handle: &IndexHandle, res: &Resources, @@ -429,6 +511,25 @@ mod tests { const N_DATAPOINTS: usize = 256; const N_FEATURES: usize = 16; + fn tiered_dense_fixture() -> ndarray::Array2 { + ndarray::Array2::from_shape_fn((1024, 32), |(r, c)| { + let value = (((r * 17 + c * 31) % 1009) as f32) / 1009.0; + // The modulo pattern repeats after 1009 rows. Distinguish those rows so + // the self-neighbor assertion cannot reject an equally close duplicate. + if c == 0 { value + r as f32 / 1024.0 } else { value } + }) + } + + #[test] + fn tiered_dense_fixture_has_unique_query_rows() { + let dataset = tiered_dense_fixture(); + for q in 0..8 { + let query = dataset.row(q); + let matches = dataset.outer_iter().filter(|row| row.iter().eq(query.iter())).count(); + assert_eq!(matches, 1, "query row {q} must have a unique exact nearest neighbor"); + } + } + /// Search the first `n_queries` rows of `dataset` against `index` and /// assert each query finds itself as the top-1 neighbor. CAGRA search /// requires queries and outputs to live in device memory. @@ -477,6 +578,109 @@ mod tests { search_and_verify_self_neighbors(&res, &index, &dataset, 4, 10); } + // FlowANN tests require a native library configured with SEARCH and BUILD. + #[test] + #[ignore = "requires native FlowANN SEARCH and BUILD support"] + fn tiered_dense_build_search_and_round_trip() { + use super::super::{GraphStorage, TieredGraphParams, TieredSearchParams}; + let res = Resources::new().unwrap(); + let dataset = tiered_dense_fixture(); + let device = DeviceTensor::from_host(&res, &dataset).unwrap(); + let params = IndexParams::builder() + .nn_descent_with_iterations(10) + .intermediate_graph_degree(32) + .graph_degree(16) + .graph_storage(GraphStorage::Tiered) + .tiered(TieredGraphParams { + device_graph_budget_bytes: 8192, + validate: true, + ..Default::default() + }) + .build() + .unwrap(); + let view = DatasetView::new(&res, &device).unwrap(); + let index = Index::build_tiered(&res, ¶ms, &view).unwrap(); + drop(view); + drop(device); + assert_eq!(index.graph_storage().unwrap(), GraphStorage::Tiered); + let search = SearchParams::builder() + .algo(super::super::SearchAlgo::MultiCta) + .itopk_size(64) + .max_iterations(32) + .tiered(TieredSearchParams { keep_pollers_running: true, ..Default::default() }) + .build() + .unwrap(); + let queries = + DeviceTensor::from_host(&res, &dataset.slice(s![0..8, ..]).to_owned()).unwrap(); + let mut neighbors = DeviceTensor::::zeros(&res, &[8, 1]).unwrap(); + let mut distances = DeviceTensor::::zeros(&res, &[8, 1]).unwrap(); + for _ in 0..2 { + index.search(&res, &search, &queries, &mut neighbors, &mut distances).unwrap(); + } + let dir = tempfile::tempdir().unwrap(); + let filename = dir.path().join("tiered.index"); + index.serialize(&res, &filename, true).unwrap(); + assert!(index.serialize(&res, dir.path().join("graph-only.index"), false).is_err()); + drop(index); + let loaded = Index::deserialize_graph_and_dataset(&res, &filename).unwrap(); + assert_eq!(loaded.graph_storage().unwrap(), GraphStorage::Tiered); + assert!(!loaded.has_dataset()); + loaded.search(&res, &search, &queries, &mut neighbors, &mut distances).unwrap(); + let mut found = ndarray::Array2::::zeros((8, 1)); + neighbors.copy_to_host(&res, &mut found).unwrap(); + assert!(found.iter().all(|&id| id < 1024)); + assert!(found.column(0).iter().enumerate().filter(|&(i, id)| *id == i as u32).count() >= 6); + } + + #[test] + #[ignore = "requires native FlowANN SEARCH and BUILD support"] + fn tiered_compressed_build_search_and_round_trip() { + use super::super::{CompressionParams, GraphStorage, TieredGraphParams}; + let res = Resources::new().unwrap(); + let dataset = ndarray::Array2::from_shape_fn((2048, 32), |(r, c)| { + (((r * 17 + c * 31) % 2017) as f32) / 2017.0 + }); + let params = IndexParams::builder() + .nn_descent_with_iterations(10) + .intermediate_graph_degree(32) + .graph_degree(16) + .graph_storage(GraphStorage::Tiered) + .tiered(TieredGraphParams { validate: true, ..Default::default() }) + .build() + .unwrap(); + let compression = CompressionParams::builder() + .pq_dim(4) + .vq_n_centers(16) + .kmeans_n_iters(5) + .build() + .unwrap(); + let view = DatasetView::standard(&res, &*dataset).unwrap(); + let index = Index::build_compressed(&res, ¶ms, &compression, &view).unwrap(); + drop(view); + let queries = + DeviceTensor::from_host(&res, &dataset.slice(s![0..8, ..]).to_owned()).unwrap(); + drop(dataset); + let search = SearchParams::builder() + .algo(super::super::SearchAlgo::MultiCta) + .itopk_size(64) + .max_iterations(32) + .build() + .unwrap(); + let mut neighbors = DeviceTensor::::zeros(&res, &[8, 1]).unwrap(); + let mut distances = DeviceTensor::::zeros(&res, &[8, 1]).unwrap(); + index.search(&res, &search, &queries, &mut neighbors, &mut distances).unwrap(); + let dir = tempfile::tempdir().unwrap(); + let filename = dir.path().join("tiered-vpq.index"); + index.serialize(&res, &filename, true).unwrap(); + drop(index); + let loaded = Index::deserialize_graph_and_dataset(&res, &filename).unwrap(); + assert_eq!(loaded.graph_storage().unwrap(), GraphStorage::Tiered); + loaded.search(&res, &search, &queries, &mut neighbors, &mut distances).unwrap(); + let mut found = ndarray::Array2::::zeros((8, 1)); + neighbors.copy_to_host(&res, &mut found).unwrap(); + assert!(found.iter().all(|&id| id < 2048)); + } + #[test] fn test_cagra_index() { let build_params = IndexParams::builder().build().unwrap(); diff --git a/rust/cuvs/src/neighbors/cagra/mod.rs b/rust/cuvs/src/neighbors/cagra/mod.rs index c8ab865e02..fe4d792baa 100644 --- a/rust/cuvs/src/neighbors/cagra/mod.rs +++ b/rust/cuvs/src/neighbors/cagra/mod.rs @@ -16,14 +16,20 @@ //! builder pattern: every setter is optional and unset values keep the cuVS C //! library defaults. Values are validated when the builder's `build()` runs, //! returning [`CagraError::Validation`] for out-of-range inputs. +//! +//! Experimental [`GraphStorage::Tiered`] uses FlowANN to store graph edges +//! across GPU and CPU memory. See `README.md` in this crate for build settings, +//! ownership, serialization, and examples. mod index; mod params; +mod tiered; pub use crate::dataset::{CuvsDataset, Dataset, DatasetKind, DatasetView, PaddedDataset}; pub use crate::neighbors::filters::{Bitset, Filter}; pub use index::{DeserializedIndex, Index}; pub use params::{IndexParams, SearchParams}; +pub use tiered::{CompressionParams, GraphStorage, TieredGraphParams, TieredSearchParams}; use crate::dlpack::DLPackError; use crate::error::LibraryError; diff --git a/rust/cuvs/src/neighbors/cagra/params.rs b/rust/cuvs/src/neighbors/cagra/params.rs index 4fcc3d18af..62780f3ecf 100644 --- a/rust/cuvs/src/neighbors/cagra/params.rs +++ b/rust/cuvs/src/neighbors/cagra/params.rs @@ -20,7 +20,10 @@ use bon::bon; use crate::distance::DistanceType; use crate::error::check_cuvs; -use super::{CagraError, GraphBuildAlgo, HashMode, SearchAlgo}; +use super::{ + CagraError, GraphBuildAlgo, GraphStorage, HashMode, SearchAlgo, TieredGraphParams, + TieredSearchParams, +}; #[derive(Debug)] enum RequestedGraphBuild { @@ -60,6 +63,8 @@ impl IndexParams { metric: Option, intermediate_graph_degree: Option, graph_degree: Option, + graph_storage: Option, + tiered: Option, #[builder(setters(vis = "", some_fn = graph_build_internal))] graph_build: Option< RequestedGraphBuild, >, @@ -96,6 +101,27 @@ impl IndexParams { } } + let storage = graph_storage.unwrap_or(GraphStorage::Device); + if tiered.is_some() && storage != GraphStorage::Tiered { + return Err(CagraError::Validation( + "tiered options require GraphStorage::Tiered".into(), + )); + } + if storage == GraphStorage::Tiered && matches!(graph_build, Some(RequestedGraphBuild::Ace)) + { + return Err(CagraError::Validation( + "tiered storage does not support ACE builds".into(), + )); + } + if let Some(options) = tiered { + options.validate()?; + unsafe { + (*params.handle).tiered = options.into(); + } + } + unsafe { + (*params.handle).graph_storage = storage.into(); + } params.apply_graph_build(graph_build); Ok(params) } @@ -247,8 +273,16 @@ impl SearchParams { persistent: Option, persistent_lifetime: Option, persistent_device_usage: Option, + tiered: Option, ) -> Result { let params = Self::create_handle()?; + if let Some(options) = tiered { + options.validate()?; + unsafe { + (*params.handle).use_tiered_params = true; + (*params.handle).tiered = options.into(); + } + } let effective_algo = algo.unwrap_or_else(|| unsafe { (*params.handle).algo.into() }); let effective_hashmap_mode = diff --git a/rust/cuvs/src/neighbors/cagra/tiered.rs b/rust/cuvs/src/neighbors/cagra/tiered.rs new file mode 100644 index 0000000000..36ed1d5710 --- /dev/null +++ b/rust/cuvs/src/neighbors/cagra/tiered.rs @@ -0,0 +1,299 @@ +/* + * 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). + */ + +//! Options for experimental CAGRA graph storage across GPU and CPU memory. + +use super::CagraError; +use crate::error::check_cuvs; + +/// Representation of the native CAGRA graph. +#[derive(Debug, Default, Copy, Clone, PartialEq, Eq)] +#[non_exhaustive] +pub enum GraphStorage { + /// Store the adjacency graph on the device. + #[default] + Device, + /// Store a compressed subset on the GPU and fetch other edges from the CPU. + Tiered, +} + +impl From for ffi::cuvsCagraGraphStorage { + fn from(value: GraphStorage) -> Self { + match value { + GraphStorage::Device => Self::CUVS_CAGRA_GRAPH_DEVICE, + GraphStorage::Tiered => Self::CUVS_CAGRA_GRAPH_TIERED, + } + } +} + +impl From for GraphStorage { + fn from(value: ffi::cuvsCagraGraphStorage) -> Self { + match value { + ffi::cuvsCagraGraphStorage::CUVS_CAGRA_GRAPH_DEVICE => Self::Device, + ffi::cuvsCagraGraphStorage::CUVS_CAGRA_GRAPH_TIERED => Self::Tiered, + } + } +} + +/// Experimental graph layout options. Unchanged defaults match the C library. +#[derive(Debug, Copy, Clone)] +pub struct TieredGraphParams { + /// Packed GPU edge budget in bytes; zero stores all edges on the host. + pub device_graph_budget_bytes: usize, + /// Nodes per packed row, from 1 to 64 with grouping or 1 to 32 without it. + pub node_per_cacheline: u16, + /// Group and reorder nodes to improve locality. + pub grouping_enabled: bool, + /// Final group count; zero chooses automatically, maximum 65,536. + pub n_groups: u32, + /// Local ID width; zero chooses automatically, otherwise from 1 to 32. + pub n_bits: u16, + /// Allowed group imbalance in the open interval (0, 1). + pub balance_tolerance: f64, + /// Maximum training rows per grouping node; zero chooses automatically. + pub training_rows: usize, + /// Assignment batch size in rows; zero chooses automatically. + pub assignment_batch_rows: usize, + /// Positive number of k-means training iterations. + pub kmeans_n_iters: u32, + /// Verify encoded edges after building the graph. + pub validate: bool, + /// Number of medoid entry points to generate; zero disables generation. + pub num_seeds: u32, + /// Seed training sample rows; zero chooses automatically. + pub seed_training_rows: usize, + /// Random seed for grouping and seed generation. + pub seed: u64, +} + +impl Default for TieredGraphParams { + fn default() -> Self { + Self { + device_graph_budget_bytes: 0, + node_per_cacheline: 2, + grouping_enabled: true, + n_groups: 0, + n_bits: 0, + balance_tolerance: 0.1, + training_rows: 0, + assignment_batch_rows: 0, + kmeans_n_iters: 20, + validate: false, + num_seeds: 0, + seed_training_rows: 0, + seed: 0x9e3779b97f4a7c15, + } + } +} + +impl TieredGraphParams { + pub(super) fn validate(&self) -> Result<(), CagraError> { + let max_nodes = if self.grouping_enabled { 64 } else { 32 }; + if self.node_per_cacheline == 0 || self.node_per_cacheline > max_nodes { + return Err(CagraError::Validation(format!( + "node_per_cacheline must be in [1, {max_nodes}]" + ))); + } + if self.n_groups > 65_536 || self.n_bits > 32 || self.kmeans_n_iters == 0 { + return Err(CagraError::Validation( + "invalid tiered grouping count, ID width, or iteration count".into(), + )); + } + if !self.balance_tolerance.is_finite() + || self.balance_tolerance <= 0.0 + || self.balance_tolerance >= 1.0 + { + return Err(CagraError::Validation("balance_tolerance must be in (0, 1)".into())); + } + Ok(()) + } +} + +impl From for ffi::cuvsCagraTieredGraphParams { + fn from(p: TieredGraphParams) -> Self { + Self { + device_graph_budget_bytes: p.device_graph_budget_bytes, + node_per_cacheline: p.node_per_cacheline, + grouping_enabled: p.grouping_enabled, + n_groups: p.n_groups, + n_bits: p.n_bits, + balance_tolerance: p.balance_tolerance, + training_rows: p.training_rows, + assignment_batch_rows: p.assignment_batch_rows, + kmeans_n_iters: p.kmeans_n_iters, + validate: p.validate, + num_seeds: p.num_seeds, + seed_training_rows: p.seed_training_rows, + seed: p.seed, + } + } +} + +/// Experimental search queue options. Searches complete before returning. +#[derive(Debug, Copy, Clone)] +pub struct TieredSearchParams { + /// Maximum stored medoid entry points used per query. + pub num_seeds: u32, + /// Adaptive synchronization window scale. + pub sync_window_scale: f32, + /// Parent position drop that forces synchronization. + pub sync_drop_threshold: u32, + /// Positive number of CPU/GPU command queues. + pub num_queues: u32, + /// CPU pause instructions after each empty poll. + pub empty_pause: u32, + /// Collect native queue statistics. + pub collect_statistics: bool, + /// Keep CPU workers active between searches, consuming CPU while idle. + pub keep_pollers_running: bool, +} + +impl Default for TieredSearchParams { + fn default() -> Self { + Self { + num_seeds: 0, + sync_window_scale: 10.0, + sync_drop_threshold: 51, + num_queues: 1, + empty_pause: 64, + collect_statistics: false, + keep_pollers_running: false, + } + } +} + +impl TieredSearchParams { + pub(super) fn validate(&self) -> Result<(), CagraError> { + if self.num_queues == 0 + || !self.sync_window_scale.is_finite() + || self.sync_window_scale <= 0.0 + { + return Err(CagraError::Validation( + "tiered search requires positive num_queues and sync_window_scale".into(), + )); + } + Ok(()) + } +} + +impl From for ffi::cuvsCagraTieredSearchParams { + fn from(p: TieredSearchParams) -> Self { + Self { + num_seeds: p.num_seeds, + sync_window_scale: p.sync_window_scale, + sync_drop_threshold: p.sync_drop_threshold, + num_queues: p.num_queues, + empty_pause: p.empty_pause, + collect_statistics: p.collect_statistics, + keep_pollers_running: p.keep_pollers_running, + } + } +} + +/// Owns the native compression parameters used by [`super::Index::build_compressed`]. +#[derive(Debug)] +pub struct CompressionParams { + handle: ffi::cuvsCagraCompressionParams_t, +} + +#[bon::bon] +impl CompressionParams { + /// Allocate native defaults and optionally set VPQ dimensions and training settings. + #[builder] + pub fn new( + pq_bits: Option, + pq_dim: Option, + vq_n_centers: Option, + kmeans_n_iters: Option, + vq_kmeans_trainset_fraction: Option, + pq_kmeans_trainset_fraction: Option, + ) -> Result { + if pq_bits.is_some_and(|bits| bits != 8) { + return Err(CagraError::Validation("tiered VPQ requires 8-bit codes".into())); + } + if kmeans_n_iters == Some(0) { + return Err(CagraError::Validation("kmeans_n_iters must be positive".into())); + } + for fraction in + [vq_kmeans_trainset_fraction, pq_kmeans_trainset_fraction].into_iter().flatten() + { + if !fraction.is_finite() || fraction <= 0.0 || fraction > 1.0 { + return Err(CagraError::Validation("training fractions must be in (0, 1]".into())); + } + } + let mut handle = std::ptr::null_mut(); + check_cuvs(unsafe { ffi::cuvsCagraCompressionParamsCreate(&mut handle) })?; + let result = Self { handle }; + unsafe { + if let Some(v) = pq_bits { + (*handle).pq_bits = v; + } + if let Some(v) = pq_dim { + (*handle).pq_dim = v; + } + if let Some(v) = vq_n_centers { + (*handle).vq_n_centers = v; + } + if let Some(v) = kmeans_n_iters { + (*handle).kmeans_n_iters = v; + } + if let Some(v) = vq_kmeans_trainset_fraction { + (*handle).vq_kmeans_trainset_fraction = v; + } + if let Some(v) = pq_kmeans_trainset_fraction { + (*handle).pq_kmeans_trainset_fraction = v; + } + } + Ok(result) + } + + pub(super) fn handle(&self) -> ffi::cuvsCagraCompressionParams_t { + self.handle + } +} + +impl Drop for CompressionParams { + fn drop(&mut self) { + let _ = unsafe { ffi::cuvsCagraCompressionParamsDestroy(self.handle) }; + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn invalid_graph_options_are_rejected() { + for p in [ + TieredGraphParams { node_per_cacheline: 0, ..Default::default() }, + TieredGraphParams { + grouping_enabled: false, + node_per_cacheline: 33, + ..Default::default() + }, + TieredGraphParams { n_groups: 65_537, ..Default::default() }, + TieredGraphParams { n_bits: 33, ..Default::default() }, + TieredGraphParams { balance_tolerance: f64::NAN, ..Default::default() }, + TieredGraphParams { kmeans_n_iters: 0, ..Default::default() }, + ] { + assert!(matches!(p.validate(), Err(CagraError::Validation(_)))); + } + assert!(TieredGraphParams::default().validate().is_ok()); + } + + #[test] + fn invalid_queue_options_are_rejected() { + assert!(TieredSearchParams { num_queues: 0, ..Default::default() }.validate().is_err()); + assert!( + TieredSearchParams { sync_window_scale: f32::NAN, ..Default::default() } + .validate() + .is_err() + ); + assert!(TieredSearchParams::default().validate().is_ok()); + } +}