Add Augmented Core Extraction Algorithm

[julianmi]

This PR introduces Augmented Core Extraction (ACE), an approach proposed by @anaruse for building CAGRA indices on very large datasets that exceed GPU memory capacity. ACE enables users to build high-quality approximate nearest neighbor search indices on datasets that would otherwise be impossible to process on a single GPU. The approach uses the host memory if large enough and falls back to the disk if required.

This work is a collaboration: @anaruse, @tfeher, @achirkin, @mfoerste4

Algorithm Description

1. Dataset Partitioning: The dataset is partitioned using balanced k-means clustering on sampled data. Each vector is assigned to its two closest partition centroids (primary and augmented). The primary partitions are non-overlapping. The augmentation ensures that cross-partition edges are captured in the final graph. Partitions smaller than a minimum threshold are automatically merged with larger partitions to ensure computational efficiency and graph quality. Vectors from small partitions are reassigned to the nearest valid partitions.
2. Per-Partition Graph Building: For each partition, a sub-index is built independently (regular build_knn_graph() flow) with its primary vectors plus augmented vectors from neighboring partitions.
3. Graph Combining: The per-partition graphs are combined into a single unified CAGRA index. Merging is not needed since the primary partitions are non-overlapping. The in-memory variant remaps the local partition IDs to global dataset IDs to create a correct index. The disk variant stores the backward index mappings (dataset_mapping.bin), the reordered dataset (reordered_dataset.bin) and the optimized CAGRA graph (cagra_graph.bin) on disk. The index is then incomplete as show by cuvs::neighbors::index::on_disk(). The files are stored in cuvs::neighbors::index::file_directory(). The HNSW index serialization was provided by @mfoerste4 in [WIP] Add disk2disk serialization foe ACE Algorithm #1410, which was merged here. This adds the serialize_to_hnsw() serialization routine that allows combination of dataset, graph, and mapping. The data will be combined on-the-fly while streamed from disk to disk while trying to minimize the required host memory. The host needs enough memory to hold the index though.

Core Components

• ace_build(): Main routine which users should call.
• ace_get_partition_labels(): Performs balanced k-means clustering to assign each vector to two closest partitions while handling small partition merging.
• ace_create_forward_and_backward_lists(): Creates bidirectional ID mappings between original dataset indices and reordered partition-local indices.
• ace_set_index_params(): Set the index parameters based on the partition and augmented dataset to ensure an efficient KNN graph building.
• ace_gather_partition_dataset(): In-memory only: gather the partition and augmented dataset.
• ace_adjust_sub_graph_ids: In-memory only: Adjust ids in sub search graph and store them into the main search graph.
• ace_adjust_final_graph_ids: In-memory only: Map graph neighbor IDs from reordered space back to original vector IDs.
• ace_reorder_and_store_dataset: Disk only: Reorder the dataset based on partitions and store to disk. Uses write buffers to improve performance.
• ace_load_partition_dataset_from_disk: Disk only: Load partition dataset and augmented dataset from disk.
• file_descriptor and ace_read_large_file() / ace_write_large_file(): RAII file handle and chunked file I/O operations.
• CAGRA index changes: Added on_disk_ flag and file_directory_ to the CAGRA index structure to support disk-backed indices.
• CAGRA parameter changes: Added ace_npartitions and ace_build_dir to the CAGRA parameters for users to specify that ACE should be used and which directory should be used if required.

Usage

C++ API

#include <cuvs/neighbors/cagra.hpp>

using namespace cuvs::neighbors;

// Configure index parameters
cagra::index_params params;
params.ace_npartitions = 10;  // Number of partitions (unset or <= 1 to disable ACE)
params.ace_build_dir = "/tmp/ace_build";  // Directory for intermediate files (should be a fast NVMe)
params.graph_degree = 64;
params.intermediate_graph_degree = 128;

// Build ACE index (dataset can be on host memory)
auto dataset = raft::make_host_matrix<float, int64_t>(n_rows, n_cols);
// ... load dataset ...

auto index = cagra::build_ace(res, params, dataset.view(), params.ace_npartitions);

// Search works identically to standard CAGRA if the host has enough memory (index.on_disk() == false)
cagra::search_params search_params;
auto neighbors = raft::make_device_matrix<uint32_t>(res, n_queries, k);
auto distances = raft::make_device_matrix<float>(res, n_queries, k);
cagra::search(res, search_params, index, queries, neighbors.view(), distances.view());

Storage Requirements

1. cagra_graph.bin: n_rows * graph_degree * sizeof(IdxT)
2. dataset_mapping.bin: n_rows * sizeof(IdxT)
3. reordered_dataset.bin: Size of the input dataset
4. augmented_dataset.bin: Size of the input dataset

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[tfeher]

/ok to test 38c03a4

[achirkin]

Thanks for pushing to implement my file-descriptor suggestions! Only now I've noticed converting from linux file descriptors to the streams is not as convenient as it should be. I have some suggestion for this below.

[tfeher]

Thanks @julianmi for addressing the issues. LGTM.

[julianmi]

Thanks for pushing to implement my file-descriptor suggestions! Only now I've noticed converting from linux file descriptors to the streams is not as convenient as it should be. I have some suggestion for this below.

Thank you for these suggestions. I have applied them and removed the on_disk() helper as you suggested.

The issues brought up by Artem were addressed. Since Artem is on leave, I am dismissing his requests for changes.

[tfeher]

/merge