diff --git a/c/include/cuvs/neighbors/cagra.h b/c/include/cuvs/neighbors/cagra.h index 602f52dbfc..94d9cb03ac 100644 --- a/c/include/cuvs/neighbors/cagra.h +++ b/c/include/cuvs/neighbors/cagra.h @@ -149,8 +149,8 @@ struct cuvsAceParams { * * Small values might improve recall but potentially degrade performance and * increase memory usage. Partitions should not be too small to prevent issues - * in KNN graph construction. The partition size is on average 2 * (n_rows / - * npartitions) * dim * sizeof(T). 2 is because of the core and augmented + * in KNN graph construction. The partition size is on average `(1 + nsubpartitions) * + * (n_rows / npartitions) * dim * sizeof(T)`, accounting for core and augmented * vectors. Please account for imbalance in the partition sizes (up to 3x in * our tests). * @@ -199,6 +199,17 @@ struct cuvsAceParams { * Useful for testing or when running alongside other memory-intensive processes. */ double max_gpu_memory_gb; + /** + * Number of additional nearest partitions assigned to each vector. Default: 1. + * Must be positive and smaller than npartitions (unless automatically selected). + * Increasing this improves overlap at the cost of build time, memory, and disk space. + */ + size_t nsubpartitions; + /** + * Build pruned partial graphs, then add reverse edges globally. Default: false. + * In disk mode, incoming edges are buffered to temporary files and merged on the CPU. + */ + bool add_global_reverse_edges; }; typedef struct cuvsAceParams* cuvsAceParams_t; diff --git a/c/src/neighbors/cagra.cpp b/c/src/neighbors/cagra.cpp index 30d356101f..6dcc21af2b 100644 --- a/c/src/neighbors/cagra.cpp +++ b/c/src/neighbors/cagra.cpp @@ -817,6 +817,8 @@ static void _set_graph_build_params( if (params.graph_build_params) { auto ace_params_c = static_cast(params.graph_build_params); ace_p.npartitions = ace_params_c->npartitions; + ace_p.nsubpartitions = ace_params_c->nsubpartitions; + ace_p.add_global_reverse_edges = ace_params_c->add_global_reverse_edges; ace_p.ef_construction = ace_params_c->ef_construction; ace_p.build_dir = std::string(ace_params_c->build_dir); ace_p.use_disk = ace_params_c->use_disk; @@ -1455,6 +1457,8 @@ static void _populate_cagra_index_params_from_cpp(cuvsCagraIndexParams_t c_param cpp_params.graph_build_params); cuvsAceParams* c_ace_params = new cuvsAceParams; c_ace_params->npartitions = ace_params.npartitions; + c_ace_params->nsubpartitions = ace_params.nsubpartitions; + c_ace_params->add_global_reverse_edges = ace_params.add_global_reverse_edges; c_ace_params->ef_construction = ace_params.ef_construction; c_ace_params->build_dir = ace_params.build_dir.empty() ? nullptr : strdup(ace_params.build_dir.c_str()); c_ace_params->use_disk = ace_params.use_disk; @@ -2342,7 +2346,9 @@ extern "C" cuvsError_t cuvsAceParamsCreate(cuvsAceParams_t* params) .build_dir = build_dir, .use_disk = ps.use_disk, .max_host_memory_gb = ps.max_host_memory_gb, - .max_gpu_memory_gb = ps.max_gpu_memory_gb}; + .max_gpu_memory_gb = ps.max_gpu_memory_gb, + .nsubpartitions = ps.nsubpartitions, + .add_global_reverse_edges = ps.add_global_reverse_edges}; }); } diff --git a/c/tests/neighbors/ann_cagra_c.cu b/c/tests/neighbors/ann_cagra_c.cu index d08357b91e..c5fca175e2 100644 --- a/c/tests/neighbors/ann_cagra_c.cu +++ b/c/tests/neighbors/ann_cagra_c.cu @@ -1186,6 +1186,10 @@ TEST(CagraC, BuildSearchACEMemory) cuvsAceParams_t ace_params; cuvsAceParamsCreate(&ace_params); + EXPECT_EQ(ace_params->nsubpartitions, 1); + EXPECT_FALSE(ace_params->add_global_reverse_edges); + ace_params->nsubpartitions = 0; + ace_params->add_global_reverse_edges = true; ace_params->npartitions = 2; ace_params->ef_construction = 120; ace_params->use_disk = false; @@ -1193,6 +1197,8 @@ TEST(CagraC, BuildSearchACEMemory) build_params->graph_build_params = ace_params; cuvsDataset_t dataset_view; ASSERT_EQ(cuvsDatasetMakeStandardView(res, &dataset_tensor, &dataset_view), CUVS_SUCCESS); + EXPECT_EQ(cuvsCagraBuild(res, build_params, dataset_view, index), CUVS_ERROR); + ace_params->nsubpartitions = 1; ASSERT_EQ(cuvsCagraBuild(res, build_params, dataset_view, index), CUVS_SUCCESS); // Host build yields a host index. Attach a caller-provided device padded dataset diff --git a/cpp/include/cuvs/neighbors/cagra.hpp b/cpp/include/cuvs/neighbors/cagra.hpp index 98df8cdadf..7d6fa0b5d8 100644 --- a/cpp/include/cuvs/neighbors/cagra.hpp +++ b/cpp/include/cuvs/neighbors/cagra.hpp @@ -180,20 +180,25 @@ struct ace_params { * * Small values might improve recall but potentially degrade performance and * increase memory usage. Partitions should not be too small to prevent issues - * in KNN graph construction. The partition size is on average 2 * (n_rows / npartitions) * dim * - * sizeof(T). 2 is because of the core and augmented vectors. Please account for imbalance in the - * partition sizes (up to 3x in our tests). + * in KNN graph construction. The partition size is on average `(1 + nsubpartitions) * (n_rows / + * npartitions) * dim * sizeof(T)`, accounting for the core and augmented vectors. Please account + * for imbalance in the partition sizes (up to 3x in our tests). * * If the specified number of partitions results in partitions that exceed * available memory, the value will be automatically increased to fit memory * constraints and a warning will be issued. */ size_t npartitions = 0; + /** Number of additional nearest partitions assigned to each vector. Must be positive and + * smaller than npartitions (unless npartitions is auto-selected). Increasing this improves + * overlap at the cost of build time, memory, and temporary disk space. Default: 1. + */ + size_t nsubpartitions = 1; /** * The index quality for the ACE build. * - * Bigger values increase the index quality. At some point, increasing this will no longer improve - * the quality. + * Bigger values increase the index quality. At some point, increasing this will no longer + * improve the quality. */ size_t ef_construction = 120; /** @@ -205,6 +210,9 @@ struct ace_params { * already exist, but ACE's named artifacts must not already exist. Simultaneous * builds must use different directories. On failure, ACE removes only artifacts * it created and never deletes unrelated directory contents. + * With `add_global_reverse_edges`, temporary incoming-edge records also require up to + * `dataset_size * graph_degree * (2 * sizeof(IdxT) + sizeof(uint32_t))` bytes, plus any + * record alignment padding. These records are removed after the merge. */ std::string build_dir = "/tmp/ace_build"; /** @@ -214,6 +222,20 @@ struct ace_params { */ bool use_disk = false; + /** + * Build pruned core-plus-argument partial graphs, then add reverse edges globally. + * + * Each partition first builds an intermediate kNN graph and prunes it without a local + * reverse-edge merge. Then build the reverse-edge graph globally and merge it into the pruned + * graph. In disk mode (explicit or automatic), incoming edges are buffered + * to temporary files while building each pruned partition, then merged on the CPU one + * destination partition at a time, retaining at most `ceil(graph_degree / 2)` incoming candidates + * per node in memory. + * The merge preserves edge-rank priority; ties use reordered source ID in disk mode + * and are GPU-scheduling-dependent in memory mode. Disabled by default. + */ + bool add_global_reverse_edges = false; + /** * Maximum host memory to use for ACE build in GiB. * @@ -330,8 +352,8 @@ struct index_params : cuvs::neighbors::index_params { * * - `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. + * 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. * @@ -417,10 +439,11 @@ struct index_params : cuvs::neighbors::index_params { * * * IMPORTANT NOTE * * - * The reference HNSW index and the corresponding from-CAGRA generated HNSW index will NOT produce - * exactly the same recalls and QPS for the same parameter `ef`. The graphs are different + * The reference HNSW index and the corresponding from-CAGRA generated HNSW index will NOT + * produce exactly the same recalls and QPS for the same parameter `ef`. The graphs are different * internally. Depending on the selected heuristics, the CAGRA-produced graph's QPS-Recall curve - * may be shifted along the curve right or left. See the heuristics descriptions for more details. + * may be shifted along the curve right or left. See the heuristics descriptions for more + * details. * * Usage example: * @code{.cpp} @@ -453,8 +476,8 @@ struct extend_params { /** The additional dataset is divided into chunks and added to the graph. This is the knob to * adjust the tradeoff between the recall and operation throughput. Large chunk sizes can result * in high throughput, but use more working memory (O(max_chunk_size*degree^2)). This can also - * degrade recall because no edges are added between the nodes in the same chunk. Auto select when - * 0. */ + * degrade recall because no edges are added between the nodes in the same chunk. Auto select + * when 0. */ uint32_t max_chunk_size = 0; }; /** @@ -724,7 +747,8 @@ struct CUVS_EXPORT index : cuvs::neighbors::index { } /** - * Replace the source indices with a new source indices taking the ownership of the passed vector. + * @brief Replace the source indices with a new source indices taking the ownership of the passed + * vector. */ void update_source_indices(raft::device_vector&& source_indices) { diff --git a/cpp/src/neighbors/cagra.cuh b/cpp/src/neighbors/cagra.cuh index 550b317552..acd8f4d6d6 100644 --- a/cpp/src/neighbors/cagra.cuh +++ b/cpp/src/neighbors/cagra.cuh @@ -7,6 +7,7 @@ #include "detail/ann_utils.cuh" #include "detail/cagra/add_nodes.cuh" +#include "detail/cagra/cagra_ace.cuh" #include "detail/cagra/cagra_build.cuh" #include "detail/cagra/cagra_merge.cuh" #include "detail/cagra/cagra_search.cuh" @@ -340,6 +341,10 @@ auto build(raft::resources const& res, const index_params& params, DatasetViewT } return idx; } else if constexpr (cuvs::neighbors::is_dense_row_major_device_dataset_view_v) { + RAFT_EXPECTS( + !std::holds_alternative(params.graph_build_params), + "cagra::build: ACE requires a dataset in host memory; device-resident datasets are not " + "supported. Pass a host dataset view when using ace_params."); auto idx = cuvs::neighbors::cagra::detail::build_from_device_matrix( res, params, dataset); if (params.attach_dataset_on_build) { diff --git a/cpp/src/neighbors/detail/cagra/cagra_ace.cuh b/cpp/src/neighbors/detail/cagra/cagra_ace.cuh index d86e6d1f66..523db08ad3 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_ace.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_ace.cuh @@ -6,6 +6,8 @@ #include "../../../core/nvtx.hpp" #include "../../../util/kvikio_io.hpp" +#include "cagra_ace_disk.hpp" +#include "cagra_build.cuh" #include #include @@ -16,11 +18,15 @@ #include #include #include +#include +#include #include +#include #include #include #include +#include #include #include #include @@ -39,6 +45,7 @@ #include #include #include +#include #include #include #include @@ -157,11 +164,29 @@ class ace_disk_workspace { bool committed_ = false; }; -template -void check_graph_degree(size_t& intermediate_degree, size_t& graph_degree, size_t dataset_size); +// Convert selected IDs and count core/subpartition assignments while they are on device. +template +__global__ void ace_store_partition_labels_kernel( + const int64_t* selected_labels, IdxT* labels, IdxT* histogram, size_t count, size_t labels_dim) +{ + for (size_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; i < count; + i += static_cast(gridDim.x) * blockDim.x) { + const auto label = static_cast(selected_labels[i]); + labels[i] = label; + const auto bin = static_cast(label) * 2 + (i % labels_dim != 0); +#if __CUDA_ARCH__ >= 700 + // Combine equal bins within a warp to reduce contention for small partition counts. + const auto peers = __match_any_sync(__activemask(), bin); + if ((threadIdx.x % warpSize) == static_cast(__ffs(peers) - 1)) { + raft::myAtomicAdd(histogram + bin, static_cast(__popc(peers))); + } +#else + raft::myAtomicAdd(histogram + bin, IdxT{1}); +#endif + } +} -// Get partition labels for partitioned approach -// TODO(julianmi): Use all neighbors APIs. +// ACE: Get partition labels for partitioned approach template void ace_get_partition_labels( raft::resources const& res, @@ -181,103 +206,116 @@ void ace_get_partition_labels( static_cast(dataset.extent(1)), dataset_dim); size_t n_partitions = partition_histogram.extent(0); - RAFT_EXPECTS(labels_dim == 2, "Labels must have 2 columns"); + RAFT_EXPECTS(labels_dim >= 2 && labels_dim <= n_partitions, + "Labels must contain a core partition and distinct augmented partitions"); RAFT_EXPECTS(partition_histogram.extent(1) == 2, "Partition histogram must have 2 columns"); cudaStream_t stream = raft::resource::get_cuda_stream(res).get(); - // Sampling vectors from dataset. Uses float conversion on host instead of - // raft::matrix::sample_rows to minimize GPU memory usage. - // TODO(julianmi): Switch to sample_rows when https://github.com/NVIDIA/cuvs/issues/1461 is - // addressed. + // Sample training rows with a fixed seed for reproducible partitioning. size_t n_samples = dataset_size * sampling_rate; const size_t min_samples = 100 * n_partitions; n_samples = std::max(n_samples, min_samples); n_samples = std::min(n_samples, dataset_size); RAFT_LOG_DEBUG("ACE build: partition labeling uses %lu sampled vectors", n_samples); - auto sample_db = raft::make_host_matrix(n_samples, dataset_dim); -#pragma omp parallel for - for (size_t i = 0; i < n_samples; i++) { - size_t j = i * dataset_size / n_samples; - for (size_t k = 0; k < dataset_dim; k++) { - sample_db(i, k) = static_cast(dataset(j, k)); + auto centroids_dev = raft::make_device_matrix(res, n_partitions, dataset_dim); + { + auto sampling_dataset = + raft::make_host_strided_matrix_view(dataset.data_handle(), + dataset.extent(0), + static_cast(dataset_dim), + dataset.stride(0)); + auto sample_db_dev = raft::matrix::sample_rows( + res, raft::random::RngState{1234}, sampling_dataset, static_cast(n_samples)); + cuvs::cluster::kmeans::balanced_params kmeans_params; + if constexpr (std::is_same_v) { + cuvs::cluster::kmeans::fit(res, kmeans_params, sample_db_dev.view(), centroids_dev.view()); + } else { + auto sample_db_float = raft::make_device_matrix(res, n_samples, dataset_dim); + raft::linalg::unaryOp( + sample_db_float.data_handle(), + sample_db_dev.data_handle(), + n_samples * dataset_dim, + [] __device__(T value) { return static_cast(value); }, + stream); + cuvs::cluster::kmeans::fit(res, kmeans_params, sample_db_float.view(), centroids_dev.view()); } } - auto sample_db_dev = raft::make_device_matrix(res, n_samples, dataset_dim); - raft::copy(res, sample_db_dev.view(), sample_db.view()); - cuvs::cluster::kmeans::balanced_params kmeans_params; - auto centroids_dev = raft::make_device_matrix(res, n_partitions, dataset_dim); - cuvs::cluster::kmeans::fit(res, kmeans_params, sample_db_dev.view(), centroids_dev.view()); - - // Compute distances between dataset and centroid vectors - // Uses float conversion on host instead of batch_load_iterator to minimize GPU memory usage. - const size_t chunk_size = 32 * 1024; - auto _sub_dataset = raft::make_host_matrix(chunk_size, dataset_dim); - auto _sub_distances = raft::make_host_matrix(chunk_size, n_partitions); - auto _sub_dataset_dev = raft::make_device_matrix(res, chunk_size, dataset_dim); - auto _sub_distances_dev = raft::make_device_matrix(res, chunk_size, n_partitions); + // Find the core and subpartition IDs on GPU; only selected IDs are copied to host. + cuvs::neighbors::brute_force::index_params centroid_index_params; + centroid_index_params.metric = cuvs::distance::DistanceType::L2Expanded; + auto centroid_index = cuvs::neighbors::brute_force::build( + res, centroid_index_params, raft::make_const_mdspan(centroids_dev.view())); + cuvs::neighbors::brute_force::search_params centroid_search_params; + + // Transfer bounded chunks in their original type, then convert on GPU when needed. + const size_t chunk_size = std::min(32 * 1024, dataset_size); + auto _sub_dataset_dev = raft::make_device_matrix(res, chunk_size, dataset_dim); + auto _sub_dataset_float = raft::make_device_matrix( + res, std::is_same_v ? 0 : chunk_size, dataset_dim); + auto _sub_output_labels = raft::make_device_matrix(res, chunk_size, labels_dim); + auto histogram_dev = raft::make_device_matrix(res, n_partitions, 2); + raft::copy(res, histogram_dev.view(), partition_histogram); + auto _sub_labels_dev = raft::make_device_matrix(res, chunk_size, labels_dim); + auto _sub_distances_dev = raft::make_device_matrix(res, chunk_size, labels_dim); size_t next_progress_percent = 10; for (size_t i_base = 0; i_base < dataset_size; i_base += chunk_size) { const size_t sub_dataset_size = std::min(chunk_size, dataset_size - i_base); - auto sub_dataset = raft::make_host_matrix_view( - _sub_dataset.data_handle(), sub_dataset_size, dataset_dim); -#pragma omp parallel for - for (size_t i_sub = 0; i_sub < sub_dataset_size; i_sub++) { - size_t i = i_base + i_sub; - for (size_t k = 0; k < dataset_dim; k++) { - sub_dataset(i_sub, k) = static_cast(dataset(i, k)); - } + RAFT_CUDA_TRY(cudaMemcpy2DAsync(_sub_dataset_dev.data_handle(), + dataset_dim * sizeof(T), + &dataset(i_base, 0), + dataset.extent(1) * sizeof(T), + dataset_dim * sizeof(T), + sub_dataset_size, + cudaMemcpyHostToDevice, + stream)); + const float* queries; + if constexpr (std::is_same_v) { + queries = _sub_dataset_dev.data_handle(); + } else { + raft::linalg::unaryOp( + _sub_dataset_float.data_handle(), + _sub_dataset_dev.data_handle(), + sub_dataset_size * dataset_dim, + [] __device__(T value) { return static_cast(value); }, + stream); + queries = _sub_dataset_float.data_handle(); } - auto sub_dataset_dev_view = raft::make_device_matrix_view( - _sub_dataset_dev.data_handle(), sub_dataset_size, dataset_dim); - raft::copy(res, sub_dataset_dev_view, sub_dataset); - auto sub_dataset_dev = raft::make_device_matrix_view( - _sub_dataset_dev.data_handle(), sub_dataset_size, dataset_dim); - - auto sub_distances = raft::make_host_matrix_view( - _sub_distances.data_handle(), sub_dataset_size, n_partitions); - auto sub_distances_dev = raft::make_device_matrix_view( - _sub_distances_dev.data_handle(), sub_dataset_size, n_partitions); - - cuvs::distance::pairwise_distance(res, - sub_dataset_dev, - centroids_dev.view(), - sub_distances_dev, - cuvs::distance::DistanceType::L2Expanded); + auto sub_dataset_dev = + raft::make_device_matrix_view(queries, sub_dataset_size, dataset_dim); - raft::copy(res, sub_distances, sub_distances_dev); + auto sub_labels_dev = raft::make_device_matrix_view( + _sub_labels_dev.data_handle(), sub_dataset_size, labels_dim); + auto sub_distances_dev = raft::make_device_matrix_view( + _sub_distances_dev.data_handle(), sub_dataset_size, labels_dim); + + cuvs::neighbors::brute_force::search(res, + centroid_search_params, + centroid_index, + sub_dataset_dev, + sub_labels_dev, + sub_distances_dev); + const size_t assignment_count = sub_dataset_size * labels_dim; + constexpr uint32_t threads = 256; + const auto blocks = static_cast( + std::min(1024, raft::div_rounding_up_safe(assignment_count, size_t{threads}))); + ace_store_partition_labels_kernel<<>>( + sub_labels_dev.data_handle(), + _sub_output_labels.data_handle(), + histogram_dev.data_handle(), + assignment_count, + labels_dim); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + // The first ID is the core partition; the rest form this node's subpartition list. + raft::copy(partition_labels.data_handle() + i_base * labels_dim, + _sub_output_labels.data_handle(), + assignment_count, + stream); raft::resource::sync_stream(res, stream); - // Find two closest partitions to each dataset vector -#pragma omp parallel for - for (size_t i_sub = 0; i_sub < sub_dataset_size; i_sub++) { - size_t core_label = 0; - size_t augmented_label = 1; - if (sub_distances(i_sub, 0) > sub_distances(i_sub, 1)) { - core_label = 1; - augmented_label = 0; - } - for (size_t c = 2; c < n_partitions; c++) { - if (sub_distances(i_sub, c) < sub_distances(i_sub, core_label)) { - augmented_label = core_label; - core_label = c; - } else if (sub_distances(i_sub, c) < sub_distances(i_sub, augmented_label)) { - augmented_label = c; - } - } - size_t i = i_base + i_sub; - partition_labels(i, 0) = core_label; - partition_labels(i, 1) = augmented_label; - -#pragma omp atomic update - partition_histogram(core_label, 0) += 1; -#pragma omp atomic update - partition_histogram(augmented_label, 1) += 1; - } - const auto completed_rows = i_base + sub_dataset_size; const auto report_percent = progress_step_10(completed_rows, dataset_size); if (report_percent >= next_progress_percent) { @@ -288,6 +326,8 @@ void ace_get_partition_labels( next_progress_percent = report_percent + 10; } } + raft::copy(res, partition_histogram, histogram_dev.view()); + raft::resource::sync_stream(res, stream); } // Check partition sizes for stable KNN graph construction @@ -295,6 +335,7 @@ template void ace_check_partition_sizes( size_t dataset_size, size_t n_partitions, + raft::host_matrix_view partition_labels, raft::host_matrix_view partition_histogram, size_t min_partition_size) { @@ -328,7 +369,8 @@ void ace_check_partition_sizes( double avg_core_vectors = static_cast(total_core_vectors) / n_partitions; double avg_augmented_vectors = static_cast(total_augmented_vectors) / n_partitions; - double expected_avg_vectors = 2.0 * static_cast(dataset_size) / n_partitions; + double expected_avg_vectors = + static_cast(partition_labels.extent(1)) * dataset_size / n_partitions; RAFT_LOG_INFO("ACE: Core vectors - Total: %lu, Avg: %.1f, Min: %lu, Max: %lu", total_core_vectors, @@ -386,6 +428,7 @@ void ace_create_forward_and_backward_lists( raft::host_vector_view core_partition_offsets, raft::host_vector_view augmented_partition_offsets) { + const size_t augmented_size = dataset_size * (partition_labels.extent(1) - 1); core_partition_offsets(0) = 0; augmented_partition_offsets(0) = 0; for (size_t c = 1; c < n_partitions; c++) { @@ -398,8 +441,8 @@ void ace_create_forward_and_backward_lists( // Memory path: both backward mappings RAFT_EXPECTS(static_cast(core_backward_mapping.extent(0)) == dataset_size, "core_backward_mapping must be of size dataset_size"); - RAFT_EXPECTS(static_cast(augmented_backward_mapping.extent(0)) == dataset_size, - "augmented_backward_mapping must be of size dataset_size"); + RAFT_EXPECTS(static_cast(augmented_backward_mapping.extent(0)) == augmented_size, + "augmented_backward_mapping must cover every augmented assignment"); #pragma omp parallel for for (size_t i = 0; i < dataset_size; i++) { size_t core_partition_id = partition_labels(i, 0); @@ -409,12 +452,15 @@ void ace_create_forward_and_backward_lists( RAFT_EXPECTS(core_id < dataset_size, "Vector ID must be smaller than dataset_size"); core_backward_mapping(core_id) = i; - size_t augmented_partition_id = partition_labels(i, 1); - size_t augmented_id; + for (size_t rank = 1; rank < static_cast(partition_labels.extent(1)); ++rank) { + size_t augmented_partition_id = partition_labels(i, rank); + size_t augmented_id; #pragma omp atomic capture - augmented_id = augmented_partition_offsets(augmented_partition_id)++; - RAFT_EXPECTS(augmented_id < dataset_size, "Vector ID must be smaller than dataset_size"); - augmented_backward_mapping(augmented_id) = i; + augmented_id = augmented_partition_offsets(augmented_partition_id)++; + RAFT_EXPECTS(augmented_id < augmented_size, + "Augmented offset must be smaller than augmented_size"); + augmented_backward_mapping(augmented_id) = i; + } } } else { // Disk path: all three mappings @@ -422,8 +468,8 @@ void ace_create_forward_and_backward_lists( "core_forward_mapping must be of size dataset_size"); RAFT_EXPECTS(static_cast(core_backward_mapping.extent(0)) == dataset_size, "core_backward_mapping must be of size dataset_size"); - RAFT_EXPECTS(static_cast(augmented_backward_mapping.extent(0)) == dataset_size, - "augmented_backward_mapping must be of size dataset_size"); + RAFT_EXPECTS(static_cast(augmented_backward_mapping.extent(0)) == augmented_size, + "augmented_backward_mapping must cover every augmented assignment"); for (size_t i = 0; i < dataset_size; i++) { size_t core_partition_id = partition_labels(i, 0); size_t core_id; @@ -432,11 +478,14 @@ void ace_create_forward_and_backward_lists( core_backward_mapping(core_id) = i; core_forward_mapping(i) = core_id; - size_t augmented_partition_id = partition_labels(i, 1); - size_t augmented_id; - augmented_id = augmented_partition_offsets(augmented_partition_id)++; - RAFT_EXPECTS(augmented_id < dataset_size, "Vector ID must be smaller than dataset_size"); - augmented_backward_mapping(augmented_id) = i; + for (size_t rank = 1; rank < static_cast(partition_labels.extent(1)); ++rank) { + size_t augmented_partition_id = partition_labels(i, rank); + size_t augmented_id; + augmented_id = augmented_partition_offsets(augmented_partition_id)++; + RAFT_EXPECTS(augmented_id < augmented_size, + "Augmented offset must be smaller than augmented_size"); + augmented_backward_mapping(augmented_id) = i; + } } } @@ -449,9 +498,9 @@ void ace_create_forward_and_backward_lists( augmented_partition_offsets(0) = 0; } -// Gather partition dataset +// ACE: Materialize the core partition followed by its argument partition. template -void ace_gather_partition_dataset( +void ace_make_core_argument_partition( size_t core_sub_dataset_size, size_t augmented_sub_dataset_size, size_t dataset_dim, @@ -472,7 +521,7 @@ void ace_gather_partition_dataset( memcpy(&sub_dataset(j, 0), &dataset(i, 0), vector_size_bytes); } - // Copy augmented partition vectors (2nd closest partition) + // Copy argument vectors whose subpartition lists include this core partition. #pragma omp parallel for for (size_t j = 0; j < augmented_sub_dataset_size; j++) { size_t i = augmented_backward_mapping(j + augmented_partition_offsets(partition_id)); @@ -480,7 +529,49 @@ void ace_gather_partition_dataset( } } -// Adjust IDs from core and augmented partitions to global reordered IDs +// ACE: Build the intermediate kNN graph for a core-plus-argument partition. +template +auto ace_build_partial_knn_graph( + raft::resources const& res, + raft::host_matrix_view core_argument_partition, + size_t intermediate_graph_degree, + cuvs::distance::DistanceType metric) -> raft::host_matrix +{ + auto const partition_size = static_cast(core_argument_partition.extent(0)); + RAFT_EXPECTS(partition_size > intermediate_graph_degree, + "ACE: partition size (%lu) must exceed the intermediate graph degree (%lu)", + partition_size, + intermediate_graph_degree); + + auto knn_graph = raft::make_host_matrix(partition_size, intermediate_graph_degree); + auto knn_params = cuvs::neighbors::cagra::graph_build_params::ivf_pq_params( + raft::make_extents(partition_size, core_argument_partition.extent(1)), metric); + knn_params.search_params.n_probes = + std::min(knn_params.search_params.n_probes * 2u, knn_params.build_params.n_lists); + build_knn_graph(res, core_argument_partition, knn_graph.view(), knn_params); + + return knn_graph; +} + +// ACE global reverse-edge mode: prune without merging local reverse edges. +template +auto ace_build_pruned_partial_graph( + raft::resources const& res, + raft::host_matrix_view core_argument_partition, + size_t intermediate_graph_degree, + size_t graph_degree, + cuvs::distance::DistanceType metric) -> raft::host_matrix +{ + auto knn_graph = ace_build_partial_knn_graph( + res, core_argument_partition, intermediate_graph_degree, metric); + auto partial_graph = + raft::make_host_matrix(core_argument_partition.extent(0), graph_degree); + graph::prune_graph_gpu(res, knn_graph.view(), partial_graph.view()); + return partial_graph; +} + +// ACE: Adjust IDs from core and augmented partitions to global reordered IDs + template void ace_adjust_sub_graph_ids( size_t core_sub_dataset_size, @@ -638,8 +729,8 @@ void ace_reorder_and_store_dataset( } RAFT_EXPECTS(total_core_vectors == dataset_size, "Total core vectors must be equal to dataset size"); - RAFT_EXPECTS(total_augmented_vectors == dataset_size, - "Total augmented vectors must be equal to dataset size"); + RAFT_EXPECTS(total_augmented_vectors == dataset_size * (partition_labels.extent(1) - 1), + "Total augmented vectors must match the number of secondary assignments"); // Pre-allocate file space for better performance const size_t vector_size = dataset_dim * sizeof(T); @@ -720,8 +811,7 @@ void ace_reorder_and_store_dataset( size_t vectors_processed = 0; size_t next_progress_percent = 10; for (size_t i = 0; i < dataset_size; i++) { - size_t core_partition = partition_labels(i, 0); - size_t secondary_partition = partition_labels(i, 1); + size_t core_partition = partition_labels(i, 0); // Add vector to core partition buffer size_t core_buffer_row = core_buffer_counts(core_partition); @@ -740,22 +830,25 @@ void ace_reorder_and_store_dataset( core_partition); } - // Add vector to augmented partition buffer - size_t augmented_buffer_row = augmented_buffer_counts(secondary_partition); - memcpy(&augmented_buffers[secondary_partition](augmented_buffer_row, 0), - &dataset(i, 0), - dataset_dim * sizeof(T)); - augmented_buffer_counts(secondary_partition)++; - - // Flush augmented buffer if full - if (augmented_buffer_counts(secondary_partition) >= vectors_per_buffer) { - flush_buffer(augmented_buffer_counts, - augmented_partition_current, - augmented_partition_starts, - augmented_buffers, - augmented_fd, - augmented_header_size, - secondary_partition); + for (size_t rank = 1; rank < static_cast(partition_labels.extent(1)); ++rank) { + size_t secondary_partition = partition_labels(i, rank); + // Add vector to augmented partition buffer + size_t augmented_buffer_row = augmented_buffer_counts(secondary_partition); + memcpy(&augmented_buffers[secondary_partition](augmented_buffer_row, 0), + &dataset(i, 0), + dataset_dim * sizeof(T)); + augmented_buffer_counts(secondary_partition)++; + + // Flush augmented buffer if full + if (augmented_buffer_counts(secondary_partition) >= vectors_per_buffer) { + flush_buffer(augmented_buffer_counts, + augmented_partition_current, + augmented_partition_starts, + augmented_buffers, + augmented_fd, + augmented_header_size, + secondary_partition); + } } vectors_processed++; @@ -960,9 +1053,6 @@ constexpr double usable_cpu_memory_fraction = 0.8; // Factor to account for imbalances in the partitions (maximum allowed is 3x the average) constexpr double imbalance_factor = 3.0; -// Current partitioning adds each vector into 2 partitions (core and augmented) -constexpr double vector_expansion_factor = 2.0; - // Use a configured GiB cap when it fits in the measured memory. Otherwise keep the measured value. inline size_t ace_resolve_available_memory(size_t actual_bytes, std::optional configured_gb, @@ -1009,9 +1099,12 @@ bool ace_check_use_disk_mode(raft::resources const& res, std::optional max_host_memory_gb, std::optional max_gpu_memory_gb, bool guarantee_connectivity, - ace_memory_requirements& mem) + bool add_global_reverse_edges, + ace_memory_requirements& mem, + size_t nsubpartitions = 1) { - const auto host_memory = cuvs::util::get_host_memory_info(); + const double vector_expansion_factor = 1.0 + nsubpartitions; + const auto host_memory = cuvs::util::get_host_memory_info(); RAFT_EXPECTS(host_memory.available > 0, "ACE: No host memory is available within the current system or cgroup limit"); if (host_memory.cgroup_limit.has_value() && host_memory.cgroup_current.has_value() && @@ -1060,6 +1153,12 @@ bool ace_check_use_disk_mode(raft::resources const& res, mem.cagra_graph_size = dataset_size * graph_degree * sizeof(IdxT); mem.total_size = mem.partition_labels_size + mem.id_mapping_size + mem.sub_dataset_size + mem.sub_graph_size + mem.cagra_graph_size + opt_host_ws_total; + if (add_global_reverse_edges) { + // The global merge keeps a full source graph copy and reserves one candidate per edge + // before deduplication, plus the per-partition vectors and their allocation counts. + mem.total_size += + 2 * mem.cagra_graph_size + n_partitions * (sizeof(std::vector) + sizeof(size_t)); + } RAFT_LOG_INFO("ACE build: host memory estimate required=%.2f GiB available=%.2f GiB", to_gib(mem.total_size), @@ -1143,8 +1242,11 @@ void ace_validate_disk_mode_partitions(raft::resources const& res, size_t intermediate_degree, size_t graph_degree, bool guarantee_connectivity, - ace_memory_requirements& mem) + bool add_global_reverse_edges, + ace_memory_requirements& mem, + size_t nsubpartitions = 1) { + const double vector_expansion_factor = 1.0 + nsubpartitions; // In disk mode, we don't need the full dataset or final graph in memory. // Host memory model for disk mode: // - Partition labels (core + augmented): vector_expansion_factor * dataset_size * sizeof(IdxT) @@ -1171,6 +1273,19 @@ void ace_validate_disk_mode_partitions(raft::resources const& res, helpers::optimize_workspace_size( sub_partition_size, graph_degree, intermediate_degree, sizeof(IdxT), guarantee_connectivity); + // Conservatively reserve the disk reverse-edge heaps in addition to build workspace. + // The record buffer is shared by all buckets and does not grow with the edge count. + size_t reverse_buffer_bytes = 0; + size_t reverse_bytes_per_row = 0; + if (add_global_reverse_edges) { + reverse_buffer_bytes = ace_disk_reverse_edges::buffer_bytes(mem.available_host_memory); + reverse_bytes_per_row = ace_disk_reverse_edges::incoming_capacity(graph_degree) * + sizeof(typename ace_disk_reverse_edges::incoming_edge) + + sizeof(uint32_t); + host_workspace_size_total += reverse_buffer_bytes + sub_partition_size * reverse_bytes_per_row; + host_workspace_size_fixed += reverse_buffer_bytes; + } + // Check host memory requirements size_t disk_mode_host_required = mem.partition_labels_size + mem.id_mapping_size + mem.sub_dataset_size + mem.sub_graph_size + @@ -1270,7 +1385,8 @@ void ace_validate_disk_mode_partitions(raft::resources const& res, mem.sub_graph_size = new_sub_partition_size * (intermediate_degree + graph_degree) * sizeof(IdxT); mem.total_size = mem.partition_labels_size + mem.id_mapping_size + mem.sub_dataset_size + - mem.sub_graph_size + mem.cagra_graph_size + new_opt_host_ws; + mem.sub_graph_size + mem.cagra_graph_size + new_opt_host_ws + + reverse_buffer_bytes + new_sub_partition_size * reverse_bytes_per_row; RAFT_LOG_INFO( "ACE build: updated partition memory estimates dataset=%.2f GiB graph=%.2f GiB host " @@ -1282,24 +1398,212 @@ void ace_validate_disk_mode_partitions(raft::resources const& res, } } -template - requires cuvs::neighbors::is_dense_row_major_device_dataset_view_v -auto build_from_device_matrix(raft::resources const& res, - const index_params& params, - DatasetViewT const& device_dataset) - -> cuvs::neighbors::cagra::index; +template +__global__ void ace_merge_partition_reverse_edges(IdxT* partition_graph, + IdxT const* reverse_graph, + uint32_t const* reverse_graph_count, + uint32_t partition_size, + uint32_t graph_degree) +{ + auto const num_protected_edges = graph_degree / 2; + for (uint32_t local_node = blockIdx.x * blockDim.x + threadIdx.x; local_node < partition_size; + local_node += blockDim.x * gridDim.x) { + auto* row = partition_graph + static_cast(local_node) * graph_degree; + auto num_reverse_edges = min(reverse_graph_count[local_node], graph_degree); + while (num_reverse_edges > 0) { + auto const reverse_node = + reverse_graph[static_cast(local_node) * graph_degree + --num_reverse_edges]; + uint32_t existing_rank = 0; + for (; existing_rank < graph_degree; ++existing_rank) { + if (row[existing_rank] == reverse_node) { break; } + } + if (existing_rank < num_protected_edges) { continue; } -template - requires cuvs::neighbors::is_host_dataset_view_v -auto build_from_host_matrix(raft::resources const& res, - const index_params& params, - DatasetViewT const& dataset) - -> cuvs::neighbors::cagra::index; + auto const shift_end = existing_rank == graph_degree ? graph_degree - 1 : existing_rank; + for (uint32_t rank = shift_end; rank > num_protected_edges; --rank) { + row[rank] = row[rank - 1]; + } + row[num_protected_edges] = reverse_node; + } + } +} + +template +void ace_merge_partition_reverse_edges(raft::resources const& res, + raft::host_matrix_view source_graph, + raft::host_matrix_view graph, + std::vector const& partition_nodes, + std::vector candidate_sources) +{ + if (partition_nodes.empty() || candidate_sources.empty()) { return; } + RAFT_EXPECTS( + source_graph.extent(0) == graph.extent(0) && source_graph.extent(1) == graph.extent(1), + "ACE: source and output graphs have different sizes"); + RAFT_EXPECTS(partition_nodes.size() <= std::numeric_limits::max(), + "ACE: partition is too large for reverse-edge processing"); + + auto const graph_degree = static_cast(graph.extent(1)); + auto const partition_size = static_cast(partition_nodes.size()); + auto const num_candidates = candidate_sources.size(); + + std::vector host_partition_graph(static_cast(partition_size) * graph_degree); + std::vector host_candidate_destinations(num_candidates); + for (uint32_t local_node = 0; local_node < partition_size; ++local_node) { + std::copy_n( + graph.data_handle() + static_cast(partition_nodes[local_node]) * graph_degree, + graph_degree, + host_partition_graph.data() + static_cast(local_node) * graph_degree); + } + + std::vector> sorted_partition_nodes; + sorted_partition_nodes.reserve(partition_size); + for (uint32_t local_node = 0; local_node < partition_size; ++local_node) { + sorted_partition_nodes.emplace_back(partition_nodes[local_node], local_node); + } + std::sort(sorted_partition_nodes.begin(), sorted_partition_nodes.end()); + std::vector host_sorted_partition_nodes(partition_size); + std::vector host_sorted_partition_local_indices(partition_size); + for (uint32_t i = 0; i < partition_size; ++i) { + host_sorted_partition_nodes[i] = sorted_partition_nodes[i].first; + host_sorted_partition_local_indices[i] = sorted_partition_nodes[i].second; + } + + auto large_workspace = raft::resource::get_large_workspace_resource_ref(res); + auto workspace = raft::resource::get_workspace_resource_ref(res); + auto d_partition_graph = raft::make_device_mdarray( + res, large_workspace, raft::make_extents(partition_size, graph_degree)); + auto d_reverse_graph = raft::make_device_mdarray( + res, large_workspace, raft::make_extents(partition_size, graph_degree)); + auto d_reverse_graph_count = raft::make_device_mdarray( + res, workspace, raft::make_extents(partition_size)); + auto d_candidate_sources = raft::make_device_mdarray( + res, workspace, raft::make_extents(num_candidates)); + auto d_candidate_destinations = raft::make_device_mdarray( + res, workspace, raft::make_extents(num_candidates, 1)); + auto d_sorted_partition_nodes = raft::make_device_mdarray( + res, workspace, raft::make_extents(partition_size)); + auto d_sorted_partition_local_indices = raft::make_device_mdarray( + res, workspace, raft::make_extents(partition_size)); + + auto const stream = raft::resource::get_cuda_stream(res); + raft::copy(d_partition_graph.data_handle(), + host_partition_graph.data(), + host_partition_graph.size(), + stream); + raft::copy(d_candidate_sources.data_handle(), candidate_sources.data(), num_candidates, stream); + raft::copy(d_sorted_partition_nodes.data_handle(), + host_sorted_partition_nodes.data(), + partition_size, + stream); + raft::copy(d_sorted_partition_local_indices.data_handle(), + host_sorted_partition_local_indices.data(), + partition_size, + stream); + RAFT_CUDA_TRY(cudaMemsetAsync( + d_reverse_graph_count.data_handle(), 0, partition_size * sizeof(uint32_t), stream.get())); + + constexpr uint32_t threads = 256; + auto const candidate_blocks = static_cast( + std::min(1024, raft::div_rounding_up_safe(num_candidates, size_t{threads}))); + auto const partition_blocks = static_cast( + std::min(1024, raft::div_rounding_up_safe(partition_size, threads))); + for (uint32_t rank = 0; rank < graph_degree; ++rank) { + for (size_t candidate_index = 0; candidate_index < num_candidates; ++candidate_index) { + host_candidate_destinations[candidate_index] = + source_graph(static_cast(candidate_sources[candidate_index]), rank); + } + raft::copy(d_candidate_destinations.data_handle(), + host_candidate_destinations.data(), + num_candidates, + stream); + graph::kern_make_rev_graph_k<<>>( + d_candidate_destinations.view(), + d_reverse_graph.view(), + d_reverse_graph_count.view(), + 0, + d_candidate_sources.data_handle(), + d_sorted_partition_nodes.data_handle(), + d_sorted_partition_local_indices.data_handle()); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + } + ace_merge_partition_reverse_edges<<>>( + d_partition_graph.data_handle(), + d_reverse_graph.data_handle(), + d_reverse_graph_count.data_handle(), + partition_size, + graph_degree); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + raft::copy(host_partition_graph.data(), + d_partition_graph.data_handle(), + host_partition_graph.size(), + stream); + raft::resource::sync_stream(res); + + for (uint32_t local_node = 0; local_node < partition_size; ++local_node) { + std::copy_n( + host_partition_graph.data() + static_cast(local_node) * graph_degree, + graph_degree, + graph.data_handle() + static_cast(partition_nodes[local_node]) * graph_degree); + } +} + +template +void ace_add_global_reverse_edges(raft::resources const& res, + raft::host_matrix_view graph, + raft::host_matrix_view partition_labels, + raft::host_vector_view core_backward_mapping, + raft::host_vector_view core_partition_offsets) +{ + auto const dataset_size = static_cast(graph.extent(0)); + auto const graph_degree = static_cast(graph.extent(1)); + auto const n_partitions = static_cast(core_partition_offsets.extent(0) - 1); + RAFT_EXPECTS(partition_labels.extent(0) == graph.extent(0), + "ACE: partition labels must match the graph size"); + RAFT_EXPECTS(core_backward_mapping.extent(0) == graph.extent(0), + "ACE: core mapping must match the graph size"); + + auto source_graph = raft::make_host_matrix(graph.extent(0), graph.extent(1)); + std::copy_n(graph.data_handle(), graph.size(), source_graph.data_handle()); + std::vector candidate_counts(n_partitions, 0); + for (size_t source = 0; source < dataset_size; ++source) { + for (uint32_t rank = 0; rank < graph_degree; ++rank) { + auto const destination = source_graph(source, rank); + RAFT_EXPECTS(destination < dataset_size, "ACE: graph contains an invalid neighbor"); + auto const destination_partition = static_cast(partition_labels(destination, 0)); + ++candidate_counts[destination_partition]; + } + } + std::vector> partition_candidates(n_partitions); + for (size_t partition_id = 0; partition_id < n_partitions; ++partition_id) { + partition_candidates[partition_id].reserve(candidate_counts[partition_id]); + } + for (size_t source = 0; source < dataset_size; ++source) { + for (uint32_t rank = 0; rank < graph_degree; ++rank) { + auto const destination = source_graph(source, rank); + auto const destination_partition = static_cast(partition_labels(destination, 0)); + partition_candidates[destination_partition].push_back(static_cast(source)); + } + } + + for (size_t partition_id = 0; partition_id < n_partitions; ++partition_id) { + auto& candidate_sources = partition_candidates[partition_id]; + if (candidate_sources.empty()) { continue; } + std::sort(candidate_sources.begin(), candidate_sources.end()); + candidate_sources.erase(std::unique(candidate_sources.begin(), candidate_sources.end()), + candidate_sources.end()); + auto const first = static_cast(core_partition_offsets(partition_id)); + auto const last = static_cast(core_partition_offsets(partition_id + 1)); + std::vector partition_nodes(core_backward_mapping.data_handle() + first, + core_backward_mapping.data_handle() + last); + ace_merge_partition_reverse_edges( + res, source_graph.view(), graph, partition_nodes, std::move(candidate_sources)); + } +} // Build CAGRA index using ACE (Augmented Core Extraction) partitioning // ACE enables building indexes for datasets too large to fit in GPU memory by: // 1. Partitioning the dataset using balanced k-means in core (non-overlapping) and augmented -// (second-closest) partitions +// (next-closest) partitions // 2. Building sub-indexes for each partition independently // 3. Concatenating sub-graphs (of core partitions) into a final unified index // Supports both in-memory and disk-based modes depending on available host memory. @@ -1341,13 +1645,21 @@ auto build_ace(raft::resources const& res, const index_params& params, DatasetVi "ACE build: intermediate graph degree must be greater than 0"); RAFT_EXPECTS(params.graph_degree > 0, "ACE build: graph degree must be greater than 0"); - size_t n_partitions = ace_resolve_partition_count(npartitions); + const size_t nsubpartitions = ace_params.nsubpartitions; + RAFT_EXPECTS(nsubpartitions > 0 && nsubpartitions < dataset_size, + "ACE: nsubpartitions must be in [1, dataset_size)"); + RAFT_EXPECTS(npartitions == 0 || nsubpartitions < ace_resolve_partition_count(npartitions), + "ACE: nsubpartitions must be smaller than npartitions"); + RAFT_EXPECTS(dataset_size <= std::numeric_limits::max() / nsubpartitions, + "ACE: augmented mapping exceeds the index type capacity"); + const size_t augmented_size = dataset_size * nsubpartitions; + size_t n_partitions = std::max(ace_resolve_partition_count(npartitions), nsubpartitions + 1); ace_validate_partition_count(n_partitions, dataset_size); size_t min_required_per_partition = 1000; if (n_partitions > dataset_size / min_required_per_partition) { - n_partitions = dataset_size / min_required_per_partition; + n_partitions = std::max(dataset_size / min_required_per_partition, nsubpartitions + 1); if (n_partitions < 2) { RAFT_LOG_WARN( "ACE build: reduced partitions to the minimum of 2 to avoid tiny partitions; regular " @@ -1360,10 +1672,12 @@ auto build_ace(raft::resources const& res, const index_params& params, DatasetVi auto total_start = std::chrono::high_resolution_clock::now(); RAFT_LOG_INFO( - "ACE build: start rows=%zu dim=%zu partitions=%zu graph_degree=%zu intermediate_degree=%zu", + "ACE build: start rows=%zu dim=%zu partitions=%zu sub_partitions=%zu graph_degree=%zu " + "intermediate_degree=%zu", dataset_size, dataset_dim, n_partitions, + nsubpartitions, static_cast(params.graph_degree), static_cast(params.intermediate_graph_degree)); @@ -1388,8 +1702,9 @@ auto build_ace(raft::resources const& res, const index_params& params, DatasetVi ace_params.max_host_memory_gb, ace_params.max_gpu_memory_gb, params.guarantee_connectivity, - mem); - + ace_params.add_global_reverse_edges, + mem, + ace_params.nsubpartitions); // Validate and adjust partitions if disk mode is enabled if (use_disk_mode) { ace_validate_disk_mode_partitions(res, @@ -1399,7 +1714,9 @@ auto build_ace(raft::resources const& res, const index_params& params, DatasetVi intermediate_degree, graph_degree, params.guarantee_connectivity, - mem); + ace_params.add_global_reverse_edges, + mem, + ace_params.nsubpartitions); ace_validate_partition_count(n_partitions, dataset_size, true); } @@ -1425,7 +1742,7 @@ auto build_ace(raft::resources const& res, const index_params& params, DatasetVi std::tie(augmented_fd, augmented_header_size) = cuvs::util::create_numpy_file( workspace.artifact_path(ace_disk_workspace::artifact::augmented_dataset), - {dataset_size, dataset_dim}, + {augmented_size, dataset_dim}, true); workspace.mark_artifact_created(ace_disk_workspace::artifact::augmented_dataset); @@ -1449,8 +1766,8 @@ auto build_ace(raft::resources const& res, const index_params& params, DatasetVi graph_header_size); } - auto partition_start = std::chrono::high_resolution_clock::now(); - auto partition_labels = raft::make_host_matrix(dataset_size, 2); + auto partition_start = std::chrono::high_resolution_clock::now(); + auto partition_labels = raft::make_host_matrix(dataset_size, nsubpartitions + 1); auto partition_histogram = raft::make_host_matrix(n_partitions, 2); for (size_t c = 0; c < n_partitions; c++) { partition_histogram(c, 0) = 0; @@ -1463,8 +1780,11 @@ auto build_ace(raft::resources const& res, const index_params& params, DatasetVi ace_get_partition_labels( res, dataset_view, dataset_dim, partition_labels.view(), partition_histogram.view()); - ace_check_partition_sizes( - dataset_size, n_partitions, partition_histogram.view(), min_partition_size); + ace_check_partition_sizes(dataset_size, + n_partitions, + partition_labels.view(), + partition_histogram.view(), + min_partition_size); auto partition_end = std::chrono::high_resolution_clock::now(); auto partition_elapsed = @@ -1479,7 +1799,7 @@ auto build_ace(raft::resources const& res, const index_params& params, DatasetVi auto core_forward_mapping = use_disk_mode ? raft::make_host_vector(dataset_size) : raft::make_host_vector(0); auto core_backward_mapping = raft::make_host_vector(dataset_size); - auto augmented_backward_mapping = raft::make_host_vector(dataset_size); + auto augmented_backward_mapping = raft::make_host_vector(augmented_size); auto core_partition_offsets = raft::make_host_vector(n_partitions + 1); auto augmented_partition_offsets = raft::make_host_vector(n_partitions + 1); @@ -1522,6 +1842,15 @@ auto build_ace(raft::resources const& res, const index_params& params, DatasetVi ? raft::make_host_matrix(0, 0) : raft::make_host_matrix(dataset_size, graph_degree); + std::optional> disk_reverse_edges; + if (use_disk_mode && ace_params.add_global_reverse_edges) { + disk_reverse_edges.emplace( + build_dir, + std::span(core_partition_offsets.data_handle(), n_partitions + 1), + graph_degree, + ace_disk_reverse_edges::buffer_bytes(mem.available_host_memory)); + } + // Process each partition auto partition_processing_start = std::chrono::high_resolution_clock::now(); for (size_t partition_id = 0; partition_id < n_partitions; partition_id++) { @@ -1558,15 +1887,28 @@ auto build_ace(raft::resources const& res, const index_params& params, DatasetVi using device_sub_index_t = cuvs::neighbors::cagra::device_padded_index; using host_sub_index_t = cuvs::neighbors::cagra::host_standard_index; using sub_index_t = std::variant; - auto build_host_sub_index = [&](auto& sub_dataset) -> sub_index_t { + std::optional> partial_sub_graph; + std::optional sub_index; + auto build_host_partition_graph = [&](auto& sub_dataset) { read_end = std::chrono::high_resolution_clock::now(); - auto sub_dataset_view = cuvs::neighbors::make_host_standard_dataset_view( - raft::make_const_mdspan(sub_dataset.view())); - auto host_index = ::cuvs::neighbors::cagra::build(res, sub_index_params, sub_dataset_view); - static_assert(std::is_same_v); - return sub_index_t{std::in_place_type, std::move(host_index)}; + auto sub_dataset_view = raft::make_const_mdspan(sub_dataset.view()); + if (ace_params.add_global_reverse_edges) { + // Keep only pruned outgoing edges. Reverse edges are added after all core + // rows have been mapped to global IDs, including cross-partition edges. + partial_sub_graph.emplace(ace_build_pruned_partial_graph( + res, sub_dataset_view, intermediate_degree, graph_degree, params.metric)); + } else { + auto host_index = ::cuvs::neighbors::cagra::build( + res, + sub_index_params, + cuvs::neighbors::make_host_standard_dataset_view(sub_dataset_view)); + static_assert(std::is_same_v); + sub_index.emplace(std::in_place_type, std::move(host_index)); + } }; - auto sub_index = [&]() -> sub_index_t { + // The prune-only builder consumes a host partition. The regular builder can + // also read directly into device memory when the partition fits. + auto build_partition_graph = [&]() { if (use_disk_mode) { const size_t current_free_gpu = rmm::available_device_memory().first; const size_t configured_gpu = @@ -1574,7 +1916,8 @@ auto build_ace(raft::resources const& res, const index_params& params, DatasetVi ? static_cast(ace_params.max_gpu_memory_gb * (1ULL << 30)) : current_free_gpu; const size_t free_gpu_bytes = std::min(current_free_gpu, configured_gpu); - if (sub_ds_bytes < static_cast(0.4 * free_gpu_bytes)) { + if (!ace_params.add_global_reverse_edges && + sub_ds_bytes < static_cast(0.4 * free_gpu_bytes)) { try { auto sub_dataset_tight = raft::make_device_matrix(res, sub_dataset_size, dataset_dim); @@ -1606,7 +1949,8 @@ auto build_ace(raft::resources const& res, const index_params& params, DatasetVi ::cuvs::neighbors::cagra::detail::build_from_device_matrix( res, sub_index_params, sub_dataset_dev); used_device_read = true; - return sub_index_t{std::in_place_type, std::move(direct_index)}; + sub_index.emplace(std::in_place_type, std::move(direct_index)); + return; } catch (const std::bad_alloc& e) { RAFT_LOG_WARN( "ACE: partition %lu did not fit in device memory for a direct (GDS) read: %s; " @@ -1627,23 +1971,23 @@ auto build_ace(raft::resources const& res, const index_params& params, DatasetVi reordered_header_size, augmented_header_size, sub_dataset.view()); - return build_host_sub_index(sub_dataset); + return build_host_partition_graph(sub_dataset); } auto sub_dataset = raft::make_host_matrix(sub_dataset_size, dataset_dim); - ace_gather_partition_dataset(core_sub_dataset_size, - augmented_sub_dataset_size, - dataset_dim, - partition_id, - dataset_view, - core_backward_mapping.view(), - augmented_backward_mapping.view(), - core_partition_offsets.view(), - augmented_partition_offsets.view(), - sub_dataset.view()); - return build_host_sub_index(sub_dataset); - }(); - auto sub_graph = std::visit([](auto const& index) { return index.graph(); }, sub_index); + ace_make_core_argument_partition(core_sub_dataset_size, + augmented_sub_dataset_size, + dataset_dim, + partition_id, + dataset_view, + core_backward_mapping.view(), + augmented_backward_mapping.view(), + core_partition_offsets.view(), + augmented_partition_offsets.view(), + sub_dataset.view()); + return build_host_partition_graph(sub_dataset); + }; + build_partition_graph(); if (used_device_read) { RAFT_LOG_DEBUG("ACE: partition %lu read directly into device memory (GDS path)", partition_id); @@ -1657,20 +2001,59 @@ auto build_ace(raft::resources const& res, const index_params& params, DatasetVi auto adjust_end = optimize_end; auto write_elapsed = 0L; const size_t graph_bytes = core_sub_dataset_size * graph_degree * sizeof(IdxT); - if (use_disk_mode) { + if (disk_reverse_edges.has_value()) { + // The pruned partial graph is already on the host. Convert only its core rows + // to disk IDs and collect incoming records before writing or releasing it. + auto adjusted_graph = + raft::make_host_vector(core_sub_dataset_size * graph_degree); + auto const core_offset = core_partition_offsets(partition_id); + auto const augmented_offset = augmented_partition_offsets(partition_id); + for (size_t edge = 0; edge < adjusted_graph.size(); ++edge) { + auto const local_id = partial_sub_graph->data_handle()[edge]; + RAFT_EXPECTS(size_t(local_id) < sub_dataset_size, + "ACE: partial graph contains an invalid neighbor"); + adjusted_graph(edge) = size_t(local_id) < core_sub_dataset_size + ? static_cast(core_offset + local_id) + : core_forward_mapping(augmented_backward_mapping( + augmented_offset + local_id - core_sub_dataset_size)); + } + disk_reverse_edges->append_partition( + partition_id, std::span(adjusted_graph.data_handle(), adjusted_graph.size())); + adjust_end = std::chrono::high_resolution_clock::now(); + auto const graph_offset = + size_t(core_offset) * graph_degree * sizeof(IdxT) + graph_header_size; + cuvs::util::write_large_file( + graph_fd, adjusted_graph.data_handle(), graph_bytes, graph_offset); + write_elapsed = std::chrono::duration_cast( + std::chrono::high_resolution_clock::now() - adjust_end) + .count(); + } else if (use_disk_mode) { auto adjusted_search_graph = raft::make_device_matrix(res, core_sub_dataset_size, graph_degree); - ace_adjust_sub_graph_ids_disk(res, - core_sub_dataset_size, - augmented_sub_dataset_size, - graph_degree, - partition_id, - sub_graph, - adjusted_search_graph.view(), - core_partition_offsets.view(), - augmented_partition_offsets.view(), - augmented_backward_mapping.view(), - core_forward_mapping.view()); + auto adjust_sub_graph = [&](auto sub_graph) { + ace_adjust_sub_graph_ids_disk(res, + core_sub_dataset_size, + augmented_sub_dataset_size, + graph_degree, + partition_id, + sub_graph, + adjusted_search_graph.view(), + core_partition_offsets.view(), + augmented_partition_offsets.view(), + augmented_backward_mapping.view(), + core_forward_mapping.view()); + }; + if (ace_params.add_global_reverse_edges) { + auto partial_sub_graph_device = + raft::make_device_matrix(res, core_sub_dataset_size, graph_degree); + raft::copy(res, + partial_sub_graph_device.view(), + raft::make_host_matrix_view( + partial_sub_graph->data_handle(), core_sub_dataset_size, graph_degree)); + adjust_sub_graph(raft::make_const_mdspan(partial_sub_graph_device.view())); + } else { + adjust_sub_graph(std::visit([](auto const& index) { return index.graph(); }, *sub_index)); + } adjust_end = std::chrono::high_resolution_clock::now(); const size_t graph_offset = @@ -1684,11 +2067,18 @@ auto build_ace(raft::resources const& res, const index_params& params, DatasetVi } else { auto sub_search_graph = raft::make_host_matrix(core_sub_dataset_size, graph_degree); - raft::copy( - res, - raft::make_host_vector_view(sub_search_graph.data_handle(), sub_search_graph.size()), - raft::make_device_vector_view(sub_graph.data_handle(), sub_search_graph.size())); - raft::resource::sync_stream(res); + if (ace_params.add_global_reverse_edges) { + std::copy_n(partial_sub_graph->data_handle(), + sub_search_graph.size(), + sub_search_graph.data_handle()); + } else { + auto sub_graph = std::visit([](auto const& index) { return index.graph(); }, *sub_index); + raft::copy( + res, + raft::make_host_vector_view(sub_search_graph.data_handle(), sub_search_graph.size()), + raft::make_device_vector_view(sub_graph.data_handle(), sub_search_graph.size())); + raft::resource::sync_stream(res); + } // Adjust IDs in sub_search_graph and save to search_graph ace_adjust_sub_graph_ids(core_sub_dataset_size, @@ -1738,6 +2128,20 @@ auto build_ace(raft::resources const& res, const index_params& params, DatasetVi partition_processing_elapsed, n_partitions); + if (ace_params.add_global_reverse_edges) { + RAFT_LOG_INFO("ACE: Adding global reverse edges partition by partition"); + if (disk_reverse_edges.has_value()) { + disk_reverse_edges->merge(graph_fd, graph_header_size); + disk_reverse_edges.reset(); + } else { + ace_add_global_reverse_edges(res, + search_graph.view(), + partition_labels.view(), + core_backward_mapping.view(), + core_partition_offsets.view()); + } + } + // Clean up augmented dataset file to save disk space (no longer needed after partitions // processed) if (use_disk_mode && diff --git a/cpp/src/neighbors/detail/cagra/cagra_ace_disk.hpp b/cpp/src/neighbors/detail/cagra/cagra_ace_disk.hpp new file mode 100644 index 0000000000..2bd2f8e9cb --- /dev/null +++ b/cpp/src/neighbors/detail/cagra/cagra_ace_disk.hpp @@ -0,0 +1,298 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace cuvs::neighbors::cagra::detail { + +// These buffers are always host memory. Use positional POSIX I/O directly to avoid +// reopening/registering the bucket with KvikIO for every buffered run. +template +void ace_disk_host_io(cuvs::util::file_descriptor const& fd, + std::conditional_t data, + size_t bytes, + size_t offset) +{ + auto constexpr max_offset = static_cast(std::numeric_limits::max()); + RAFT_EXPECTS(offset <= max_offset && bytes <= max_offset - offset, + "ACE: reverse-edge file offset overflow"); + auto* buffer = static_cast>(data); + while (bytes > 0) { + auto const chunk = std::min(bytes, size_t{1} << 30); + ssize_t transferred; + if constexpr (Write) { + transferred = ::pwrite(fd.get(), buffer, chunk, static_cast(offset)); + } else { + transferred = ::pread(fd.get(), buffer, chunk, static_cast(offset)); + } + if (transferred < 0 && errno == EINTR) { continue; } + RAFT_EXPECTS(transferred >= 0, + "ACE: reverse-edge %s failed: %s", + Write ? "write" : "read", + std::strerror(errno)); + RAFT_EXPECTS(transferred > 0, "ACE: incomplete reverse-edge %s", Write ? "write" : "read"); + buffer += transferred; + offset += transferred; + bytes -= transferred; + } +} + +// Spool all edges while the pruned core rows are still in memory. A single +// bounded buffer is sorted into destination-partition runs before writing; only one bucket +// file is open at a time. Neither candidate sets nor a second full graph are kept in RAM. +template +class ace_disk_reverse_edges { + public: + struct record { + IdxT destination; + IdxT source; + uint32_t rank; + }; + + struct incoming_edge { + IdxT source; + uint32_t rank; + + bool operator<(incoming_edge const& other) const + { + return rank < other.rank || (rank == other.rank && source < other.source); + } + }; + + // The first floor(degree / 2) graph neighbors are protected during the merge. + static uint32_t incoming_capacity(uint32_t degree) { return degree - degree / 2; } + + static size_t buffer_bytes(size_t available_host_memory) + { + return std::max(sizeof(record), std::min(size_t{8} << 20, available_host_memory / 100)); + } + + ace_disk_reverse_edges(std::string const& build_dir, + std::span offsets, + uint32_t degree, + size_t buffer_size) + : offsets_(raft::make_host_vector(offsets.size())), + degree_(degree), + capacity_(std::max(size_t{1}, buffer_size / sizeof(record))), + sizes_(raft::make_host_vector(offsets.empty() ? 0 : offsets.size() - 1)), + buffer_(raft::make_host_vector(capacity_)) + { + RAFT_EXPECTS(offsets.size() >= 2 && offsets.front() == 0 && degree_ > 0, + "ACE: invalid reverse-edge partition layout"); + for (size_t i = 0; i < offsets.size(); ++i) { + RAFT_EXPECTS(i == 0 || offsets[i - 1] <= offsets[i], + "ACE: invalid reverse-edge partition layout"); + offsets_(i) = offsets[i]; + if (i + 1 < offsets.size()) { sizes_(i) = 0; } + } + directory_ = build_dir + "/reverse_edges_XXXXXX"; + RAFT_EXPECTS(::mkdtemp(directory_.data()) != nullptr, + "ACE: failed to create reverse-edge workspace: %s", + std::strerror(errno)); + } + + ace_disk_reverse_edges(ace_disk_reverse_edges const&) = delete; + ace_disk_reverse_edges& operator=(ace_disk_reverse_edges const&) = delete; + + ~ace_disk_reverse_edges() + { + // This uniquely created directory contains only this build's temporary records. + std::error_code error; + std::filesystem::remove_all(directory_, error); + if (error) { + RAFT_LOG_WARN("ACE: failed to remove reverse-edge workspace %s: %s", + directory_.c_str(), + error.message().c_str()); + } + } + + // Rows and neighbors both use core-reordered IDs, as in the on-disk graph. + void append_partition(size_t partition, std::span graph) + { + RAFT_EXPECTS(partition < sizes_.size(), "ACE: invalid reverse-edge partition"); + auto const first = offsets_(partition); + auto const last = offsets_(partition + 1); + RAFT_EXPECTS(graph.size() == size_t(last - first) * degree_, + "ACE: invalid pruned core graph size"); + for (size_t row = 0; row < size_t(last - first); ++row) { + for (uint32_t rank = 0; rank < degree_; ++rank) { + auto const destination = graph[row * degree_ + rank]; + RAFT_EXPECTS( + static_cast(destination) < static_cast(offsets_(offsets_.size() - 1)), + "ACE: graph contains an invalid neighbor"); + buffer_(buffer_count_++) = {destination, static_cast(first + row), rank}; + if (buffer_count_ == capacity_) { flush(); } + } + } + } + + void merge(cuvs::util::file_descriptor const& graph_fd, size_t graph_header_size) + { + flush(); + auto const candidate_capacity = incoming_capacity(degree_); + for (size_t partition = 0; partition < sizes_.size(); ++partition) { + if (sizes_(partition) == 0) { continue; } + auto const first = offsets_(partition); + auto const rows = size_t(offsets_(partition + 1) - first); + + // Keep only enough incoming candidates to fill the replaceable half of each row. + // Buckets can be arbitrarily skewed: their size does not determine RAM usage. + // Max-heaps preserve rank priority across buffer flushes and source-partition build order. + // GPU rank ties are scheduling-dependent; the disk path breaks them by source ID. + auto incoming = raft::make_host_vector(rows * candidate_capacity); + auto counts = raft::make_host_vector(rows); + for (size_t row = 0; row < rows; ++row) { + counts(row) = 0; + } + cuvs::util::file_descriptor bucket(bucket_path(partition), O_RDONLY); + for (size_t offset = 0; offset < sizes_(partition);) { + auto const count = std::min(capacity_, sizes_(partition) - offset); + ace_disk_host_io( + bucket, buffer_.data_handle(), count * sizeof(record), offset * sizeof(record)); + for (size_t i = 0; i < count; ++i) { + auto const& edge = buffer_(i); + RAFT_EXPECTS(edge.destination >= first && edge.destination < offsets_(partition + 1) && + edge.rank < degree_ && + static_cast(edge.source) < + static_cast(offsets_(offsets_.size() - 1)), + "ACE: invalid reverse-edge record"); + auto const row = size_t(edge.destination - first); + auto* heap = incoming.data_handle() + row * candidate_capacity; + auto& n = counts(row); + incoming_edge candidate{edge.source, edge.rank}; + if (n < candidate_capacity) { + auto index = n++; + while (index > 0) { + auto const parent = (index - 1) / 2; + if (!(heap[parent] < candidate)) { break; } + heap[index] = heap[parent]; + index = parent; + } + heap[index] = candidate; + } else if (candidate < heap[0]) { + heap[0] = candidate; + sift_down(heap, n); + } + } + offset += count; + } + + auto graph = raft::make_host_vector(rows * degree_); + auto const file_offset = graph_header_size + size_t(first) * degree_ * sizeof(IdxT); + ace_disk_host_io( + graph_fd, graph.data_handle(), graph.size() * sizeof(IdxT), file_offset); + auto const protected_edges = degree_ / 2; + for (size_t row = 0; row < rows; ++row) { + auto* edges = incoming.data_handle() + row * candidate_capacity; + auto n = counts(row); + auto* neighbors = graph.data_handle() + row * degree_; + // Match ace_merge_partition_reverse_edges: insert in reverse priority order, + // protect the first half of the pruned row, and move rather than duplicate edges. + while (n > 0) { + auto const source = edges[0].source; + edges[0] = edges[--n]; + sift_down(edges, n); + uint32_t existing = 0; + while (existing < degree_ && neighbors[existing] != source) { + ++existing; + } + if (existing < protected_edges) { continue; } + auto const end = existing == degree_ ? degree_ - 1 : existing; + for (auto rank = end; rank > protected_edges; --rank) { + neighbors[rank] = neighbors[rank - 1]; + } + neighbors[protected_edges] = source; + } + } + ace_disk_host_io( + graph_fd, graph.data_handle(), graph.size() * sizeof(IdxT), file_offset); + bucket.close(); + std::filesystem::remove(bucket_path(partition)); + } + buffer_count_ = 0; + } + + private: + std::string bucket_path(size_t partition) const + { + return directory_ + "/" + std::to_string(partition); + } + + void flush() + { + // Sort the fixed POD buffer without STL containers or algorithms. + ::qsort(buffer_.data_handle(), buffer_count_, sizeof(record), [](void const* a, void const* b) { + auto const lhs = static_cast(a)->destination; + auto const rhs = static_cast(b)->destination; + return lhs < rhs ? -1 : lhs > rhs ? 1 : 0; + }); + size_t begin = 0; + while (begin < buffer_count_) { + // Upper-bound search also handles empty core partitions (repeated offsets). + size_t low = 0; + size_t high = offsets_.size(); + while (low < high) { + auto const middle = low + (high - low) / 2; + if (offsets_(middle) <= buffer_(begin).destination) { + low = middle + 1; + } else { + high = middle; + } + } + auto const partition = low - 1; + size_t end = begin + 1; + while (end < buffer_count_ && buffer_(end).destination < offsets_(partition + 1)) { + ++end; + } + cuvs::util::file_descriptor bucket( + bucket_path(partition), sizes_(partition) == 0 ? O_CREAT | O_EXCL | O_WRONLY : O_WRONLY); + ace_disk_host_io(bucket, + buffer_.data_handle() + begin, + (end - begin) * sizeof(record), + sizes_(partition) * sizeof(record)); + sizes_(partition) += end - begin; + begin = end; + } + buffer_count_ = 0; + } + + // Restore a max-heap after replacing its root. No allocation occurs per candidate. + static void sift_down(incoming_edge* heap, uint32_t count) + { + if (count == 0) { return; } + auto const value = heap[0]; + uint32_t index = 0; + while (index < count / 2) { + auto child = index * 2 + 1; + if (child + 1 < count && heap[child] < heap[child + 1]) { ++child; } + if (!(value < heap[child])) { break; } + heap[index] = heap[child]; + index = child; + } + heap[index] = value; + } + + raft::host_vector offsets_; + uint32_t degree_; + size_t capacity_; + raft::host_vector sizes_; + raft::host_vector buffer_; + size_t buffer_count_ = 0; + std::string directory_; +}; + +} // namespace cuvs::neighbors::cagra::detail diff --git a/cpp/src/neighbors/detail/cagra/cagra_build.cuh b/cpp/src/neighbors/detail/cagra/cagra_build.cuh index 8c8766cbd5..66a331739d 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_build.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_build.cuh @@ -5,9 +5,7 @@ #pragma once #include "../../../core/nvtx.hpp" -#include "../../../util/kvikio_io.hpp" #include "../../ivf_pq/ivf_pq_fp16_overflow.cuh" -#include "cagra_ace.cuh" #include "cagra_search.cuh" #include "graph_core.cuh" #include @@ -27,34 +25,25 @@ #include #include -#include #include #include #include #include #include #include -#include - -#include // TODO: This shouldn't be calling spatial/knn APIs #include "../ann_utils.cuh" #include -#include -#include -#include #include #include #include -#include -#include +#include #include #include #include -#include #include #include #include @@ -87,6 +76,13 @@ void check_graph_degree(size_t& intermediate_degree, size_t& graph_degree, size_ } } +template +void build_knn_graph( + raft::resources const& res, + raft::mdspan, raft::row_major, accessor> dataset, + raft::host_matrix_view knn_graph, + cuvs::neighbors::cagra::graph_build_params::ivf_pq_params pq); + template void write_to_graph(raft::host_matrix_view knn_graph, raft::host_matrix_view neighbors_host_view, diff --git a/cpp/src/neighbors/detail/cagra/graph_core.cuh b/cpp/src/neighbors/detail/cagra/graph_core.cuh index 6138412414..dd96685aa1 100644 --- a/cpp/src/neighbors/detail/cagra/graph_core.cuh +++ b/cpp/src/neighbors/detail/cagra/graph_core.cuh @@ -54,26 +54,51 @@ inline double cur_time(void) return ((double)tv.tv_sec + (double)tv.tv_usec * 1e-6); } +// When index_conversion_table is provided, it maps an output-graph row to the node ID stored in +// the reverse graph. sorted_destination_nodes/local_indices optionally maps global destination +// IDs into compact reverse-graph rows; omitting both preserves the standard dense-ID path. template __global__ void kern_make_rev_graph_k( OutputMatrixView output_graph, // [graph_size, degree] raft::device_matrix_view rev_graph, // [graph_size, degree] raft::device_vector_view rev_graph_count, // [graph_size] - uint64_t k) + uint64_t k, + IdxT const* index_conversion_table = nullptr, + IdxT const* sorted_destination_nodes = nullptr, + uint32_t const* sorted_destination_local_indices = nullptr) { const uint64_t tid = threadIdx.x + (blockDim.x * blockIdx.x); const uint64_t tnum = blockDim.x * gridDim.x; - const uint64_t graph_size = rev_graph.extent(0); - const uint32_t rev_graph_degree = rev_graph.extent(1); - const uint32_t output_graph_degree = output_graph.extent(1); - - for (uint64_t src_id = tid; src_id < graph_size; src_id += tnum) { - IdxT dest_id = output_graph(src_id, k); - if (dest_id >= graph_size) continue; + const uint64_t graph_size = rev_graph.extent(0); + const uint32_t rev_graph_degree = rev_graph.extent(1); + + const uint64_t num_sources = output_graph.extent(0); + for (uint64_t src_id = tid; src_id < num_sources; src_id += tnum) { + const auto dest_id = output_graph(src_id, k); + auto local_dest_id = dest_id; + if (sorted_destination_nodes != nullptr) { + uint64_t lower = 0; + uint64_t upper = graph_size; + while (lower < upper) { + const auto middle = lower + (upper - lower) / 2; + if (sorted_destination_nodes[middle] < dest_id) { + lower = middle + 1; + } else { + upper = middle; + } + } + if (lower == graph_size || sorted_destination_nodes[lower] != dest_id) { continue; } + local_dest_id = sorted_destination_local_indices[lower]; + } + if (local_dest_id >= graph_size) continue; - const uint32_t pos = atomicAdd(&rev_graph_count(dest_id), 1); - if (pos < rev_graph_degree) { rev_graph(dest_id, pos) = static_cast(src_id); } + const uint32_t pos = atomicAdd(&rev_graph_count(local_dest_id), 1); + if (pos < rev_graph_degree) { + rev_graph(local_dest_id, pos) = index_conversion_table == nullptr + ? static_cast(src_id) + : index_conversion_table[src_id]; + } } } diff --git a/cpp/tests/neighbors/ann_cagra/test_float_uint32_t.cu b/cpp/tests/neighbors/ann_cagra/test_float_uint32_t.cu index 2b445df7bc..841c7ad4f4 100644 --- a/cpp/tests/neighbors/ann_cagra/test_float_uint32_t.cu +++ b/cpp/tests/neighbors/ann_cagra/test_float_uint32_t.cu @@ -6,6 +6,7 @@ #include #include "../ann_cagra.cuh" +#include "../ann_cagra_ace.cuh" #include #include @@ -14,6 +15,13 @@ namespace cuvs::neighbors::cagra { +typedef AnnCagraAceTest AnnCagraAceTestF_U32; +TEST_P(AnnCagraAceTestF_U32, BuildSearch) { this->testCagraAce(); } +INSTANTIATE_TEST_CASE_P(AnnCagraAceTest, + AnnCagraAceTestF_U32, + ::testing::ValuesIn(inputs_ace), + ace_case_name); + typedef AnnCagraTest AnnCagraTestF_U32; TEST_P(AnnCagraTestF_U32, AnnCagra_U32) { this->testCagra(); } TEST_P(AnnCagraTestF_U32, AnnCagra_I64) { this->testCagra(); } diff --git a/cpp/tests/neighbors/ann_cagra/test_half_uint32_t.cu b/cpp/tests/neighbors/ann_cagra/test_half_uint32_t.cu index b4906ae4ff..901f59aef1 100644 --- a/cpp/tests/neighbors/ann_cagra/test_half_uint32_t.cu +++ b/cpp/tests/neighbors/ann_cagra/test_half_uint32_t.cu @@ -6,9 +6,17 @@ #include #include "../ann_cagra.cuh" +#include "../ann_cagra_ace.cuh" namespace cuvs::neighbors::cagra { +typedef AnnCagraAceTest AnnCagraAceTestF16_U32; +TEST_P(AnnCagraAceTestF16_U32, BuildSearch) { this->testCagraAce(); } +INSTANTIATE_TEST_CASE_P(AnnCagraAceTest, + AnnCagraAceTestF16_U32, + ::testing::ValuesIn(inputs_ace), + ace_case_name); + typedef AnnCagraTest AnnCagraTestF16_U32; TEST_P(AnnCagraTestF16_U32, AnnCagra_U32) { this->testCagra(); } TEST_P(AnnCagraTestF16_U32, AnnCagra_I64) { this->testCagra(); } diff --git a/cpp/tests/neighbors/ann_cagra_ace.cuh b/cpp/tests/neighbors/ann_cagra_ace.cuh new file mode 100644 index 0000000000..fb54f04cfd --- /dev/null +++ b/cpp/tests/neighbors/ann_cagra_ace.cuh @@ -0,0 +1,206 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include "ann_cagra.cuh" + +#include +#include +#include + +#include +#include + +namespace cuvs::neighbors::cagra { +namespace { + +struct AnnCagraAceInputs { + bool use_disk; + bool add_global_reverse_edges; + cuvs::distance::DistanceType metric; + int npartitions; + size_t nsubpartitions = 1; +}; + +// Own the workspace even if a build throws or a fatal assertion returns early. +struct ace_test_workspace { + std::string path = (std::filesystem::temp_directory_path() / "cuvs_cagra_ace_XXXXXX").string(); + ace_test_workspace() + { + RAFT_EXPECTS(::mkdtemp(path.data()) != nullptr, "Cannot create ACE test workspace"); + } + ~ace_test_workspace() + { + std::error_code error; + std::filesystem::remove_all(path, error); + } +}; + +template +void read_ace_array(cuvs::util::file_descriptor const& fd, + std::vector& values, + std::vector const& shape) +{ + auto stream = fd.make_istream(); + stream.seekg(0); + auto header = raft::numpy_serializer::read_header(stream); + ASSERT_EQ(header.shape, shape); + ASSERT_EQ(header.dtype, raft::numpy_serializer::get_numpy_dtype()); + stream.read(reinterpret_cast(values.data()), values.size() * sizeof(T)); + ASSERT_TRUE(stream.good()); +} + +template +class AnnCagraAceTest : public ::testing::TestWithParam { + public: + void testCagraAce() + { + auto const [use_disk, global_reverse_edges, metric, npartitions, nsubpartitions] = + this->GetParam(); + // Keep the average core partition size fixed while varying the partition count. + const uint32_t rows = 5000 * npartitions; + constexpr int dim = 16, queries = 100, k = 10, degree = 64; + ace_test_workspace workspace; + raft::resources res; + auto stream = raft::resource::get_cuda_stream(res); + auto dataset = raft::make_host_matrix(rows, dim); + auto device_data = raft::make_device_matrix(res, rows, dim); + auto device_queries = raft::make_device_matrix(res, queries, dim); + raft::random::RngState rng(1234ULL); + InitDataset(res, device_data.data_handle(), rows, dim, metric, rng); + InitDataset(res, device_queries.data_handle(), queries, dim, metric, rng); + raft::copy(dataset.data_handle(), device_data.data_handle(), dataset.size(), stream); + raft::resource::sync_stream(res); + + index_params params; + params.metric = metric; + params.graph_degree = degree; + params.intermediate_graph_degree = 2 * degree; + params.attach_dataset_on_build = false; + graph_build_params::ace_params ace; + ace.npartitions = npartitions; + ace.nsubpartitions = nsubpartitions; + ace.ef_construction = 100; + ace.use_disk = use_disk; + ace.add_global_reverse_edges = global_reverse_edges; + ace.build_dir = workspace.path; + params.graph_build_params = ace; + auto host_view = + cuvs::neighbors::make_host_standard_dataset_view(raft::make_const_mdspan(dataset.view())); + auto built = cagra::build(res, params, host_view); + ASSERT_EQ(built.size(), rows); + ASSERT_EQ(built.graph_degree(), degree); + ASSERT_EQ(built.graph_fd().has_value(), use_disk); + ASSERT_EQ(built.dataset_fd().has_value(), use_disk); + ASSERT_EQ(built.mapping_fd().has_value(), use_disk); + + auto graph = raft::make_host_matrix(rows, degree); + if (use_disk) { + std::vector disk_graph(rows * degree), mapping(rows); + std::vector reordered(rows * dim); + ASSERT_NO_FATAL_FAILURE( + read_ace_array(*built.graph_fd(), disk_graph, {size_t(rows), degree})); + ASSERT_NO_FATAL_FAILURE(read_ace_array(*built.mapping_fd(), mapping, {size_t(rows)})); + ASSERT_NO_FATAL_FAILURE(read_ace_array(*built.dataset_fd(), reordered, {size_t(rows), dim})); + auto sorted_mapping = mapping; + std::sort(sorted_mapping.begin(), sorted_mapping.end()); + for (uint32_t row = 0; row < rows; ++row) { + ASSERT_EQ(sorted_mapping[row], row); + } + // Restore original IDs to search the persisted graph against the original dataset. + for (uint32_t row = 0; row < rows; ++row) { + for (int col = 0; col < dim; ++col) { + ASSERT_EQ(reordered[row * dim + col], dataset(mapping[row], col)); + } + for (int rank = 0; rank < degree; ++rank) { + auto neighbor = disk_graph[row * degree + rank]; + ASSERT_LT(neighbor, rows); + graph(mapping[row], rank) = mapping[neighbor]; + } + } + // The augmented vectors and reverse-edge scratch files must be removed on success. + EXPECT_EQ(std::distance(std::filesystem::directory_iterator(workspace.path), + std::filesystem::directory_iterator{}), + 3); + } else { + raft::copy(graph.data_handle(), built.graph().data_handle(), graph.size(), stream); + raft::resource::sync_stream(res); + } + for (uint32_t row = 0; row < rows; ++row) { + auto* first = graph.data_handle() + row * degree; + std::vector neighbors(first, first + degree); + std::sort(neighbors.begin(), neighbors.end()); + ASSERT_LT(neighbors.back(), rows); + EXPECT_EQ(std::find(neighbors.begin(), neighbors.end(), row), neighbors.end()); + EXPECT_EQ(std::adjacent_find(neighbors.begin(), neighbors.end()), neighbors.end()); + } + + cuvs::neighbors::test::padded_device_matrix_for_cagra padded( + res, raft::make_const_mdspan(device_data.view())); + device_padded_index searchable( + res, metric, padded.view, raft::make_const_mdspan(graph.view())); + auto expected_ids = raft::make_device_matrix(res, queries, k); + auto expected_distances = raft::make_device_matrix(res, queries, k); + cuvs::neighbors::naive_knn(res, + expected_distances.data_handle(), + expected_ids.data_handle(), + device_queries.data_handle(), + device_data.data_handle(), + queries, + rows, + dim, + k, + metric); + auto actual_ids = raft::make_device_matrix(res, queries, k); + auto actual_distances = raft::make_device_matrix(res, queries, k); + search_params search; + search.itopk_size = 128; + cagra::search(res, + search, + searchable, + raft::make_const_mdspan(device_queries.view()), + actual_ids.view(), + actual_distances.view()); + std::vector expected(queries * k), actual(queries * k); + std::vector expected_d(queries * k), actual_d(queries * k); + raft::copy(expected.data(), expected_ids.data_handle(), expected.size(), stream); + raft::copy(actual.data(), actual_ids.data_handle(), actual.size(), stream); + raft::copy(expected_d.data(), expected_distances.data_handle(), expected_d.size(), stream); + raft::copy(actual_d.data(), actual_distances.data_handle(), actual_d.size(), stream); + raft::resource::sync_stream(res); + EXPECT_TRUE(cuvs::neighbors::eval_neighbours( + expected, actual, expected_d, actual_d, queries, k, 0.003, 0.9)); + } +}; + +const std::vector inputs_ace = [] { + auto inputs = raft::util::itertools::product( + {false, true}, // use_disk + {false, true}, // add_global_reverse_edges + {cuvs::distance::DistanceType::L2Expanded, cuvs::distance::DistanceType::InnerProduct}, + {2, 8, 32, 64}, // npartitions + {size_t{1}}); + for (bool disk : {false, true}) { + for (bool reverse : {false, true}) { + for (size_t secondary : {size_t{2}, size_t{3}}) { + inputs.push_back({disk, reverse, cuvs::distance::DistanceType::L2Expanded, 4, secondary}); + } + } + } + return inputs; +}(); + +std::string ace_case_name(::testing::TestParamInfo const& info) +{ + auto const [disk, reverse, metric, npartitions, nsubpartitions] = info.param; + return std::string(disk ? "Disk" : "Memory") + (reverse ? "GlobalReverse" : "LocalReverse") + + (metric == cuvs::distance::DistanceType::L2Expanded ? "L2" : "InnerProduct") + + "Partitions" + std::to_string(npartitions) + "Subpartitions" + + std::to_string(nsubpartitions); +} + +} // namespace +} // namespace cuvs::neighbors::cagra diff --git a/fern/pages/c_api/c-api-neighbors-cagra.md b/fern/pages/c_api/c-api-neighbors-cagra.md index c6ceb153d5..faedc849c4 100644 --- a/fern/pages/c_api/c-api-neighbors-cagra.md +++ b/fern/pages/c_api/c-api-neighbors-cagra.md @@ -102,6 +102,8 @@ struct cuvsAceParams { bool use_disk; double max_host_memory_gb; double max_gpu_memory_gb; + size_t nsubpartitions; + bool add_global_reverse_edges; }; ``` @@ -109,12 +111,14 @@ struct cuvsAceParams { | Name | Type | Description | | --- | --- | --- | -| `npartitions` | `size_t` | Number of partitions for ACE (Augmented Core Extraction) partitioned build.

