From 5a9c71cc43914c94a706321885d4d4cedf0521a1 Mon Sep 17 00:00:00 2001 From: enp1s0 Date: Fri, 11 Sep 2026 13:40:27 +0900 Subject: [PATCH 01/18] cagra-improve-ace-revedge --- cpp/include/cuvs/neighbors/cagra.hpp | 105 ++- cpp/src/neighbors/cagra.cuh | 14 + cpp/src/neighbors/cagra_build_inst.cu.in | 16 + .../neighbors/detail/cagra/cagra_build.cuh | 599 +++++++++++++++--- cpp/src/neighbors/detail/cagra/graph_core.cuh | 45 +- fern/pages/cpp_api/cpp-api-neighbors-cagra.md | 27 +- 6 files changed, 693 insertions(+), 113 deletions(-) diff --git a/cpp/include/cuvs/neighbors/cagra.hpp b/cpp/include/cuvs/neighbors/cagra.hpp index 84eaa7ac69..d25e40cca1 100644 --- a/cpp/include/cuvs/neighbors/cagra.hpp +++ b/cpp/include/cuvs/neighbors/cagra.hpp @@ -8,6 +8,7 @@ #include #include #include +#include #include #include #include @@ -37,6 +38,12 @@ using iterative_search_params = cuvs::neighbors::search_params; /** Specialized parameters for ACE (Augmented Core Extraction) graph build */ struct ace_params { + /** Algorithm used to build each ACE partial kNN graph. */ + enum class knn_graph_build_algo : uint32_t { + IVF_PQ, + BRUTE_FORCE, + }; + /** * Number of partitions for ACE (Augmented Core Extraction) partitioned build. * @@ -58,8 +65,8 @@ struct ace_params { /** * 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; /** @@ -80,6 +87,28 @@ 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. The final pass preserves the first half of every adjacency row and + * inserts incoming edges originating in other core partitions into the remaining entries. + * Reverse edges are staged and merged one target partition at a time, so GPU workspace scales + * with a partition and its incoming-edge candidates rather than the full graph. This mode does + * not support `use_disk`. + */ + bool add_global_reverse_edges = false; + + /** + * Algorithm used to construct each intermediate partition kNN graph. + * + * Applies to both the original ACE partial CAGRA build with local reverse edges and the global + * reverse-edge mode. IVF_PQ is the default approximate builder. BRUTE_FORCE computes exact + * neighbors for float and half datasets and is intended for smaller partitions or quality + * evaluation. + */ + knn_graph_build_algo partial_graph_knn_build_algo = knn_graph_build_algo::IVF_PQ; + /** * Maximum host memory to use for ACE build in GiB. * @@ -120,7 +149,8 @@ using graph_build_params_t = std::variant; + graph_build_params::iterative_search_params, + graph_build_params::brute_force_params>; /** * @brief A strategy for selecting the graph build parameters based on similar HNSW index @@ -183,6 +213,9 @@ struct index_params : cuvs::neighbors::index_params { * // 4. Choose iterative graph building using CAGRA's search() and optimize() [Experimental] * params.graph_build_params = * cagra::graph_build_params::iterative_search_params(); + * + * // 5. Choose exact brute-force kNN construction (float and half datasets) + * params.graph_build_params = cagra::graph_build_params::brute_force_params{}; * @endcode */ graph_build_params_t graph_build_params; @@ -196,8 +229,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. * @@ -283,10 +316,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} @@ -423,8 +457,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; }; /** @@ -694,7 +728,8 @@ struct CUVS_EXPORT index : cuvs::neighbors::index { } /** - * Replace the source indices with a new source indices taking the ownership of the passed vector. + * 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) { @@ -4360,6 +4395,54 @@ void distribute(const raft::resources& clique, const std::string& filename, cuvs::neighbors::mg_index, T, IdxT>* index); +/** + * @brief Build an exact, distance-ordered kNN graph, excluding self neighbors. + * @param res RAFT resources. + * @param dataset Row-major host dataset of float or half values. + * @param knn_graph Host output with matching rows and degree in [1, n_rows - 1]. + * @param build_params Brute-force index and search parameters. + */ +void build_knn_graph(raft::resources const& res, + raft::host_matrix_view dataset, + raft::host_matrix_view knn_graph, + graph_build_params::brute_force_params build_params); + +/** + * @brief Build an exact, distance-ordered kNN graph, excluding self neighbors. + * @param res RAFT resources. + * @param dataset Row-major device dataset of float or half values. + * @param knn_graph Host output with matching rows and degree in [1, n_rows - 1]. + * @param build_params Brute-force index and search parameters. + */ +void build_knn_graph(raft::resources const& res, + raft::device_matrix_view dataset, + raft::host_matrix_view knn_graph, + graph_build_params::brute_force_params build_params); + +/** + * @brief Build an exact, distance-ordered kNN graph, excluding self neighbors. + * @param res RAFT resources. + * @param dataset Row-major host dataset of float or half values. + * @param knn_graph Host output with matching rows and degree in [1, n_rows - 1]. + * @param build_params Brute-force index and search parameters. + */ +void build_knn_graph(raft::resources const& res, + raft::host_matrix_view dataset, + raft::host_matrix_view knn_graph, + graph_build_params::brute_force_params build_params); + +/** + * @brief Build an exact, distance-ordered kNN graph, excluding self neighbors. + * @param res RAFT resources. + * @param dataset Row-major device dataset of float or half values. + * @param knn_graph Host output with matching rows and degree in [1, n_rows - 1]. + * @param build_params Brute-force index and search parameters. + */ +void build_knn_graph(raft::resources const& res, + raft::device_matrix_view dataset, + raft::host_matrix_view knn_graph, + graph_build_params::brute_force_params build_params); + /** * @brief Build a kNN graph using IVF-PQ. * diff --git a/cpp/src/neighbors/cagra.cuh b/cpp/src/neighbors/cagra.cuh index b23b264dcb..f13ca5d3a0 100644 --- a/cpp/src/neighbors/cagra.cuh +++ b/cpp/src/neighbors/cagra.cuh @@ -189,6 +189,16 @@ void build_knn_graph( detail::build_knn_graph(res, dataset, knn_graph, build_params); } +template +void build_knn_graph( + raft::resources const& res, + raft::mdspan, raft::row_major, Accessor> dataset, + raft::host_matrix_view knn_graph, + graph_build_params::brute_force_params params) +{ + detail::build_knn_graph(res, dataset, knn_graph, params); +} + /** * @brief Sort a KNN graph index. * Preprocessing step for `cagra::optimize`: If a KNN graph is not built using @@ -302,6 +312,10 @@ auto build(raft::resources const& res, const index_params& params, DatasetViewT if constexpr (cuvs::neighbors::is_device_vpq_dataset_view_v) { RAFT_FAIL("cagra::build: VPQ-compressed dataset cannot be used for dense graph construction."); } 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/cagra_build_inst.cu.in b/cpp/src/neighbors/cagra_build_inst.cu.in index e81f8cae95..5406316748 100644 --- a/cpp/src/neighbors/cagra_build_inst.cu.in +++ b/cpp/src/neighbors/cagra_build_inst.cu.in @@ -40,6 +40,22 @@ void build_knn_graph(raft::resources const& handle, cuvs::neighbors::cagra::build_knn_graph(handle, dataset, knn_graph, params); } +void build_knn_graph(raft::resources const& res, + raft::host_matrix_view dataset, + raft::host_matrix_view knn_graph, + graph_build_params::brute_force_params params) +{ + detail::build_knn_graph(res, dataset, knn_graph, params); +} + +void build_knn_graph(raft::resources const& res, + raft::device_matrix_view dataset, + raft::host_matrix_view knn_graph, + graph_build_params::brute_force_params params) +{ + detail::build_knn_graph(res, dataset, knn_graph, params); +} + #define CUVS_DEFINE_CAGRA_BUILD_OVERLOAD(DatasetViewT, ...) \ auto build(raft::resources const& res, \ const cuvs::neighbors::cagra::index_params& params, \ diff --git a/cpp/src/neighbors/detail/cagra/cagra_build.cuh b/cpp/src/neighbors/detail/cagra/cagra_build.cuh index 24fb9776e4..0b7465c9c8 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_build.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_build.cuh @@ -10,6 +10,7 @@ #include "graph_core.cuh" #include +#include #include #include #include @@ -47,6 +48,7 @@ #include #include #include +#include #include #include #include @@ -461,9 +463,9 @@ void ace_create_forward_and_backward_lists( augmented_partition_offsets(0) = 0; } -// ACE: 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, @@ -492,7 +494,137 @@ void ace_gather_partition_dataset( } } +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); + +// Exact kNN construction shared by ordinary CAGRA and ACE. Request one extra +// neighbor and remove self by ID, including when duplicate vectors tie with self. +template +void build_knn_graph(raft::resources const& res, + raft::mdspan, Layout, Accessor> dataset, + raft::host_matrix_view knn_graph, + graph_build_params::brute_force_params params) +{ + auto const partition_size = static_cast(dataset.extent(0)); + auto const intermediate_graph_degree = static_cast(knn_graph.extent(1)); + RAFT_EXPECTS(knn_graph.extent(0) == dataset.extent(0), + "CAGRA: kNN graph and dataset must have the same number of rows"); + RAFT_EXPECTS(intermediate_graph_degree > 0 && partition_size > intermediate_graph_degree, + "CAGRA: kNN graph degree must be positive and smaller than the dataset size"); + if constexpr (std::is_same_v || std::is_same_v) { + auto device_queries = + raft::make_device_matrix(res, partition_size, dataset.extent(1)); + raft::copy(res, device_queries.view(), dataset); + auto brute_force_index = cuvs::neighbors::brute_force::build( + res, params.build_params, raft::make_const_mdspan(device_queries.view())); + auto device_neighbors = raft::make_device_matrix( + res, partition_size, intermediate_graph_degree + 1); + auto device_distances = + raft::make_device_matrix(res, partition_size, intermediate_graph_degree + 1); + cuvs::neighbors::brute_force::search(res, + params.search_params, + brute_force_index, + raft::make_const_mdspan(device_queries.view()), + device_neighbors.view(), + device_distances.view()); + + auto host_neighbors = + raft::make_host_matrix(partition_size, intermediate_graph_degree + 1); + raft::copy(res, host_neighbors.view(), device_neighbors.view()); + raft::resource::sync_stream(res); + for (size_t row = 0; row < partition_size; ++row) { + size_t output_rank = 0; + for (size_t rank = 0; rank < intermediate_graph_degree + 1; ++rank) { + auto const neighbor = host_neighbors(row, rank); + if (neighbor == static_cast(row)) { continue; } + RAFT_EXPECTS(neighbor >= 0 && static_cast(neighbor) < partition_size, + "CAGRA: brute-force kNN returned an invalid neighbor"); + knn_graph(row, output_rank++) = static_cast(neighbor); + if (output_rank == intermediate_graph_degree) { break; } + } + RAFT_EXPECTS(output_rank == intermediate_graph_degree, + "CAGRA: brute-force kNN did not return enough non-self neighbors"); + } + } else { + RAFT_FAIL("CAGRA: brute-force kNN build supports only float and half datasets"); + } +} + +// 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, + graph_build_params::ace_params::knn_graph_build_algo knn_build_algo) + -> 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); + if (knn_build_algo == graph_build_params::ace_params::knn_graph_build_algo::IVF_PQ) { + 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); + } else { + auto knn_params = graph_build_params::brute_force_params{}; + knn_params.build_params.metric = metric; + 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, + graph_build_params::ace_params::knn_graph_build_algo knn_build_algo) + -> raft::host_matrix +{ + auto knn_graph = ace_build_partial_knn_graph( + res, core_argument_partition, intermediate_graph_degree, metric, knn_build_algo); + 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; +} + +// Original ACE mode: optimize each partial graph, including its local reverse-edge merge. +template +auto ace_build_optimized_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, + graph_build_params::ace_params::knn_graph_build_algo knn_build_algo, + bool guarantee_connectivity) -> raft::host_matrix +{ + auto knn_graph = ace_build_partial_knn_graph( + res, core_argument_partition, intermediate_graph_degree, metric, knn_build_algo); + auto partial_graph = + raft::make_host_matrix(core_argument_partition.extent(0), graph_degree); + graph::optimize(res, knn_graph.view(), partial_graph.view(), guarantee_connectivity); + 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, @@ -1310,6 +1442,199 @@ auto build_from_host_matrix(raft::resources const& res, DatasetViewT const& 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; } + + 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::sort(candidate_sources.begin(), candidate_sources.end()); + + 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)); + + 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_sets(n_partitions); + for (size_t source = 0; source < dataset_size; ++source) { + auto const source_partition = static_cast(partition_labels(source, 0)); + 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)); + if (source_partition != destination_partition) { + candidate_sets[destination_partition].insert(static_cast(source)); + } + } + } + + for (size_t partition_id = 0; partition_id < n_partitions; ++partition_id) { + if (candidate_sets[partition_id].empty()) { continue; } + 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); + std::vector candidate_sources(candidate_sets[partition_id].begin(), + candidate_sets[partition_id].end()); + 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 @@ -1398,6 +1723,8 @@ auto build_ace(raft::resources const& res, const index_params& params, DatasetVi ace_params.max_gpu_memory_gb, params.guarantee_connectivity, mem); + RAFT_EXPECTS(!ace_params.add_global_reverse_edges || !use_disk_mode, + "ACE: partitioned global reverse-edge addition does not support disk mode"); // Validate and adjust partitions if disk mode is enabled if (use_disk_mode) { @@ -1579,56 +1906,16 @@ 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 sub_index = [&]() -> sub_index_t { + std::optional> partial_sub_graph; + std::optional sub_index; + const bool build_partial_graph_on_host = + ace_params.add_global_reverse_edges || + ace_params.partial_graph_knn_build_algo == + graph_build_params::ace_params::knn_graph_build_algo::BRUTE_FORCE; + if (build_partial_graph_on_host) { + auto core_argument_partition = + raft::make_host_matrix(sub_dataset_size, dataset_dim); if (use_disk_mode) { - const size_t current_free_gpu = rmm::available_device_memory().first; - const size_t configured_gpu = - ace_params.max_gpu_memory_gb > 0 - ? 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)) { - try { - auto sub_dataset_tight = - raft::make_device_matrix(res, sub_dataset_size, dataset_dim); - raft::resource::sync_stream(res); - ace_load_partition_dataset_from_disk(partition_id, - dataset_dim, - partition_histogram.view(), - core_partition_offsets.view(), - augmented_partition_offsets.view(), - reordered_fd, - augmented_fd, - reordered_header_size, - augmented_header_size, - sub_dataset_tight.view()); - read_end = std::chrono::high_resolution_clock::now(); - auto sub_dataset_view = raft::make_const_mdspan(sub_dataset_tight.view()); - std::unique_ptr> - sub_dataset_padded; - auto sub_dataset_dev = [&]() { - if (cuvs::neighbors::matrix_row_width_matches_cagra_required(sub_dataset_view)) { - return cuvs::neighbors::make_device_padded_dataset_view(res, sub_dataset_view); - } - sub_dataset_padded = - cuvs::neighbors::make_device_padded_dataset(res, sub_dataset_view); - return sub_dataset_padded->as_dataset_view(); - }(); - auto direct_index = - ::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)}; - } catch (const std::bad_alloc& e) { - RAFT_LOG_WARN( - "ACE: partition %lu did not fit in device memory for a direct (GDS) read: %s; " - "falling back to a host read", - partition_id, - e.what()); - } - } - - auto sub_dataset = raft::make_host_matrix(sub_dataset_size, dataset_dim); ace_load_partition_dataset_from_disk(partition_id, dataset_dim, partition_histogram.view(), @@ -1638,7 +1925,121 @@ auto build_ace(raft::resources const& res, const index_params& params, DatasetVi augmented_fd, reordered_header_size, augmented_header_size, - sub_dataset.view()); + core_argument_partition.view()); + } else { + 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(), + core_argument_partition.view()); + } + read_end = std::chrono::high_resolution_clock::now(); + auto core_argument_partition_view = raft::make_host_matrix_view( + core_argument_partition.data_handle(), sub_dataset_size, dataset_dim); + if (ace_params.add_global_reverse_edges) { + partial_sub_graph.emplace( + ace_build_pruned_partial_graph(res, + core_argument_partition_view, + intermediate_degree, + graph_degree, + params.metric, + ace_params.partial_graph_knn_build_algo)); + } else { + partial_sub_graph.emplace( + ace_build_optimized_partial_graph(res, + core_argument_partition_view, + intermediate_degree, + graph_degree, + params.metric, + ace_params.partial_graph_knn_build_algo, + params.guarantee_connectivity)); + } + } else { + sub_index.emplace([&]() -> sub_index_t { + if (use_disk_mode) { + const size_t current_free_gpu = rmm::available_device_memory().first; + const size_t configured_gpu = + ace_params.max_gpu_memory_gb > 0 + ? 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)) { + try { + auto sub_dataset_tight = + raft::make_device_matrix(res, sub_dataset_size, dataset_dim); + raft::resource::sync_stream(res); + ace_load_partition_dataset_from_disk(partition_id, + dataset_dim, + partition_histogram.view(), + core_partition_offsets.view(), + augmented_partition_offsets.view(), + reordered_fd, + augmented_fd, + reordered_header_size, + augmented_header_size, + sub_dataset_tight.view()); + read_end = std::chrono::high_resolution_clock::now(); + auto sub_dataset_view = raft::make_const_mdspan(sub_dataset_tight.view()); + std::unique_ptr> + sub_dataset_padded; + auto sub_dataset_dev = [&]() { + if (cuvs::neighbors::matrix_row_width_matches_cagra_required(sub_dataset_view)) { + return cuvs::neighbors::make_device_padded_dataset_view(res, sub_dataset_view); + } + sub_dataset_padded = + cuvs::neighbors::make_device_padded_dataset(res, sub_dataset_view); + return sub_dataset_padded->as_dataset_view(); + }(); + auto direct_index = + ::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)}; + } catch (const std::bad_alloc& e) { + RAFT_LOG_WARN( + "ACE: partition %lu did not fit in device memory for a direct (GDS) read: %s; " + "falling back to a host read", + partition_id, + e.what()); + } + } + + auto sub_dataset = raft::make_host_matrix(sub_dataset_size, dataset_dim); + ace_load_partition_dataset_from_disk(partition_id, + dataset_dim, + partition_histogram.view(), + core_partition_offsets.view(), + augmented_partition_offsets.view(), + reordered_fd, + augmented_fd, + reordered_header_size, + augmented_header_size, + sub_dataset.view()); + 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 = raft::make_host_matrix(sub_dataset_size, dataset_dim); + 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()); 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())); @@ -1646,27 +2047,8 @@ auto build_ace(raft::resources const& res, const index_params& params, DatasetVi ::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 = 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()); - 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_graph = std::visit([](auto const& index) { return index.graph(); }, sub_index); + }()); + } if (used_device_read) { RAFT_LOG_DEBUG("ACE: partition %lu read directly into device memory (GDS path)", partition_id); @@ -1683,17 +2065,30 @@ auto build_ace(raft::resources const& res, const index_params& params, DatasetVi 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 (build_partial_graph_on_host) { + 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 = @@ -1707,11 +2102,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 (build_partial_graph_on_host) { + 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, @@ -1761,6 +2163,15 @@ 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 cross-partition reverse edges partition by partition"); + 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) { @@ -2516,11 +2927,12 @@ inline void validate_cagra_knn_graph_build_constraints(index_params const& param std::holds_alternative( knn_build_params) || std::holds_alternative(knn_build_params), - "IVF_PQ for CAGRA graph build does not support BitwiseHamming as a metric. Please " + "The selected CAGRA graph builder does not support BitwiseHamming as a metric. Please " "use nn-descent or the iterative CAGRA search build."); RAFT_EXPECTS( params.metric != cuvs::distance::DistanceType::CosineExpanded || std::holds_alternative(knn_build_params) || + std::holds_alternative(knn_build_params) || std::holds_alternative(knn_build_params), "CosineExpanded distance is not supported for iterative CAGRA graph build."); @@ -2531,7 +2943,8 @@ inline void validate_cagra_knn_graph_build_constraints(index_params const& param } /** - * Iterative / IVF-PQ / NN-descent KNN graph construction and `optimize` → final host CAGRA graph. + * Iterative / IVF-PQ / NN-descent / brute-force KNN graph construction and `optimize` → final host + * CAGRA graph. * * @param knn_graph_dataset mdspan passed to IVF-PQ / NN-descent `build_knn_graph` (any stride). */ @@ -2560,6 +2973,10 @@ auto build_cagra_host_graph_from_knn_params(raft::resources const& res, ivf_pq_params.build_params.metric = params.metric; } build_knn_graph(res, knn_graph_dataset, knn_graph->view(), ivf_pq_params); + } else if (std::holds_alternative(knn_build_params)) { + auto brute_force_params = std::get(knn_build_params); + brute_force_params.build_params.metric = params.metric; + build_knn_graph(res, knn_graph_dataset, knn_graph->view(), brute_force_params); } else { auto nn_descent_params = std::get(knn_build_params); 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/fern/pages/cpp_api/cpp-api-neighbors-cagra.md b/fern/pages/cpp_api/cpp-api-neighbors-cagra.md index 5f1865967d..cf1d492b25 100644 --- a/fern/pages/cpp_api/cpp-api-neighbors-cagra.md +++ b/fern/pages/cpp_api/cpp-api-neighbors-cagra.md @@ -19,6 +19,8 @@ struct ace_params { size_t ef_construction; std::string build_dir; bool use_disk; + bool add_global_reverse_edges; + knn_graph_build_algo partial_graph_knn_build_algo; double max_host_memory_gb; double max_gpu_memory_gb; }; @@ -32,9 +34,30 @@ struct ace_params { | `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. | | `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. The final pass preserves the first half of every adjacency row and inserts incoming edges originating in other core partitions into the remaining entries. It requires two graph-sized device workspace allocations. With `use_disk`, partition data is loaded from the reordered dataset artifacts and the final graph is updated in place on disk. | +| `partial_graph_knn_build_algo` | [`knn_graph_build_algo`](/api-reference/cpp-api-neighbors-cagra#neighbors-graph-build-params-ace-params-knn-graph-build-algo) | Algorithm used to construct each intermediate partition kNN graph.

