From 818326cfd54599e1557003393b18d3648432245e Mon Sep 17 00:00:00 2001 From: Julian Miller Date: Wed, 5 Aug 2026 11:58:46 +0200 Subject: [PATCH 1/2] Fix ACE regressions from #1846 --- cpp/src/neighbors/cagra.cuh | 3 +- .../neighbors/detail/cagra/cagra_build.cuh | 64 +++++++++++++------ 2 files changed, 47 insertions(+), 20 deletions(-) diff --git a/cpp/src/neighbors/cagra.cuh b/cpp/src/neighbors/cagra.cuh index 80e2f2a07e..138b899fc5 100644 --- a/cpp/src/neighbors/cagra.cuh +++ b/cpp/src/neighbors/cagra.cuh @@ -310,8 +310,7 @@ auto build(raft::resources const& res, const index_params& params, DatasetViewT return idx; } else { if (std::holds_alternative(params.graph_build_params)) { - return cuvs::neighbors::cagra::detail::build_ace( - res, params, dataset.view()); + return cuvs::neighbors::cagra::detail::build_ace(res, params, dataset); } return cuvs::neighbors::cagra::detail::build_from_host_matrix( res, params, dataset); diff --git a/cpp/src/neighbors/detail/cagra/cagra_build.cuh b/cpp/src/neighbors/detail/cagra/cagra_build.cuh index c06f9b12e3..64d4b2a002 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_build.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_build.cuh @@ -80,16 +80,21 @@ 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 dataset_dim = dataset.extent(1); 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"); @@ -468,6 +473,7 @@ 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, @@ -483,8 +489,12 @@ void ace_reorder_and_store_dataset( auto start = std::chrono::high_resolution_clock::now(); size_t dataset_size = dataset.extent(0); - size_t dataset_dim = dataset.extent(1); 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)", @@ -1120,6 +1130,13 @@ auto build_from_device_matrix(raft::resources const& res, 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; + // 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 @@ -1131,9 +1148,7 @@ auto build_from_device_matrix(raft::resources const& res, // 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, - raft::host_matrix_view dataset) +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 @@ -1153,8 +1168,9 @@ auto build_ace(raft::resources const& res, params.graph_degree, npartitions); - size_t dataset_size = dataset.extent(0); - size_t dataset_dim = dataset.extent(1); + 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) { @@ -1274,8 +1290,12 @@ auto build_ace(raft::resources const& res, // 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, partition_labels.view(), partition_histogram.view(), min_partition_size); + 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, @@ -1323,7 +1343,8 @@ auto build_ace(raft::resources const& res, if (use_disk_mode) { ace_reorder_and_store_dataset(res, build_dir, - dataset, + dataset_view, + dataset_dim, partition_labels.view(), partition_histogram.view(), core_backward_mapping.view(), @@ -1382,7 +1403,7 @@ auto build_ace(raft::resources const& res, augmented_sub_dataset_size, dataset_dim, partition_id, - dataset, + dataset_view, core_backward_mapping.view(), augmented_backward_mapping.view(), core_partition_offsets.view(), @@ -1402,12 +1423,14 @@ auto build_ace(raft::resources const& res, sub_index_params.attach_dataset_on_build = false; sub_index_params.guarantee_connectivity = params.guarantee_connectivity; - // Copy host partition to device with padding; build_from_device_matrix accepts - // device_padded_dataset_view. - auto sub_dataset_dev = cuvs::neighbors::make_device_padded_dataset( - res, raft::make_const_mdspan(sub_dataset.view())); - auto sub_index = ::cuvs::neighbors::cagra::detail::build_from_device_matrix( - res, sub_index_params, sub_dataset_dev->as_dataset_view()); + // Keep the partition host-resident so IVF-PQ and NN-descent can consume it in batches. + // Iterative CAGRA uploads and pads the partition inside build_from_host_matrix. + auto sub_dataset_view = cuvs::neighbors::make_host_standard_dataset_view( + raft::make_const_mdspan(sub_dataset.view())); + auto sub_index = ::cuvs::neighbors::cagra::detail::build_from_host_matrix( + res, sub_index_params, sub_dataset_view); + static_assert( + std::is_same_v>); auto optimize_end = std::chrono::high_resolution_clock::now(); auto optimize_elapsed = @@ -1509,7 +1532,12 @@ auto build_ace(raft::resources const& res, auto index_creation_start = std::chrono::high_resolution_clock::now(); cuvs::neighbors::cagra::index idx(res, params.metric); if (!use_disk_mode) { - idx.update_graph(res, raft::make_const_mdspan(search_graph.view())); + 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)); From b67a87eb1d691093d014c4b52cafbb58ca2fea55 Mon Sep 17 00:00:00 2001 From: Julian Miller Date: Thu, 6 Aug 2026 09:59:37 +0200 Subject: [PATCH 2/2] Use public cagra::build API --- cpp/src/neighbors/detail/cagra/cagra_build.cuh | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/cagra_build.cuh b/cpp/src/neighbors/detail/cagra/cagra_build.cuh index 64d4b2a002..9d4bfc422b 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_build.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_build.cuh @@ -1427,8 +1427,7 @@ auto build_ace(raft::resources const& res, const index_params& params, DatasetVi // Iterative CAGRA uploads and pads the partition inside build_from_host_matrix. auto sub_dataset_view = cuvs::neighbors::make_host_standard_dataset_view( raft::make_const_mdspan(sub_dataset.view())); - auto sub_index = ::cuvs::neighbors::cagra::detail::build_from_host_matrix( - res, sub_index_params, sub_dataset_view); + auto sub_index = ::cuvs::neighbors::cagra::build(res, sub_index_params, sub_dataset_view); static_assert( std::is_same_v>);