When set to 0 (default), the number of partitions is automatically derived based on available host and GPU memory to maximize partition size while ensuring the build fits in memory.

Small values might improve recall but potentially degrade performance and increase memory usage. Partitions should not be too small to prevent issues in KNN graph construction. The partition size is on average 2 * (n_rows / npartitions) * dim * sizeof(T). 2 is because of the core and augmented vectors. Please account for imbalance in the partition sizes (up to 3x in our tests).

If the specified number of partitions results in partitions that exceed available memory, the value will be automatically increased to fit memory constraints and a warning will be issued. | +| `npartitions` | `size_t` | Number of partitions for ACE (Augmented Core Extraction) partitioned build.

When set to 0 (default), the number of partitions is automatically derived based on available host and GPU memory to maximize partition size while ensuring the build fits in memory.

Small values might improve recall but potentially degrade performance and increase memory usage. Partitions should not be too small to prevent issues in KNN graph construction. The partition size is on average `(1 + nsubpartitions) * (n_rows / npartitions) * dim * sizeof(T)`, accounting for core and augmented vectors. Please account for imbalance in the partition sizes (up to 3x in our tests).

If the specified number of partitions results in partitions that exceed available memory, the value will be automatically increased to fit memory constraints and a warning will be issued. | | `ef_construction` | `size_t` | The index quality for the ACE build.