Applies to both the original ACE partial CAGRA build with local reverse edges and the global reverse-edge mode. IVF_PQ is the default approximate builder. BRUTE_FORCE computes exact neighbors for float and half datasets and is intended for smaller partitions or quality evaluation. | | `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. | + +### neighbors::graph_build_params::ace_params::knn_graph_build_algo + +Algorithm used to build each ACE partial kNN graph. + +```cpp +enum class knn_graph_build_algo : uint32_t { + IVF_PQ, + BRUTE_FORCE +}; +``` + +**Values** + +| Name | Value | +| --- | --- | +| `IVF_PQ` | `` | +| `BRUTE_FORCE` | `` | + ## CAGRA index build parameters @@ -551,12 +574,14 @@ We create a copy of the graph on the device. The index manages the lifetime of t ### neighbors::cagra::index::update_source_indices -Replace the source indices with a new source indices taking the ownership of the passed vector. +Replace the source indices with a new source indices taking the ownership of the passed ```cpp void update_source_indices(raft::device_vector&& source_indices); ``` +vector. + **Parameters** | Name | Direction | Type | Description | From a9a3c5b3d3f70a965a1c9a9404fa7370f6dcdeb0 Mon Sep 17 00:00:00 2001 From: enp1s0 Date: Fri, 11 Sep 2026 14:39:22 +0900 Subject: [PATCH 02/18] Separate ACE code --- cpp/src/neighbors/cagra.cuh | 1 + cpp/src/neighbors/detail/cagra/cagra_ace.cuh | 2097 ++++++++++++++++ .../neighbors/detail/cagra/cagra_build.cuh | 2098 ----------------- 3 files changed, 2098 insertions(+), 2098 deletions(-) create mode 100644 cpp/src/neighbors/detail/cagra/cagra_ace.cuh diff --git a/cpp/src/neighbors/cagra.cuh b/cpp/src/neighbors/cagra.cuh index f13ca5d3a0..7f347bd02f 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" diff --git a/cpp/src/neighbors/detail/cagra/cagra_ace.cuh b/cpp/src/neighbors/detail/cagra/cagra_ace.cuh new file mode 100644 index 0000000000..4b1d270ff5 --- /dev/null +++ b/cpp/src/neighbors/detail/cagra/cagra_ace.cuh @@ -0,0 +1,2097 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include "../../../util/kvikio_io.hpp" +#include "cagra_build.cuh" + +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +namespace cuvs::neighbors::cagra::detail { + +// Helpers to convert bytes to MiB and GiB +constexpr double to_mib(size_t bytes) { return static_cast(bytes) / (1 << 20); } +constexpr double to_gib(size_t bytes) { return static_cast(bytes) / (1 << 30); } + +class ace_disk_workspace { + public: + enum class artifact : size_t { + reordered_dataset, + augmented_dataset, + dataset_mapping, + cagra_graph, + }; + + explicit ace_disk_workspace(std::string build_dir) + : build_dir_(std::move(build_dir)), + artifacts_{build_dir_ / "reordered_dataset.npy", + build_dir_ / "augmented_dataset.npy", + build_dir_ / "dataset_mapping.npy", + build_dir_ / "cagra_graph.npy"} + { + } + + void initialize() + { + if (mkdir(build_dir_.c_str(), 0755) == 0) { + directory_created_by_this_build_ = true; + return; + } + + if (errno != EEXIST) { + RAFT_FAIL("Failed to create ACE build directory: %s (errno: %d, %s)", + build_dir_.c_str(), + errno, + strerror(errno)); + } + + std::error_code error; + const bool is_directory = std::filesystem::is_directory(build_dir_, error); + RAFT_EXPECTS(!error, + "Failed to inspect ACE build directory: %s (%s)", + build_dir_.c_str(), + error.message().c_str()); + RAFT_EXPECTS(is_directory, "ACE build path is not a directory: %s", build_dir_.c_str()); + } + + [[nodiscard]] std::string artifact_path(artifact which) const + { + return artifacts_[static_cast(which)].string(); + } + + void mark_artifact_created(artifact which) noexcept + { + artifacts_created_[static_cast(which)] = true; + } + + void commit() noexcept { committed_ = true; } + + void cleanup() noexcept + { + if (committed_) { return; } + + for (size_t i = artifacts_.size(); i > 0; --i) { + if (!artifacts_created_[i - 1]) { continue; } + + std::error_code error; + std::filesystem::remove(artifacts_[i - 1], error); + if (error) { + RAFT_LOG_WARN("ACE: Failed to remove build artifact %s: %s", + artifacts_[i - 1].c_str(), + error.message().c_str()); + } + } + + if (directory_created_by_this_build_) { + std::error_code error; + std::filesystem::remove(build_dir_, error); + if (error) { + RAFT_LOG_WARN("ACE: Failed to remove empty build directory %s: %s", + build_dir_.c_str(), + error.message().c_str()); + } + } + } + + private: + std::filesystem::path build_dir_; + std::array artifacts_; + std::array artifacts_created_{}; + bool directory_created_by_this_build_ = false; + bool committed_ = false; +}; + +// ACE: Get partition labels for partitioned approach +// TODO(julianmi): Use all neighbors APIs. +template +void ace_get_partition_labels( + raft::resources const& res, + raft::host_matrix_view dataset, + size_t dataset_dim, + raft::host_matrix_view partition_labels, + raft::host_matrix_view partition_histogram, + size_t min_partition_size, + double sampling_rate = 0.01) +{ + size_t dataset_size = dataset.extent(0); + size_t labels_size = partition_labels.extent(0); + size_t labels_dim = partition_labels.extent(1); + RAFT_EXPECTS(dataset_size == labels_size, "Dataset size must match partition labels extent"); + RAFT_EXPECTS(dataset_dim > 0, "Dataset dimension must be greater than 0"); + RAFT_EXPECTS(static_cast(dataset.extent(1)) >= dataset_dim, + "Dataset row extent (%zu) must be >= logical dimension (%zu)", + 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(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. + 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: n_samples: %lu", 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 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); + size_t report_interval = dataset_size / 10; + report_interval = (report_interval / chunk_size) * chunk_size; + report_interval = std::max(report_interval, chunk_size); + + 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); + if (i_base % report_interval == 0) { + RAFT_LOG_INFO("ACE: Processing chunk %lu / %lu (%.1f%%)", + i_base, + dataset_size, + static_cast(100 * i_base) / dataset_size); + } + + 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)); + } + } + 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); + + raft::copy(res, sub_distances, sub_distances_dev); + 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; + } + } +} + +// ACE: Check partition sizes for stable KNN graph construction +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) +{ + // Collect partition histogram statistics + size_t total_core_vectors = 0; + size_t total_augmented_vectors = 0; + size_t min_core_vectors = dataset_size; + size_t max_core_vectors = 0; + size_t min_augmented_vectors = dataset_size; + size_t max_augmented_vectors = 0; + size_t min_total_vectors = dataset_size; + size_t max_total_vectors = 0; + + for (size_t c = 0; c < n_partitions; c++) { + size_t core_count = partition_histogram(c, 0); + size_t augmented_count = partition_histogram(c, 1); + size_t total_count = core_count + augmented_count; + + if (total_count > 0) { + total_core_vectors += core_count; + total_augmented_vectors += augmented_count; + + min_core_vectors = std::min(min_core_vectors, core_count); + max_core_vectors = std::max(max_core_vectors, core_count); + min_augmented_vectors = std::min(min_augmented_vectors, augmented_count); + max_augmented_vectors = std::max(max_augmented_vectors, augmented_count); + min_total_vectors = std::min(min_total_vectors, total_count); + max_total_vectors = std::max(max_total_vectors, total_count); + } + } + + double avg_core_vectors = static_cast(total_core_vectors) / n_partitions; + double avg_augmented_vectors = static_cast(total_augmented_vectors) / n_partitions; + double avg_total_vectors = 2.0 * static_cast(dataset_size) / n_partitions; + double expected_avg_vectors = 2.0 * static_cast(dataset_size) / n_partitions; + + RAFT_LOG_INFO("ACE: Core vectors - Total: %lu, Avg: %.1f, Min: %lu, Max: %lu", + total_core_vectors, + avg_core_vectors, + min_core_vectors, + max_core_vectors); + RAFT_LOG_INFO("ACE: Augmented vectors - Total: %lu, Avg: %.1f, Min: %lu, Max: %lu", + total_augmented_vectors, + avg_augmented_vectors, + min_augmented_vectors, + max_augmented_vectors); + RAFT_LOG_INFO("ACE: Total per partition - Total: %lu, Avg: %.1f, Min: %lu, Max: %lu", + total_core_vectors + total_augmented_vectors, + avg_total_vectors, + min_total_vectors, + max_total_vectors); + + // Check for partition imbalance and issue warnings + size_t very_small_threshold = min_partition_size; + size_t very_large_threshold = static_cast(5.0 * expected_avg_vectors); + + for (size_t c = 0; c < n_partitions; c++) { + size_t total_count = partition_histogram(c, 0) + partition_histogram(c, 1); + + if (total_count > 0 && total_count < very_small_threshold) { + RAFT_LOG_WARN( + "ACE: Partition %lu is very small (%lu vectors, expected ~%.1f). This may affect graph " + "quality.", + c, + total_count, + expected_avg_vectors); + } else if (total_count > very_large_threshold) { + RAFT_LOG_WARN( + "ACE: Partition %lu is very large (%lu vectors, expected ~%.1f, threshold: %lu). This may " + "indicate imbalance and can lead to memory issues in restricted environments.", + c, + total_count, + expected_avg_vectors, + very_large_threshold); + } + } +} + +// ACE: Create forward/backward mappings between original and reordered vector IDs +// The in-memory path can be parallelized but the disk path requires ordering. +template +void ace_create_forward_and_backward_lists( + size_t dataset_size, + size_t n_partitions, + raft::host_matrix_view partition_labels, + raft::host_matrix_view partition_histogram, + raft::host_vector_view core_forward_mapping, + raft::host_vector_view core_backward_mapping, + raft::host_vector_view augmented_backward_mapping, + raft::host_vector_view core_partition_offsets, + raft::host_vector_view augmented_partition_offsets) +{ + core_partition_offsets(0) = 0; + augmented_partition_offsets(0) = 0; + for (size_t c = 1; c < n_partitions; c++) { + core_partition_offsets(c) = core_partition_offsets(c - 1) + partition_histogram(c - 1, 0); + augmented_partition_offsets(c) = + augmented_partition_offsets(c - 1) + partition_histogram(c - 1, 1); + } + + if (static_cast(core_forward_mapping.extent(0)) == 0) { + // 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"); +#pragma omp parallel for + for (size_t i = 0; i < dataset_size; i++) { + size_t core_partition_id = partition_labels(i, 0); + size_t core_id; +#pragma omp atomic capture + core_id = core_partition_offsets(core_partition_id)++; + 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; +#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; + } + } else { + // Disk path: all three mappings + RAFT_EXPECTS(static_cast(core_forward_mapping.extent(0)) == dataset_size, + "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"); + for (size_t i = 0; i < dataset_size; i++) { + size_t core_partition_id = partition_labels(i, 0); + size_t core_id; + core_id = core_partition_offsets(core_partition_id)++; + RAFT_EXPECTS(core_id < dataset_size, "Vector ID must be smaller than dataset_size"); + 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; + } + } + + // Restore idxptr arrays + for (size_t c = n_partitions; c > 0; c--) { + core_partition_offsets(c) = core_partition_offsets(c - 1); + augmented_partition_offsets(c) = augmented_partition_offsets(c - 1); + } + core_partition_offsets(0) = 0; + augmented_partition_offsets(0) = 0; +} + +// ACE: Materialize the core partition followed by its argument partition. +template +void ace_make_core_argument_partition( + size_t core_sub_dataset_size, + size_t augmented_sub_dataset_size, + size_t dataset_dim, + size_t partition_id, + raft::host_matrix_view dataset, + raft::host_vector_view core_backward_mapping, + raft::host_vector_view augmented_backward_mapping, + raft::host_vector_view core_partition_offsets, + raft::host_vector_view augmented_partition_offsets, + raft::host_matrix_view sub_dataset) +{ + const size_t vector_size_bytes = dataset_dim * sizeof(T); + + // Copy core partition vectors +#pragma omp parallel for + for (size_t j = 0; j < core_sub_dataset_size; j++) { + size_t i = core_backward_mapping(j + core_partition_offsets(partition_id)); + memcpy(&sub_dataset(j, 0), &dataset(i, 0), vector_size_bytes); + } + + // Copy augmented partition vectors (2nd closest 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)); + memcpy(&sub_dataset(j + core_sub_dataset_size, 0), &dataset(i, 0), vector_size_bytes); + } +} + +// 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, + graph_build_params::ace_params::knn_graph_build_algo knn_build_algo) + -> 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); + if (knn_build_algo == graph_build_params::ace_params::knn_graph_build_algo::IVF_PQ) { + 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); + } else { + auto knn_params = graph_build_params::brute_force_params{}; + knn_params.build_params.metric = metric; + 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, + graph_build_params::ace_params::knn_graph_build_algo knn_build_algo) + -> raft::host_matrix +{ + auto knn_graph = ace_build_partial_knn_graph( + res, core_argument_partition, intermediate_graph_degree, metric, knn_build_algo); + 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; +} + +// Original ACE mode: optimize each partial graph, including its local reverse-edge merge. +template +auto ace_build_optimized_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, + graph_build_params::ace_params::knn_graph_build_algo knn_build_algo, + bool guarantee_connectivity) -> raft::host_matrix +{ + auto knn_graph = ace_build_partial_knn_graph( + res, core_argument_partition, intermediate_graph_degree, metric, knn_build_algo); + auto partial_graph = + raft::make_host_matrix(core_argument_partition.extent(0), graph_degree); + graph::optimize(res, knn_graph.view(), partial_graph.view(), guarantee_connectivity); + 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, + size_t augmented_sub_dataset_size, + size_t graph_degree, + size_t partition_id, + raft::host_matrix_view sub_search_graph, + raft::host_matrix_view search_graph, + raft::host_vector_view core_partition_offsets, + raft::host_vector_view augmented_partition_offsets, + raft::host_vector_view core_backward_mapping, + raft::host_vector_view augmented_backward_mapping) +{ +#pragma omp parallel for + for (size_t i = 0; i < core_sub_dataset_size; i++) { + // Map row index from local → reordered → original + size_t i_reordered = i + core_partition_offsets(partition_id); + size_t i_original = core_backward_mapping(i_reordered); + + for (size_t k = 0; k < graph_degree; k++) { + size_t j = sub_search_graph(i, k); + size_t j_original; + + if (j < core_sub_dataset_size) { + // core partition neighbor: local → core reordered → original + size_t j_reordered = j + core_partition_offsets(partition_id); + j_original = core_backward_mapping(j_reordered); + } else { + // Augmented partition neighbor: local → augmented reordered → original + size_t j_augmented = j - core_sub_dataset_size; + j_original = + augmented_backward_mapping(j_augmented + augmented_partition_offsets(partition_id)); + } + search_graph(i_original, k) = j_original; + } + } +} + +// ACE: Adjust IDs into global reordered IDs on the device for the disk version. +template +__global__ void ace_adjust_sub_graph_ids_disk_kernel(const GraphIdxT* sub_search_graph, + IdxT* adjusted_search_graph, + size_t graph_edges, + size_t core_sub_dataset_size, + IdxT core_partition_offset, + const IdxT* augmented_reordered_ids) +{ + const size_t edge_id = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (edge_id >= graph_edges) { return; } + + const GraphIdxT neighbor = sub_search_graph[edge_id]; + if (static_cast(neighbor) < core_sub_dataset_size) { + adjusted_search_graph[edge_id] = static_cast(neighbor) + core_partition_offset; + } else { + adjusted_search_graph[edge_id] = + augmented_reordered_ids[static_cast(neighbor) - core_sub_dataset_size]; + } +} + +template +void ace_adjust_sub_graph_ids_disk( + raft::resources const& res, + size_t core_sub_dataset_size, + size_t augmented_sub_dataset_size, + size_t graph_degree, + size_t partition_id, + raft::device_matrix_view sub_search_graph, + raft::device_matrix_view adjusted_search_graph, + raft::host_vector_view core_partition_offsets, + raft::host_vector_view augmented_partition_offsets, + raft::host_vector_view augmented_backward_mapping, + raft::host_vector_view core_forward_mapping) +{ + RAFT_EXPECTS(static_cast(sub_search_graph.extent(0)) >= core_sub_dataset_size, + "ACE: source graph has fewer rows than the core partition"); + RAFT_EXPECTS(static_cast(sub_search_graph.extent(1)) == graph_degree, + "ACE: source graph degree does not match the requested graph degree"); + RAFT_EXPECTS(static_cast(adjusted_search_graph.extent(0)) == core_sub_dataset_size && + static_cast(adjusted_search_graph.extent(1)) == graph_degree, + "ACE: adjusted graph shape does not match the core partition"); + + const size_t augmented_offset = augmented_partition_offsets(partition_id); + RAFT_EXPECTS(augmented_offset + augmented_sub_dataset_size <= + static_cast(augmented_backward_mapping.extent(0)), + "ACE: augmented partition exceeds the backward mapping"); + + auto augmented_reordered_ids = raft::make_host_vector(augmented_sub_dataset_size); +#pragma omp parallel for + for (size_t i = 0; i < augmented_sub_dataset_size; i++) { + const size_t original_id = augmented_backward_mapping(augmented_offset + i); + augmented_reordered_ids(i) = core_forward_mapping(original_id); + } + + auto augmented_reordered_ids_dev = + raft::make_device_vector(res, augmented_sub_dataset_size); + raft::copy(res, augmented_reordered_ids_dev.view(), augmented_reordered_ids.view()); + + const size_t graph_edges = core_sub_dataset_size * graph_degree; + if (graph_edges == 0) { return; } + + constexpr uint32_t block_size = 256; + const dim3 grid_size(raft::ceildiv(graph_edges, static_cast(block_size))); + ace_adjust_sub_graph_ids_disk_kernel<<>>( + sub_search_graph.data_handle(), + adjusted_search_graph.data_handle(), + graph_edges, + core_sub_dataset_size, + core_partition_offsets(partition_id), + augmented_reordered_ids_dev.data_handle()); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + raft::resource::sync_stream(res); +} + +// ACE: Reorder dataset based on partition assignments and store to disk +// Writes two files: reordered_dataset.npy (core partitions) and augmented_dataset.npy (secondary +// partitions). Uses buffered writes optimized for NVMe storage. +template +void ace_reorder_and_store_dataset( + raft::resources const& res, + const std::string& build_dir, + raft::host_matrix_view dataset, + size_t dataset_dim, + raft::host_matrix_view partition_labels, + raft::host_matrix_view partition_histogram, + raft::host_vector_view core_backward_mapping, + raft::host_vector_view core_partition_offsets, + raft::host_vector_view augmented_partition_offsets, + cuvs::util::file_descriptor& reordered_fd, + cuvs::util::file_descriptor& augmented_fd, + cuvs::util::file_descriptor& mapping_fd, + size_t reordered_header_size, + size_t augmented_header_size, + size_t mapping_header_size) +{ + auto start = std::chrono::high_resolution_clock::now(); + + size_t dataset_size = dataset.extent(0); + size_t n_partitions = partition_histogram.extent(0); + RAFT_EXPECTS(dataset_dim > 0, "Dataset dimension must be greater than 0"); + RAFT_EXPECTS(static_cast(dataset.extent(1)) >= dataset_dim, + "Dataset row extent (%zu) must be >= logical dimension (%zu)", + static_cast(dataset.extent(1)), + dataset_dim); + + RAFT_LOG_DEBUG( + "ACE: Reordering and storing dataset to disk (%lu vectors, %lu dimensions, %lu partitions)", + dataset_size, + dataset_dim, + n_partitions); + + // Calculate total sizes for pre-allocation + size_t total_core_vectors = 0; + size_t total_augmented_vectors = 0; + size_t max_core_vectors = 0; + size_t max_augmented_vectors = 0; + for (size_t p = 0; p < n_partitions; p++) { + total_core_vectors += partition_histogram(p, 0); + total_augmented_vectors += partition_histogram(p, 1); + max_core_vectors = std::max(max_core_vectors, partition_histogram(p, 0)); + max_augmented_vectors = std::max(max_augmented_vectors, partition_histogram(p, 1)); + } + 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"); + + // Pre-allocate file space for better performance + const size_t vector_size = dataset_dim * sizeof(T); + size_t reordered_file_size = total_core_vectors * vector_size; + size_t augmented_file_size = total_augmented_vectors * vector_size; + + RAFT_LOG_DEBUG("ACE: Reordered dataset: %lu core vectors (%.2f GiB)", + total_core_vectors, + reordered_file_size / (1024.0 * 1024.0 * 1024.0)); + RAFT_LOG_DEBUG("ACE: Augmented dataset: %lu secondary vectors (%.2f GiB)", + total_augmented_vectors, + augmented_file_size / (1024.0 * 1024.0 * 1024.0)); + + // Calculate partition start offsets for reordered and augmented datasets + auto core_partition_starts = raft::make_host_vector(n_partitions + 1); + memset(core_partition_starts.data_handle(), 0, (n_partitions + 1) * sizeof(size_t)); + auto augmented_partition_starts = raft::make_host_vector(n_partitions + 1); + memset(augmented_partition_starts.data_handle(), 0, (n_partitions + 1) * sizeof(size_t)); + auto core_partition_current = raft::make_host_vector(n_partitions); + memset(core_partition_current.data_handle(), 0, n_partitions * sizeof(size_t)); + auto augmented_partition_current = raft::make_host_vector(n_partitions); + memset(augmented_partition_current.data_handle(), 0, n_partitions * sizeof(size_t)); + + for (size_t p = 0; p < n_partitions; p++) { + core_partition_starts(p + 1) = core_partition_starts(p) + partition_histogram(p, 0); + augmented_partition_starts(p + 1) = augmented_partition_starts(p) + partition_histogram(p, 1); + } + + const size_t free_memory = cuvs::util::get_free_host_memory(); + // Conservatively allocate 50% of free memory per partition. Accounts for internal buffers and + // overhead. + // TODO: Adjust overhead if needed. + const size_t memory_per_partition = 0.5 * free_memory / (n_partitions * 2); + size_t disk_write_size = raft::bound_by_power_of_two(memory_per_partition); + // 64MB should be enough to saturate typical NVMe SSDs. + disk_write_size = std::min(disk_write_size, 64 * 1024 * 1024); + size_t vectors_per_buffer = std::max(64, disk_write_size / vector_size); + + RAFT_LOG_DEBUG("ACE: Reorder buffers: %lu vectors per buffer (%.2f MiB)", + vectors_per_buffer, + to_mib(vectors_per_buffer * vector_size)); + + std::vector> core_buffers; + std::vector> augmented_buffers; + auto core_buffer_counts = raft::make_host_vector(n_partitions); + auto augmented_buffer_counts = raft::make_host_vector(n_partitions); + + core_buffers.reserve(n_partitions); + augmented_buffers.reserve(n_partitions); + + for (size_t p = 0; p < n_partitions; p++) { + core_buffers.emplace_back(raft::make_host_matrix(vectors_per_buffer, dataset_dim)); + augmented_buffers.emplace_back( + raft::make_host_matrix(vectors_per_buffer, dataset_dim)); + core_buffer_counts(p) = 0; + augmented_buffer_counts(p) = 0; + } + auto flush_core_buffer = [&](size_t partition_id) { + const size_t count = core_buffer_counts(partition_id); + if (count > 0) { + const size_t bytes_to_write = count * vector_size; + const size_t file_offset = + (core_partition_starts(partition_id) + core_partition_current(partition_id)) * vector_size + + reordered_header_size; + + cuvs::util::write_large_file( + reordered_fd, core_buffers[partition_id].data_handle(), bytes_to_write, file_offset); + + core_partition_current(partition_id) += count; + core_buffer_counts(partition_id) = 0; + } + }; + + auto flush_augmented_buffer = [&](size_t partition_id) { + const size_t count = augmented_buffer_counts(partition_id); + if (count > 0) { + const size_t bytes_to_write = count * vector_size; + const size_t file_offset = + (augmented_partition_starts(partition_id) + augmented_partition_current(partition_id)) * + vector_size + + augmented_header_size; + + cuvs::util::write_large_file( + augmented_fd, augmented_buffers[partition_id].data_handle(), bytes_to_write, file_offset); + + augmented_partition_current(partition_id) += count; + augmented_buffer_counts(partition_id) = 0; + } + }; + + size_t vectors_processed = 0; + const size_t log_interval = std::max(dataset_size / 10, size_t(1)); + 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); + + // Add vector to core partition buffer + size_t core_buffer_row = core_buffer_counts(core_partition); + memcpy( + &core_buffers[core_partition](core_buffer_row, 0), &dataset(i, 0), dataset_dim * sizeof(T)); + core_buffer_counts(core_partition)++; + + // Flush core buffer if full + if (core_buffer_counts(core_partition) >= vectors_per_buffer) { + flush_core_buffer(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_augmented_buffer(secondary_partition); + } + + vectors_processed++; + if (vectors_processed % log_interval == 0) { + RAFT_LOG_INFO("ACE: Processed %lu/%lu vectors (%.1f%%)", + vectors_processed, + dataset_size, + 100.0 * vectors_processed / dataset_size); + } + } + + // Flush all remaining buffers + RAFT_LOG_DEBUG("ACE: Flushing remaining buffers..."); +#pragma omp parallel sections + { +#pragma omp section + { + for (size_t p = 0; p < n_partitions; p++) { + flush_core_buffer(p); + } + } +#pragma omp section + { + for (size_t p = 0; p < n_partitions; p++) { + flush_augmented_buffer(p); + } + } + } + + const size_t mapping_file_size = dataset_size * sizeof(IdxT); + cuvs::util::write_large_file( + mapping_fd, core_backward_mapping.data_handle(), mapping_file_size, mapping_header_size); + + auto end = std::chrono::high_resolution_clock::now(); + auto elapsed_ms = std::chrono::duration_cast(end - start).count(); + + // Calculate total bytes written + size_t total_bytes_written = reordered_file_size + augmented_file_size + mapping_file_size; + double throughput_mb_s = + elapsed_ms > 0 ? to_mib(total_bytes_written) / (elapsed_ms / 1000.0) : 0.0; + + RAFT_LOG_INFO( + "ACE: Dataset (%.2f GiB reordered, %.2f GiB augmented, %.2f GiB mapping) reordering completed " + "in %ld ms (%.1f MiB/s)", + reordered_file_size / (1024.0 * 1024.0 * 1024.0), + augmented_file_size / (1024.0 * 1024.0 * 1024.0), + mapping_file_size / (1024.0 * 1024.0 * 1024.0), + elapsed_ms, + throughput_mb_s); +} + +// ACE: Load partition dataset and augmented dataset from disk +template +void ace_load_partition_dataset_from_disk( + size_t partition_id, + size_t dataset_dim, + raft::host_matrix_view partition_histogram, + raft::host_vector_view core_partition_offsets, + raft::host_vector_view augmented_partition_offsets, + const cuvs::util::file_descriptor& reordered_fd, + const cuvs::util::file_descriptor& augmented_fd, + size_t reordered_header_size, + size_t augmented_header_size, + raft::mdspan, raft::row_major, Accessor> sub_dataset) +{ + RAFT_LOG_DEBUG("ACE: Loading partition %lu dataset from disk", partition_id); + + const std::string reordered_path = reordered_fd.get_path(); + const std::string augmented_path = augmented_fd.get_path(); + T* sub_dataset_ptr = sub_dataset.data_handle(); + RAFT_EXPECTS(sub_dataset_ptr != nullptr, "ACE: sub-dataset destination must not be null"); + RAFT_EXPECTS(reordered_fd.is_valid() && !reordered_path.empty(), + "ACE: reordered dataset file descriptor is not valid"); + RAFT_EXPECTS(augmented_fd.is_valid() && !augmented_path.empty(), + "ACE: augmented dataset file descriptor is not valid"); + RAFT_EXPECTS(partition_id < static_cast(partition_histogram.extent(0)), + "ACE: partition id is out of range"); + RAFT_EXPECTS(partition_id < static_cast(core_partition_offsets.extent(0)), + "ACE: core partition offset is out of range"); + RAFT_EXPECTS(partition_id < static_cast(augmented_partition_offsets.extent(0)), + "ACE: augmented partition offset is out of range"); + + size_t core_size = partition_histogram(partition_id, 0); + size_t augmented_size = partition_histogram(partition_id, 1); + size_t total_size = core_size + augmented_size; + + RAFT_EXPECTS(static_cast(sub_dataset.extent(0)) == total_size, + "ACE: sub-dataset rows (%zu) must match partition size (%zu)", + static_cast(sub_dataset.extent(0)), + total_size); + RAFT_EXPECTS(static_cast(sub_dataset.extent(1)) == dataset_dim, + "ACE: sub-dataset columns (%zu) must match dataset dimensions (%zu)", + static_cast(sub_dataset.extent(1)), + dataset_dim); + + RAFT_LOG_DEBUG("ACE: Partition %lu: %lu core + %lu augmented = %lu total vectors", + partition_id, + core_size, + augmented_size, + core_size + augmented_size); + + const size_t vector_size = dataset_dim * sizeof(T); + const size_t core_file_offset = + reordered_header_size + static_cast(core_partition_offsets(partition_id)) * vector_size; + const size_t augmented_file_offset = + augmented_header_size + + static_cast(augmented_partition_offsets(partition_id)) * vector_size; + + RAFT_LOG_DEBUG("ACE: Core file offset: %lu bytes, Augmented file offset: %lu bytes", + core_file_offset, + augmented_file_offset); + + const size_t core_bytes = core_size * vector_size; + const size_t augmented_bytes = augmented_size * vector_size; + T* augmented_dest = sub_dataset_ptr + (core_size * dataset_dim); + + auto expect_complete_read = [partition_id](const char* name, size_t expected, size_t actual) { + RAFT_EXPECTS(actual == expected, + "ACE: Short %s read for partition %lu: expected %zu bytes, got %zu", + name, + partition_id, + expected, + actual); + }; + + if (core_bytes > 0 && augmented_bytes > 0) { + RAFT_LOG_DEBUG("ACE: Reading %lu core vectors from offset %lu", core_size, core_file_offset); + RAFT_LOG_DEBUG( + "ACE: Reading %lu augmented vectors from offset %lu", augmented_size, augmented_file_offset); + auto reordered_handle = + cuvs::util::detail::open_kvikio_file_for_ace_io(reordered_path, "r", sub_dataset_ptr); + auto augmented_handle = + cuvs::util::detail::open_kvikio_file_for_ace_io(augmented_path, "r", augmented_dest); + auto core_future = reordered_handle.pread(sub_dataset_ptr, core_bytes, core_file_offset); + auto augmented_future = + augmented_handle.pread(augmented_dest, augmented_bytes, augmented_file_offset); + std::exception_ptr read_exception = nullptr; + size_t core_read = 0; + size_t augmented_read = 0; + try { + core_read = core_future.get(); + } catch (...) { + read_exception = std::current_exception(); + } + try { + augmented_read = augmented_future.get(); + } catch (...) { + if (!read_exception) { read_exception = std::current_exception(); } + } + if (read_exception) { std::rethrow_exception(read_exception); } + expect_complete_read("core", core_bytes, core_read); + expect_complete_read("augmented", augmented_bytes, augmented_read); + } else if (core_bytes > 0) { + RAFT_LOG_DEBUG("ACE: Reading %lu core vectors from offset %lu", core_size, core_file_offset); + auto reordered_handle = + cuvs::util::detail::open_kvikio_file_for_ace_io(reordered_path, "r", sub_dataset_ptr); + expect_complete_read( + "core", + core_bytes, + reordered_handle.pread(sub_dataset_ptr, core_bytes, core_file_offset).get()); + } else if (augmented_bytes > 0) { + RAFT_LOG_DEBUG( + "ACE: Reading %lu augmented vectors from offset %lu", augmented_size, augmented_file_offset); + auto augmented_handle = + cuvs::util::detail::open_kvikio_file_for_ace_io(augmented_path, "r", augmented_dest); + expect_complete_read( + "augmented", + augmented_bytes, + augmented_handle.pread(augmented_dest, augmented_bytes, augmented_file_offset).get()); + } +} + +// Memory requirements for ACE operation +struct ace_memory_requirements { + size_t partition_labels_size; + size_t id_mapping_size; + size_t sub_dataset_size; + size_t sub_graph_size; + size_t cagra_graph_size; + size_t total_size; + size_t available_host_memory; + size_t available_gpu_memory; +}; + +// Amount of host memory that can be used for the build +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; + +// Check if disk mode should be used for ACE based on memory constraints +template +bool ace_check_use_disk_mode(raft::resources const& res, + bool use_disk, + const std::string& build_dir, + size_t dataset_size, + size_t dataset_dim, + size_t n_partitions, + size_t intermediate_degree, + size_t graph_degree, + std::optional max_host_memory_gb, + std::optional max_gpu_memory_gb, + bool guarantee_connectivity, + ace_memory_requirements& mem) +{ + 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() && + host_memory.cgroup_reclaimable_file.has_value() && + host_memory.cgroup_working_set.has_value()) { + const size_t cgroup_headroom = *host_memory.cgroup_working_set < *host_memory.cgroup_limit + ? *host_memory.cgroup_limit - *host_memory.cgroup_working_set + : 0; + RAFT_LOG_INFO( + "ACE: Cgroup host memory limit: %.2f GiB, current: %.2f GiB, reclaimable file cache: %.2f " + "GiB, working set: %.2f GiB, headroom: %.2f GiB; system available: %.2f GiB, effective " + "available: %.2f GiB", + to_gib(*host_memory.cgroup_limit), + to_gib(*host_memory.cgroup_current), + to_gib(*host_memory.cgroup_reclaimable_file), + to_gib(*host_memory.cgroup_working_set), + to_gib(cgroup_headroom), + to_gib(host_memory.system_available), + to_gib(host_memory.available)); + } + + // Use overridden memory limits if provided (> 0), otherwise query actual system memory + if (max_host_memory_gb.has_value() && max_host_memory_gb.value() > 0) { + auto actual_available_host_memory = host_memory.available; + auto configured_host_memory = static_cast(max_host_memory_gb.value() * (1ULL << 30)); + if (actual_available_host_memory < configured_host_memory) { + RAFT_LOG_WARN( + "ACE: Actual host memory (%.2f GiB) is less than configured limit (%.2f GiB). " + "Using actual host memory.", + to_gib(actual_available_host_memory), + to_gib(configured_host_memory)); + mem.available_host_memory = actual_available_host_memory; + } else { + RAFT_LOG_INFO("ACE: Using overridden host memory limit: %.2f GiB", + max_host_memory_gb.value()); + mem.available_host_memory = configured_host_memory; + } + } else { + mem.available_host_memory = host_memory.available; + } + size_t sub_partition_size = + static_cast(imbalance_factor * vector_expansion_factor * + raft::div_rounding_up_safe(dataset_size, n_partitions)); + auto [opt_host_ws_total, opt_dev_ws_total, opt_host_ws_fixed, opt_dev_ws_fixed] = + helpers::optimize_workspace_size( + sub_partition_size, graph_degree, intermediate_degree, sizeof(IdxT), guarantee_connectivity); + + // Optimistic memory model: focus on largest arrays, assumes all partitions are of equal size + // For memory path: + // - Partition labels (core + augmented): vector_expansion_factor * dataset_size * sizeof(IdxT) + // - Backward ID mapping arrays (core + augmented): vector_expansion_factor * dataset_size * + // sizeof(IdxT) + // - Avg. per-partition dataset: vector_expansion_factor * (dataset_size / n_partitions) * + // dataset_dim * sizeof(T) + // - Avg. per-partition graph during build: vector_expansion_factor * (dataset_size / + // n_partitions) * (intermediate + final) * sizeof(IdxT) + // - Final assembled graph: dataset_size * graph_degree * sizeof(IdxT) + mem.partition_labels_size = vector_expansion_factor * dataset_size * sizeof(IdxT); + mem.id_mapping_size = vector_expansion_factor * dataset_size * sizeof(IdxT); + mem.sub_dataset_size = sub_partition_size * dataset_dim * sizeof(T); + mem.sub_graph_size = sub_partition_size * (intermediate_degree + graph_degree) * sizeof(IdxT); + 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; + + RAFT_LOG_INFO("ACE: Estimated host memory required: %.2f GiB, available: %.2f GiB", + to_gib(mem.total_size), + to_gib(mem.available_host_memory)); + + bool host_memory_limited = + static_cast(usable_cpu_memory_fraction * mem.available_host_memory) < mem.total_size; + + // GPU is mostly limited by the index size (update_graph() in the end of this routine). + // Check if GPU has enough memory for the final graph or use disk mode instead. + // TODO: Extend model or use managed memory if running out of GPU memory. + if (max_gpu_memory_gb.has_value() && max_gpu_memory_gb.value() > 0) { + auto actual_available_gpu_memory = rmm::available_device_memory().second; + auto configured_gpu_memory = static_cast(max_gpu_memory_gb.value() * (1ULL << 30)); + if (actual_available_gpu_memory < configured_gpu_memory) { + RAFT_LOG_WARN( + "ACE: Actual GPU memory (%zu GiB) is less than configured limit (%zu GiB). " + "Using actual GPU memory.", + to_gib(actual_available_gpu_memory), + to_gib(configured_gpu_memory)); + mem.available_gpu_memory = actual_available_gpu_memory; + } else { + RAFT_LOG_INFO("ACE: Using overridden GPU memory limit: %.2f GiB", max_gpu_memory_gb.value()); + mem.available_gpu_memory = configured_gpu_memory; + } + } else { + mem.available_gpu_memory = rmm::available_device_memory().second; + } + + // what we need is maximum of: + // * IVF-PQ on partition (sub_dataset_size, uncompressed upper bound) + // * optimize workspace (opt_dev_ws_total) + // + some extra workspace (IVF-PQ search, ...) + size_t extra_gpu_workspace_size = raft::resource::get_workspace_total_bytes(res); + size_t gpu_memory_required = + std::max(mem.sub_dataset_size, opt_dev_ws_total) + extra_gpu_workspace_size; + + bool gpu_memory_limited = mem.available_gpu_memory < gpu_memory_required; + + RAFT_LOG_INFO("ACE: Estimated GPU memory required: %.2f GiB, available: %.2f GiB", + to_gib(gpu_memory_required), + to_gib(mem.available_gpu_memory)); + + bool use_disk_mode = use_disk || host_memory_limited || gpu_memory_limited; + if (use_disk_mode) { RAFT_EXPECTS(!build_dir.empty(), "ACE build directory must not be empty"); } + + if (host_memory_limited && gpu_memory_limited) { + RAFT_LOG_INFO( + "ACE: Graph does not fit in host and GPU memory. Using disk-mode with temporary storage %s", + build_dir.c_str()); + } else if (host_memory_limited) { + RAFT_LOG_INFO( + "ACE: Graph does not fit in host memory. Using disk-mode with temporary storage %s", + build_dir.c_str()); + } else if (gpu_memory_limited) { + RAFT_LOG_INFO( + "ACE: Graph does not fit in GPU memory. Using disk-mode with temporary storage %s", + build_dir.c_str()); + } else if (use_disk) { + RAFT_LOG_INFO( + "ACE: Graph fits in host and GPU memory but disk mode is forced. Using disk-mode with " + "temporary storage %s", + build_dir.c_str()); + } else { + RAFT_LOG_INFO("ACE: Graph fits in host and GPU memory. Using in-memory mode."); + } + + return use_disk_mode; +} + +// Resolve the ACE partition count while preserving 0 as the auto-selection sentinel. +inline size_t ace_resolve_partition_count(size_t n_partitions) +{ + if (n_partitions == 0) { return 2; } + if (n_partitions == 1) { + RAFT_LOG_WARN( + "ACE: Requested 1 partition; adjusted to 2 before applying partitioning heuristics"); + return 2; + } + return n_partitions; +} + +// Validate the structural ACE partition-count invariants required by the labeler. +inline void ace_validate_partition_count(size_t n_partitions, + size_t dataset_size, + bool adjusted_for_memory = false) +{ + RAFT_EXPECTS(n_partitions <= dataset_size, + adjusted_for_memory + ? "ACE: configured memory limit is unsatisfiable because the requested partition " + "count cannot exceed dataset size" + : "ACE: number of partitions cannot exceed dataset size"); +} + +// Validate and adjust partitions for disk mode memory requirements +template +void ace_validate_disk_mode_partitions(raft::resources const& res, + size_t& n_partitions, + size_t dataset_size, + size_t dataset_dim, + size_t intermediate_degree, + size_t graph_degree, + bool guarantee_connectivity, + ace_memory_requirements& mem) +{ + // 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) + // - ID mapping arrays (core + augmented): vector_expansion_factor * dataset_size * sizeof(IdxT) + // - Avg. per-partition dataset during processing: vector_expansion_factor * (dataset_size / + // n_partitions) * dataset_dim * sizeof(T) + // - Avg. per-partition graph during build: vector_expansion_factor * (dataset_size / + // n_partitions) * (intermediate + final) * sizeof(IdxT) + + size_t original_n_partitions = n_partitions; + size_t host_suggested_partitions = n_partitions; + size_t gpu_suggested_partitions = n_partitions; + bool host_memory_insufficient = false; + bool gpu_memory_insufficient = false; + + // Compute optimize workspace requirements + size_t sub_partition_size = + static_cast(imbalance_factor * vector_expansion_factor * + raft::div_rounding_up_safe(dataset_size, n_partitions)); + auto [host_workspace_size_total, + gpu_workspace_size_total, + host_workspace_size_fixed, + gpu_workspace_size_fixed] = + helpers::optimize_workspace_size( + sub_partition_size, graph_degree, intermediate_degree, sizeof(IdxT), guarantee_connectivity); + + // 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 + + host_workspace_size_total; + + if (static_cast(usable_cpu_memory_fraction * mem.available_host_memory) < + disk_mode_host_required) { + host_memory_insufficient = true; + RAFT_LOG_WARN( + "ACE: Host memory insufficient for disk mode. Required: %.2f GiB, available: %.2f GiB. " + "Per-partition breakdown: dataset %.2f GiB, graph %.2f GiB, workspace %.2f GiB", + to_gib(disk_mode_host_required), + to_gib(mem.available_host_memory), + to_gib(mem.sub_dataset_size), + to_gib(mem.sub_graph_size), + to_gib(host_workspace_size_total)); + + // Calculate suggested number of partitions for host memory + size_t disk_mode_host_static = + mem.partition_labels_size + mem.id_mapping_size + host_workspace_size_fixed; + size_t disk_mode_host_dynamic = disk_mode_host_required - disk_mode_host_static; + double available_for_scaling = + usable_cpu_memory_fraction * mem.available_host_memory - disk_mode_host_static; + + RAFT_EXPECTS(available_for_scaling > 0, + "ACE: Host memory insufficient even for constant overhead (labels + id_mapping + " + "static workspace). " + "Required: %.2f GiB, available: %.2f GiB", + to_gib(disk_mode_host_static), + to_gib(usable_cpu_memory_fraction * mem.available_host_memory)); + host_suggested_partitions = + static_cast(std::ceil(disk_mode_host_dynamic * n_partitions / available_for_scaling)); + // Ensure we always increase partitions (current count is insufficient by definition) + host_suggested_partitions = std::max(host_suggested_partitions, n_partitions + 1); + } + + // Check GPU memory requirements in disk mode + // GPU memory model for disk mode (per-partition processing): + // * IVF-PQ on partition (mem.sub_dataset_size) (compressed?) + // * optimize workspace (gpu_workspace_size_total) + // + some extra workspace (IVF-PQ search, ...) + size_t extra_gpu_workspace_size = raft::resource::get_workspace_total_bytes(res); + size_t disk_mode_gpu_required = + std::max(mem.sub_dataset_size, gpu_workspace_size_total) + extra_gpu_workspace_size; + + if (mem.available_gpu_memory < disk_mode_gpu_required) { + gpu_memory_insufficient = true; + RAFT_LOG_WARN( + "ACE: GPU memory insufficient for per-partition processing. Required: %.2f GiB, " + "available: %.2f GiB. Per-partition breakdown: dataset %.2f GiB, workspace %.2f GiB", + to_gib(disk_mode_gpu_required), + to_gib(mem.available_gpu_memory), + to_gib(mem.sub_dataset_size), + to_gib(gpu_workspace_size_total + extra_gpu_workspace_size)); + + size_t disk_mode_gpu_static = gpu_workspace_size_fixed + extra_gpu_workspace_size; + size_t disk_mode_gpu_dynamic = disk_mode_gpu_required - disk_mode_gpu_static; + double available_for_scaling = mem.available_gpu_memory - disk_mode_gpu_static; + + RAFT_EXPECTS(available_for_scaling > 0, + "ACE: GPU memory insufficient even for constant overhead. Required: %.2f GiB, " + "available: %.2f GiB", + to_gib(disk_mode_gpu_static), + to_gib(mem.available_gpu_memory)); + + gpu_suggested_partitions = + static_cast(std::ceil(disk_mode_gpu_dynamic * n_partitions / available_for_scaling)); + gpu_suggested_partitions = std::max(gpu_suggested_partitions, n_partitions + 1); + } + + // Auto-adjust to the maximum of host and GPU requirements + if (host_memory_insufficient || gpu_memory_insufficient) { + size_t new_n_partitions = std::max(host_suggested_partitions, gpu_suggested_partitions); + + RAFT_LOG_WARN( + "ACE: Automatically increasing number of partitions from %zu to %zu to satisfy memory " + "constraints.%s%s", + original_n_partitions, + new_n_partitions, + host_memory_insufficient + ? " Host memory requires >= " + std::to_string(host_suggested_partitions) + " partitions." + : "", + gpu_memory_insufficient + ? " GPU memory requires >= " + std::to_string(gpu_suggested_partitions) + " partitions." + : ""); + + n_partitions = new_n_partitions; + + size_t new_sub_partition_size = + static_cast(imbalance_factor * vector_expansion_factor * + raft::div_rounding_up_safe(dataset_size, n_partitions)); + auto [new_opt_host_ws, new_opt_dev_ws, new_opt_host_ws_fixed, new_opt_dev_ws_fixed] = + helpers::optimize_workspace_size(new_sub_partition_size, + graph_degree, + intermediate_degree, + sizeof(IdxT), + guarantee_connectivity); + + mem.sub_dataset_size = new_sub_partition_size * dataset_dim * sizeof(T); + 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; + + RAFT_LOG_INFO( + "ACE: Updated per-partition memory estimates: dataset %.2f GiB, graph %.2f GiB, " + "host workspace %.2f GiB, GPU workspace %.2f GiB", + to_gib(mem.sub_dataset_size), + to_gib(mem.sub_graph_size), + to_gib(new_opt_host_ws), + to_gib(new_opt_dev_ws + extra_gpu_workspace_size)); + } +} + +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; } + + 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::sort(candidate_sources.begin(), candidate_sources.end()); + + 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)); + + 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_sets(n_partitions); + for (size_t source = 0; source < dataset_size; ++source) { + auto const source_partition = static_cast(partition_labels(source, 0)); + 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)); + if (source_partition != destination_partition) { + candidate_sets[destination_partition].insert(static_cast(source)); + } + } + } + + for (size_t partition_id = 0; partition_id < n_partitions; ++partition_id) { + if (candidate_sets[partition_id].empty()) { continue; } + 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); + std::vector candidate_sources(candidate_sets[partition_id].begin(), + candidate_sets[partition_id].end()); + 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 +// 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. +// In disk mode, the graph is stored in build_dir and dataset is reordered on disk. +// The returned index is not usable for search. Use the created files for search instead. +template + requires cuvs::neighbors::is_host_dataset_view_v +auto build_ace(raft::resources const& res, const index_params& params, DatasetViewT const& dataset) + -> cuvs::neighbors::cagra::index +{ + // Extract ACE parameters from graph_build_params + RAFT_EXPECTS( + std::holds_alternative(params.graph_build_params), + "ACE build requires graph_build_params to be set to ace_params"); + + auto ace_params = std::get(params.graph_build_params); + size_t npartitions = ace_params.npartitions; + size_t ef_construction = ace_params.ef_construction; + std::string build_dir = ace_params.build_dir; + bool use_disk = ace_params.use_disk; + + common::nvtx::range function_scope( + "cagra::detail::build_ace(%zu, %zu, %zu)", + params.intermediate_graph_degree, + params.graph_degree, + npartitions); + + auto dataset_view = dataset.view(); + size_t dataset_size = dataset.n_rows(); + size_t dataset_dim = dataset.dim(); + + RAFT_EXPECTS(dataset_size > 0, "ACE: Dataset must not be empty"); + if (dataset_size < 1000) { + RAFT_LOG_WARN("ACE: Very small dataset size (%zu), consider using regular CAGRA build instead.", + dataset_size); + } + RAFT_EXPECTS(dataset_dim > 0, "ACE: Dataset dimension must be greater than 0"); + RAFT_EXPECTS(params.intermediate_graph_degree > 0, + "ACE: Intermediate graph degree must be greater than 0"); + RAFT_EXPECTS(params.graph_degree > 0, "ACE: Graph degree must be greater than 0"); + + size_t n_partitions = ace_resolve_partition_count(npartitions); + + 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; + if (n_partitions < 2) { + RAFT_LOG_WARN( + "ACE: Reduced number of partitions to the minimum of 2 to avoid tiny partitions. Consider " + "using regular CAGRA build instead."); + n_partitions = 2; + } else { + RAFT_LOG_WARN("ACE: Reduced number of partitions to %zu to avoid tiny partitions", + n_partitions); + } + } + + auto total_start = std::chrono::high_resolution_clock::now(); + RAFT_LOG_INFO("ACE: Starting partitioned CAGRA build with %zu partitions", n_partitions); + + size_t intermediate_degree = params.intermediate_graph_degree; + size_t graph_degree = params.graph_degree; + + ace_disk_workspace workspace(build_dir); + + try { + check_graph_degree(intermediate_degree, graph_degree, dataset_size); + + // Check if disk mode should be used based on memory constraints + ace_memory_requirements mem; + bool use_disk_mode = ace_check_use_disk_mode(res, + use_disk, + build_dir, + dataset_size, + dataset_dim, + n_partitions, + intermediate_degree, + graph_degree, + ace_params.max_host_memory_gb, + ace_params.max_gpu_memory_gb, + params.guarantee_connectivity, + mem); + RAFT_EXPECTS(!ace_params.add_global_reverse_edges || !use_disk_mode, + "ACE: partitioned global reverse-edge addition does not support disk mode"); + + // Validate and adjust partitions if disk mode is enabled + if (use_disk_mode) { + ace_validate_disk_mode_partitions(res, + n_partitions, + dataset_size, + dataset_dim, + intermediate_degree, + graph_degree, + params.guarantee_connectivity, + mem); + ace_validate_partition_count(n_partitions, dataset_size, true); + } + + // Preallocate space for files for better performance and fail early if not enough space. + cuvs::util::file_descriptor reordered_fd; + cuvs::util::file_descriptor augmented_fd; + cuvs::util::file_descriptor mapping_fd; + cuvs::util::file_descriptor graph_fd; + size_t reordered_header_size = 0; + size_t augmented_header_size = 0; + size_t mapping_header_size = 0; + size_t graph_header_size = 0; + + if (use_disk_mode) { + workspace.initialize(); + + // Create numpy files with pre-allocated space + std::tie(reordered_fd, reordered_header_size) = cuvs::util::create_numpy_file( + workspace.artifact_path(ace_disk_workspace::artifact::reordered_dataset), + {dataset_size, dataset_dim}, + true); + workspace.mark_artifact_created(ace_disk_workspace::artifact::reordered_dataset); + + 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}, + true); + workspace.mark_artifact_created(ace_disk_workspace::artifact::augmented_dataset); + + std::tie(mapping_fd, mapping_header_size) = cuvs::util::create_numpy_file( + workspace.artifact_path(ace_disk_workspace::artifact::dataset_mapping), + {dataset_size}, + true); + workspace.mark_artifact_created(ace_disk_workspace::artifact::dataset_mapping); + + std::tie(graph_fd, graph_header_size) = cuvs::util::create_numpy_file( + workspace.artifact_path(ace_disk_workspace::artifact::cagra_graph), + {dataset_size, graph_degree}, + true); + workspace.mark_artifact_created(ace_disk_workspace::artifact::cagra_graph); + + RAFT_LOG_DEBUG( + "ACE: Wrote numpy headers (reordered: %zu, augmented: %zu, mapping: %zu, graph: %zu bytes)", + reordered_header_size, + augmented_header_size, + mapping_header_size, + graph_header_size); + } + + auto partition_start = std::chrono::high_resolution_clock::now(); + auto partition_labels = raft::make_host_matrix(dataset_size, 2); + auto partition_histogram = raft::make_host_matrix(n_partitions, 2); + for (size_t c = 0; c < n_partitions; c++) { + partition_histogram(c, 0) = 0; + partition_histogram(c, 1) = 0; + } + + // Determine minimum partition size for stable KNN graph construction + size_t min_partition_size = std::max(1000ULL, dataset_size / n_partitions * 0.1); + + ace_get_partition_labels(res, + dataset_view, + dataset_dim, + partition_labels.view(), + 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 = + std::chrono::duration_cast(partition_end - partition_start) + .count(); + RAFT_LOG_INFO( + "ACE: Partition labeling completed in %ld ms (min_partition_size: " + "%lu)", + partition_elapsed, + min_partition_size); + + // Create vector lists for each partition + auto vectorlist_start = std::chrono::high_resolution_clock::now(); + 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 core_partition_offsets = raft::make_host_vector(n_partitions + 1); + auto augmented_partition_offsets = raft::make_host_vector(n_partitions + 1); + + ace_create_forward_and_backward_lists(dataset_size, + n_partitions, + partition_labels.view(), + partition_histogram.view(), + core_forward_mapping.view(), + core_backward_mapping.view(), + augmented_backward_mapping.view(), + core_partition_offsets.view(), + augmented_partition_offsets.view()); + + auto vectorlist_end = std::chrono::high_resolution_clock::now(); + auto vectorlist_elapsed = + std::chrono::duration_cast(vectorlist_end - vectorlist_start) + .count(); + RAFT_LOG_INFO("ACE: Vector list creation completed in %ld ms", vectorlist_elapsed); + + // Reorder the dataset based on partitions and store to disk. Uses write buffers to improve + // performance. + if (use_disk_mode) { + ace_reorder_and_store_dataset(res, + build_dir, + dataset_view, + dataset_dim, + partition_labels.view(), + partition_histogram.view(), + core_backward_mapping.view(), + core_partition_offsets.view(), + augmented_partition_offsets.view(), + reordered_fd, + augmented_fd, + mapping_fd, + reordered_header_size, + augmented_header_size, + mapping_header_size); + // core_backward_mapping is not needed anymore. + core_backward_mapping = raft::make_host_vector(0); + } + + // Placeholder search graph for in-memory version + auto search_graph = use_disk_mode + ? raft::make_host_matrix(0, 0) + : raft::make_host_matrix(dataset_size, graph_degree); + + // 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++) { + RAFT_LOG_DEBUG("ACE: Processing partition %lu/%lu", partition_id + 1, n_partitions); + auto start = std::chrono::high_resolution_clock::now(); + + // Extract vectors for this partition + size_t core_sub_dataset_size = partition_histogram(partition_id, 0); + size_t augmented_sub_dataset_size = partition_histogram(partition_id, 1); + size_t sub_dataset_size = core_sub_dataset_size + augmented_sub_dataset_size; + + if (sub_dataset_size == 0) { + RAFT_LOG_WARN("ACE: Skipping empty partition %lu", partition_id); + continue; + } + RAFT_LOG_DEBUG("ACE: Sub-dataset size: %lu (%lu + %lu)", + sub_dataset_size, + core_sub_dataset_size, + augmented_sub_dataset_size); + + // Create index for this partition + auto sub_index_params = cuvs::neighbors::cagra::index_params::from_dataset( + raft::make_extents(sub_dataset_size, dataset_dim), + graph_degree, + params.metric, + ef_construction / 16); + sub_index_params.attach_dataset_on_build = false; + sub_index_params.guarantee_connectivity = params.guarantee_connectivity; + + auto read_end = start; + bool used_device_read = false; + const size_t sub_ds_bytes = sub_dataset_size * dataset_dim * sizeof(T); + 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; + std::optional> partial_sub_graph; + std::optional sub_index; + const bool build_partial_graph_on_host = + ace_params.add_global_reverse_edges || + ace_params.partial_graph_knn_build_algo == + graph_build_params::ace_params::knn_graph_build_algo::BRUTE_FORCE; + if (build_partial_graph_on_host) { + auto core_argument_partition = + raft::make_host_matrix(sub_dataset_size, dataset_dim); + if (use_disk_mode) { + ace_load_partition_dataset_from_disk(partition_id, + dataset_dim, + partition_histogram.view(), + core_partition_offsets.view(), + augmented_partition_offsets.view(), + reordered_fd, + augmented_fd, + reordered_header_size, + augmented_header_size, + core_argument_partition.view()); + } else { + 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(), + core_argument_partition.view()); + } + read_end = std::chrono::high_resolution_clock::now(); + auto core_argument_partition_view = raft::make_host_matrix_view( + core_argument_partition.data_handle(), sub_dataset_size, dataset_dim); + if (ace_params.add_global_reverse_edges) { + partial_sub_graph.emplace( + ace_build_pruned_partial_graph(res, + core_argument_partition_view, + intermediate_degree, + graph_degree, + params.metric, + ace_params.partial_graph_knn_build_algo)); + } else { + partial_sub_graph.emplace( + ace_build_optimized_partial_graph(res, + core_argument_partition_view, + intermediate_degree, + graph_degree, + params.metric, + ace_params.partial_graph_knn_build_algo, + params.guarantee_connectivity)); + } + } else { + sub_index.emplace([&]() -> sub_index_t { + if (use_disk_mode) { + const size_t current_free_gpu = rmm::available_device_memory().first; + const size_t configured_gpu = + ace_params.max_gpu_memory_gb > 0 + ? 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)) { + try { + auto sub_dataset_tight = + raft::make_device_matrix(res, sub_dataset_size, dataset_dim); + raft::resource::sync_stream(res); + ace_load_partition_dataset_from_disk(partition_id, + dataset_dim, + partition_histogram.view(), + core_partition_offsets.view(), + augmented_partition_offsets.view(), + reordered_fd, + augmented_fd, + reordered_header_size, + augmented_header_size, + sub_dataset_tight.view()); + read_end = std::chrono::high_resolution_clock::now(); + auto sub_dataset_view = raft::make_const_mdspan(sub_dataset_tight.view()); + std::unique_ptr> + sub_dataset_padded; + auto sub_dataset_dev = [&]() { + if (cuvs::neighbors::matrix_row_width_matches_cagra_required(sub_dataset_view)) { + return cuvs::neighbors::make_device_padded_dataset_view(res, sub_dataset_view); + } + sub_dataset_padded = + cuvs::neighbors::make_device_padded_dataset(res, sub_dataset_view); + return sub_dataset_padded->as_dataset_view(); + }(); + auto direct_index = + ::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)}; + } catch (const std::bad_alloc& e) { + RAFT_LOG_WARN( + "ACE: partition %lu did not fit in device memory for a direct (GDS) read: %s; " + "falling back to a host read", + partition_id, + e.what()); + } + } + + auto sub_dataset = raft::make_host_matrix(sub_dataset_size, dataset_dim); + ace_load_partition_dataset_from_disk(partition_id, + dataset_dim, + partition_histogram.view(), + core_partition_offsets.view(), + augmented_partition_offsets.view(), + reordered_fd, + augmented_fd, + reordered_header_size, + augmented_header_size, + sub_dataset.view()); + 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 = raft::make_host_matrix(sub_dataset_size, dataset_dim); + 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()); + 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)}; + }()); + } + if (used_device_read) { + RAFT_LOG_DEBUG("ACE: partition %lu read directly into device memory (GDS path)", + partition_id); + } + auto read_elapsed = + std::chrono::duration_cast(read_end - start).count(); + auto optimize_end = std::chrono::high_resolution_clock::now(); + auto optimize_elapsed = + std::chrono::duration_cast(optimize_end - read_end).count(); + + 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) { + auto adjusted_search_graph = + raft::make_device_matrix(res, core_sub_dataset_size, graph_degree); + 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 (build_partial_graph_on_host) { + 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 = + static_cast(core_partition_offsets(partition_id)) * graph_degree * sizeof(IdxT) + + graph_header_size; + cuvs::util::write_large_file( + graph_fd, adjusted_search_graph.data_handle(), graph_bytes, graph_offset); + auto write_end = std::chrono::high_resolution_clock::now(); + write_elapsed = + std::chrono::duration_cast(write_end - adjust_end).count(); + } else { + auto sub_search_graph = + raft::make_host_matrix(core_sub_dataset_size, graph_degree); + if (build_partial_graph_on_host) { + 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, + augmented_sub_dataset_size, + graph_degree, + partition_id, + sub_search_graph.view(), + search_graph.view(), + core_partition_offsets.view(), + augmented_partition_offsets.view(), + core_backward_mapping.view(), + augmented_backward_mapping.view()); + adjust_end = std::chrono::high_resolution_clock::now(); + } + + auto adjust_elapsed = + std::chrono::duration_cast(adjust_end - optimize_end).count(); + + auto end = std::chrono::high_resolution_clock::now(); + auto elapsed_ms = std::chrono::duration_cast(end - start).count(); + double read_throughput = + read_elapsed > 0 + ? to_mib(sub_dataset_size * dataset_dim * sizeof(T)) / (read_elapsed / 1000.0) + : 0.0; + double write_throughput = + write_elapsed > 0 ? to_mib(graph_bytes) / (write_elapsed / 1000.0) : 0.0; + RAFT_LOG_INFO( + "ACE: Partition %4lu (%8lu + %8lu) completed in %6ld ms: read %6ld ms (%7.1f MiB/s), " + "optimize %6ld ms, adjust %6ld ms, write %6ld ms (%7.1f MiB/s)", + partition_id, + core_sub_dataset_size, + augmented_sub_dataset_size, + elapsed_ms, + read_elapsed, + read_throughput, + optimize_elapsed, + adjust_elapsed, + write_elapsed, + write_throughput); + } + + auto partition_processing_end = std::chrono::high_resolution_clock::now(); + auto partition_processing_elapsed = std::chrono::duration_cast( + partition_processing_end - partition_processing_start) + .count(); + RAFT_LOG_INFO("ACE: All partition processing completed in %ld ms (%zu partitions)", + partition_processing_elapsed, + n_partitions); + + if (ace_params.add_global_reverse_edges) { + RAFT_LOG_INFO("ACE: Adding global cross-partition reverse edges partition by partition"); + 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) { + const std::string augmented_dataset_path = build_dir + "/augmented_dataset.npy"; + if (std::filesystem::exists(augmented_dataset_path)) { + std::filesystem::remove(augmented_dataset_path); + RAFT_LOG_INFO("ACE: Removed augmented dataset file to save disk space"); + } + } + + auto index_creation_start = std::chrono::high_resolution_clock::now(); + cuvs::neighbors::cagra::index idx(res, params.metric); + if (!use_disk_mode) { + if (params.attach_dataset_on_build) { + idx = cuvs::neighbors::cagra::index( + res, params.metric, dataset, raft::make_const_mdspan(search_graph.view())); + } else { + idx.update_graph(res, raft::make_const_mdspan(search_graph.view())); + } + } else { + idx.update_dataset(res, std::move(reordered_fd)); + idx.update_graph(res, std::move(graph_fd)); + idx.update_mapping(res, std::move(mapping_fd)); + + RAFT_LOG_INFO( + "ACE: Set disk storage at %s (dataset shape [%zu, %zu], graph shape [%zu, %zu])", + build_dir.c_str(), + idx.size(), + idx.dim(), + idx.size(), + idx.graph_degree()); + } + + auto index_creation_end = std::chrono::high_resolution_clock::now(); + auto index_creation_elapsed = std::chrono::duration_cast( + index_creation_end - index_creation_start) + .count(); + RAFT_LOG_INFO("ACE: Final index creation completed in %ld ms", index_creation_elapsed); + + auto total_end = std::chrono::high_resolution_clock::now(); + auto total_elapsed = + std::chrono::duration_cast(total_end - total_start).count(); + RAFT_LOG_INFO("ACE: Partitioned CAGRA build completed in %ld ms total", total_elapsed); + + workspace.commit(); + return std::move(idx); + } catch (const std::exception& e) { + RAFT_LOG_ERROR("ACE: Build failed with exception: %s", e.what()); + workspace.cleanup(); + throw; + } catch (...) { + workspace.cleanup(); + throw; + } +} + +} // 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 0b7465c9c8..7122e56478 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_build.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_build.cuh @@ -5,7 +5,6 @@ #pragma once #include "../../../core/nvtx.hpp" -#include "../../../util/kvikio_io.hpp" #include "../../ivf_pq/ivf_pq_fp16_overflow.cuh" #include "graph_core.cuh" #include @@ -25,135 +24,33 @@ #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 -#include namespace cuvs::neighbors::cagra::detail { -// Helpers to convert bytes to MiB and GiB -constexpr double to_mib(size_t bytes) { return static_cast(bytes) / (1 << 20); } -constexpr double to_gib(size_t bytes) { return static_cast(bytes) / (1 << 30); } - -class ace_disk_workspace { - public: - enum class artifact : size_t { - reordered_dataset, - augmented_dataset, - dataset_mapping, - cagra_graph, - }; - - explicit ace_disk_workspace(std::string build_dir) - : build_dir_(std::move(build_dir)), - artifacts_{build_dir_ / "reordered_dataset.npy", - build_dir_ / "augmented_dataset.npy", - build_dir_ / "dataset_mapping.npy", - build_dir_ / "cagra_graph.npy"} - { - } - - void initialize() - { - if (mkdir(build_dir_.c_str(), 0755) == 0) { - directory_created_by_this_build_ = true; - return; - } - - if (errno != EEXIST) { - RAFT_FAIL("Failed to create ACE build directory: %s (errno: %d, %s)", - build_dir_.c_str(), - errno, - strerror(errno)); - } - - std::error_code error; - const bool is_directory = std::filesystem::is_directory(build_dir_, error); - RAFT_EXPECTS(!error, - "Failed to inspect ACE build directory: %s (%s)", - build_dir_.c_str(), - error.message().c_str()); - RAFT_EXPECTS(is_directory, "ACE build path is not a directory: %s", build_dir_.c_str()); - } - - [[nodiscard]] std::string artifact_path(artifact which) const - { - return artifacts_[static_cast(which)].string(); - } - - void mark_artifact_created(artifact which) noexcept - { - artifacts_created_[static_cast(which)] = true; - } - - void commit() noexcept { committed_ = true; } - - void cleanup() noexcept - { - if (committed_) { return; } - - for (size_t i = artifacts_.size(); i > 0; --i) { - if (!artifacts_created_[i - 1]) { continue; } - - std::error_code error; - std::filesystem::remove(artifacts_[i - 1], error); - if (error) { - RAFT_LOG_WARN("ACE: Failed to remove build artifact %s: %s", - artifacts_[i - 1].c_str(), - error.message().c_str()); - } - } - - if (directory_created_by_this_build_) { - std::error_code error; - std::filesystem::remove(build_dir_, error); - if (error) { - RAFT_LOG_WARN("ACE: Failed to remove empty build directory %s: %s", - build_dir_.c_str(), - error.message().c_str()); - } - } - } - - private: - std::filesystem::path build_dir_; - std::array artifacts_; - std::array artifacts_created_{}; - bool directory_created_by_this_build_ = false; - bool committed_ = false; -}; - template void check_graph_degree(size_t& intermediate_degree, size_t& graph_degree, size_t dataset_size) { @@ -173,327 +70,6 @@ void check_graph_degree(size_t& intermediate_degree, size_t& graph_degree, size_ } } -// ACE: Get partition labels for partitioned approach -// TODO(julianmi): Use all neighbors APIs. -template -void ace_get_partition_labels( - raft::resources const& res, - raft::host_matrix_view dataset, - size_t dataset_dim, - raft::host_matrix_view partition_labels, - raft::host_matrix_view partition_histogram, - size_t min_partition_size, - double sampling_rate = 0.01) -{ - size_t dataset_size = dataset.extent(0); - size_t labels_size = partition_labels.extent(0); - size_t labels_dim = partition_labels.extent(1); - RAFT_EXPECTS(dataset_size == labels_size, "Dataset size must match partition labels extent"); - RAFT_EXPECTS(dataset_dim > 0, "Dataset dimension must be greater than 0"); - RAFT_EXPECTS(static_cast(dataset.extent(1)) >= dataset_dim, - "Dataset row extent (%zu) must be >= logical dimension (%zu)", - 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(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. - 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: n_samples: %lu", 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 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); - size_t report_interval = dataset_size / 10; - report_interval = (report_interval / chunk_size) * chunk_size; - report_interval = std::max(report_interval, chunk_size); - - 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); - if (i_base % report_interval == 0) { - RAFT_LOG_INFO("ACE: Processing chunk %lu / %lu (%.1f%%)", - i_base, - dataset_size, - static_cast(100 * i_base) / dataset_size); - } - - 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)); - } - } - 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); - - raft::copy(res, sub_distances, sub_distances_dev); - 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; - } - } -} - -// ACE: Check partition sizes for stable KNN graph construction -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) -{ - // Collect partition histogram statistics - size_t total_core_vectors = 0; - size_t total_augmented_vectors = 0; - size_t min_core_vectors = dataset_size; - size_t max_core_vectors = 0; - size_t min_augmented_vectors = dataset_size; - size_t max_augmented_vectors = 0; - size_t min_total_vectors = dataset_size; - size_t max_total_vectors = 0; - - for (size_t c = 0; c < n_partitions; c++) { - size_t core_count = partition_histogram(c, 0); - size_t augmented_count = partition_histogram(c, 1); - size_t total_count = core_count + augmented_count; - - if (total_count > 0) { - total_core_vectors += core_count; - total_augmented_vectors += augmented_count; - - min_core_vectors = std::min(min_core_vectors, core_count); - max_core_vectors = std::max(max_core_vectors, core_count); - min_augmented_vectors = std::min(min_augmented_vectors, augmented_count); - max_augmented_vectors = std::max(max_augmented_vectors, augmented_count); - min_total_vectors = std::min(min_total_vectors, total_count); - max_total_vectors = std::max(max_total_vectors, total_count); - } - } - - double avg_core_vectors = static_cast(total_core_vectors) / n_partitions; - double avg_augmented_vectors = static_cast(total_augmented_vectors) / n_partitions; - double avg_total_vectors = 2.0 * static_cast(dataset_size) / n_partitions; - double expected_avg_vectors = 2.0 * static_cast(dataset_size) / n_partitions; - - RAFT_LOG_INFO("ACE: Core vectors - Total: %lu, Avg: %.1f, Min: %lu, Max: %lu", - total_core_vectors, - avg_core_vectors, - min_core_vectors, - max_core_vectors); - RAFT_LOG_INFO("ACE: Augmented vectors - Total: %lu, Avg: %.1f, Min: %lu, Max: %lu", - total_augmented_vectors, - avg_augmented_vectors, - min_augmented_vectors, - max_augmented_vectors); - RAFT_LOG_INFO("ACE: Total per partition - Total: %lu, Avg: %.1f, Min: %lu, Max: %lu", - total_core_vectors + total_augmented_vectors, - avg_total_vectors, - min_total_vectors, - max_total_vectors); - - // Check for partition imbalance and issue warnings - size_t very_small_threshold = min_partition_size; - size_t very_large_threshold = static_cast(5.0 * expected_avg_vectors); - - for (size_t c = 0; c < n_partitions; c++) { - size_t total_count = partition_histogram(c, 0) + partition_histogram(c, 1); - - if (total_count > 0 && total_count < very_small_threshold) { - RAFT_LOG_WARN( - "ACE: Partition %lu is very small (%lu vectors, expected ~%.1f). This may affect graph " - "quality.", - c, - total_count, - expected_avg_vectors); - } else if (total_count > very_large_threshold) { - RAFT_LOG_WARN( - "ACE: Partition %lu is very large (%lu vectors, expected ~%.1f, threshold: %lu). This may " - "indicate imbalance and can lead to memory issues in restricted environments.", - c, - total_count, - expected_avg_vectors, - very_large_threshold); - } - } -} - -// ACE: Create forward/backward mappings between original and reordered vector IDs -// The in-memory path can be parallelized but the disk path requires ordering. -template -void ace_create_forward_and_backward_lists( - size_t dataset_size, - size_t n_partitions, - raft::host_matrix_view partition_labels, - raft::host_matrix_view partition_histogram, - raft::host_vector_view core_forward_mapping, - raft::host_vector_view core_backward_mapping, - raft::host_vector_view augmented_backward_mapping, - raft::host_vector_view core_partition_offsets, - raft::host_vector_view augmented_partition_offsets) -{ - core_partition_offsets(0) = 0; - augmented_partition_offsets(0) = 0; - for (size_t c = 1; c < n_partitions; c++) { - core_partition_offsets(c) = core_partition_offsets(c - 1) + partition_histogram(c - 1, 0); - augmented_partition_offsets(c) = - augmented_partition_offsets(c - 1) + partition_histogram(c - 1, 1); - } - - if (static_cast(core_forward_mapping.extent(0)) == 0) { - // 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"); -#pragma omp parallel for - for (size_t i = 0; i < dataset_size; i++) { - size_t core_partition_id = partition_labels(i, 0); - size_t core_id; -#pragma omp atomic capture - core_id = core_partition_offsets(core_partition_id)++; - 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; -#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; - } - } else { - // Disk path: all three mappings - RAFT_EXPECTS(static_cast(core_forward_mapping.extent(0)) == dataset_size, - "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"); - for (size_t i = 0; i < dataset_size; i++) { - size_t core_partition_id = partition_labels(i, 0); - size_t core_id; - core_id = core_partition_offsets(core_partition_id)++; - RAFT_EXPECTS(core_id < dataset_size, "Vector ID must be smaller than dataset_size"); - 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; - } - } - - // Restore idxptr arrays - for (size_t c = n_partitions; c > 0; c--) { - core_partition_offsets(c) = core_partition_offsets(c - 1); - augmented_partition_offsets(c) = augmented_partition_offsets(c - 1); - } - core_partition_offsets(0) = 0; - augmented_partition_offsets(0) = 0; -} - -// ACE: Materialize the core partition followed by its argument partition. -template -void ace_make_core_argument_partition( - size_t core_sub_dataset_size, - size_t augmented_sub_dataset_size, - size_t dataset_dim, - size_t partition_id, - raft::host_matrix_view dataset, - raft::host_vector_view core_backward_mapping, - raft::host_vector_view augmented_backward_mapping, - raft::host_vector_view core_partition_offsets, - raft::host_vector_view augmented_partition_offsets, - raft::host_matrix_view sub_dataset) -{ - const size_t vector_size_bytes = dataset_dim * sizeof(T); - - // Copy core partition vectors -#pragma omp parallel for - for (size_t j = 0; j < core_sub_dataset_size; j++) { - size_t i = core_backward_mapping(j + core_partition_offsets(partition_id)); - memcpy(&sub_dataset(j, 0), &dataset(i, 0), vector_size_bytes); - } - - // Copy augmented partition vectors (2nd closest 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)); - memcpy(&sub_dataset(j + core_sub_dataset_size, 0), &dataset(i, 0), vector_size_bytes); - } -} - template void build_knn_graph( raft::resources const& res, @@ -554,1680 +130,6 @@ void build_knn_graph(raft::resources const& res, } } -// 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, - graph_build_params::ace_params::knn_graph_build_algo knn_build_algo) - -> 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); - if (knn_build_algo == graph_build_params::ace_params::knn_graph_build_algo::IVF_PQ) { - 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); - } else { - auto knn_params = graph_build_params::brute_force_params{}; - knn_params.build_params.metric = metric; - 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, - graph_build_params::ace_params::knn_graph_build_algo knn_build_algo) - -> raft::host_matrix -{ - auto knn_graph = ace_build_partial_knn_graph( - res, core_argument_partition, intermediate_graph_degree, metric, knn_build_algo); - 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; -} - -// Original ACE mode: optimize each partial graph, including its local reverse-edge merge. -template -auto ace_build_optimized_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, - graph_build_params::ace_params::knn_graph_build_algo knn_build_algo, - bool guarantee_connectivity) -> raft::host_matrix -{ - auto knn_graph = ace_build_partial_knn_graph( - res, core_argument_partition, intermediate_graph_degree, metric, knn_build_algo); - auto partial_graph = - raft::make_host_matrix(core_argument_partition.extent(0), graph_degree); - graph::optimize(res, knn_graph.view(), partial_graph.view(), guarantee_connectivity); - 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, - size_t augmented_sub_dataset_size, - size_t graph_degree, - size_t partition_id, - raft::host_matrix_view sub_search_graph, - raft::host_matrix_view search_graph, - raft::host_vector_view core_partition_offsets, - raft::host_vector_view augmented_partition_offsets, - raft::host_vector_view core_backward_mapping, - raft::host_vector_view augmented_backward_mapping) -{ -#pragma omp parallel for - for (size_t i = 0; i < core_sub_dataset_size; i++) { - // Map row index from local → reordered → original - size_t i_reordered = i + core_partition_offsets(partition_id); - size_t i_original = core_backward_mapping(i_reordered); - - for (size_t k = 0; k < graph_degree; k++) { - size_t j = sub_search_graph(i, k); - size_t j_original; - - if (j < core_sub_dataset_size) { - // core partition neighbor: local → core reordered → original - size_t j_reordered = j + core_partition_offsets(partition_id); - j_original = core_backward_mapping(j_reordered); - } else { - // Augmented partition neighbor: local → augmented reordered → original - size_t j_augmented = j - core_sub_dataset_size; - j_original = - augmented_backward_mapping(j_augmented + augmented_partition_offsets(partition_id)); - } - search_graph(i_original, k) = j_original; - } - } -} - -// ACE: Adjust IDs into global reordered IDs on the device for the disk version. -template -__global__ void ace_adjust_sub_graph_ids_disk_kernel(const GraphIdxT* sub_search_graph, - IdxT* adjusted_search_graph, - size_t graph_edges, - size_t core_sub_dataset_size, - IdxT core_partition_offset, - const IdxT* augmented_reordered_ids) -{ - const size_t edge_id = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - if (edge_id >= graph_edges) { return; } - - const GraphIdxT neighbor = sub_search_graph[edge_id]; - if (static_cast(neighbor) < core_sub_dataset_size) { - adjusted_search_graph[edge_id] = static_cast(neighbor) + core_partition_offset; - } else { - adjusted_search_graph[edge_id] = - augmented_reordered_ids[static_cast(neighbor) - core_sub_dataset_size]; - } -} - -template -void ace_adjust_sub_graph_ids_disk( - raft::resources const& res, - size_t core_sub_dataset_size, - size_t augmented_sub_dataset_size, - size_t graph_degree, - size_t partition_id, - raft::device_matrix_view sub_search_graph, - raft::device_matrix_view adjusted_search_graph, - raft::host_vector_view core_partition_offsets, - raft::host_vector_view augmented_partition_offsets, - raft::host_vector_view augmented_backward_mapping, - raft::host_vector_view core_forward_mapping) -{ - RAFT_EXPECTS(static_cast(sub_search_graph.extent(0)) >= core_sub_dataset_size, - "ACE: source graph has fewer rows than the core partition"); - RAFT_EXPECTS(static_cast(sub_search_graph.extent(1)) == graph_degree, - "ACE: source graph degree does not match the requested graph degree"); - RAFT_EXPECTS(static_cast(adjusted_search_graph.extent(0)) == core_sub_dataset_size && - static_cast(adjusted_search_graph.extent(1)) == graph_degree, - "ACE: adjusted graph shape does not match the core partition"); - - const size_t augmented_offset = augmented_partition_offsets(partition_id); - RAFT_EXPECTS(augmented_offset + augmented_sub_dataset_size <= - static_cast(augmented_backward_mapping.extent(0)), - "ACE: augmented partition exceeds the backward mapping"); - - auto augmented_reordered_ids = raft::make_host_vector(augmented_sub_dataset_size); -#pragma omp parallel for - for (size_t i = 0; i < augmented_sub_dataset_size; i++) { - const size_t original_id = augmented_backward_mapping(augmented_offset + i); - augmented_reordered_ids(i) = core_forward_mapping(original_id); - } - - auto augmented_reordered_ids_dev = - raft::make_device_vector(res, augmented_sub_dataset_size); - raft::copy(res, augmented_reordered_ids_dev.view(), augmented_reordered_ids.view()); - - const size_t graph_edges = core_sub_dataset_size * graph_degree; - if (graph_edges == 0) { return; } - - constexpr uint32_t block_size = 256; - const dim3 grid_size(raft::ceildiv(graph_edges, static_cast(block_size))); - ace_adjust_sub_graph_ids_disk_kernel<<>>( - sub_search_graph.data_handle(), - adjusted_search_graph.data_handle(), - graph_edges, - core_sub_dataset_size, - core_partition_offsets(partition_id), - augmented_reordered_ids_dev.data_handle()); - RAFT_CUDA_TRY(cudaPeekAtLastError()); - raft::resource::sync_stream(res); -} - -// ACE: Reorder dataset based on partition assignments and store to disk -// Writes two files: reordered_dataset.npy (core partitions) and augmented_dataset.npy (secondary -// partitions). Uses buffered writes optimized for NVMe storage. -template -void ace_reorder_and_store_dataset( - raft::resources const& res, - const std::string& build_dir, - raft::host_matrix_view dataset, - size_t dataset_dim, - raft::host_matrix_view partition_labels, - raft::host_matrix_view partition_histogram, - raft::host_vector_view core_backward_mapping, - raft::host_vector_view core_partition_offsets, - raft::host_vector_view augmented_partition_offsets, - cuvs::util::file_descriptor& reordered_fd, - cuvs::util::file_descriptor& augmented_fd, - cuvs::util::file_descriptor& mapping_fd, - size_t reordered_header_size, - size_t augmented_header_size, - size_t mapping_header_size) -{ - auto start = std::chrono::high_resolution_clock::now(); - - size_t dataset_size = dataset.extent(0); - size_t n_partitions = partition_histogram.extent(0); - RAFT_EXPECTS(dataset_dim > 0, "Dataset dimension must be greater than 0"); - RAFT_EXPECTS(static_cast(dataset.extent(1)) >= dataset_dim, - "Dataset row extent (%zu) must be >= logical dimension (%zu)", - static_cast(dataset.extent(1)), - dataset_dim); - - RAFT_LOG_DEBUG( - "ACE: Reordering and storing dataset to disk (%lu vectors, %lu dimensions, %lu partitions)", - dataset_size, - dataset_dim, - n_partitions); - - // Calculate total sizes for pre-allocation - size_t total_core_vectors = 0; - size_t total_augmented_vectors = 0; - size_t max_core_vectors = 0; - size_t max_augmented_vectors = 0; - for (size_t p = 0; p < n_partitions; p++) { - total_core_vectors += partition_histogram(p, 0); - total_augmented_vectors += partition_histogram(p, 1); - max_core_vectors = std::max(max_core_vectors, partition_histogram(p, 0)); - max_augmented_vectors = std::max(max_augmented_vectors, partition_histogram(p, 1)); - } - 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"); - - // Pre-allocate file space for better performance - const size_t vector_size = dataset_dim * sizeof(T); - size_t reordered_file_size = total_core_vectors * vector_size; - size_t augmented_file_size = total_augmented_vectors * vector_size; - - RAFT_LOG_DEBUG("ACE: Reordered dataset: %lu core vectors (%.2f GiB)", - total_core_vectors, - reordered_file_size / (1024.0 * 1024.0 * 1024.0)); - RAFT_LOG_DEBUG("ACE: Augmented dataset: %lu secondary vectors (%.2f GiB)", - total_augmented_vectors, - augmented_file_size / (1024.0 * 1024.0 * 1024.0)); - - // Calculate partition start offsets for reordered and augmented datasets - auto core_partition_starts = raft::make_host_vector(n_partitions + 1); - memset(core_partition_starts.data_handle(), 0, (n_partitions + 1) * sizeof(size_t)); - auto augmented_partition_starts = raft::make_host_vector(n_partitions + 1); - memset(augmented_partition_starts.data_handle(), 0, (n_partitions + 1) * sizeof(size_t)); - auto core_partition_current = raft::make_host_vector(n_partitions); - memset(core_partition_current.data_handle(), 0, n_partitions * sizeof(size_t)); - auto augmented_partition_current = raft::make_host_vector(n_partitions); - memset(augmented_partition_current.data_handle(), 0, n_partitions * sizeof(size_t)); - - for (size_t p = 0; p < n_partitions; p++) { - core_partition_starts(p + 1) = core_partition_starts(p) + partition_histogram(p, 0); - augmented_partition_starts(p + 1) = augmented_partition_starts(p) + partition_histogram(p, 1); - } - - const size_t free_memory = cuvs::util::get_free_host_memory(); - // Conservatively allocate 50% of free memory per partition. Accounts for internal buffers and - // overhead. - // TODO: Adjust overhead if needed. - const size_t memory_per_partition = 0.5 * free_memory / (n_partitions * 2); - size_t disk_write_size = raft::bound_by_power_of_two(memory_per_partition); - // 64MB should be enough to saturate typical NVMe SSDs. - disk_write_size = std::min(disk_write_size, 64 * 1024 * 1024); - size_t vectors_per_buffer = std::max(64, disk_write_size / vector_size); - - RAFT_LOG_DEBUG("ACE: Reorder buffers: %lu vectors per buffer (%.2f MiB)", - vectors_per_buffer, - to_mib(vectors_per_buffer * vector_size)); - - std::vector> core_buffers; - std::vector> augmented_buffers; - auto core_buffer_counts = raft::make_host_vector(n_partitions); - auto augmented_buffer_counts = raft::make_host_vector(n_partitions); - - core_buffers.reserve(n_partitions); - augmented_buffers.reserve(n_partitions); - - for (size_t p = 0; p < n_partitions; p++) { - core_buffers.emplace_back(raft::make_host_matrix(vectors_per_buffer, dataset_dim)); - augmented_buffers.emplace_back( - raft::make_host_matrix(vectors_per_buffer, dataset_dim)); - core_buffer_counts(p) = 0; - augmented_buffer_counts(p) = 0; - } - auto flush_core_buffer = [&](size_t partition_id) { - const size_t count = core_buffer_counts(partition_id); - if (count > 0) { - const size_t bytes_to_write = count * vector_size; - const size_t file_offset = - (core_partition_starts(partition_id) + core_partition_current(partition_id)) * vector_size + - reordered_header_size; - - cuvs::util::write_large_file( - reordered_fd, core_buffers[partition_id].data_handle(), bytes_to_write, file_offset); - - core_partition_current(partition_id) += count; - core_buffer_counts(partition_id) = 0; - } - }; - - auto flush_augmented_buffer = [&](size_t partition_id) { - const size_t count = augmented_buffer_counts(partition_id); - if (count > 0) { - const size_t bytes_to_write = count * vector_size; - const size_t file_offset = - (augmented_partition_starts(partition_id) + augmented_partition_current(partition_id)) * - vector_size + - augmented_header_size; - - cuvs::util::write_large_file( - augmented_fd, augmented_buffers[partition_id].data_handle(), bytes_to_write, file_offset); - - augmented_partition_current(partition_id) += count; - augmented_buffer_counts(partition_id) = 0; - } - }; - - size_t vectors_processed = 0; - const size_t log_interval = std::max(dataset_size / 10, size_t(1)); - 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); - - // Add vector to core partition buffer - size_t core_buffer_row = core_buffer_counts(core_partition); - memcpy( - &core_buffers[core_partition](core_buffer_row, 0), &dataset(i, 0), dataset_dim * sizeof(T)); - core_buffer_counts(core_partition)++; - - // Flush core buffer if full - if (core_buffer_counts(core_partition) >= vectors_per_buffer) { - flush_core_buffer(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_augmented_buffer(secondary_partition); - } - - vectors_processed++; - if (vectors_processed % log_interval == 0) { - RAFT_LOG_INFO("ACE: Processed %lu/%lu vectors (%.1f%%)", - vectors_processed, - dataset_size, - 100.0 * vectors_processed / dataset_size); - } - } - - // Flush all remaining buffers - RAFT_LOG_DEBUG("ACE: Flushing remaining buffers..."); -#pragma omp parallel sections - { -#pragma omp section - { - for (size_t p = 0; p < n_partitions; p++) { - flush_core_buffer(p); - } - } -#pragma omp section - { - for (size_t p = 0; p < n_partitions; p++) { - flush_augmented_buffer(p); - } - } - } - - const size_t mapping_file_size = dataset_size * sizeof(IdxT); - cuvs::util::write_large_file( - mapping_fd, core_backward_mapping.data_handle(), mapping_file_size, mapping_header_size); - - auto end = std::chrono::high_resolution_clock::now(); - auto elapsed_ms = std::chrono::duration_cast(end - start).count(); - - // Calculate total bytes written - size_t total_bytes_written = reordered_file_size + augmented_file_size + mapping_file_size; - double throughput_mb_s = - elapsed_ms > 0 ? to_mib(total_bytes_written) / (elapsed_ms / 1000.0) : 0.0; - - RAFT_LOG_INFO( - "ACE: Dataset (%.2f GiB reordered, %.2f GiB augmented, %.2f GiB mapping) reordering completed " - "in %ld ms (%.1f MiB/s)", - reordered_file_size / (1024.0 * 1024.0 * 1024.0), - augmented_file_size / (1024.0 * 1024.0 * 1024.0), - mapping_file_size / (1024.0 * 1024.0 * 1024.0), - elapsed_ms, - throughput_mb_s); -} - -// ACE: Load partition dataset and augmented dataset from disk -template -void ace_load_partition_dataset_from_disk( - size_t partition_id, - size_t dataset_dim, - raft::host_matrix_view partition_histogram, - raft::host_vector_view core_partition_offsets, - raft::host_vector_view augmented_partition_offsets, - const cuvs::util::file_descriptor& reordered_fd, - const cuvs::util::file_descriptor& augmented_fd, - size_t reordered_header_size, - size_t augmented_header_size, - raft::mdspan, raft::row_major, Accessor> sub_dataset) -{ - RAFT_LOG_DEBUG("ACE: Loading partition %lu dataset from disk", partition_id); - - const std::string reordered_path = reordered_fd.get_path(); - const std::string augmented_path = augmented_fd.get_path(); - T* sub_dataset_ptr = sub_dataset.data_handle(); - RAFT_EXPECTS(sub_dataset_ptr != nullptr, "ACE: sub-dataset destination must not be null"); - RAFT_EXPECTS(reordered_fd.is_valid() && !reordered_path.empty(), - "ACE: reordered dataset file descriptor is not valid"); - RAFT_EXPECTS(augmented_fd.is_valid() && !augmented_path.empty(), - "ACE: augmented dataset file descriptor is not valid"); - RAFT_EXPECTS(partition_id < static_cast(partition_histogram.extent(0)), - "ACE: partition id is out of range"); - RAFT_EXPECTS(partition_id < static_cast(core_partition_offsets.extent(0)), - "ACE: core partition offset is out of range"); - RAFT_EXPECTS(partition_id < static_cast(augmented_partition_offsets.extent(0)), - "ACE: augmented partition offset is out of range"); - - size_t core_size = partition_histogram(partition_id, 0); - size_t augmented_size = partition_histogram(partition_id, 1); - size_t total_size = core_size + augmented_size; - - RAFT_EXPECTS(static_cast(sub_dataset.extent(0)) == total_size, - "ACE: sub-dataset rows (%zu) must match partition size (%zu)", - static_cast(sub_dataset.extent(0)), - total_size); - RAFT_EXPECTS(static_cast(sub_dataset.extent(1)) == dataset_dim, - "ACE: sub-dataset columns (%zu) must match dataset dimensions (%zu)", - static_cast(sub_dataset.extent(1)), - dataset_dim); - - RAFT_LOG_DEBUG("ACE: Partition %lu: %lu core + %lu augmented = %lu total vectors", - partition_id, - core_size, - augmented_size, - core_size + augmented_size); - - const size_t vector_size = dataset_dim * sizeof(T); - const size_t core_file_offset = - reordered_header_size + static_cast(core_partition_offsets(partition_id)) * vector_size; - const size_t augmented_file_offset = - augmented_header_size + - static_cast(augmented_partition_offsets(partition_id)) * vector_size; - - RAFT_LOG_DEBUG("ACE: Core file offset: %lu bytes, Augmented file offset: %lu bytes", - core_file_offset, - augmented_file_offset); - - const size_t core_bytes = core_size * vector_size; - const size_t augmented_bytes = augmented_size * vector_size; - T* augmented_dest = sub_dataset_ptr + (core_size * dataset_dim); - - auto expect_complete_read = [partition_id](const char* name, size_t expected, size_t actual) { - RAFT_EXPECTS(actual == expected, - "ACE: Short %s read for partition %lu: expected %zu bytes, got %zu", - name, - partition_id, - expected, - actual); - }; - - if (core_bytes > 0 && augmented_bytes > 0) { - RAFT_LOG_DEBUG("ACE: Reading %lu core vectors from offset %lu", core_size, core_file_offset); - RAFT_LOG_DEBUG( - "ACE: Reading %lu augmented vectors from offset %lu", augmented_size, augmented_file_offset); - auto reordered_handle = - cuvs::util::detail::open_kvikio_file_for_ace_io(reordered_path, "r", sub_dataset_ptr); - auto augmented_handle = - cuvs::util::detail::open_kvikio_file_for_ace_io(augmented_path, "r", augmented_dest); - auto core_future = reordered_handle.pread(sub_dataset_ptr, core_bytes, core_file_offset); - auto augmented_future = - augmented_handle.pread(augmented_dest, augmented_bytes, augmented_file_offset); - std::exception_ptr read_exception = nullptr; - size_t core_read = 0; - size_t augmented_read = 0; - try { - core_read = core_future.get(); - } catch (...) { - read_exception = std::current_exception(); - } - try { - augmented_read = augmented_future.get(); - } catch (...) { - if (!read_exception) { read_exception = std::current_exception(); } - } - if (read_exception) { std::rethrow_exception(read_exception); } - expect_complete_read("core", core_bytes, core_read); - expect_complete_read("augmented", augmented_bytes, augmented_read); - } else if (core_bytes > 0) { - RAFT_LOG_DEBUG("ACE: Reading %lu core vectors from offset %lu", core_size, core_file_offset); - auto reordered_handle = - cuvs::util::detail::open_kvikio_file_for_ace_io(reordered_path, "r", sub_dataset_ptr); - expect_complete_read( - "core", - core_bytes, - reordered_handle.pread(sub_dataset_ptr, core_bytes, core_file_offset).get()); - } else if (augmented_bytes > 0) { - RAFT_LOG_DEBUG( - "ACE: Reading %lu augmented vectors from offset %lu", augmented_size, augmented_file_offset); - auto augmented_handle = - cuvs::util::detail::open_kvikio_file_for_ace_io(augmented_path, "r", augmented_dest); - expect_complete_read( - "augmented", - augmented_bytes, - augmented_handle.pread(augmented_dest, augmented_bytes, augmented_file_offset).get()); - } -} - -// Memory requirements for ACE operation -struct ace_memory_requirements { - size_t partition_labels_size; - size_t id_mapping_size; - size_t sub_dataset_size; - size_t sub_graph_size; - size_t cagra_graph_size; - size_t total_size; - size_t available_host_memory; - size_t available_gpu_memory; -}; - -// Amount of host memory that can be used for the build -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; - -// Check if disk mode should be used for ACE based on memory constraints -template -bool ace_check_use_disk_mode(raft::resources const& res, - bool use_disk, - const std::string& build_dir, - size_t dataset_size, - size_t dataset_dim, - size_t n_partitions, - size_t intermediate_degree, - size_t graph_degree, - std::optional max_host_memory_gb, - std::optional max_gpu_memory_gb, - bool guarantee_connectivity, - ace_memory_requirements& mem) -{ - 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() && - host_memory.cgroup_reclaimable_file.has_value() && - host_memory.cgroup_working_set.has_value()) { - const size_t cgroup_headroom = *host_memory.cgroup_working_set < *host_memory.cgroup_limit - ? *host_memory.cgroup_limit - *host_memory.cgroup_working_set - : 0; - RAFT_LOG_INFO( - "ACE: Cgroup host memory limit: %.2f GiB, current: %.2f GiB, reclaimable file cache: %.2f " - "GiB, working set: %.2f GiB, headroom: %.2f GiB; system available: %.2f GiB, effective " - "available: %.2f GiB", - to_gib(*host_memory.cgroup_limit), - to_gib(*host_memory.cgroup_current), - to_gib(*host_memory.cgroup_reclaimable_file), - to_gib(*host_memory.cgroup_working_set), - to_gib(cgroup_headroom), - to_gib(host_memory.system_available), - to_gib(host_memory.available)); - } - - // Use overridden memory limits if provided (> 0), otherwise query actual system memory - if (max_host_memory_gb.has_value() && max_host_memory_gb.value() > 0) { - auto actual_available_host_memory = host_memory.available; - auto configured_host_memory = static_cast(max_host_memory_gb.value() * (1ULL << 30)); - if (actual_available_host_memory < configured_host_memory) { - RAFT_LOG_WARN( - "ACE: Actual host memory (%.2f GiB) is less than configured limit (%.2f GiB). " - "Using actual host memory.", - to_gib(actual_available_host_memory), - to_gib(configured_host_memory)); - mem.available_host_memory = actual_available_host_memory; - } else { - RAFT_LOG_INFO("ACE: Using overridden host memory limit: %.2f GiB", - max_host_memory_gb.value()); - mem.available_host_memory = configured_host_memory; - } - } else { - mem.available_host_memory = host_memory.available; - } - size_t sub_partition_size = - static_cast(imbalance_factor * vector_expansion_factor * - raft::div_rounding_up_safe(dataset_size, n_partitions)); - auto [opt_host_ws_total, opt_dev_ws_total, opt_host_ws_fixed, opt_dev_ws_fixed] = - helpers::optimize_workspace_size( - sub_partition_size, graph_degree, intermediate_degree, sizeof(IdxT), guarantee_connectivity); - - // Optimistic memory model: focus on largest arrays, assumes all partitions are of equal size - // For memory path: - // - Partition labels (core + augmented): vector_expansion_factor * dataset_size * sizeof(IdxT) - // - Backward ID mapping arrays (core + augmented): vector_expansion_factor * dataset_size * - // sizeof(IdxT) - // - Avg. per-partition dataset: vector_expansion_factor * (dataset_size / n_partitions) * - // dataset_dim * sizeof(T) - // - Avg. per-partition graph during build: vector_expansion_factor * (dataset_size / - // n_partitions) * (intermediate + final) * sizeof(IdxT) - // - Final assembled graph: dataset_size * graph_degree * sizeof(IdxT) - mem.partition_labels_size = vector_expansion_factor * dataset_size * sizeof(IdxT); - mem.id_mapping_size = vector_expansion_factor * dataset_size * sizeof(IdxT); - mem.sub_dataset_size = sub_partition_size * dataset_dim * sizeof(T); - mem.sub_graph_size = sub_partition_size * (intermediate_degree + graph_degree) * sizeof(IdxT); - 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; - - RAFT_LOG_INFO("ACE: Estimated host memory required: %.2f GiB, available: %.2f GiB", - to_gib(mem.total_size), - to_gib(mem.available_host_memory)); - - bool host_memory_limited = - static_cast(usable_cpu_memory_fraction * mem.available_host_memory) < mem.total_size; - - // GPU is mostly limited by the index size (update_graph() in the end of this routine). - // Check if GPU has enough memory for the final graph or use disk mode instead. - // TODO: Extend model or use managed memory if running out of GPU memory. - if (max_gpu_memory_gb.has_value() && max_gpu_memory_gb.value() > 0) { - auto actual_available_gpu_memory = rmm::available_device_memory().second; - auto configured_gpu_memory = static_cast(max_gpu_memory_gb.value() * (1ULL << 30)); - if (actual_available_gpu_memory < configured_gpu_memory) { - RAFT_LOG_WARN( - "ACE: Actual GPU memory (%zu GiB) is less than configured limit (%zu GiB). " - "Using actual GPU memory.", - to_gib(actual_available_gpu_memory), - to_gib(configured_gpu_memory)); - mem.available_gpu_memory = actual_available_gpu_memory; - } else { - RAFT_LOG_INFO("ACE: Using overridden GPU memory limit: %.2f GiB", max_gpu_memory_gb.value()); - mem.available_gpu_memory = configured_gpu_memory; - } - } else { - mem.available_gpu_memory = rmm::available_device_memory().second; - } - - // what we need is maximum of: - // * IVF-PQ on partition (sub_dataset_size, uncompressed upper bound) - // * optimize workspace (opt_dev_ws_total) - // + some extra workspace (IVF-PQ search, ...) - size_t extra_gpu_workspace_size = raft::resource::get_workspace_total_bytes(res); - size_t gpu_memory_required = - std::max(mem.sub_dataset_size, opt_dev_ws_total) + extra_gpu_workspace_size; - - bool gpu_memory_limited = mem.available_gpu_memory < gpu_memory_required; - - RAFT_LOG_INFO("ACE: Estimated GPU memory required: %.2f GiB, available: %.2f GiB", - to_gib(gpu_memory_required), - to_gib(mem.available_gpu_memory)); - - bool use_disk_mode = use_disk || host_memory_limited || gpu_memory_limited; - if (use_disk_mode) { RAFT_EXPECTS(!build_dir.empty(), "ACE build directory must not be empty"); } - - if (host_memory_limited && gpu_memory_limited) { - RAFT_LOG_INFO( - "ACE: Graph does not fit in host and GPU memory. Using disk-mode with temporary storage %s", - build_dir.c_str()); - } else if (host_memory_limited) { - RAFT_LOG_INFO( - "ACE: Graph does not fit in host memory. Using disk-mode with temporary storage %s", - build_dir.c_str()); - } else if (gpu_memory_limited) { - RAFT_LOG_INFO( - "ACE: Graph does not fit in GPU memory. Using disk-mode with temporary storage %s", - build_dir.c_str()); - } else if (use_disk) { - RAFT_LOG_INFO( - "ACE: Graph fits in host and GPU memory but disk mode is forced. Using disk-mode with " - "temporary storage %s", - build_dir.c_str()); - } else { - RAFT_LOG_INFO("ACE: Graph fits in host and GPU memory. Using in-memory mode."); - } - - return use_disk_mode; -} - -// Resolve the ACE partition count while preserving 0 as the auto-selection sentinel. -inline size_t ace_resolve_partition_count(size_t n_partitions) -{ - if (n_partitions == 0) { return 2; } - if (n_partitions == 1) { - RAFT_LOG_WARN( - "ACE: Requested 1 partition; adjusted to 2 before applying partitioning heuristics"); - return 2; - } - return n_partitions; -} - -// Validate the structural ACE partition-count invariants required by the labeler. -inline void ace_validate_partition_count(size_t n_partitions, - size_t dataset_size, - bool adjusted_for_memory = false) -{ - RAFT_EXPECTS(n_partitions <= dataset_size, - adjusted_for_memory - ? "ACE: configured memory limit is unsatisfiable because the requested partition " - "count cannot exceed dataset size" - : "ACE: number of partitions cannot exceed dataset size"); -} - -// Validate and adjust partitions for disk mode memory requirements -template -void ace_validate_disk_mode_partitions(raft::resources const& res, - size_t& n_partitions, - size_t dataset_size, - size_t dataset_dim, - size_t intermediate_degree, - size_t graph_degree, - bool guarantee_connectivity, - ace_memory_requirements& mem) -{ - // 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) - // - ID mapping arrays (core + augmented): vector_expansion_factor * dataset_size * sizeof(IdxT) - // - Avg. per-partition dataset during processing: vector_expansion_factor * (dataset_size / - // n_partitions) * dataset_dim * sizeof(T) - // - Avg. per-partition graph during build: vector_expansion_factor * (dataset_size / - // n_partitions) * (intermediate + final) * sizeof(IdxT) - - size_t original_n_partitions = n_partitions; - size_t host_suggested_partitions = n_partitions; - size_t gpu_suggested_partitions = n_partitions; - bool host_memory_insufficient = false; - bool gpu_memory_insufficient = false; - - // Compute optimize workspace requirements - size_t sub_partition_size = - static_cast(imbalance_factor * vector_expansion_factor * - raft::div_rounding_up_safe(dataset_size, n_partitions)); - auto [host_workspace_size_total, - gpu_workspace_size_total, - host_workspace_size_fixed, - gpu_workspace_size_fixed] = - helpers::optimize_workspace_size( - sub_partition_size, graph_degree, intermediate_degree, sizeof(IdxT), guarantee_connectivity); - - // 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 + - host_workspace_size_total; - - if (static_cast(usable_cpu_memory_fraction * mem.available_host_memory) < - disk_mode_host_required) { - host_memory_insufficient = true; - RAFT_LOG_WARN( - "ACE: Host memory insufficient for disk mode. Required: %.2f GiB, available: %.2f GiB. " - "Per-partition breakdown: dataset %.2f GiB, graph %.2f GiB, workspace %.2f GiB", - to_gib(disk_mode_host_required), - to_gib(mem.available_host_memory), - to_gib(mem.sub_dataset_size), - to_gib(mem.sub_graph_size), - to_gib(host_workspace_size_total)); - - // Calculate suggested number of partitions for host memory - size_t disk_mode_host_static = - mem.partition_labels_size + mem.id_mapping_size + host_workspace_size_fixed; - size_t disk_mode_host_dynamic = disk_mode_host_required - disk_mode_host_static; - double available_for_scaling = - usable_cpu_memory_fraction * mem.available_host_memory - disk_mode_host_static; - - RAFT_EXPECTS(available_for_scaling > 0, - "ACE: Host memory insufficient even for constant overhead (labels + id_mapping + " - "static workspace). " - "Required: %.2f GiB, available: %.2f GiB", - to_gib(disk_mode_host_static), - to_gib(usable_cpu_memory_fraction * mem.available_host_memory)); - host_suggested_partitions = - static_cast(std::ceil(disk_mode_host_dynamic * n_partitions / available_for_scaling)); - // Ensure we always increase partitions (current count is insufficient by definition) - host_suggested_partitions = std::max(host_suggested_partitions, n_partitions + 1); - } - - // Check GPU memory requirements in disk mode - // GPU memory model for disk mode (per-partition processing): - // * IVF-PQ on partition (mem.sub_dataset_size) (compressed?) - // * optimize workspace (gpu_workspace_size_total) - // + some extra workspace (IVF-PQ search, ...) - size_t extra_gpu_workspace_size = raft::resource::get_workspace_total_bytes(res); - size_t disk_mode_gpu_required = - std::max(mem.sub_dataset_size, gpu_workspace_size_total) + extra_gpu_workspace_size; - - if (mem.available_gpu_memory < disk_mode_gpu_required) { - gpu_memory_insufficient = true; - RAFT_LOG_WARN( - "ACE: GPU memory insufficient for per-partition processing. Required: %.2f GiB, " - "available: %.2f GiB. Per-partition breakdown: dataset %.2f GiB, workspace %.2f GiB", - to_gib(disk_mode_gpu_required), - to_gib(mem.available_gpu_memory), - to_gib(mem.sub_dataset_size), - to_gib(gpu_workspace_size_total + extra_gpu_workspace_size)); - - size_t disk_mode_gpu_static = gpu_workspace_size_fixed + extra_gpu_workspace_size; - size_t disk_mode_gpu_dynamic = disk_mode_gpu_required - disk_mode_gpu_static; - double available_for_scaling = mem.available_gpu_memory - disk_mode_gpu_static; - - RAFT_EXPECTS(available_for_scaling > 0, - "ACE: GPU memory insufficient even for constant overhead. Required: %.2f GiB, " - "available: %.2f GiB", - to_gib(disk_mode_gpu_static), - to_gib(mem.available_gpu_memory)); - - gpu_suggested_partitions = - static_cast(std::ceil(disk_mode_gpu_dynamic * n_partitions / available_for_scaling)); - gpu_suggested_partitions = std::max(gpu_suggested_partitions, n_partitions + 1); - } - - // Auto-adjust to the maximum of host and GPU requirements - if (host_memory_insufficient || gpu_memory_insufficient) { - size_t new_n_partitions = std::max(host_suggested_partitions, gpu_suggested_partitions); - - RAFT_LOG_WARN( - "ACE: Automatically increasing number of partitions from %zu to %zu to satisfy memory " - "constraints.%s%s", - original_n_partitions, - new_n_partitions, - host_memory_insufficient - ? " Host memory requires >= " + std::to_string(host_suggested_partitions) + " partitions." - : "", - gpu_memory_insufficient - ? " GPU memory requires >= " + std::to_string(gpu_suggested_partitions) + " partitions." - : ""); - - n_partitions = new_n_partitions; - - size_t new_sub_partition_size = - static_cast(imbalance_factor * vector_expansion_factor * - raft::div_rounding_up_safe(dataset_size, n_partitions)); - auto [new_opt_host_ws, new_opt_dev_ws, new_opt_host_ws_fixed, new_opt_dev_ws_fixed] = - helpers::optimize_workspace_size(new_sub_partition_size, - graph_degree, - intermediate_degree, - sizeof(IdxT), - guarantee_connectivity); - - mem.sub_dataset_size = new_sub_partition_size * dataset_dim * sizeof(T); - 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; - - RAFT_LOG_INFO( - "ACE: Updated per-partition memory estimates: dataset %.2f GiB, graph %.2f GiB, " - "host workspace %.2f GiB, GPU workspace %.2f GiB", - to_gib(mem.sub_dataset_size), - to_gib(mem.sub_graph_size), - to_gib(new_opt_host_ws), - to_gib(new_opt_dev_ws + extra_gpu_workspace_size)); - } -} - -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 - 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; - -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; } - - 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::sort(candidate_sources.begin(), candidate_sources.end()); - - 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)); - - 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_sets(n_partitions); - for (size_t source = 0; source < dataset_size; ++source) { - auto const source_partition = static_cast(partition_labels(source, 0)); - 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)); - if (source_partition != destination_partition) { - candidate_sets[destination_partition].insert(static_cast(source)); - } - } - } - - for (size_t partition_id = 0; partition_id < n_partitions; ++partition_id) { - if (candidate_sets[partition_id].empty()) { continue; } - 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); - std::vector candidate_sources(candidate_sets[partition_id].begin(), - candidate_sets[partition_id].end()); - 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 -// 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. -// In disk mode, the graph is stored in build_dir and dataset is reordered on disk. -// The returned index is not usable for search. Use the created files for search instead. -template - requires cuvs::neighbors::is_host_dataset_view_v -auto build_ace(raft::resources const& res, const index_params& params, DatasetViewT const& dataset) - -> cuvs::neighbors::cagra::index -{ - // Extract ACE parameters from graph_build_params - RAFT_EXPECTS( - std::holds_alternative(params.graph_build_params), - "ACE build requires graph_build_params to be set to ace_params"); - - auto ace_params = std::get(params.graph_build_params); - size_t npartitions = ace_params.npartitions; - size_t ef_construction = ace_params.ef_construction; - std::string build_dir = ace_params.build_dir; - bool use_disk = ace_params.use_disk; - - common::nvtx::range function_scope( - "cagra::detail::build_ace(%zu, %zu, %zu)", - params.intermediate_graph_degree, - params.graph_degree, - npartitions); - - auto dataset_view = dataset.view(); - size_t dataset_size = dataset.n_rows(); - size_t dataset_dim = dataset.dim(); - - RAFT_EXPECTS(dataset_size > 0, "ACE: Dataset must not be empty"); - if (dataset_size < 1000) { - RAFT_LOG_WARN("ACE: Very small dataset size (%zu), consider using regular CAGRA build instead.", - dataset_size); - } - RAFT_EXPECTS(dataset_dim > 0, "ACE: Dataset dimension must be greater than 0"); - RAFT_EXPECTS(params.intermediate_graph_degree > 0, - "ACE: Intermediate graph degree must be greater than 0"); - RAFT_EXPECTS(params.graph_degree > 0, "ACE: Graph degree must be greater than 0"); - - size_t n_partitions = ace_resolve_partition_count(npartitions); - - 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; - if (n_partitions < 2) { - RAFT_LOG_WARN( - "ACE: Reduced number of partitions to the minimum of 2 to avoid tiny partitions. Consider " - "using regular CAGRA build instead."); - n_partitions = 2; - } else { - RAFT_LOG_WARN("ACE: Reduced number of partitions to %zu to avoid tiny partitions", - n_partitions); - } - } - - auto total_start = std::chrono::high_resolution_clock::now(); - RAFT_LOG_INFO("ACE: Starting partitioned CAGRA build with %zu partitions", n_partitions); - - size_t intermediate_degree = params.intermediate_graph_degree; - size_t graph_degree = params.graph_degree; - - ace_disk_workspace workspace(build_dir); - - try { - check_graph_degree(intermediate_degree, graph_degree, dataset_size); - - // Check if disk mode should be used based on memory constraints - ace_memory_requirements mem; - bool use_disk_mode = ace_check_use_disk_mode(res, - use_disk, - build_dir, - dataset_size, - dataset_dim, - n_partitions, - intermediate_degree, - graph_degree, - ace_params.max_host_memory_gb, - ace_params.max_gpu_memory_gb, - params.guarantee_connectivity, - mem); - RAFT_EXPECTS(!ace_params.add_global_reverse_edges || !use_disk_mode, - "ACE: partitioned global reverse-edge addition does not support disk mode"); - - // Validate and adjust partitions if disk mode is enabled - if (use_disk_mode) { - ace_validate_disk_mode_partitions(res, - n_partitions, - dataset_size, - dataset_dim, - intermediate_degree, - graph_degree, - params.guarantee_connectivity, - mem); - ace_validate_partition_count(n_partitions, dataset_size, true); - } - - // Preallocate space for files for better performance and fail early if not enough space. - cuvs::util::file_descriptor reordered_fd; - cuvs::util::file_descriptor augmented_fd; - cuvs::util::file_descriptor mapping_fd; - cuvs::util::file_descriptor graph_fd; - size_t reordered_header_size = 0; - size_t augmented_header_size = 0; - size_t mapping_header_size = 0; - size_t graph_header_size = 0; - - if (use_disk_mode) { - workspace.initialize(); - - // Create numpy files with pre-allocated space - std::tie(reordered_fd, reordered_header_size) = cuvs::util::create_numpy_file( - workspace.artifact_path(ace_disk_workspace::artifact::reordered_dataset), - {dataset_size, dataset_dim}, - true); - workspace.mark_artifact_created(ace_disk_workspace::artifact::reordered_dataset); - - 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}, - true); - workspace.mark_artifact_created(ace_disk_workspace::artifact::augmented_dataset); - - std::tie(mapping_fd, mapping_header_size) = cuvs::util::create_numpy_file( - workspace.artifact_path(ace_disk_workspace::artifact::dataset_mapping), - {dataset_size}, - true); - workspace.mark_artifact_created(ace_disk_workspace::artifact::dataset_mapping); - - std::tie(graph_fd, graph_header_size) = cuvs::util::create_numpy_file( - workspace.artifact_path(ace_disk_workspace::artifact::cagra_graph), - {dataset_size, graph_degree}, - true); - workspace.mark_artifact_created(ace_disk_workspace::artifact::cagra_graph); - - RAFT_LOG_DEBUG( - "ACE: Wrote numpy headers (reordered: %zu, augmented: %zu, mapping: %zu, graph: %zu bytes)", - reordered_header_size, - augmented_header_size, - mapping_header_size, - graph_header_size); - } - - auto partition_start = std::chrono::high_resolution_clock::now(); - auto partition_labels = raft::make_host_matrix(dataset_size, 2); - auto partition_histogram = raft::make_host_matrix(n_partitions, 2); - for (size_t c = 0; c < n_partitions; c++) { - partition_histogram(c, 0) = 0; - partition_histogram(c, 1) = 0; - } - - // Determine minimum partition size for stable KNN graph construction - size_t min_partition_size = std::max(1000ULL, dataset_size / n_partitions * 0.1); - - ace_get_partition_labels(res, - dataset_view, - dataset_dim, - partition_labels.view(), - 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 = - std::chrono::duration_cast(partition_end - partition_start) - .count(); - RAFT_LOG_INFO( - "ACE: Partition labeling completed in %ld ms (min_partition_size: " - "%lu)", - partition_elapsed, - min_partition_size); - - // Create vector lists for each partition - auto vectorlist_start = std::chrono::high_resolution_clock::now(); - 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 core_partition_offsets = raft::make_host_vector(n_partitions + 1); - auto augmented_partition_offsets = raft::make_host_vector(n_partitions + 1); - - ace_create_forward_and_backward_lists(dataset_size, - n_partitions, - partition_labels.view(), - partition_histogram.view(), - core_forward_mapping.view(), - core_backward_mapping.view(), - augmented_backward_mapping.view(), - core_partition_offsets.view(), - augmented_partition_offsets.view()); - - auto vectorlist_end = std::chrono::high_resolution_clock::now(); - auto vectorlist_elapsed = - std::chrono::duration_cast(vectorlist_end - vectorlist_start) - .count(); - RAFT_LOG_INFO("ACE: Vector list creation completed in %ld ms", vectorlist_elapsed); - - // Reorder the dataset based on partitions and store to disk. Uses write buffers to improve - // performance. - if (use_disk_mode) { - ace_reorder_and_store_dataset(res, - build_dir, - dataset_view, - dataset_dim, - partition_labels.view(), - partition_histogram.view(), - core_backward_mapping.view(), - core_partition_offsets.view(), - augmented_partition_offsets.view(), - reordered_fd, - augmented_fd, - mapping_fd, - reordered_header_size, - augmented_header_size, - mapping_header_size); - // core_backward_mapping is not needed anymore. - core_backward_mapping = raft::make_host_vector(0); - } - - // Placeholder search graph for in-memory version - auto search_graph = use_disk_mode - ? raft::make_host_matrix(0, 0) - : raft::make_host_matrix(dataset_size, graph_degree); - - // 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++) { - RAFT_LOG_DEBUG("ACE: Processing partition %lu/%lu", partition_id + 1, n_partitions); - auto start = std::chrono::high_resolution_clock::now(); - - // Extract vectors for this partition - size_t core_sub_dataset_size = partition_histogram(partition_id, 0); - size_t augmented_sub_dataset_size = partition_histogram(partition_id, 1); - size_t sub_dataset_size = core_sub_dataset_size + augmented_sub_dataset_size; - - if (sub_dataset_size == 0) { - RAFT_LOG_WARN("ACE: Skipping empty partition %lu", partition_id); - continue; - } - RAFT_LOG_DEBUG("ACE: Sub-dataset size: %lu (%lu + %lu)", - sub_dataset_size, - core_sub_dataset_size, - augmented_sub_dataset_size); - - // Create index for this partition - auto sub_index_params = cuvs::neighbors::cagra::index_params::from_dataset( - raft::make_extents(sub_dataset_size, dataset_dim), - graph_degree, - params.metric, - ef_construction / 16); - sub_index_params.attach_dataset_on_build = false; - sub_index_params.guarantee_connectivity = params.guarantee_connectivity; - - auto read_end = start; - bool used_device_read = false; - const size_t sub_ds_bytes = sub_dataset_size * dataset_dim * sizeof(T); - 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; - std::optional> partial_sub_graph; - std::optional sub_index; - const bool build_partial_graph_on_host = - ace_params.add_global_reverse_edges || - ace_params.partial_graph_knn_build_algo == - graph_build_params::ace_params::knn_graph_build_algo::BRUTE_FORCE; - if (build_partial_graph_on_host) { - auto core_argument_partition = - raft::make_host_matrix(sub_dataset_size, dataset_dim); - if (use_disk_mode) { - ace_load_partition_dataset_from_disk(partition_id, - dataset_dim, - partition_histogram.view(), - core_partition_offsets.view(), - augmented_partition_offsets.view(), - reordered_fd, - augmented_fd, - reordered_header_size, - augmented_header_size, - core_argument_partition.view()); - } else { - 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(), - core_argument_partition.view()); - } - read_end = std::chrono::high_resolution_clock::now(); - auto core_argument_partition_view = raft::make_host_matrix_view( - core_argument_partition.data_handle(), sub_dataset_size, dataset_dim); - if (ace_params.add_global_reverse_edges) { - partial_sub_graph.emplace( - ace_build_pruned_partial_graph(res, - core_argument_partition_view, - intermediate_degree, - graph_degree, - params.metric, - ace_params.partial_graph_knn_build_algo)); - } else { - partial_sub_graph.emplace( - ace_build_optimized_partial_graph(res, - core_argument_partition_view, - intermediate_degree, - graph_degree, - params.metric, - ace_params.partial_graph_knn_build_algo, - params.guarantee_connectivity)); - } - } else { - sub_index.emplace([&]() -> sub_index_t { - if (use_disk_mode) { - const size_t current_free_gpu = rmm::available_device_memory().first; - const size_t configured_gpu = - ace_params.max_gpu_memory_gb > 0 - ? 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)) { - try { - auto sub_dataset_tight = - raft::make_device_matrix(res, sub_dataset_size, dataset_dim); - raft::resource::sync_stream(res); - ace_load_partition_dataset_from_disk(partition_id, - dataset_dim, - partition_histogram.view(), - core_partition_offsets.view(), - augmented_partition_offsets.view(), - reordered_fd, - augmented_fd, - reordered_header_size, - augmented_header_size, - sub_dataset_tight.view()); - read_end = std::chrono::high_resolution_clock::now(); - auto sub_dataset_view = raft::make_const_mdspan(sub_dataset_tight.view()); - std::unique_ptr> - sub_dataset_padded; - auto sub_dataset_dev = [&]() { - if (cuvs::neighbors::matrix_row_width_matches_cagra_required(sub_dataset_view)) { - return cuvs::neighbors::make_device_padded_dataset_view(res, sub_dataset_view); - } - sub_dataset_padded = - cuvs::neighbors::make_device_padded_dataset(res, sub_dataset_view); - return sub_dataset_padded->as_dataset_view(); - }(); - auto direct_index = - ::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)}; - } catch (const std::bad_alloc& e) { - RAFT_LOG_WARN( - "ACE: partition %lu did not fit in device memory for a direct (GDS) read: %s; " - "falling back to a host read", - partition_id, - e.what()); - } - } - - auto sub_dataset = raft::make_host_matrix(sub_dataset_size, dataset_dim); - ace_load_partition_dataset_from_disk(partition_id, - dataset_dim, - partition_histogram.view(), - core_partition_offsets.view(), - augmented_partition_offsets.view(), - reordered_fd, - augmented_fd, - reordered_header_size, - augmented_header_size, - sub_dataset.view()); - 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 = raft::make_host_matrix(sub_dataset_size, dataset_dim); - 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()); - 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)}; - }()); - } - if (used_device_read) { - RAFT_LOG_DEBUG("ACE: partition %lu read directly into device memory (GDS path)", - partition_id); - } - auto read_elapsed = - std::chrono::duration_cast(read_end - start).count(); - auto optimize_end = std::chrono::high_resolution_clock::now(); - auto optimize_elapsed = - std::chrono::duration_cast(optimize_end - read_end).count(); - - 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) { - auto adjusted_search_graph = - raft::make_device_matrix(res, core_sub_dataset_size, graph_degree); - 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 (build_partial_graph_on_host) { - 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 = - static_cast(core_partition_offsets(partition_id)) * graph_degree * sizeof(IdxT) + - graph_header_size; - cuvs::util::write_large_file( - graph_fd, adjusted_search_graph.data_handle(), graph_bytes, graph_offset); - auto write_end = std::chrono::high_resolution_clock::now(); - write_elapsed = - std::chrono::duration_cast(write_end - adjust_end).count(); - } else { - auto sub_search_graph = - raft::make_host_matrix(core_sub_dataset_size, graph_degree); - if (build_partial_graph_on_host) { - 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, - augmented_sub_dataset_size, - graph_degree, - partition_id, - sub_search_graph.view(), - search_graph.view(), - core_partition_offsets.view(), - augmented_partition_offsets.view(), - core_backward_mapping.view(), - augmented_backward_mapping.view()); - adjust_end = std::chrono::high_resolution_clock::now(); - } - - auto adjust_elapsed = - std::chrono::duration_cast(adjust_end - optimize_end).count(); - - auto end = std::chrono::high_resolution_clock::now(); - auto elapsed_ms = std::chrono::duration_cast(end - start).count(); - double read_throughput = - read_elapsed > 0 - ? to_mib(sub_dataset_size * dataset_dim * sizeof(T)) / (read_elapsed / 1000.0) - : 0.0; - double write_throughput = - write_elapsed > 0 ? to_mib(graph_bytes) / (write_elapsed / 1000.0) : 0.0; - RAFT_LOG_INFO( - "ACE: Partition %4lu (%8lu + %8lu) completed in %6ld ms: read %6ld ms (%7.1f MiB/s), " - "optimize %6ld ms, adjust %6ld ms, write %6ld ms (%7.1f MiB/s)", - partition_id, - core_sub_dataset_size, - augmented_sub_dataset_size, - elapsed_ms, - read_elapsed, - read_throughput, - optimize_elapsed, - adjust_elapsed, - write_elapsed, - write_throughput); - } - - auto partition_processing_end = std::chrono::high_resolution_clock::now(); - auto partition_processing_elapsed = std::chrono::duration_cast( - partition_processing_end - partition_processing_start) - .count(); - RAFT_LOG_INFO("ACE: All partition processing completed in %ld ms (%zu partitions)", - partition_processing_elapsed, - n_partitions); - - if (ace_params.add_global_reverse_edges) { - RAFT_LOG_INFO("ACE: Adding global cross-partition reverse edges partition by partition"); - 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) { - const std::string augmented_dataset_path = build_dir + "/augmented_dataset.npy"; - if (std::filesystem::exists(augmented_dataset_path)) { - std::filesystem::remove(augmented_dataset_path); - RAFT_LOG_INFO("ACE: Removed augmented dataset file to save disk space"); - } - } - - auto index_creation_start = std::chrono::high_resolution_clock::now(); - cuvs::neighbors::cagra::index idx(res, params.metric); - if (!use_disk_mode) { - if (params.attach_dataset_on_build) { - idx = cuvs::neighbors::cagra::index( - res, params.metric, dataset, raft::make_const_mdspan(search_graph.view())); - } else { - idx.update_graph(res, raft::make_const_mdspan(search_graph.view())); - } - } else { - idx.update_dataset(res, std::move(reordered_fd)); - idx.update_graph(res, std::move(graph_fd)); - idx.update_mapping(res, std::move(mapping_fd)); - - RAFT_LOG_INFO( - "ACE: Set disk storage at %s (dataset shape [%zu, %zu], graph shape [%zu, %zu])", - build_dir.c_str(), - idx.size(), - idx.dim(), - idx.size(), - idx.graph_degree()); - } - - auto index_creation_end = std::chrono::high_resolution_clock::now(); - auto index_creation_elapsed = std::chrono::duration_cast( - index_creation_end - index_creation_start) - .count(); - RAFT_LOG_INFO("ACE: Final index creation completed in %ld ms", index_creation_elapsed); - - auto total_end = std::chrono::high_resolution_clock::now(); - auto total_elapsed = - std::chrono::duration_cast(total_end - total_start).count(); - RAFT_LOG_INFO("ACE: Partitioned CAGRA build completed in %ld ms total", total_elapsed); - - workspace.commit(); - return std::move(idx); - } catch (const std::exception& e) { - RAFT_LOG_ERROR("ACE: Build failed with exception: %s", e.what()); - workspace.cleanup(); - throw; - } catch (...) { - workspace.cleanup(); - throw; - } -} - template void write_to_graph(raft::host_matrix_view knn_graph, raft::host_matrix_view neighbors_host_view, From 180ab16296f02f657bb67444865e9510ccafffb0 Mon Sep 17 00:00:00 2001 From: enp1s0 Date: Fri, 11 Sep 2026 15:24:19 +0900 Subject: [PATCH 03/18] Fix docs --- cpp/include/cuvs/neighbors/cagra.hpp | 7 ++----- fern/pages/cpp_api/cpp-api-neighbors-cagra.md | 5 +++-- 2 files changed, 5 insertions(+), 7 deletions(-) diff --git a/cpp/include/cuvs/neighbors/cagra.hpp b/cpp/include/cuvs/neighbors/cagra.hpp index d25e40cca1..b94b8447ab 100644 --- a/cpp/include/cuvs/neighbors/cagra.hpp +++ b/cpp/include/cuvs/neighbors/cagra.hpp @@ -91,11 +91,8 @@ struct ace_params { * 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. The final pass preserves the first half of every adjacency row and - * inserts incoming edges originating in other core partitions into the remaining entries. - * Reverse edges are staged and merged one target partition at a time, so GPU workspace scales - * with a partition and its incoming-edge candidates rather than the full graph. This mode does - * not support `use_disk`. + * reverse-edge merge. Then build the reverse-edge graph globally and merge it into the pruned + * graph. This mode does not support `use_disk`. */ bool add_global_reverse_edges = false; diff --git a/fern/pages/cpp_api/cpp-api-neighbors-cagra.md b/fern/pages/cpp_api/cpp-api-neighbors-cagra.md index cf1d492b25..a2318ba23f 100644 --- a/fern/pages/cpp_api/cpp-api-neighbors-cagra.md +++ b/fern/pages/cpp_api/cpp-api-neighbors-cagra.md @@ -34,7 +34,7 @@ struct ace_params { | `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. | | `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. The final pass preserves the first half of every adjacency row and inserts incoming edges originating in other core partitions into the remaining entries. It requires two graph-sized device workspace allocations. With `use_disk`, partition data is loaded from the reordered dataset artifacts and the final graph is updated in place on disk. | +| `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. This mode does not support `use_disk`. | | `partial_graph_knn_build_algo` | [`knn_graph_build_algo`](/api-reference/cpp-api-neighbors-cagra#neighbors-graph-build-params-ace-params-knn-graph-build-algo) | Algorithm used to construct each intermediate partition kNN graph.

Applies to both the original ACE partial CAGRA build with local reverse edges and the global reverse-edge mode. IVF_PQ is the default approximate builder. BRUTE_FORCE computes exact neighbors for float and half datasets and is intended for smaller partitions or quality evaluation. | | `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. | @@ -103,7 +103,8 @@ using graph_build_params_t = std::variant; +graph_build_params::iterative_search_params, +graph_build_params::brute_force_params>; ``` From 25ac2d5094a1e8e9b9e712830cc8db3afe173928 Mon Sep 17 00:00:00 2001 From: enp1s0 Date: Mon, 14 Sep 2026 15:56:50 +0900 Subject: [PATCH 04/18] Add disk global rev edge support --- cpp/src/neighbors/detail/cagra/cagra_ace.cuh | 91 +++++-- .../neighbors/detail/cagra/cagra_ace_disk.hpp | 252 ++++++++++++++++++ 2 files changed, 323 insertions(+), 20 deletions(-) create mode 100644 cpp/src/neighbors/detail/cagra/cagra_ace_disk.hpp diff --git a/cpp/src/neighbors/detail/cagra/cagra_ace.cuh b/cpp/src/neighbors/detail/cagra/cagra_ace.cuh index 4b1d270ff5..2fcb3f6bf9 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_ace.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_ace.cuh @@ -6,6 +6,7 @@ #pragma once #include "../../../util/kvikio_io.hpp" +#include "cagra_ace_disk.hpp" #include "cagra_build.cuh" #include @@ -1166,6 +1167,7 @@ void ace_validate_disk_mode_partitions(raft::resources const& res, size_t intermediate_degree, size_t graph_degree, bool guarantee_connectivity, + bool add_global_reverse_edges, ace_memory_requirements& mem) { // In disk mode, we don't need the full dataset or final graph in memory. @@ -1194,6 +1196,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 = + 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 + @@ -1296,7 +1311,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: Updated per-partition memory estimates: dataset %.2f GiB, graph %.2f GiB, " @@ -1411,7 +1427,7 @@ void ace_merge_partition_reverse_edges(raft::resources const& res, partition_size, stream); RAFT_CUDA_TRY(cudaMemsetAsync( - d_reverse_graph_count.data_handle(), 0, partition_size * sizeof(uint32_t), stream)); + 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( @@ -1427,17 +1443,17 @@ void ace_merge_partition_reverse_edges(raft::resources const& res, 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()); + 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<<>>( + ace_merge_partition_reverse_edges<<>>( d_partition_graph.data_handle(), d_reverse_graph.data_handle(), d_reverse_graph_count.data_handle(), @@ -1589,9 +1605,6 @@ auto build_ace(raft::resources const& res, const index_params& params, DatasetVi ace_params.max_gpu_memory_gb, params.guarantee_connectivity, mem); - RAFT_EXPECTS(!ace_params.add_global_reverse_edges || !use_disk_mode, - "ACE: partitioned global reverse-edge addition does not support disk mode"); - // Validate and adjust partitions if disk mode is enabled if (use_disk_mode) { ace_validate_disk_mode_partitions(res, @@ -1601,6 +1614,7 @@ auto build_ace(raft::resources const& res, const index_params& params, DatasetVi intermediate_degree, graph_degree, params.guarantee_connectivity, + ace_params.add_global_reverse_edges, mem); ace_validate_partition_count(n_partitions, dataset_size, true); } @@ -1737,6 +1751,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++) { @@ -1928,7 +1951,30 @@ 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. + std::vector adjusted_graph(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, adjusted_graph); + 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(), 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); auto adjust_sub_graph = [&](auto sub_graph) { @@ -2031,11 +2077,16 @@ auto build_ace(raft::resources const& res, const index_params& params, DatasetVi if (ace_params.add_global_reverse_edges) { RAFT_LOG_INFO("ACE: Adding global cross-partition reverse edges partition by partition"); - ace_add_global_reverse_edges(res, - search_graph.view(), - partition_labels.view(), - core_backward_mapping.view(), - core_partition_offsets.view()); + 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 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..58d96b2d1a --- /dev/null +++ b/cpp/src/neighbors/detail/cagra/cagra_ace_disk.hpp @@ -0,0 +1,252 @@ +/* + * 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 cross-partition 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; + + // Avoid the generic device-only detail::swap overload pulled in by CAGRA's top-k code. + friend void swap(record& a, record& b) noexcept + { + std::swap(a.destination, b.destination); + std::swap(a.source, b.source); + std::swap(a.rank, b.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); + } + }; + + 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_(offsets.begin(), offsets.end()), + degree_(degree), + capacity_(std::max(size_t{1}, buffer_size / sizeof(record))) + { + RAFT_EXPECTS(offsets_.size() >= 2 && offsets_.front() == 0 && + std::is_sorted(offsets_.begin(), offsets_.end()) && degree_ > 0, + "ACE: invalid reverse-edge partition layout"); + sizes_.resize(offsets_.size() - 1); + buffer_.reserve(capacity_); + 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) + { + auto const first = offsets_.at(partition); + auto const last = offsets_.at(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_.back()), + "ACE: graph contains an invalid neighbor"); + if (destination >= first && destination < last) { continue; } + buffer_.push_back({destination, static_cast(first + row), rank}); + if (buffer_.size() == capacity_) { flush(); } + } + } + } + + void merge(cuvs::util::file_descriptor const& graph_fd, size_t graph_header_size) + { + flush(); + buffer_.resize(capacity_); + 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 the best degree incoming edges per 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 deterministically by source ID. + std::vector incoming(rows * degree_); + std::vector counts(rows, 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(), 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_.back()), + "ACE: invalid reverse-edge record"); + auto const row = size_t(edge.destination - first); + auto* heap = incoming.data() + row * degree_; + auto& n = counts[row]; + incoming_edge candidate{edge.source, edge.rank}; + if (n < degree_) { + heap[n++] = candidate; + std::push_heap(heap, heap + n); + } else if (candidate < heap[0]) { + std::pop_heap(heap, heap + n); + heap[n - 1] = candidate; + std::push_heap(heap, heap + n); + } + } + offset += count; + } + + std::vector graph(rows * degree_); + auto const file_offset = graph_header_size + size_t(first) * degree_ * sizeof(IdxT); + ace_disk_host_io(graph_fd, graph.data(), graph.size() * sizeof(IdxT), file_offset); + auto const protected_edges = degree_ / 2; + for (size_t row = 0; row < rows; ++row) { + auto* edges = incoming.data() + row * degree_; + auto n = counts[row]; + std::sort_heap(edges, edges + n); + auto* neighbors = graph.data() + 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[--n].source; + auto* existing = std::find(neighbors, neighbors + degree_, source); + if (existing < neighbors + protected_edges) { continue; } + auto* end = existing == neighbors + degree_ ? existing - 1 : existing; + std::move_backward(neighbors + protected_edges, end, end + 1); + neighbors[protected_edges] = source; + } + } + ace_disk_host_io(graph_fd, graph.data(), graph.size() * sizeof(IdxT), file_offset); + bucket.close(); + std::filesystem::remove(bucket_path(partition)); + } + buffer_.clear(); + } + + private: + std::string bucket_path(size_t partition) const + { + return directory_ + "/" + std::to_string(partition); + } + + void flush() + { + std::sort(buffer_.begin(), buffer_.end(), [](record const& a, record const& b) { + return a.destination < b.destination; + }); + size_t begin = 0; + while (begin < buffer_.size()) { + auto const partition = + size_t(std::upper_bound(offsets_.begin(), offsets_.end(), buffer_[begin].destination) - + offsets_.begin() - 1); + size_t end = begin + 1; + while (end < buffer_.size() && 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() + begin, + (end - begin) * sizeof(record), + sizes_[partition] * sizeof(record)); + sizes_[partition] += end - begin; + begin = end; + } + buffer_.clear(); + } + + std::vector offsets_; + uint32_t degree_; + size_t capacity_; + std::vector sizes_; + std::vector buffer_; + std::string directory_; +}; + +} // namespace cuvs::neighbors::cagra::detail From 6cd66ca434f9263f4f4bd837a389c450ea8dd4bd Mon Sep 17 00:00:00 2001 From: enp1s0 Date: Tue, 15 Sep 2026 19:36:17 +0900 Subject: [PATCH 05/18] Update docs --- cpp/include/cuvs/neighbors/cagra.hpp | 10 +++++++++- fern/pages/cpp_api/cpp-api-neighbors-cagra.md | 4 ++-- 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/cpp/include/cuvs/neighbors/cagra.hpp b/cpp/include/cuvs/neighbors/cagra.hpp index b94b8447ab..53d1a76f7c 100644 --- a/cpp/include/cuvs/neighbors/cagra.hpp +++ b/cpp/include/cuvs/neighbors/cagra.hpp @@ -78,6 +78,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"; /** @@ -92,7 +95,12 @@ struct ace_params { * * 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. This mode does not support `use_disk`. + * graph. In disk mode (explicit or automatic), cross-partition 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 `graph_degree` 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; diff --git a/fern/pages/cpp_api/cpp-api-neighbors-cagra.md b/fern/pages/cpp_api/cpp-api-neighbors-cagra.md index a2318ba23f..6bbf36bbf0 100644 --- a/fern/pages/cpp_api/cpp-api-neighbors-cagra.md +++ b/fern/pages/cpp_api/cpp-api-neighbors-cagra.md @@ -32,9 +32,9 @@ struct ace_params { | --- | --- | --- | | `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. | | `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. This mode does not support `use_disk`. | +| `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), cross-partition 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 `graph_degree` 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. | | `partial_graph_knn_build_algo` | [`knn_graph_build_algo`](/api-reference/cpp-api-neighbors-cagra#neighbors-graph-build-params-ace-params-knn-graph-build-algo) | Algorithm used to construct each intermediate partition kNN graph.

Applies to both the original ACE partial CAGRA build with local reverse edges and the global reverse-edge mode. IVF_PQ is the default approximate builder. BRUTE_FORCE computes exact neighbors for float and half datasets and is intended for smaller partitions or quality evaluation. | | `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. | From aa09b9f6a8f5e9413326150b964cf075a8bf32f3 Mon Sep 17 00:00:00 2001 From: enp1s0 Date: Wed, 16 Sep 2026 17:56:57 +0900 Subject: [PATCH 06/18] Fix stream --- cpp/src/neighbors/detail/cagra/cagra_ace.cuh | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/cagra_ace.cuh b/cpp/src/neighbors/detail/cagra/cagra_ace.cuh index 4b1d270ff5..d011e0a758 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_ace.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_ace.cuh @@ -1411,7 +1411,7 @@ void ace_merge_partition_reverse_edges(raft::resources const& res, partition_size, stream); RAFT_CUDA_TRY(cudaMemsetAsync( - d_reverse_graph_count.data_handle(), 0, partition_size * sizeof(uint32_t), stream)); + 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( @@ -1427,17 +1427,17 @@ void ace_merge_partition_reverse_edges(raft::resources const& res, 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()); + 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<<>>( + ace_merge_partition_reverse_edges<<>>( d_partition_graph.data_handle(), d_reverse_graph.data_handle(), d_reverse_graph_count.data_handle(), From 007a63840be56343bae790a332b27f2fd7549108 Mon Sep 17 00:00:00 2001 From: enp1s0 Date: Tue, 22 Sep 2026 16:29:31 +0900 Subject: [PATCH 07/18] Add ACE test --- .../ann_cagra/test_float_uint32_t.cu | 8 + .../neighbors/ann_cagra/test_half_uint32_t.cu | 8 + cpp/tests/neighbors/ann_cagra_ace.cuh | 191 ++++++++++++++++++ 3 files changed, 207 insertions(+) create mode 100644 cpp/tests/neighbors/ann_cagra_ace.cuh 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 5dccce21da..b9a5b4fa39 100644 --- a/cpp/tests/neighbors/ann_cagra/test_float_uint32_t.cu +++ b/cpp/tests/neighbors/ann_cagra/test_float_uint32_t.cu @@ -6,12 +6,20 @@ #include #include "../ann_cagra.cuh" +#include "../ann_cagra_ace.cuh" #include #include 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..855859b58e --- /dev/null +++ b/cpp/tests/neighbors/ann_cagra_ace.cuh @@ -0,0 +1,191 @@ +/* + * 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; +}; + +// 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] = 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.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 = 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 + +std::string ace_case_name(::testing::TestParamInfo const& info) +{ + auto const [disk, reverse, metric, npartitions] = info.param; + return std::string(disk ? "Disk" : "Memory") + (reverse ? "GlobalReverse" : "LocalReverse") + + (metric == cuvs::distance::DistanceType::L2Expanded ? "L2" : "InnerProduct") + + "Partitions" + std::to_string(npartitions); +} + +} // namespace +} // namespace cuvs::neighbors::cagra From eccded6b821b268010776eb706fe100513077bcb Mon Sep 17 00:00:00 2001 From: enp1s0 Date: Tue, 22 Sep 2026 23:34:56 +0900 Subject: [PATCH 08/18] Add support for subpartitions > 1 --- cpp/include/cuvs/neighbors/cagra.hpp | 11 +- cpp/src/neighbors/detail/cagra/cagra_ace.cuh | 309 +++++++++++------- cpp/tests/neighbors/ann_cagra_ace.cuh | 31 +- fern/pages/cpp_api/cpp-api-neighbors-cagra.md | 4 +- 4 files changed, 220 insertions(+), 135 deletions(-) diff --git a/cpp/include/cuvs/neighbors/cagra.hpp b/cpp/include/cuvs/neighbors/cagra.hpp index 5be67003fc..6330dfa05f 100644 --- a/cpp/include/cuvs/neighbors/cagra.hpp +++ b/cpp/include/cuvs/neighbors/cagra.hpp @@ -187,15 +187,20 @@ 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. * diff --git a/cpp/src/neighbors/detail/cagra/cagra_ace.cuh b/cpp/src/neighbors/detail/cagra/cagra_ace.cuh index c95b01231a..91d053a328 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_ace.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_ace.cuh @@ -10,14 +10,19 @@ #include "cagra_build.cuh" #include +#include #include #include +#include +#include +#include #include #include #include #include #include +#include #include #include @@ -122,8 +127,29 @@ class ace_disk_workspace { bool committed_ = false; }; +// 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 + } +} + // ACE: Get partition labels for partitioned approach -// TODO(julianmi): Use all neighbors APIs. template void ace_get_partition_labels( raft::resources const& res, @@ -144,103 +170,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) { @@ -251,6 +290,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); } // ACE: Check partition sizes for stable KNN graph construction @@ -292,8 +333,10 @@ 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 avg_total_vectors = 2.0 * static_cast(dataset_size) / n_partitions; - double expected_avg_vectors = 2.0 * static_cast(dataset_size) / n_partitions; + double avg_total_vectors = + static_cast(partition_labels.extent(1)) * 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, @@ -351,6 +394,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++) { @@ -363,8 +407,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); @@ -374,12 +418,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 @@ -387,8 +434,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; @@ -397,11 +444,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; + } } } @@ -437,7 +487,7 @@ void ace_make_core_argument_partition( 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)); @@ -681,8 +731,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); @@ -776,8 +826,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); @@ -790,16 +839,19 @@ void ace_reorder_and_store_dataset( flush_core_buffer(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_augmented_buffer(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_augmented_buffer(secondary_partition); + } } vectors_processed++; @@ -992,9 +1044,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; - // Check if disk mode should be used for ACE based on memory constraints template bool ace_check_use_disk_mode(raft::resources const& res, @@ -1008,9 +1057,11 @@ 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) + 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() && @@ -1174,8 +1225,10 @@ void ace_validate_disk_mode_partitions(raft::resources const& res, size_t graph_degree, bool guarantee_connectivity, bool add_global_reverse_edges, - ace_memory_requirements& mem) + 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) @@ -1525,7 +1578,7 @@ void ace_add_global_reverse_edges(raft::resources const& res, // 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. @@ -1567,13 +1620,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 " @@ -1614,7 +1675,8 @@ 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); + mem, + ace_params.nsubpartitions); // Validate and adjust partitions if disk mode is enabled if (use_disk_mode) { ace_validate_disk_mode_partitions(res, @@ -1625,7 +1687,8 @@ auto build_ace(raft::resources const& res, const index_params& params, DatasetVi graph_degree, params.guarantee_connectivity, ace_params.add_global_reverse_edges, - mem); + mem, + ace_params.nsubpartitions); ace_validate_partition_count(n_partitions, dataset_size, true); } @@ -1651,7 +1714,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); @@ -1675,8 +1738,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; @@ -1712,7 +1775,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); diff --git a/cpp/tests/neighbors/ann_cagra_ace.cuh b/cpp/tests/neighbors/ann_cagra_ace.cuh index 855859b58e..fb54f04cfd 100644 --- a/cpp/tests/neighbors/ann_cagra_ace.cuh +++ b/cpp/tests/neighbors/ann_cagra_ace.cuh @@ -22,6 +22,7 @@ struct AnnCagraAceInputs { 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. @@ -57,7 +58,8 @@ class AnnCagraAceTest : public ::testing::TestWithParam { public: void testCagraAce() { - auto const [use_disk, global_reverse_edges, metric, npartitions] = this->GetParam(); + 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; @@ -80,6 +82,7 @@ class AnnCagraAceTest : public ::testing::TestWithParam { 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; @@ -173,18 +176,30 @@ class AnnCagraAceTest : public ::testing::TestWithParam { } }; -const std::vector inputs_ace = 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 +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] = info.param; + 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); + "Partitions" + std::to_string(npartitions) + "Subpartitions" + + std::to_string(nsubpartitions); } } // namespace diff --git a/fern/pages/cpp_api/cpp-api-neighbors-cagra.md b/fern/pages/cpp_api/cpp-api-neighbors-cagra.md index 97c45cb6ab..57fcf70371 100644 --- a/fern/pages/cpp_api/cpp-api-neighbors-cagra.md +++ b/fern/pages/cpp_api/cpp-api-neighbors-cagra.md @@ -204,6 +204,7 @@ 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; @@ -218,7 +219,8 @@ 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. 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. | From bba9a4459c4ffe9e5e5fb5f4d54194ab4132ad3c Mon Sep 17 00:00:00 2001 From: enp1s0 Date: Wed, 23 Sep 2026 13:14:27 +0900 Subject: [PATCH 09/18] Update ACE log --- cpp/src/neighbors/detail/cagra/cagra_ace.cuh | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/cpp/src/neighbors/detail/cagra/cagra_ace.cuh b/cpp/src/neighbors/detail/cagra/cagra_ace.cuh index 91d053a328..b3d8956848 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_ace.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_ace.cuh @@ -1647,10 +1647,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)); From 1e8703c9d7608f7961f507280c4850ed2136f25a Mon Sep 17 00:00:00 2001 From: enp1s0 Date: Wed, 7 Oct 2026 18:05:47 +0900 Subject: [PATCH 10/18] Update global rev edge addition not to use STL --- cpp/include/cuvs/neighbors/cagra.hpp | 2 +- cpp/src/neighbors/detail/cagra/cagra_ace.cuh | 19 +- .../neighbors/detail/cagra/cagra_ace_disk.hpp | 203 +++++++++++------- fern/pages/cpp_api/cpp-api-neighbors-cagra.md | 2 +- 4 files changed, 138 insertions(+), 88 deletions(-) diff --git a/cpp/include/cuvs/neighbors/cagra.hpp b/cpp/include/cuvs/neighbors/cagra.hpp index a27befbaef..ecad75c486 100644 --- a/cpp/include/cuvs/neighbors/cagra.hpp +++ b/cpp/include/cuvs/neighbors/cagra.hpp @@ -236,7 +236,7 @@ struct ace_params { * reverse-edge merge. Then build the reverse-edge graph globally and merge it into the pruned * graph. In disk mode (explicit or automatic), cross-partition 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 `graph_degree` incoming candidates + * 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. diff --git a/cpp/src/neighbors/detail/cagra/cagra_ace.cuh b/cpp/src/neighbors/detail/cagra/cagra_ace.cuh index b3d8956848..5a25d1e8a9 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_ace.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_ace.cuh @@ -1260,10 +1260,10 @@ void ace_validate_disk_mode_partitions(raft::resources const& res, 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 = - graph_degree * sizeof(typename ace_disk_reverse_edges::incoming_edge) + - sizeof(uint32_t); + 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; } @@ -2028,23 +2028,26 @@ auto build_ace(raft::resources const& res, const index_params& params, DatasetVi 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. - std::vector adjusted_graph(core_sub_dataset_size * graph_degree); + 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 + 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, adjusted_graph); + 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(), graph_bytes, graph_offset); + 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(); diff --git a/cpp/src/neighbors/detail/cagra/cagra_ace_disk.hpp b/cpp/src/neighbors/detail/cagra/cagra_ace_disk.hpp index 58d96b2d1a..63f432e064 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_ace_disk.hpp +++ b/cpp/src/neighbors/detail/cagra/cagra_ace_disk.hpp @@ -6,6 +6,7 @@ #pragma once #include +#include #include #include @@ -15,7 +16,6 @@ #include #include #include -#include namespace cuvs::neighbors::cagra::detail { @@ -61,14 +61,6 @@ class ace_disk_reverse_edges { IdxT destination; IdxT source; uint32_t rank; - - // Avoid the generic device-only detail::swap overload pulled in by CAGRA's top-k code. - friend void swap(record& a, record& b) noexcept - { - std::swap(a.destination, b.destination); - std::swap(a.source, b.source); - std::swap(a.rank, b.rank); - } }; struct incoming_edge { @@ -81,6 +73,9 @@ class ace_disk_reverse_edges { } }; + // 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)); @@ -90,15 +85,20 @@ class ace_disk_reverse_edges { std::span offsets, uint32_t degree, size_t buffer_size) - : offsets_(offsets.begin(), offsets.end()), + : offsets_(raft::make_host_vector(offsets.size())), degree_(degree), - capacity_(std::max(size_t{1}, buffer_size / sizeof(record))) + 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 && - std::is_sorted(offsets_.begin(), offsets_.end()) && degree_ > 0, + RAFT_EXPECTS(offsets.size() >= 2 && offsets.front() == 0 && degree_ > 0, "ACE: invalid reverse-edge partition layout"); - sizes_.resize(offsets_.size() - 1); - buffer_.reserve(capacity_); + 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", @@ -123,18 +123,20 @@ class ace_disk_reverse_edges { // Rows and neighbors both use core-reordered IDs, as in the on-disk graph. void append_partition(size_t partition, std::span graph) { - auto const first = offsets_.at(partition); - auto const last = offsets_.at(partition + 1); + 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_.back()), - "ACE: graph contains an invalid neighbor"); + RAFT_EXPECTS( + static_cast(destination) < static_cast(offsets_(offsets_.size() - 1)), + "ACE: graph contains an invalid neighbor"); if (destination >= first && destination < last) { continue; } - buffer_.push_back({destination, static_cast(first + row), rank}); - if (buffer_.size() == capacity_) { flush(); } + buffer_(buffer_count_++) = {destination, static_cast(first + row), rank}; + if (buffer_count_ == capacity_) { flush(); } } } } @@ -142,71 +144,87 @@ class ace_disk_reverse_edges { void merge(cuvs::util::file_descriptor const& graph_fd, size_t graph_header_size) { flush(); - buffer_.resize(capacity_); + 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 the best degree incoming edges per 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 deterministically by source ID. - std::vector incoming(rows * degree_); - std::vector counts(rows, 0); + 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); + for (size_t offset = 0; offset < sizes_(partition);) { + auto const count = std::min(capacity_, sizes_(partition) - offset); ace_disk_host_io( - bucket, buffer_.data(), count * sizeof(record), offset * sizeof(record)); + 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_.back()), - "ACE: invalid reverse-edge record"); + 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() + row * degree_; - auto& n = counts[row]; + auto* heap = incoming.data_handle() + row * candidate_capacity; + auto& n = counts(row); incoming_edge candidate{edge.source, edge.rank}; - if (n < degree_) { - heap[n++] = candidate; - std::push_heap(heap, heap + n); + 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]) { - std::pop_heap(heap, heap + n); - heap[n - 1] = candidate; - std::push_heap(heap, heap + n); + heap[0] = candidate; + sift_down(heap, n); } } offset += count; } - std::vector graph(rows * degree_); + 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(), graph.size() * sizeof(IdxT), file_offset); + 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() + row * degree_; - auto n = counts[row]; - std::sort_heap(edges, edges + n); - auto* neighbors = graph.data() + row * degree_; + 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[--n].source; - auto* existing = std::find(neighbors, neighbors + degree_, source); - if (existing < neighbors + protected_edges) { continue; } - auto* end = existing == neighbors + degree_ ? existing - 1 : existing; - std::move_backward(neighbors + protected_edges, end, end + 1); + 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(), graph.size() * sizeof(IdxT), file_offset); + ace_disk_host_io( + graph_fd, graph.data_handle(), graph.size() * sizeof(IdxT), file_offset); bucket.close(); std::filesystem::remove(bucket_path(partition)); } - buffer_.clear(); + buffer_count_ = 0; } private: @@ -217,35 +235,64 @@ class ace_disk_reverse_edges { void flush() { - std::sort(buffer_.begin(), buffer_.end(), [](record const& a, record const& b) { - return a.destination < b.destination; + // 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_.size()) { - auto const partition = - size_t(std::upper_bound(offsets_.begin(), offsets_.end(), buffer_[begin].destination) - - offsets_.begin() - 1); - size_t end = begin + 1; - while (end < buffer_.size() && buffer_[end].destination < offsets_[partition + 1]) { + 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); + bucket_path(partition), sizes_(partition) == 0 ? O_CREAT | O_EXCL | O_WRONLY : O_WRONLY); ace_disk_host_io(bucket, - buffer_.data() + begin, + buffer_.data_handle() + begin, (end - begin) * sizeof(record), - sizes_[partition] * sizeof(record)); - sizes_[partition] += end - begin; + sizes_(partition) * sizeof(record)); + sizes_(partition) += end - begin; begin = end; } - buffer_.clear(); + 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; } - std::vector offsets_; + raft::host_vector offsets_; uint32_t degree_; size_t capacity_; - std::vector sizes_; - std::vector buffer_; + raft::host_vector sizes_; + raft::host_vector buffer_; + size_t buffer_count_ = 0; std::string directory_; }; diff --git a/fern/pages/cpp_api/cpp-api-neighbors-cagra.md b/fern/pages/cpp_api/cpp-api-neighbors-cagra.md index 57fcf70371..3dab3a7f26 100644 --- a/fern/pages/cpp_api/cpp-api-neighbors-cagra.md +++ b/fern/pages/cpp_api/cpp-api-neighbors-cagra.md @@ -224,7 +224,7 @@ struct ace_params { | `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. 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), cross-partition 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 `graph_degree` 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. | +| `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), cross-partition 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. | | `partial_graph_knn_build_algo` | [`knn_graph_build_algo`](/api-reference/cpp-api-neighbors-cagra#neighbors-graph-build-params-ace-params-knn-graph-build-algo) | Algorithm used to construct each intermediate partition kNN graph.

Applies to both the original ACE partial CAGRA build with local reverse edges and the global reverse-edge mode. IVF_PQ is the default approximate builder. BRUTE_FORCE computes exact neighbors for float and half datasets and is intended for smaller partitions or quality evaluation. | | `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. | From 02c7cae25ec4bb7ce36891e6046e46961fdc67b1 Mon Sep 17 00:00:00 2001 From: enp1s0 Date: Tue, 20 Oct 2026 16:21:31 +0900 Subject: [PATCH 11/18] Fix C interface --- c/src/neighbors/cagra.cpp | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/c/src/neighbors/cagra.cpp b/c/src/neighbors/cagra.cpp index 97cc9d8bac..3ff1e8d623 100644 --- a/c/src/neighbors/cagra.cpp +++ b/c/src/neighbors/cagra.cpp @@ -762,11 +762,7 @@ static void update_dataset(raft::resources* res_ptr, } static void _set_graph_build_params( - std::variant& out_params, + cuvs::neighbors::cagra::graph_build_params_t& out_params, cuvsCagraIndexParams& params, cuvsCagraGraphBuildAlgo algo, int64_t n_rows, From 6536752209c868b0315f8fa4eb43e508d5f284ba Mon Sep 17 00:00:00 2001 From: enp1s0 Date: Tue, 20 Oct 2026 17:01:47 +0900 Subject: [PATCH 12/18] Remove brute-force kNN graph build from ACE reverse-edge branch --- c/src/neighbors/cagra.cpp | 6 +- cpp/include/cuvs/neighbors/cagra.hpp | 71 +----------------- cpp/src/neighbors/cagra.cuh | 10 --- cpp/src/neighbors/cagra_build_inst.cu.in | 16 ---- cpp/src/neighbors/detail/cagra/cagra_ace.cuh | 75 ++++--------------- .../neighbors/detail/cagra/cagra_build.cuh | 64 +--------------- fern/pages/cpp_api/cpp-api-neighbors-cagra.md | 24 +----- 7 files changed, 23 insertions(+), 243 deletions(-) diff --git a/c/src/neighbors/cagra.cpp b/c/src/neighbors/cagra.cpp index 3ff1e8d623..97cc9d8bac 100644 --- a/c/src/neighbors/cagra.cpp +++ b/c/src/neighbors/cagra.cpp @@ -762,7 +762,11 @@ static void update_dataset(raft::resources* res_ptr, } static void _set_graph_build_params( - cuvs::neighbors::cagra::graph_build_params_t& out_params, + std::variant& out_params, cuvsCagraIndexParams& params, cuvsCagraGraphBuildAlgo algo, int64_t n_rows, diff --git a/cpp/include/cuvs/neighbors/cagra.hpp b/cpp/include/cuvs/neighbors/cagra.hpp index ecad75c486..164f56a80f 100644 --- a/cpp/include/cuvs/neighbors/cagra.hpp +++ b/cpp/include/cuvs/neighbors/cagra.hpp @@ -8,7 +8,6 @@ #include #include #include -#include #include #include #include @@ -172,12 +171,6 @@ struct iterative_search_params : cuvs::neighbors::cagra::search_params { /** Specialized parameters for ACE (Augmented Core Extraction) graph build */ struct ace_params { - /** Algorithm used to build each ACE partial kNN graph. */ - enum class knn_graph_build_algo : uint32_t { - IVF_PQ, - BRUTE_FORCE, - }; - /** * Number of partitions for ACE (Augmented Core Extraction) partitioned build. * @@ -243,16 +236,6 @@ struct ace_params { */ bool add_global_reverse_edges = false; - /** - * Algorithm used to construct each intermediate partition kNN graph. - * - * Applies to both the original ACE partial CAGRA build with local reverse edges and the global - * reverse-edge mode. IVF_PQ is the default approximate builder. BRUTE_FORCE computes exact - * neighbors for float and half datasets and is intended for smaller partitions or quality - * evaluation. - */ - knn_graph_build_algo partial_graph_knn_build_algo = knn_graph_build_algo::IVF_PQ; - /** * Maximum host memory to use for ACE build in GiB. * @@ -293,8 +276,7 @@ using graph_build_params_t = std::variant; + graph_build_params::iterative_search_params>; /** * @brief A strategy for selecting the graph build parameters based on similar HNSW index @@ -357,9 +339,6 @@ struct index_params : cuvs::neighbors::index_params { * // 4. Choose iterative graph building using CAGRA's search() and optimize() [Experimental] * params.graph_build_params = * cagra::graph_build_params::iterative_search_params(); - * - * // 5. Choose exact brute-force kNN construction (float and half datasets) - * params.graph_build_params = cagra::graph_build_params::brute_force_params{}; * @endcode */ graph_build_params_t graph_build_params; @@ -4419,54 +4398,6 @@ void distribute(const raft::resources& clique, const std::string& filename, cuvs::neighbors::mg_index, T, IdxT>* index); -/** - * @brief Build an exact, distance-ordered kNN graph, excluding self neighbors. - * @param res RAFT resources. - * @param dataset Row-major host dataset of float or half values. - * @param knn_graph Host output with matching rows and degree in [1, n_rows - 1]. - * @param build_params Brute-force index and search parameters. - */ -void build_knn_graph(raft::resources const& res, - raft::host_matrix_view dataset, - raft::host_matrix_view knn_graph, - graph_build_params::brute_force_params build_params); - -/** - * @brief Build an exact, distance-ordered kNN graph, excluding self neighbors. - * @param res RAFT resources. - * @param dataset Row-major device dataset of float or half values. - * @param knn_graph Host output with matching rows and degree in [1, n_rows - 1]. - * @param build_params Brute-force index and search parameters. - */ -void build_knn_graph(raft::resources const& res, - raft::device_matrix_view dataset, - raft::host_matrix_view knn_graph, - graph_build_params::brute_force_params build_params); - -/** - * @brief Build an exact, distance-ordered kNN graph, excluding self neighbors. - * @param res RAFT resources. - * @param dataset Row-major host dataset of float or half values. - * @param knn_graph Host output with matching rows and degree in [1, n_rows - 1]. - * @param build_params Brute-force index and search parameters. - */ -void build_knn_graph(raft::resources const& res, - raft::host_matrix_view dataset, - raft::host_matrix_view knn_graph, - graph_build_params::brute_force_params build_params); - -/** - * @brief Build an exact, distance-ordered kNN graph, excluding self neighbors. - * @param res RAFT resources. - * @param dataset Row-major device dataset of float or half values. - * @param knn_graph Host output with matching rows and degree in [1, n_rows - 1]. - * @param build_params Brute-force index and search parameters. - */ -void build_knn_graph(raft::resources const& res, - raft::device_matrix_view dataset, - raft::host_matrix_view knn_graph, - graph_build_params::brute_force_params build_params); - /** * @brief Build a kNN graph using IVF-PQ. * diff --git a/cpp/src/neighbors/cagra.cuh b/cpp/src/neighbors/cagra.cuh index c471f1b857..856950f292 100644 --- a/cpp/src/neighbors/cagra.cuh +++ b/cpp/src/neighbors/cagra.cuh @@ -192,16 +192,6 @@ void build_knn_graph( detail::build_knn_graph(res, dataset, knn_graph, build_params); } -template -void build_knn_graph( - raft::resources const& res, - raft::mdspan, raft::row_major, Accessor> dataset, - raft::host_matrix_view knn_graph, - graph_build_params::brute_force_params params) -{ - detail::build_knn_graph(res, dataset, knn_graph, params); -} - /** * @brief Sort a KNN graph index. * Preprocessing step for `cagra::optimize`: If a KNN graph is not built using diff --git a/cpp/src/neighbors/cagra_build_inst.cu.in b/cpp/src/neighbors/cagra_build_inst.cu.in index 80e0e33fa3..2a3e2b43c1 100644 --- a/cpp/src/neighbors/cagra_build_inst.cu.in +++ b/cpp/src/neighbors/cagra_build_inst.cu.in @@ -43,22 +43,6 @@ void build_knn_graph(raft::resources const& handle, cuvs::neighbors::cagra::build_knn_graph(handle, dataset, knn_graph, params); } -void build_knn_graph(raft::resources const& res, - raft::host_matrix_view dataset, - raft::host_matrix_view knn_graph, - graph_build_params::brute_force_params params) -{ - detail::build_knn_graph(res, dataset, knn_graph, params); -} - -void build_knn_graph(raft::resources const& res, - raft::device_matrix_view dataset, - raft::host_matrix_view knn_graph, - graph_build_params::brute_force_params params) -{ - detail::build_knn_graph(res, dataset, knn_graph, params); -} - #define CUVS_DEFINE_CAGRA_BUILD_OVERLOAD(DatasetViewT, ...) \ auto build(raft::resources const& res, \ const cuvs::neighbors::cagra::index_params& params, \ diff --git a/cpp/src/neighbors/detail/cagra/cagra_ace.cuh b/cpp/src/neighbors/detail/cagra/cagra_ace.cuh index 6696aaf24d..abdbf027c7 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_ace.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_ace.cuh @@ -26,6 +26,7 @@ #include #include +#include #include #include #include @@ -535,9 +536,7 @@ 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, - graph_build_params::ace_params::knn_graph_build_algo knn_build_algo) - -> raft::host_matrix + 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, @@ -546,17 +545,11 @@ auto ace_build_partial_knn_graph( intermediate_graph_degree); auto knn_graph = raft::make_host_matrix(partition_size, intermediate_graph_degree); - if (knn_build_algo == graph_build_params::ace_params::knn_graph_build_algo::IVF_PQ) { - 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); - } else { - auto knn_params = graph_build_params::brute_force_params{}; - knn_params.build_params.metric = metric; - build_knn_graph(res, core_argument_partition, knn_graph.view(), knn_params); - } + 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; } @@ -568,36 +561,16 @@ auto ace_build_pruned_partial_graph( raft::host_matrix_view core_argument_partition, size_t intermediate_graph_degree, size_t graph_degree, - cuvs::distance::DistanceType metric, - graph_build_params::ace_params::knn_graph_build_algo knn_build_algo) - -> raft::host_matrix + cuvs::distance::DistanceType metric) -> raft::host_matrix { auto knn_graph = ace_build_partial_knn_graph( - res, core_argument_partition, intermediate_graph_degree, metric, knn_build_algo); + 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; } -// Original ACE mode: optimize each partial graph, including its local reverse-edge merge. -template -auto ace_build_optimized_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, - graph_build_params::ace_params::knn_graph_build_algo knn_build_algo, - bool guarantee_connectivity) -> raft::host_matrix -{ - auto knn_graph = ace_build_partial_knn_graph( - res, core_argument_partition, intermediate_graph_degree, metric, knn_build_algo); - auto partial_graph = - raft::make_host_matrix(core_argument_partition.extent(0), graph_degree); - graph::optimize(res, knn_graph.view(), partial_graph.view(), guarantee_connectivity); - return partial_graph; -} // ACE: Adjust IDs from core and augmented partitions to global reordered IDs template @@ -1908,11 +1881,7 @@ auto build_ace(raft::resources const& res, const index_params& params, DatasetVi }; std::optional> partial_sub_graph; std::optional sub_index; - const bool build_partial_graph_on_host = - ace_params.add_global_reverse_edges || - ace_params.partial_graph_knn_build_algo == - graph_build_params::ace_params::knn_graph_build_algo::BRUTE_FORCE; - if (build_partial_graph_on_host) { + if (ace_params.add_global_reverse_edges) { auto core_argument_partition = raft::make_host_matrix(sub_dataset_size, dataset_dim); if (use_disk_mode) { @@ -1941,24 +1910,8 @@ auto build_ace(raft::resources const& res, const index_params& params, DatasetVi read_end = std::chrono::high_resolution_clock::now(); auto core_argument_partition_view = raft::make_host_matrix_view( core_argument_partition.data_handle(), sub_dataset_size, dataset_dim); - if (ace_params.add_global_reverse_edges) { - partial_sub_graph.emplace( - ace_build_pruned_partial_graph(res, - core_argument_partition_view, - intermediate_degree, - graph_degree, - params.metric, - ace_params.partial_graph_knn_build_algo)); - } else { - partial_sub_graph.emplace( - ace_build_optimized_partial_graph(res, - core_argument_partition_view, - intermediate_degree, - graph_degree, - params.metric, - ace_params.partial_graph_knn_build_algo, - params.guarantee_connectivity)); - } + partial_sub_graph.emplace(ace_build_pruned_partial_graph( + res, core_argument_partition_view, intermediate_degree, graph_degree, params.metric)); } else { sub_index.emplace([&]() -> sub_index_t { if (use_disk_mode) { @@ -2092,7 +2045,7 @@ auto build_ace(raft::resources const& res, const index_params& params, DatasetVi augmented_backward_mapping.view(), core_forward_mapping.view()); }; - if (build_partial_graph_on_host) { + 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, @@ -2116,7 +2069,7 @@ 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); - if (build_partial_graph_on_host) { + if (ace_params.add_global_reverse_edges) { std::copy_n(partial_sub_graph->data_handle(), sub_search_graph.size(), sub_search_graph.data_handle()); diff --git a/cpp/src/neighbors/detail/cagra/cagra_build.cuh b/cpp/src/neighbors/detail/cagra/cagra_build.cuh index 3e84f23096..9012174488 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_build.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_build.cuh @@ -10,7 +10,6 @@ #include "graph_core.cuh" #include -#include #include #include #include @@ -79,59 +78,6 @@ void build_knn_graph( raft::host_matrix_view knn_graph, cuvs::neighbors::cagra::graph_build_params::ivf_pq_params pq); -// Exact kNN construction shared by ordinary CAGRA and ACE. Request one extra -// neighbor and remove self by ID, including when duplicate vectors tie with self. -template -void build_knn_graph(raft::resources const& res, - raft::mdspan, Layout, Accessor> dataset, - raft::host_matrix_view knn_graph, - graph_build_params::brute_force_params params) -{ - auto const partition_size = static_cast(dataset.extent(0)); - auto const intermediate_graph_degree = static_cast(knn_graph.extent(1)); - RAFT_EXPECTS(knn_graph.extent(0) == dataset.extent(0), - "CAGRA: kNN graph and dataset must have the same number of rows"); - RAFT_EXPECTS(intermediate_graph_degree > 0 && partition_size > intermediate_graph_degree, - "CAGRA: kNN graph degree must be positive and smaller than the dataset size"); - if constexpr (std::is_same_v || std::is_same_v) { - auto device_queries = - raft::make_device_matrix(res, partition_size, dataset.extent(1)); - raft::copy(res, device_queries.view(), dataset); - auto brute_force_index = cuvs::neighbors::brute_force::build( - res, params.build_params, raft::make_const_mdspan(device_queries.view())); - auto device_neighbors = raft::make_device_matrix( - res, partition_size, intermediate_graph_degree + 1); - auto device_distances = - raft::make_device_matrix(res, partition_size, intermediate_graph_degree + 1); - cuvs::neighbors::brute_force::search(res, - params.search_params, - brute_force_index, - raft::make_const_mdspan(device_queries.view()), - device_neighbors.view(), - device_distances.view()); - - auto host_neighbors = - raft::make_host_matrix(partition_size, intermediate_graph_degree + 1); - raft::copy(res, host_neighbors.view(), device_neighbors.view()); - raft::resource::sync_stream(res); - for (size_t row = 0; row < partition_size; ++row) { - size_t output_rank = 0; - for (size_t rank = 0; rank < intermediate_graph_degree + 1; ++rank) { - auto const neighbor = host_neighbors(row, rank); - if (neighbor == static_cast(row)) { continue; } - RAFT_EXPECTS(neighbor >= 0 && static_cast(neighbor) < partition_size, - "CAGRA: brute-force kNN returned an invalid neighbor"); - knn_graph(row, output_rank++) = static_cast(neighbor); - if (output_rank == intermediate_graph_degree) { break; } - } - RAFT_EXPECTS(output_rank == intermediate_graph_degree, - "CAGRA: brute-force kNN did not return enough non-self neighbors"); - } - } else { - RAFT_FAIL("CAGRA: brute-force kNN build supports only float and half datasets"); - } -} - template void write_to_graph(raft::host_matrix_view knn_graph, raft::host_matrix_view neighbors_host_view, @@ -1022,12 +968,11 @@ inline void validate_cagra_knn_graph_build_constraints(index_params const& param std::holds_alternative( knn_build_params) || std::holds_alternative(knn_build_params), - "The selected CAGRA graph builder does not support BitwiseHamming as a metric. Please " + "IVF_PQ for CAGRA graph build does not support BitwiseHamming as a metric. Please " "use nn-descent or the iterative CAGRA search build."); RAFT_EXPECTS( params.metric != cuvs::distance::DistanceType::CosineExpanded || std::holds_alternative(knn_build_params) || - std::holds_alternative(knn_build_params) || std::holds_alternative(knn_build_params), "CosineExpanded distance is not supported for iterative CAGRA graph build."); @@ -1038,8 +983,7 @@ inline void validate_cagra_knn_graph_build_constraints(index_params const& param } /** - * Iterative / IVF-PQ / NN-descent / brute-force KNN graph construction and `optimize` → final host - * CAGRA graph. + * Iterative / IVF-PQ / NN-descent KNN graph construction and `optimize` → final host CAGRA graph. * * @param knn_graph_dataset mdspan passed to IVF-PQ / NN-descent `build_knn_graph` (any stride). */ @@ -1068,10 +1012,6 @@ auto build_cagra_host_graph_from_knn_params(raft::resources const& res, ivf_pq_params.build_params.metric = params.metric; } build_knn_graph(res, knn_graph_dataset, knn_graph->view(), ivf_pq_params); - } else if (std::holds_alternative(knn_build_params)) { - auto brute_force_params = std::get(knn_build_params); - brute_force_params.build_params.metric = params.metric; - build_knn_graph(res, knn_graph_dataset, knn_graph->view(), brute_force_params); } else { auto nn_descent_params = std::get(knn_build_params); diff --git a/fern/pages/cpp_api/cpp-api-neighbors-cagra.md b/fern/pages/cpp_api/cpp-api-neighbors-cagra.md index cfa7dfbe85..9eab93844b 100644 --- a/fern/pages/cpp_api/cpp-api-neighbors-cagra.md +++ b/fern/pages/cpp_api/cpp-api-neighbors-cagra.md @@ -76,8 +76,7 @@ using graph_build_params_t = std::variant; +graph_build_params::iterative_search_params>; ``` @@ -219,7 +218,6 @@ struct ace_params { std::string build_dir; bool use_disk; bool add_global_reverse_edges; - knn_graph_build_algo partial_graph_knn_build_algo; double max_host_memory_gb; double max_gpu_memory_gb; }; @@ -235,29 +233,9 @@ struct ace_params { | `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), cross-partition 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. | -| `partial_graph_knn_build_algo` | [`knn_graph_build_algo`](/api-reference/cpp-api-neighbors-cagra#neighbors-graph-build-params-ace-params-knn-graph-build-algo) | Algorithm used to construct each intermediate partition kNN graph.

Applies to both the original ACE partial CAGRA build with local reverse edges and the global reverse-edge mode. IVF_PQ is the default approximate builder. BRUTE_FORCE computes exact neighbors for float and half datasets and is intended for smaller partitions or quality evaluation. | | `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. | - -### neighbors::graph_build_params::ace_params::knn_graph_build_algo - -Algorithm used to build each ACE partial kNN graph. - -```cpp -enum class knn_graph_build_algo : uint32_t { - IVF_PQ, - BRUTE_FORCE -}; -``` - -**Values** - -| Name | Value | -| --- | --- | -| `IVF_PQ` | `` | -| `BRUTE_FORCE` | `` | - ## CAGRA index extend parameters From e1a8a38a0cad279c7818c2a8790be6877aba83c1 Mon Sep 17 00:00:00 2001 From: enp1s0 Date: Wed, 21 Oct 2026 01:38:54 +0900 Subject: [PATCH 13/18] Fix rev edge addition --- cpp/include/cuvs/neighbors/cagra.hpp | 2 +- cpp/src/neighbors/detail/cagra/cagra_ace.cuh | 7 ++----- cpp/src/neighbors/detail/cagra/cagra_ace_disk.hpp | 3 +-- fern/pages/cpp_api/cpp-api-neighbors-cagra.md | 2 +- 4 files changed, 5 insertions(+), 9 deletions(-) diff --git a/cpp/include/cuvs/neighbors/cagra.hpp b/cpp/include/cuvs/neighbors/cagra.hpp index 164f56a80f..6c1ff7822e 100644 --- a/cpp/include/cuvs/neighbors/cagra.hpp +++ b/cpp/include/cuvs/neighbors/cagra.hpp @@ -227,7 +227,7 @@ struct ace_params { * * 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), cross-partition incoming edges are buffered + * 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. diff --git a/cpp/src/neighbors/detail/cagra/cagra_ace.cuh b/cpp/src/neighbors/detail/cagra/cagra_ace.cuh index abdbf027c7..07a3882eeb 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_ace.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_ace.cuh @@ -1561,14 +1561,11 @@ void ace_add_global_reverse_edges(raft::resources const& res, std::copy_n(graph.data_handle(), graph.size(), source_graph.data_handle()); std::vector> candidate_sets(n_partitions); for (size_t source = 0; source < dataset_size; ++source) { - auto const source_partition = static_cast(partition_labels(source, 0)); 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)); - if (source_partition != destination_partition) { - candidate_sets[destination_partition].insert(static_cast(source)); - } + candidate_sets[destination_partition].insert(static_cast(source)); } } @@ -2131,7 +2128,7 @@ auto build_ace(raft::resources const& res, const index_params& params, DatasetVi n_partitions); if (ace_params.add_global_reverse_edges) { - RAFT_LOG_INFO("ACE: Adding global cross-partition reverse edges partition by partition"); + 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(); diff --git a/cpp/src/neighbors/detail/cagra/cagra_ace_disk.hpp b/cpp/src/neighbors/detail/cagra/cagra_ace_disk.hpp index 63f432e064..2bd2f8e9cb 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_ace_disk.hpp +++ b/cpp/src/neighbors/detail/cagra/cagra_ace_disk.hpp @@ -51,7 +51,7 @@ void ace_disk_host_io(cuvs::util::file_descriptor const& fd, } } -// Spool cross-partition edges while the pruned core rows are still in memory. A single +// 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 @@ -134,7 +134,6 @@ class ace_disk_reverse_edges { RAFT_EXPECTS( static_cast(destination) < static_cast(offsets_(offsets_.size() - 1)), "ACE: graph contains an invalid neighbor"); - if (destination >= first && destination < last) { continue; } buffer_(buffer_count_++) = {destination, static_cast(first + row), rank}; if (buffer_count_ == capacity_) { flush(); } } diff --git a/fern/pages/cpp_api/cpp-api-neighbors-cagra.md b/fern/pages/cpp_api/cpp-api-neighbors-cagra.md index 9eab93844b..a51e17e205 100644 --- a/fern/pages/cpp_api/cpp-api-neighbors-cagra.md +++ b/fern/pages/cpp_api/cpp-api-neighbors-cagra.md @@ -232,7 +232,7 @@ struct ace_params { | `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. 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), cross-partition 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. | +| `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. | From 9b8f9236acf77c53272a858fc386ad5352de7e14 Mon Sep 17 00:00:00 2001 From: enp1s0 Date: Wed, 21 Oct 2026 01:45:41 +0900 Subject: [PATCH 14/18] Fix host working memory calculation --- cpp/src/neighbors/detail/cagra/cagra_ace.cuh | 33 +++++++++++++++----- 1 file changed, 26 insertions(+), 7 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/cagra_ace.cuh b/cpp/src/neighbors/detail/cagra/cagra_ace.cuh index 07a3882eeb..4e4dd21c29 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_ace.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_ace.cuh @@ -51,7 +51,6 @@ #include #include #include -#include #include #include #include @@ -1100,6 +1099,7 @@ 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, + bool add_global_reverse_edges, ace_memory_requirements& mem, size_t nsubpartitions = 1) { @@ -1153,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), @@ -1439,7 +1445,6 @@ void ace_merge_partition_reverse_edges(raft::resources const& res, 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::sort(candidate_sources.begin(), candidate_sources.end()); std::vector host_partition_graph(static_cast(partition_size) * graph_degree); std::vector host_candidate_destinations(num_candidates); @@ -1559,24 +1564,37 @@ void ace_add_global_reverse_edges(raft::resources const& res, 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_sets(n_partitions); + 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_sets[destination_partition].insert(static_cast(source)); + ++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) { - if (candidate_sets[partition_id].empty()) { continue; } + 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); - std::vector candidate_sources(candidate_sets[partition_id].begin(), - candidate_sets[partition_id].end()); ace_merge_partition_reverse_edges( res, source_graph.view(), graph, partition_nodes, std::move(candidate_sources)); } @@ -1684,6 +1702,7 @@ 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, + ace_params.add_global_reverse_edges, mem, ace_params.nsubpartitions); // Validate and adjust partitions if disk mode is enabled From 253dc2ad3d430a5b9617a6b75b0dff117eb779c3 Mon Sep 17 00:00:00 2001 From: enp1s0 Date: Wed, 21 Oct 2026 16:25:41 +0900 Subject: [PATCH 15/18] Fix comment --- cpp/include/cuvs/neighbors/cagra.hpp | 4 ++-- fern/pages/cpp_api/cpp-api-neighbors-cagra.md | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/cpp/include/cuvs/neighbors/cagra.hpp b/cpp/include/cuvs/neighbors/cagra.hpp index 6c1ff7822e..e852d5f3dd 100644 --- a/cpp/include/cuvs/neighbors/cagra.hpp +++ b/cpp/include/cuvs/neighbors/cagra.hpp @@ -180,8 +180,8 @@ 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 (1 + nsubpartitions) * (n_rows / - * npartitions) * dim * sizeof(T), accounting for the core and augmented vectors. Please account + * 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 diff --git a/fern/pages/cpp_api/cpp-api-neighbors-cagra.md b/fern/pages/cpp_api/cpp-api-neighbors-cagra.md index a51e17e205..5764e25952 100644 --- a/fern/pages/cpp_api/cpp-api-neighbors-cagra.md +++ b/fern/pages/cpp_api/cpp-api-neighbors-cagra.md @@ -227,7 +227,7 @@ 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 (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. | +| `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. 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. | From 4dc7c05a1717f68d0d1d6827c8ef36e439226ddd Mon Sep 17 00:00:00 2001 From: enp1s0 Date: Wed, 21 Oct 2026 16:38:07 +0900 Subject: [PATCH 16/18] Update C interface --- c/include/cuvs/neighbors/cagra.h | 15 +++++++++++++-- c/src/neighbors/cagra.cpp | 8 +++++++- c/tests/neighbors/ann_cagra_c.cu | 6 ++++++ fern/pages/c_api/c-api-neighbors-cagra.md | 6 +++++- 4 files changed, 31 insertions(+), 4 deletions(-) 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 97cc9d8bac..9a59e6ee3a 100644 --- a/c/src/neighbors/cagra.cpp +++ b/c/src/neighbors/cagra.cpp @@ -822,6 +822,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; @@ -1460,6 +1462,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; @@ -2341,7 +2345,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/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 From e6abf6af15e3b19ae98de7acd9c139a3ec002395 Mon Sep 17 00:00:00 2001 From: enp1s0 Date: Wed, 21 Oct 2026 16:42:11 +0900 Subject: [PATCH 17/18] Add Python interface --- .../python_api/python-api-neighbors-cagra.md | 20 +++++++++-- python/cuvs/cuvs/neighbors/cagra/cagra.pxd | 2 ++ python/cuvs/cuvs/neighbors/cagra/cagra.pyx | 34 ++++++++++++++++--- python/cuvs/cuvs/tests/test_cagra_ace.py | 22 ++++++++++++ 4 files changed, 71 insertions(+), 7 deletions(-) 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, From 3a7b5b9b90b1baf311ebd115d987d300dd6ec550 Mon Sep 17 00:00:00 2001 From: enp1s0 Date: Wed, 21 Oct 2026 16:49:01 +0900 Subject: [PATCH 18/18] Fix fern --- cpp/include/cuvs/neighbors/cagra.hpp | 2 +- fern/pages/cpp_api/cpp-api-neighbors-cagra.md | 4 +--- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/cpp/include/cuvs/neighbors/cagra.hpp b/cpp/include/cuvs/neighbors/cagra.hpp index e852d5f3dd..55d1c8830b 100644 --- a/cpp/include/cuvs/neighbors/cagra.hpp +++ b/cpp/include/cuvs/neighbors/cagra.hpp @@ -747,7 +747,7 @@ struct CUVS_EXPORT index : cuvs::neighbors::index { } /** - * Replace the source indices with a new source indices taking the ownership of the passed + * @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/fern/pages/cpp_api/cpp-api-neighbors-cagra.md b/fern/pages/cpp_api/cpp-api-neighbors-cagra.md index 5764e25952..b86a1cd5f8 100644 --- a/fern/pages/cpp_api/cpp-api-neighbors-cagra.md +++ b/fern/pages/cpp_api/cpp-api-neighbors-cagra.md @@ -576,14 +576,12 @@ We create a copy of the graph on the device. The index manages the lifetime of t ### neighbors::cagra::index::update_source_indices -Replace the source indices with a new source indices taking the ownership of the passed +Replace the source indices with a new source indices taking the ownership of the passed vector. ```cpp void update_source_indices(raft::device_vector&& source_indices); ``` -vector. - **Parameters** | Name | Direction | Type | Description |