Bigger values increase the index quality. At some point, increasing this will no longer improve the quality. | | `build_dir` | `const char*` | Directory to store ACE build artifacts (e.g., KNN graph, optimized graph).

Used when `use_disk` is true or when the graph does not fit in host and GPU memory. This should be the fastest disk in the system and hold enough space for twice the dataset, final graph, and label mapping. The directory may already exist, but ACE's named artifacts must not already exist. Simultaneous builds must use different directories. On failure, ACE removes only artifacts it created and never deletes unrelated directory contents. | | `use_disk` | `bool` | Whether to use disk-based storage for ACE build.

When true, enables disk-based operations for memory-efficient graph construction. | | `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. | +| `nsubpartitions` | `size_t` | Number of additional nearest partitions assigned to each vector. Default: 1. Must be positive and smaller than npartitions (unless automatically selected). Increasing this improves overlap at the cost of build time, memory, and disk space. | +| `add_global_reverse_edges` | `bool` | Build pruned partial graphs, then add reverse edges globally. Default: false. In disk mode, incoming edges are buffered to temporary files and merged on the CPU. | ### cuvsCagraIndexParams diff --git a/fern/pages/cpp_api/cpp-api-neighbors-cagra.md b/fern/pages/cpp_api/cpp-api-neighbors-cagra.md index 9c3f67dec8..6ff9f3d69b 100644 --- a/fern/pages/cpp_api/cpp-api-neighbors-cagra.md +++ b/fern/pages/cpp_api/cpp-api-neighbors-cagra.md @@ -213,9 +213,11 @@ Specialized parameters for ACE (Augmented Core Extraction) graph build ```cpp struct ace_params { size_t npartitions; + size_t nsubpartitions; size_t ef_construction; std::string build_dir; bool use_disk; + bool add_global_reverse_edges; double max_host_memory_gb; double max_gpu_memory_gb; }; @@ -225,10 +227,12 @@ struct ace_params { | Name | Type | Description | | --- | --- | --- | -| `npartitions` | `size_t` | Number of partitions for ACE (Augmented Core Extraction) partitioned build.

When set to 0 (default), the number of partitions is automatically derived based on available host and GPU memory to maximize partition size while ensuring the build fits in memory.

Small values might improve recall but potentially degrade performance and increase memory usage. Partitions should not be too small to prevent issues in KNN graph construction. The partition size is on average 2 * (n_rows / npartitions) * dim * sizeof(T). 2 is because of the core and augmented vectors. Please account for imbalance in the partition sizes (up to 3x in our tests).

If the specified number of partitions results in partitions that exceed available memory, the value will be automatically increased to fit memory constraints and a warning will be issued. | +| `npartitions` | `size_t` | Number of partitions for ACE (Augmented Core Extraction) partitioned build.

When set to 0 (default), the number of partitions is automatically derived based on available host and GPU memory to maximize partition size while ensuring the build fits in memory.

Small values might improve recall but potentially degrade performance and increase memory usage. Partitions should not be too small to prevent issues in KNN graph construction. The partition size is on average `(1 + nsubpartitions) * (n_rows / npartitions) * dim * sizeof(T)`, accounting for the core and augmented vectors. Please account for imbalance in the partition sizes (up to 3x in our tests).

If the specified number of partitions results in partitions that exceed available memory, the value will be automatically increased to fit memory constraints and a warning will be issued. | +| `nsubpartitions` | `size_t` | Number of additional nearest partitions assigned to each vector. Must be positive and smaller than npartitions (unless npartitions is auto-selected). Increasing this improves overlap at the cost of build time, memory, and temporary disk space. Default: 1. | | `ef_construction` | `size_t` | The index quality for the ACE build.

Bigger values increase the index quality. At some point, increasing this will no longer improve the quality. | -| `build_dir` | `std::string` | Directory to store ACE build artifacts (e.g., KNN graph, optimized graph).

Used when `use_disk` is true or when the graph does not fit in host and GPU memory. This should be the fastest disk in the system and hold enough space for twice the dataset, final graph, and label mapping. The directory may already exist, but ACE's named artifacts must not already exist. Simultaneous builds must use different directories. On failure, ACE removes only artifacts it created and never deletes unrelated directory contents. | +| `build_dir` | `std::string` | Directory to store ACE build artifacts (e.g., KNN graph, optimized graph).

Used when `use_disk` is true or when the graph does not fit in host and GPU memory. This should be the fastest disk in the system and hold enough space for twice the dataset, final graph, and label mapping. The directory may already exist, but ACE's named artifacts must not already exist. Simultaneous builds must use different directories. On failure, ACE removes only artifacts it created and never deletes unrelated directory contents. With `add_global_reverse_edges`, temporary incoming-edge records also require up to `dataset_size * graph_degree * (2 * sizeof(IdxT) + sizeof(uint32_t))` bytes, plus any record alignment padding. These records are removed after the merge. | | `use_disk` | `bool` | Whether to use disk-based storage for ACE build.

When true, enables disk-based operations for memory-efficient graph construction. | +| `add_global_reverse_edges` | `bool` | Build pruned core-plus-argument partial graphs, then add reverse edges globally.

Each partition first builds an intermediate kNN graph and prunes it without a local reverse-edge merge. Then build the reverse-edge graph globally and merge it into the pruned graph. In disk mode (explicit or automatic), incoming edges are buffered to temporary files while building each pruned partition, then merged on the CPU one destination partition at a time, retaining at most `ceil(graph_degree / 2)` incoming candidates per node in memory. The merge preserves edge-rank priority; ties use reordered source ID in disk mode and are GPU-scheduling-dependent in memory mode. Disabled by default. | | `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. | diff --git a/fern/pages/python_api/python-api-neighbors-cagra.md b/fern/pages/python_api/python-api-neighbors-cagra.md index 9b1a3e3924..d3e38629ae 100644 --- a/fern/pages/python_api/python-api-neighbors-cagra.md +++ b/fern/pages/python_api/python-api-neighbors-cagra.md @@ -22,7 +22,9 @@ for each partition independently. | Name | Type | Description | | --- | --- | --- | -| `npartitions` | `int, default = 0` | Number of partitions for ACE partitioned build. When set to 0 (default), the number of partitions is automatically derived based on available host and GPU memory to maximize partition size while ensuring the build fits in memory.

Small values might improve recall but potentially degrade performance and increase memory usage. Partitions should not be too small to prevent issues in KNN graph construction. The partition size is on average 2 * (n_rows / npartitions) * dim * sizeof(T). 2 is because of the core and augmented vectors. Please account for imbalance in the partition sizes (up to 3x in our tests).

If the specified number of partitions results in partitions that exceed available memory, the value will be automatically increased to fit memory constraints and a warning will be issued. | +| `npartitions` | `int, default = 0` | Number of partitions for ACE partitioned build. When set to 0 (default), the number of partitions is automatically derived based on available host and GPU memory to maximize partition size while ensuring the build fits in memory.

Small values might improve recall but potentially degrade performance and increase memory usage. Partitions should not be too small to prevent issues in KNN graph construction. The partition size is on average ``(1 + nsubpartitions) * (n_rows / npartitions) * dim * sizeof(T)``, accounting for core and augmented vectors. Please account for imbalance in the partition sizes (up to 3x in our tests).

If the specified number of partitions results in partitions that exceed available memory, the value will be automatically increased to fit memory constraints and a warning will be issued. | +| `nsubpartitions` | `int, default = 1` | Number of additional nearest partitions assigned to each vector. Must be positive and smaller than npartitions (unless automatically selected). Increasing this improves overlap at the cost of build time, memory, and temporary disk space. | +| `add_global_reverse_edges` | `bool, default = False` | Build pruned partial graphs, then add reverse edges globally. In disk mode, incoming edges are buffered to temporary files and merged on the CPU. | | `ef_construction` | `int, default = 120` | The index quality for the ACE build. Bigger values increase the index quality. At some point, increasing this will no longer improve the quality. | | `build_dir` | `str, default = "/tmp/ace_build"` | Directory to store ACE build artifacts (e.g., KNN graph, optimized graph). Used when `use_disk` is true or when the graph does not fit in host and GPU memory. This should be the fastest disk in the system and hold enough space for twice the dataset, final graph, and label mapping. The directory may already exist, but ACE's named artifacts must not already exist. Simultaneous builds must use different directories. On failure, ACE removes only artifacts it created and never deletes unrelated directory contents. | | `use_disk` | `bool, default = False` | Whether to use disk-based storage for ACE build. When true, enables disk-based operations for memory-efficient graph construction. | @@ -32,7 +34,7 @@ for each partition independently. **Constructor** ```python -def __init__(self, *, npartitions=0, ef_construction=120, build_dir="/tmp/ace_build", use_disk=False, max_host_memory_gb=0, max_gpu_memory_gb=0) +def __init__(self, *, npartitions=0, ef_construction=120, build_dir="/tmp/ace_build", use_disk=False, max_host_memory_gb=0, max_gpu_memory_gb=0, nsubpartitions=1, add_global_reverse_edges=False) ``` **Members** @@ -40,6 +42,8 @@ def __init__(self, *, npartitions=0, ef_construction=120, build_dir="/tmp/ace_bu | Name | Kind | | --- | --- | | `npartitions` | property | +| `nsubpartitions` | property | +| `add_global_reverse_edges` | property | | `ef_construction` | property | | `build_dir` | property | | `use_disk` | property | @@ -53,6 +57,18 @@ def __init__(self, *, npartitions=0, ef_construction=120, build_dir="/tmp/ace_bu def npartitions(self) ``` +### nsubpartitions + +```python +def nsubpartitions(self) +``` + +### add_global_reverse_edges + +```python +def add_global_reverse_edges(self) +``` + ### ef_construction ```python diff --git a/python/cuvs/cuvs/neighbors/cagra/cagra.pxd b/python/cuvs/cuvs/neighbors/cagra/cagra.pxd index cfde71855e..8e9e01367c 100644 --- a/python/cuvs/cuvs/neighbors/cagra/cagra.pxd +++ b/python/cuvs/cuvs/neighbors/cagra/cagra.pxd @@ -70,6 +70,8 @@ cdef extern from "cuvs/neighbors/cagra.h" nogil: bool use_disk double max_host_memory_gb double max_gpu_memory_gb + size_t nsubpartitions + bool add_global_reverse_edges ctypedef cuvsAceParams* cuvsAceParams_t ctypedef struct cuvsCagraIndexParams: diff --git a/python/cuvs/cuvs/neighbors/cagra/cagra.pyx b/python/cuvs/cuvs/neighbors/cagra/cagra.pyx index 27880c8dd6..97f1807806 100644 --- a/python/cuvs/cuvs/neighbors/cagra/cagra.pyx +++ b/python/cuvs/cuvs/neighbors/cagra/cagra.pyx @@ -74,14 +74,22 @@ cdef class AceParams: Small values might improve recall but potentially degrade performance and increase memory usage. Partitions should not be too small to prevent - issues in KNN graph construction. The partition size is on average 2 * - (n_rows / npartitions) * dim * sizeof(T). 2 is because of the core and - augmented vectors. Please account for imbalance in the partition sizes - (up to 3x in our tests). + issues in KNN graph construction. The partition size is on average + ``(1 + nsubpartitions) * (n_rows / npartitions) * dim * sizeof(T)``, + accounting for core and augmented vectors. Please account for imbalance + in the partition sizes (up to 3x in our tests). If the specified number of partitions results in partitions that exceed available memory, the value will be automatically increased to fit memory constraints and a warning will be issued. + nsubpartitions : int, default = 1 + Number of additional nearest partitions assigned to each vector. Must + be positive and smaller than npartitions (unless automatically selected). + Increasing this improves overlap at the cost of build time, memory, + and temporary disk space. + add_global_reverse_edges : bool, default = False + Build pruned partial graphs, then add reverse edges globally. In disk + mode, incoming edges are buffered to temporary files and merged on the CPU. ef_construction : int, default = 120 The index quality for the ACE build. Bigger values increase the index quality. At some point, increasing this will no longer improve the @@ -124,8 +132,12 @@ cdef class AceParams: build_dir="/tmp/ace_build", use_disk=False, max_host_memory_gb=0, - max_gpu_memory_gb=0): + max_gpu_memory_gb=0, + nsubpartitions=1, + add_global_reverse_edges=False): self.params.npartitions = npartitions + self.params.nsubpartitions = nsubpartitions + self.params.add_global_reverse_edges = add_global_reverse_edges self.params.ef_construction = ef_construction self.params.use_disk = use_disk self.params.max_host_memory_gb = max_host_memory_gb @@ -146,6 +158,14 @@ cdef class AceParams: def npartitions(self): return self.params.npartitions + @property + def nsubpartitions(self): + return self.params.nsubpartitions + + @property + def add_global_reverse_edges(self): + return self.params.add_global_reverse_edges + @property def ef_construction(self): return self.params.ef_construction @@ -273,6 +293,10 @@ cdef class IndexParams: # Copy values from Python object to new C struct new_ace_params.npartitions = ace_params.params.npartitions + new_ace_params.nsubpartitions = ace_params.params.nsubpartitions + new_ace_params.add_global_reverse_edges = ( + ace_params.params.add_global_reverse_edges + ) new_ace_params.ef_construction = ace_params.params.ef_construction new_ace_params.use_disk = ace_params.params.use_disk new_ace_params.max_host_memory_gb = ace_params.params.max_host_memory_gb diff --git a/python/cuvs/cuvs/tests/test_cagra_ace.py b/python/cuvs/cuvs/tests/test_cagra_ace.py index dd5f323974..b5efb021bc 100644 --- a/python/cuvs/cuvs/tests/test_cagra_ace.py +++ b/python/cuvs/cuvs/tests/test_cagra_ace.py @@ -22,6 +22,28 @@ ) +@pytest.mark.parametrize( + "nsubpartitions,add_global_reverse_edges", [(1, False), (2, True)] +) +def test_ace_overlap_params(nsubpartitions, add_global_reverse_edges): + defaults = cagra.AceParams() + assert defaults.nsubpartitions == 1 + assert not defaults.add_global_reverse_edges + params = cagra.AceParams( + npartitions=4, + nsubpartitions=nsubpartitions, + add_global_reverse_edges=add_global_reverse_edges, + ) + assert params.nsubpartitions == nsubpartitions + assert params.add_global_reverse_edges == add_global_reverse_edges + index_params = cagra.IndexParams(build_algo="ace", ace_params=params) + assert index_params.ace_params.nsubpartitions == nsubpartitions + assert ( + index_params.ace_params.add_global_reverse_edges + == add_global_reverse_edges + ) + + def run_cagra_ace_build_search_test( n_rows=10000, n_cols=10,