From 1ff9f1746aa98c3da2ed8cda4176fac2f40c8071 Mon Sep 17 00:00:00 2001 From: jinsolp Date: Thu, 2 Jul 2026 06:36:32 +0000 Subject: [PATCH 1/4] cpu side all neighbors graph impl --- cpp/include/cuvs/neighbors/all_neighbors.hpp | 42 +- .../neighbors/all_neighbors/all_neighbors.cu | 14 +- .../neighbors/all_neighbors/all_neighbors.cuh | 293 ++++++++--- .../all_neighbors/all_neighbors_batched.cuh | 229 +++++---- .../all_neighbors/all_neighbors_builder.cuh | 456 ++++++++++-------- .../all_neighbors/all_neighbors_merge.cuh | 85 +++- cpp/tests/neighbors/all_neighbors.cuh | 43 +- 7 files changed, 778 insertions(+), 384 deletions(-) diff --git a/cpp/include/cuvs/neighbors/all_neighbors.hpp b/cpp/include/cuvs/neighbors/all_neighbors.hpp index 8f0a98d8a7..f88399eda3 100644 --- a/cpp/include/cuvs/neighbors/all_neighbors.hpp +++ b/cpp/include/cuvs/neighbors/all_neighbors.hpp @@ -119,8 +119,8 @@ struct all_neighbors_params { * to build all-neighbors knn graph * @param[in] dataset raft::host_matrix_view input dataset expected to be located * in host memory - * @param[out] indices nearest neighbor indices of shape [n_row x k] - * @param[out] distances nearest neighbor distances [n_row x k] + * @param[out] indices nearest neighbor indices of shape [n_row x k] on device memory + * @param[out] distances nearest neighbor distances [n_row x k] on device memory * @param[out] core_distances array for core distances of size [n_row]. Requires distances matrix to * compute core_distances. If core_distances is given, the resulting indices and distances will be * mutual reachability space. @@ -135,6 +135,40 @@ void build( std::optional> core_distances = std::nullopt, float alpha = 1.0); +/** + * @brief Builds an approximate all-neighbors knn graph (find nearest neighbors for all the + * training vectors) + * + * Usage example: + * @code{.cpp} + * using namespace cuvs::neighbors; + * all_neighbors::all_neighbors_params params; + * params.n_clusters = 4; + * auto indices = raft::make_host_matrix(n_row, k); + * auto distances = raft::make_host_matrix(n_row, k); + * all_neighbors::build(res, params, dataset, indices.view(), distances.view()); + * @endcode + * + * @param[in] handle raft::resources is an object managing resources + * @param[in] params an instance of all_neighbors::all_neighbors_params that are parameters + * to build all-neighbors knn graph + * @param[in] dataset raft::host_matrix_view input dataset expected to be located in host memory + * @param[out] indices nearest neighbor indices of shape [n_row x k] on host memory + * @param[out] distances nearest neighbor distances [n_row x k] on host memory + * @param[out] core_distances array for core distances of size [n_row] on host memory. Requires + * distances matrix to compute core_distances. If core_distances is given, the resulting indices and + * distances will be mutual reachability space. + * @param[in] alpha distance scaling parameter as used in robust single linkage. + */ +void build( + const raft::resources& handle, + const all_neighbors_params& params, + raft::host_matrix_view dataset, + raft::host_matrix_view indices, + std::optional> distances = std::nullopt, + std::optional> core_distances = std::nullopt, + float alpha = 1.0); + /** * @brief Builds an approximate all-neighbors knn graph (find nearest neighbors for all the training * vectors) params.n_clusters should be 1 for data on device. To use a larger params.n_clusters for @@ -155,8 +189,8 @@ void build( * to build all-neighbors knn graph * @param[in] dataset raft::device_matrix_view input dataset expected to be located * in device memory - * @param[out] indices nearest neighbor indices of shape [n_row x k] - * @param[out] distances nearest neighbor distances [n_row x k] + * @param[out] indices nearest neighbor indices of shape [n_row x k] on device memory + * @param[out] distances nearest neighbor distances [n_row x k] on device memory * @param[out] core_distances array for core distances of size [n_row]. Requires distances matrix to * compute core_distances. If core_distances is given, the resulting indices and distances will be * mutual reachability space. diff --git a/cpp/src/neighbors/all_neighbors/all_neighbors.cu b/cpp/src/neighbors/all_neighbors/all_neighbors.cu index d44c8cb7a0..a9288a83f5 100644 --- a/cpp/src/neighbors/all_neighbors/all_neighbors.cu +++ b/cpp/src/neighbors/all_neighbors/all_neighbors.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ @@ -30,6 +30,18 @@ namespace cuvs::neighbors::all_neighbors { { \ return all_neighbors::detail::build( \ handle, params, dataset, indices, distances, core_distances, alpha); \ + } \ + \ + void build(const raft::resources& handle, \ + const all_neighbors_params& params, \ + raft::host_matrix_view dataset, \ + raft::host_matrix_view indices, \ + std::optional> distances, \ + std::optional> core_distances, \ + T alpha) \ + { \ + return all_neighbors::detail::build( \ + handle, params, dataset, indices, distances, core_distances, alpha); \ } CUVS_INST_ALL_NEIGHBORS(float, int64_t); diff --git a/cpp/src/neighbors/all_neighbors/all_neighbors.cuh b/cpp/src/neighbors/all_neighbors/all_neighbors.cuh index 02de3ea799..df9f86c7f1 100644 --- a/cpp/src/neighbors/all_neighbors/all_neighbors.cuh +++ b/cpp/src/neighbors/all_neighbors/all_neighbors.cuh @@ -6,18 +6,63 @@ #pragma once #include "../detail/reachability.cuh" #include "all_neighbors_batched.cuh" +#include #include #include #include +#include #include namespace cuvs::neighbors::all_neighbors::detail { using namespace cuvs::neighbors; +template GRAPH_BUILD_ALGO check_params_validity(const all_neighbors_params& params, - bool do_mutual_reachability_dist) + DatasetView dataset, + IndicesView indices, + const std::optional& distances, + const std::optional& core_distances) { - using DT = cuvs::distance::DistanceType; + // 1. Check output memory-space consistency + constexpr bool host_output = + std::is_same_v>; + constexpr bool device_output = + std::is_same_v>; + static_assert(host_output || device_output, + "indices must be a host_matrix_view or device_matrix_view"); + if constexpr (host_output) { + static_assert(std::is_same_v>, + "distances must be a host_matrix_view when indices are on host"); + static_assert(std::is_same_v>, + "core_distances must be a host_vector_view when indices are on host"); + } else { + static_assert(std::is_same_v>, + "distances must be a device_matrix_view when indices are on device"); + static_assert(std::is_same_v>, + "core_distances must be a device_vector_view when indices are on device"); + } + + // 2. Check shape consistency + RAFT_EXPECTS(dataset.extent(0) == indices.extent(0), + "number of rows in dataset should be the same as number of rows in indices matrix"); + if (distances.has_value()) { + RAFT_EXPECTS(indices.extent(0) == distances.value().extent(0) && + indices.extent(1) == distances.value().extent(1), + "indices matrix and distances matrix has to be the same shape."); + } + if (core_distances.has_value()) { + RAFT_EXPECTS(distances.has_value(), + "distances matrix should be allocated to get mutual reachability distance."); + } + + // 3. Check metric / algorithm validity + const bool do_mutual_reachability_dist = core_distances.has_value(); + using DT = cuvs::distance::DistanceType; // InnerProduct is not supported for mutual reachability distance, because mutual reachability // distance takes "max" of core distances and pairwise distance. @@ -78,9 +123,9 @@ GRAPH_BUILD_ALGO check_params_validity(const all_neighbors_params& params, } } -// Single build (i.e. no batching) supports both host and device datasets +// Full build (i.e. no batching) for output indices/distances on device memory. template -void single_build( +void full_build( const raft::resources& handle, const all_neighbors_params& params, mdspan, row_major, Accessor> dataset, @@ -91,42 +136,147 @@ void single_build( size_t num_rows = static_cast(dataset.extent(0)); size_t num_cols = static_cast(dataset.extent(1)); - auto knn_builder = get_knn_builder( - handle, params, num_rows, num_rows, indices.extent(1), indices, distances, dist_epilogue); + auto knn_builder = + get_knn_builder(handle, params, num_rows, num_rows, indices.extent(1), dist_epilogue); knn_builder->prepare_build(dataset); - knn_builder->build_knn(dataset); + knn_builder->build_knn(dataset, std::nullopt, global_graph_view{indices, distances}); } -template -void build( +// Full build (i.e. no batching) for output indices/distances on host memory. +template +void full_build_host( const raft::resources& handle, const all_neighbors_params& params, - raft::host_matrix_view dataset, - raft::device_matrix_view indices, - std::optional> distances = std::nullopt, - std::optional> core_distances = std::nullopt, - T alpha = 1.0) + mdspan, row_major, Accessor> dataset, + raft::host_matrix_view indices, + std::optional> distances = std::nullopt, + DistEpilogueT dist_epilogue = DistEpilogueT{}) { - auto build_algo = check_params_validity(params, core_distances.has_value()); - - RAFT_EXPECTS(dataset.extent(0) == indices.extent(0), - "number of rows in dataset should be the same as number of rows in indices matrix"); + size_t num_rows = static_cast(dataset.extent(0)); + size_t k = static_cast(indices.extent(1)); + // kNN algorithms require device-resident output graph + auto indices_d = raft::make_device_matrix(handle, num_rows, k); + std::optional> distances_d; + std::optional> distances_d_view; if (distances.has_value()) { - RAFT_EXPECTS(indices.extent(0) == distances.value().extent(0) && - indices.extent(1) == distances.value().extent(1), - "indices matrix and distances matrix has to be the same shape."); + distances_d.emplace(raft::make_device_matrix(handle, num_rows, k)); + distances_d_view = distances_d.value().view(); } - if (core_distances.has_value()) { - RAFT_EXPECTS(distances.has_value(), - "distances matrix should be allocated to get mutual reachability distance."); + full_build(handle, params, dataset, indices_d.view(), distances_d_view, dist_epilogue); + + raft::copy(handle, indices, indices_d.view()); + if (distances.has_value()) { raft::copy(handle, distances.value(), distances_d.value().view()); } + raft::resource::sync_stream(handle); +} + +// Host counterparts of raft::matrix::shift used to insert the self-references. +template +void host_shift_self_indices(raft::host_matrix_view indices) +{ + size_t num_rows = static_cast(indices.extent(0)); + size_t k = static_cast(indices.extent(1)); +#pragma omp parallel for + for (size_t i = 0; i < num_rows; i++) { + IdxT* row = indices.data_handle() + i * k; + std::shift_right(row, row + k, 1); + row[0] = static_cast(i); } +} + +// Mirrors use cases of raft::matrix::shift, the distance side adjustments for inserting +// self-references. Shift each row of distances right by 1. Column 0 is filled from col0(i) when +// provided (e.g. the core distances), otherwise with fill_val. +template +void host_shift_distances( + raft::host_matrix_view distances, + std::optional fill_val, + std::optional> col0 = std::nullopt) +{ + size_t num_rows = static_cast(distances.extent(0)); + size_t k = static_cast(distances.extent(1)); +#pragma omp parallel for + for (size_t i = 0; i < num_rows; i++) { + T* row = distances.data_handle() + i * k; + std::shift_right(row, row + k, 1); + row[0] = col0.has_value() ? col0.value()(i) : fill_val.value(); + } +} + +// Inserts the self-reference for the kNN graph +template > +void shift_indices_distances(const raft::resources& handle, + IndicesView indices, + std::optional distances = std::nullopt, + std::optional core_distances = std::nullopt) +{ + constexpr bool host_output = + std::is_same_v>; + + if constexpr (host_output) { + host_shift_self_indices(indices); + if (distances.has_value()) { + if (core_distances.has_value()) { + host_shift_distances( + distances.value(), std::nullopt, raft::make_const_mdspan(core_distances.value())); + } else { + host_shift_distances(distances.value(), std::make_optional(0.0)); + } + } + } else { + raft::matrix::shift(handle, indices, 1); + if (distances.has_value()) { + if (core_distances.has_value()) { + raft::matrix::shift(handle, + distances.value(), + raft::make_device_matrix_view( + core_distances.value().data_handle(), indices.extent(0), IdxT{1})); + } else { + raft::matrix::shift(handle, distances.value(), 1, std::make_optional(0.0)); + } + } + } +} + +// Builds an all-neighbors knn graph with the dataset on host. The output (indices/distances/ +// core_distances) can be either on device or host. +template +void build(const raft::resources& handle, + const all_neighbors_params& params, + raft::host_matrix_view dataset, + IndicesView indices, + std::optional distances = std::nullopt, + std::optional core_distances = std::nullopt, + T alpha = 1.0) +{ + constexpr bool host_output = + std::is_same_v>; + + auto build_algo = + check_params_validity(params, dataset, indices, distances, core_distances); + + // Runs the no-batching build into the (device or host) output, selected at compile time. + auto run_full_build = [&](auto epilogue) { + if constexpr (host_output) { + full_build_host(handle, params, dataset, indices, distances, epilogue); + } else { + full_build(handle, params, dataset, indices, distances, epilogue); + } + }; std::unique_ptr> aux_vectors; if (params.n_clusters == 1) { - single_build(handle, params, dataset, indices, distances); + run_full_build(raft::identity_op{}); } else { if (core_distances.has_value()) { aux_vectors = std::make_unique>( @@ -140,39 +290,55 @@ void build( // NN Descent doesn't include self loops. Shifted to keep it consistent with brute force and ivfpq bool need_shift = (build_algo == GRAPH_BUILD_ALGO::NN_DESCENT) && (params.metric != cuvs::distance::DistanceType::InnerProduct); - - if (need_shift) { - raft::matrix::shift(handle, indices, 1); - if (distances.has_value()) { - raft::matrix::shift(handle, distances.value(), 1, std::make_optional(0.0)); - } - } + if (need_shift) { shift_indices_distances(handle, indices, distances); } if (core_distances.has_value()) { // calculate mutual reachability distances size_t k = indices.extent(1); size_t num_rows = core_distances.value().size(); - cuvs::neighbors::detail::reachability::core_distances( - handle, - distances.value().data_handle(), - k, - k, - num_rows, - core_distances.value().data_handle()); + + std::optional> core_dists_d; + const T* core_dists_ptr = nullptr; + + if constexpr (host_output) { + auto core_dists = core_distances.value(); + auto distances_value = distances.value(); +#pragma omp parallel for + for (size_t r = 0; r < num_rows; r++) { + core_dists(r) = distances_value(r, k - 1); + } + + if (params.n_clusters > 1 && raft::resource::is_multi_gpu(handle)) { + // Multi-GPU batched builds redistribute core distances to each GPU through host inside + // multi_gpu_batch_build + core_dists_ptr = core_dists.data_handle(); + } else { + core_dists_d.emplace(raft::make_device_vector(handle, num_rows)); + raft::copy(handle, + core_dists_d.value().view(), + raft::make_host_vector_view(core_dists.data_handle(), num_rows)); + core_dists_ptr = core_dists_d.value().data_handle(); + } + } else { + cuvs::neighbors::detail::reachability::core_distances( + handle, + distances.value().data_handle(), + k, + k, + num_rows, + core_distances.value().data_handle()); + core_dists_ptr = core_distances.value().data_handle(); + } using ReachabilityPP = cuvs::neighbors::detail::reachability::ReachabilityPostProcess; - auto dist_epilogue = ReachabilityPP{core_distances.value().data_handle(), alpha, num_rows}; + auto dist_epilogue = ReachabilityPP{core_dists_ptr, alpha, num_rows}; if (params.n_clusters == 1) { - single_build(handle, params, dataset, indices, distances, dist_epilogue); + run_full_build(dist_epilogue); } else { batch_build(handle, params, dataset, indices, distances, aux_vectors.get(), dist_epilogue); } if (need_shift) { - raft::matrix::shift(handle, indices, 1); - raft::matrix::shift(handle, - distances.value(), - raft::make_device_matrix_view( - core_distances.value().data_handle(), num_rows, 1)); + shift_indices_distances(handle, indices, distances, core_distances); } } } @@ -187,40 +353,21 @@ void build( std::optional> core_distances = std::nullopt, T alpha = 1.0) { - auto build_algo = check_params_validity(params, core_distances.has_value()); - - RAFT_EXPECTS(dataset.extent(0) == indices.extent(0), - "number of rows in dataset should be the same as number of rows in indices matrix"); - - if (distances.has_value()) { - RAFT_EXPECTS(indices.extent(0) == distances.value().extent(0) && - indices.extent(1) == distances.value().extent(1), - "indices matrix and distances matrix has to be the same shape."); - } - - if (core_distances.has_value()) { - RAFT_EXPECTS(distances.has_value(), - "distances matrix should be allocated to get mutual reachability distance."); - } + auto build_algo = + check_params_validity(params, dataset, indices, distances, core_distances); if (params.n_clusters > 1) { RAFT_FAIL( "Batched all-neighbors build is not supported with data on device. Put data on host for " "batch build."); } else { - single_build(handle, params, dataset, indices, distances); + full_build(handle, params, dataset, indices, distances); } // NN Descent doesn't include self loops. Shifted to keep it consistent with brute force and ivfpq bool need_shift = (build_algo == GRAPH_BUILD_ALGO::NN_DESCENT) && (params.metric != cuvs::distance::DistanceType::InnerProduct); - - if (need_shift) { - raft::matrix::shift(handle, indices, 1); - if (distances.has_value()) { - raft::matrix::shift(handle, distances.value(), 1, std::make_optional(0.0)); - } - } + if (need_shift) { shift_indices_distances(handle, indices, distances); } if (core_distances.has_value()) { size_t k = indices.extent(1); @@ -235,14 +382,10 @@ void build( using ReachabilityPP = cuvs::neighbors::detail::reachability::ReachabilityPostProcess; auto dist_epilogue = ReachabilityPP{core_distances.value().data_handle(), alpha, num_rows}; - single_build(handle, params, dataset, indices, distances, dist_epilogue); + full_build(handle, params, dataset, indices, distances, dist_epilogue); if (need_shift) { - raft::matrix::shift(handle, indices, 1); - raft::matrix::shift(handle, - distances.value(), - raft::make_device_matrix_view( - core_distances.value().data_handle(), num_rows, 1)); + shift_indices_distances(handle, indices, distances, core_distances); } } } diff --git a/cpp/src/neighbors/all_neighbors/all_neighbors_batched.cuh b/cpp/src/neighbors/all_neighbors/all_neighbors_batched.cuh index 7e4ff748a4..33d9b3b41f 100644 --- a/cpp/src/neighbors/all_neighbors/all_neighbors_batched.cuh +++ b/cpp/src/neighbors/all_neighbors/all_neighbors_batched.cuh @@ -14,6 +14,7 @@ #include #include #include +#include #include #include #include @@ -26,16 +27,17 @@ #include #include #include +#include #include namespace cuvs::neighbors::all_neighbors::detail { using namespace cuvs::neighbors; -template +template void reset_global_matrices(raft::resources const& res, cuvs::distance::DistanceType metric, - raft::managed_matrix_view global_neighbors, - raft::managed_matrix_view global_distances) + NbrView global_neighbors, + DistView global_distances) { size_t num_rows = static_cast(global_neighbors.extent(0)); size_t k = static_cast(global_neighbors.extent(1)); @@ -46,14 +48,24 @@ void reset_global_matrices(raft::resources const& res, T global_distances_fill_value = select_min ? std::numeric_limits::max() : std::numeric_limits::min(); - raft::matrix::fill( - res, - raft::make_device_matrix_view(global_neighbors.data_handle(), num_rows, k), - global_neighbors_fill_value); - raft::matrix::fill( - res, - raft::make_device_matrix_view(global_distances.data_handle(), num_rows, k), - global_distances_fill_value); + if constexpr (std::is_same_v>) { + std::fill(global_neighbors.data_handle(), + global_neighbors.data_handle() + num_rows * k, + global_neighbors_fill_value); + std::fill(global_distances.data_handle(), + global_distances.data_handle() + num_rows * k, + global_distances_fill_value); + } else { + // device or managed memory: fill on the GPU (managed is device-accessible) + raft::matrix::fill( + res, + raft::make_device_matrix_view(global_neighbors.data_handle(), num_rows, k), + global_neighbors_fill_value); + raft::matrix::fill( + res, + raft::make_device_matrix_view(global_distances.data_handle(), num_rows, k), + global_distances_fill_value); + } } /** @@ -291,8 +303,7 @@ void single_gpu_batch_build(const raft::resources& handle, raft::host_matrix_view dataset, detail::all_neighbors_builder& knn_builder, size_t n_clusters, - raft::managed_matrix_view global_neighbors, - raft::managed_matrix_view global_distances, + global_graph_view global, raft::host_vector_view cluster_sizes, raft::host_vector_view cluster_offsets, raft::host_vector_view inverted_indices) @@ -330,24 +341,28 @@ void single_gpu_batch_build(const raft::resources& handle, auto inverted_indices_view = raft::make_host_vector_view( inverted_indices.data_handle() + offset, num_data_in_cluster); - knn_builder.build_knn( - cluster_data_view, inverted_indices_view, global_neighbors, global_distances); + knn_builder.build_knn(cluster_data_view, inverted_indices_view, global); } } -template +template void multi_gpu_batch_build(const raft::resources& handle, const all_neighbors_params& params, raft::host_matrix_view dataset, - raft::managed_matrix_view global_neighbors, - raft::managed_matrix_view global_distances, + GNbrView global_neighbors, + GDistView global_distances, raft::host_vector_view cluster_sizes, raft::host_vector_view cluster_offsets_c, raft::host_vector_view inverted_indices, DistEpilogueT dist_epilogue = DistEpilogueT{}) { + constexpr bool host_output = std::is_same_v>; + size_t num_rows = dataset.extent(0); - size_t num_cols = dataset.extent(1); size_t k = global_neighbors.extent(1); int num_ranks = raft::resource::get_num_ranks(handle); @@ -358,16 +373,31 @@ void multi_gpu_batch_build(const raft::resources& handle, auto cluster_offsets = raft::make_host_vector(cluster_offsets_c.size()); raft::copy(handle, cluster_offsets.view(), cluster_offsets_c); + reset_global_matrices(handle, params.metric, global_neighbors, global_distances); + using ReachabilityPP = cuvs::neighbors::detail::reachability::ReachabilityPostProcess; const bool mutual_reach_dist = std::is_same_v; - std::optional> core_distances_h; + + // Mutual-reachability core distances are staged on host to redistribute to each rank's device. + std::optional> core_distances_h_owner; + std::optional> core_distances_h; if constexpr (mutual_reach_dist) { - core_distances_h.emplace(raft::make_host_vector(num_rows)); - raft::copy(handle, - core_distances_h.value().view(), - raft::make_device_vector_view(dist_epilogue.core_dists, num_rows)); + if constexpr (host_output) { + core_distances_h = + raft::make_host_vector_view(dist_epilogue.core_dists, num_rows); + } else { + core_distances_h_owner.emplace(raft::make_host_vector(num_rows)); + raft::copy(handle, + core_distances_h_owner.value().view(), + raft::make_device_vector_view(dist_epilogue.core_dists, num_rows)); + core_distances_h = raft::make_const_mdspan(core_distances_h_owner.value().view()); + } } + // We only have at max num_ranks contentions + constexpr size_t max_row_locks = 8192; + std::vector row_locks(host_output ? std::min(num_rows, max_row_locks) : size_t{0}); + // Ensure all async copies complete before starting parallel region raft::resource::sync_stream(handle); @@ -418,9 +448,7 @@ void multi_gpu_batch_build(const raft::resources& handle, auto dist_epilgogue_for_rank = [&]() { if constexpr (mutual_reach_dist) { core_distances_d_for_rank.emplace(raft::make_device_vector(dev_res, num_rows)); - raft::copy(dev_res, - core_distances_d_for_rank.value().view(), - raft::make_const_mdspan(core_distances_h.value().view())); + raft::copy(dev_res, core_distances_d_for_rank.value().view(), core_distances_h.value()); return ReachabilityPP{ core_distances_d_for_rank.value().data_handle(), dist_epilogue.alpha, num_rows}; } else { @@ -428,22 +456,20 @@ void multi_gpu_batch_build(const raft::resources& handle, } }(); - std::unique_ptr> knn_builder = - get_knn_builder(dev_res, - params, - min_cluster_size, - max_cluster_size, - k, - std::nullopt, - std::nullopt, - dist_epilgogue_for_rank); + std::unique_ptr> knn_builder = get_knn_builder( + dev_res, params, min_cluster_size, max_cluster_size, k, dist_epilgogue_for_rank); + + auto global = global_graph_view{global_neighbors, global_distances}; + if constexpr (host_output) { + global.row_locks = row_locks.data(); + global.num_row_locks = row_locks.size(); + } single_gpu_batch_build(dev_res, dataset, *knn_builder, num_clusters_for_this_rank, - global_neighbors, - global_distances, + global, cluster_sizes_for_this_rank, cluster_offsets_for_this_rank, inverted_indices_for_this_rank); @@ -466,21 +492,27 @@ struct BatchBuildAux { } }; -template -void batch_build( - const raft::resources& handle, - const all_neighbors_params& params, - raft::host_matrix_view dataset, - raft::device_matrix_view indices, - std::optional> distances = std::nullopt, - BatchBuildAux* aux_vectors = nullptr, - DistEpilogueT dist_epilogue = DistEpilogueT{}) +template +void batch_build(const raft::resources& handle, + const all_neighbors_params& params, + raft::host_matrix_view dataset, + IndicesView indices, + std::optional distances = std::nullopt, + BatchBuildAux* aux_vectors = nullptr, + DistEpilogueT dist_epilogue = DistEpilogueT{}) { if (raft::resource::is_multi_gpu(handle)) { // For efficient CPU-computation of omp parallel for regions per GPU cuvs::core::omp::set_nested(1); } + constexpr bool host_output = + std::is_same_v>; + size_t num_rows = static_cast(dataset.extent(0)); size_t num_cols = static_cast(dataset.extent(1)); size_t k = indices.extent(1); @@ -528,10 +560,64 @@ void batch_build( if (aux_vectors != nullptr) { aux_vectors->is_computed = true; } } - auto global_neighbors = raft::make_managed_matrix(handle, num_rows, k); - auto global_distances = raft::make_managed_matrix(handle, num_rows, k); + auto run_single_gpu = [&](global_graph_view global) { + size_t max_cluster_size, min_cluster_size; + get_min_max_cluster_size(k, max_cluster_size, min_cluster_size, cluster_sizes_view); + std::unique_ptr> knn_builder = get_knn_builder( + handle, params, min_cluster_size, max_cluster_size, k, dist_epilogue); + single_gpu_batch_build(handle, + dataset, + *knn_builder, + params.n_clusters, + global, + cluster_sizes_view, + cluster_offsets_view, + inverted_indices_view); + }; + + if constexpr (host_output) { + // Host output: the caller's host indices array is the global neighbors buffer. We need a + // distances matrix for the merge, so use the caller's when provided, otherwise a host buffer. + std::optional> global_distances_buffer; + auto global_neighbors = indices; + auto global_distances = + distances.has_value() + ? distances.value() + : global_distances_buffer.emplace(raft::make_host_matrix(num_rows, k)).view(); + + if (raft::resource::is_multi_gpu(handle)) { + multi_gpu_batch_build(handle, + params, + dataset, + global_neighbors, + global_distances, + cluster_sizes_view, + raft::make_const_mdspan(cluster_offsets_view), + inverted_indices_view, + dist_epilogue); + } else { + reset_global_matrices(handle, params.metric, global_neighbors, global_distances); + run_single_gpu(global_graph_view{global_neighbors, global_distances}); + } + } else if (!raft::resource::is_multi_gpu(handle)) { + // Single-GPU device output: the caller's device output is the global graph. We need a distances + // matrix for the merge, so use the caller's when provided, otherwise a device buffer. + std::optional> global_distances_buffer; + auto global_distances = + distances.has_value() + ? distances.value() + : global_distances_buffer.emplace(raft::make_device_matrix(handle, num_rows, k)) + .view(); + + reset_global_matrices(handle, params.metric, indices, global_distances); + run_single_gpu(global_graph_view{indices, global_distances}); + } else { + // Multi-GPU device output: the shared global graph must be reachable from every GPU, so it is + // accumulated in managed memory (viewed as device) and copied to the caller's output at the + // end. + auto global_neighbors = raft::make_managed_matrix(handle, num_rows, k); + auto global_distances = raft::make_managed_matrix(handle, num_rows, k); - if (raft::resource::is_multi_gpu(handle)) { // Check if any GPU is Turing (SM 7.5) or older. These architectures have issues with // multi-GPU managed memory coherence for concurrent writes. Force CPU-resident memory // to ensure all GPUs access through host memory, avoiding page migration issues. @@ -565,14 +651,6 @@ void batch_build( cudaMemAdviseSetPreferredLocation, cpu_location)); } - } - - reset_global_matrices(handle, params.metric, global_neighbors.view(), global_distances.view()); - - if (raft::resource::is_multi_gpu(handle)) { - // For multi-GPU: sync the stream to ensure fill completes before other GPUs access - // the managed memory. - raft::resource::sync_stream(handle); multi_gpu_batch_build(handle, params, @@ -583,38 +661,11 @@ void batch_build( raft::make_const_mdspan(cluster_offsets_view), inverted_indices_view, dist_epilogue); - } else { - size_t max_cluster_size, min_cluster_size; - get_min_max_cluster_size(k, max_cluster_size, min_cluster_size, cluster_sizes_view); - std::unique_ptr> knn_builder = - get_knn_builder(handle, - params, - min_cluster_size, - max_cluster_size, - k, - std::nullopt, - std::nullopt, - dist_epilogue); - single_gpu_batch_build(handle, - dataset, - *knn_builder, - params.n_clusters, - global_neighbors.view(), - global_distances.view(), - cluster_sizes_view, - cluster_offsets_view, - inverted_indices_view); - } - raft::copy( - handle, - raft::make_device_vector_view(indices.data_handle(), num_rows * k), - raft::make_device_vector_view(global_neighbors.data_handle(), num_rows * k)); - if (distances.has_value()) { - raft::copy( - handle, - raft::make_device_vector_view(distances.value().data_handle(), num_rows * k), - raft::make_device_vector_view(global_distances.data_handle(), num_rows * k)); + raft::copy(handle, indices, raft::make_const_mdspan(global_neighbors.view())); + if (distances.has_value()) { + raft::copy(handle, distances.value(), raft::make_const_mdspan(global_distances.view())); + } } } diff --git a/cpp/src/neighbors/all_neighbors/all_neighbors_builder.cuh b/cpp/src/neighbors/all_neighbors/all_neighbors_builder.cuh index 8f54ea0b24..8e56735cc2 100644 --- a/cpp/src/neighbors/all_neighbors/all_neighbors_builder.cuh +++ b/cpp/src/neighbors/all_neighbors/all_neighbors_builder.cuh @@ -28,30 +28,78 @@ #include #include +#include + namespace cuvs::neighbors::all_neighbors::detail { using namespace cuvs::neighbors; +/** + * The final destination for the all-neighbors graph a builder produces. Its memory location selects + * the strategy: + * - direct_t: full build (no batching). The builder writes the full [num_rows x k] result + * straight into device-resident output (optional distances). + * - merge_device: batched, plain-device global graph (single-GPU device output, merged into the + * caller's device arrays). Merge happens on the GPU. + * - merge_managed: batched, managed global graph for multi-GPU setting. Merge happens on the GPU. + * - merge_host: batched, host-resident global graph. Merged happens on the host. + */ +template +struct global_graph_view { + using direct_t = std::pair, + std::optional>>; + using merge_device = + std::pair, raft::device_matrix_view>; + using merge_managed = + std::pair, raft::managed_matrix_view>; + using merge_host = std::pair, raft::host_matrix_view>; + + std::variant dest; + + // Used for merge_host type. + std::mutex* row_locks = nullptr; + size_t num_row_locks = 0; + + // full-build (no batching): write directly to device output + global_graph_view(raft::device_matrix_view indices, + std::optional> distances) + : dest{direct_t{indices, distances}} + { + } + // batched, device global graph (single-GPU device output) + global_graph_view(raft::device_matrix_view neighbors, + raft::device_matrix_view distances) + : dest{merge_device{neighbors, distances}} + { + } + // batched, managed global graph (multi-GPU device output) + global_graph_view(raft::managed_matrix_view neighbors, + raft::managed_matrix_view distances) + : dest{merge_managed{neighbors, distances}} + { + } + // batched, host-resident global graph + global_graph_view(raft::host_matrix_view neighbors, + raft::host_matrix_view distances) + : dest{merge_host{neighbors, distances}} + { + } + + [[nodiscard]] bool is_full_build() const { return std::holds_alternative(dest); } +}; + template struct all_neighbors_builder { - all_neighbors_builder( - raft::resources const& res, - size_t n_clusters, - size_t min_cluster_size, - size_t max_cluster_size, - size_t k, - std::optional> indices = std::nullopt, - std::optional> distances = std::nullopt) + all_neighbors_builder(raft::resources const& res, + size_t n_clusters, + size_t min_cluster_size, + size_t max_cluster_size, + size_t k) : res{res}, n_clusters{n_clusters}, min_cluster_size{min_cluster_size}, max_cluster_size{max_cluster_size}, - k{k}, - indices_{indices}, - distances_{distances} + k{k} { - RAFT_EXPECTS(this->n_clusters > 1 || indices.has_value(), - "indices should be preallocated to create knn builder for n_clusters == 1 (no " - "batching mode)"); if (n_clusters > 1) { // allocating additional space needed for batching inverted_indices_d.emplace(raft::make_device_vector(res, max_cluster_size)); batch_neighbors_h.emplace(raft::make_host_matrix(max_cluster_size, k)); @@ -60,6 +108,74 @@ struct all_neighbors_builder { } } + /** + * Merge this cluster's knn results into the global graph. The merge strategy is selected by the + * memory location of the global_graph_view: + * - device / managed: the global [num_rows x k] graph is device-accessible; merge runs on the + * GPU (merge_subgraphs_kernel). + * - host: the global graph lives in host memory; the per-cluster distances are staged + * device->host and the merge runs on the host. + */ + template + void do_merge(raft::host_matrix_view indices_for_remap_h, + raft::host_vector_view inverted_indices, + const global_graph_view& global, + size_t num_data_in_cluster, + bool select_min) + { + using merge_device = typename global_graph_view::merge_device; + using merge_managed = typename global_graph_view::merge_managed; + using merge_host = typename global_graph_view::merge_host; + + // GPU merge for a device-accessible global graph (device or managed). + auto gpu_merge = [&](auto global_neighbors, auto global_distances) { + remap_and_merge_subgraphs(res, + inverted_indices_d.value().view(), + inverted_indices, + indices_for_remap_h, + batch_neighbors_h.value().view(), + batch_neighbors_d.value().view(), + batch_distances_d.value().view(), + global_neighbors, + global_distances, + num_data_in_cluster, + k, + select_min); + }; + + if (const auto* g = std::get_if(&global.dest)) { + gpu_merge(g->first, g->second); + } else if (const auto* g = std::get_if(&global.dest)) { + gpu_merge(g->first, g->second); + } else if (const auto* g = std::get_if(&global.dest)) { + if (!batch_distances_h.has_value()) { + batch_distances_h.emplace(raft::make_host_matrix(max_cluster_size, k)); + } + // stage this cluster's distances from device to host + auto batch_distances_h_view = raft::make_host_matrix_view( + batch_distances_h.value().data_handle(), num_data_in_cluster, k); + raft::copy(res, + batch_distances_h_view, + raft::make_device_matrix_view( + batch_distances_d.value().data_handle(), num_data_in_cluster, k)); + raft::resource::sync_stream(res); + + remap_and_merge_subgraphs(res, + inverted_indices, + indices_for_remap_h, + batch_distances_h_view, + g->first, + g->second, + num_data_in_cluster, + k, + select_min, + global.row_locks, + global.num_row_locks); + } else { + RAFT_FAIL("do_merge requires a device, managed, or host global graph destination"); + } + } + /** * Some memory-heavy allocations that can be used over multiple clusters should be allocated here * Arguments: @@ -74,20 +190,20 @@ struct all_neighbors_builder { * - [in] dataset: host_matrix_view or device_matrix_view of the cluster dataset * - [in] inverted_indices (optional): global data indices for the data points in the current * cluster of size [num_data_in_cluster]. Only needed when using the batching algorithm. - * - [out] global_neighbors (optional): raft::managed_matrix_view type of [total_num_rows, k] for - * final all-neighbors graph indices. Only needed when using the batching algorithm. - * - [out] global_distances (optional): raft::managed_matrix_view type of [total_num_rows, k] for - * final all-neighbors graph distances. Only needed when using the batching algorithm. + * - [out] global: the destination for the all-neighbors graph (global_graph_view). */ - virtual void build_knn( - raft::host_matrix_view dataset, - std::optional> inverted_indices = std::nullopt, - std::optional> global_neighbors = std::nullopt, - std::optional> global_distances = std::nullopt) + virtual void build_knn(raft::host_matrix_view dataset, + std::optional> inverted_indices, + global_graph_view global) { } - virtual void build_knn(raft::device_matrix_view dataset) {} + // device dataset is full-build only (no batching), so inverted_indices is always empty. + virtual void build_knn(raft::device_matrix_view dataset, + std::optional> inverted_indices, + global_graph_view global) + { + } virtual ~all_neighbors_builder() = default; @@ -100,27 +216,19 @@ struct all_neighbors_builder { std::optional> batch_neighbors_h; std::optional> batch_neighbors_d; std::optional> batch_distances_d; - - // optional indices and distances, only used when n_clusters=1 - // when n_clusters > 1 (i.e. doing batching), we write to the global_neighbors and - // global_distances managed memory arrays. - std::optional> indices_; - std::optional> distances_; + // host staging buffer for per-cluster distances feeding the host-side merge + std::optional> batch_distances_h; }; template struct all_neighbors_builder_ivfpq : public all_neighbors_builder { - all_neighbors_builder_ivfpq( - raft::resources const& res, - size_t n_clusters, - size_t min_cluster_size, - size_t max_cluster_size, - size_t k, - graph_build_params::ivf_pq_params& params, - std::optional> indices = std::nullopt, - std::optional> distances = std::nullopt) - : all_neighbors_builder( - res, n_clusters, min_cluster_size, max_cluster_size, k, indices, distances), + all_neighbors_builder_ivfpq(raft::resources const& res, + size_t n_clusters, + size_t min_cluster_size, + size_t max_cluster_size, + size_t k, + graph_build_params::ivf_pq_params& params) + : all_neighbors_builder(res, n_clusters, min_cluster_size, max_cluster_size, k), all_ivf_pq_params{params} { } @@ -162,18 +270,14 @@ struct all_neighbors_builder_ivfpq : public all_neighbors_builder { // Actual build logic using ivfpq. // need device and host views of the dataset because ivfpq build and search uses the device view, // and refine uses the host view - void build_knn_common( - raft::device_matrix_view dataset_d, - raft::host_matrix_view dataset_h, - std::optional> inverted_indices = std::nullopt, - std::optional> global_neighbors = std::nullopt, - std::optional> global_distances = std::nullopt) + void build_knn_common(raft::device_matrix_view dataset_d, + raft::host_matrix_view dataset_h, + std::optional> inverted_indices, + global_graph_view global) { RAFT_EXPECTS( - this->n_clusters <= 1 || (inverted_indices.has_value() && global_neighbors.has_value() && - global_distances.has_value()), - "need valid inverted_indices, global_neighbors, and global_distances for " - "build_knn if doing batching."); + global.is_full_build() || inverted_indices.has_value(), + "need valid inverted_indices for a batched (managed/host) global graph destination"); size_t num_data_in_cluster = dataset_d.extent(0); size_t num_cols = dataset_d.extent(1); @@ -212,49 +316,39 @@ struct all_neighbors_builder_ivfpq : public all_neighbors_builder { refined_distances_h_view, all_ivf_pq_params.build_params.metric); - if (this->n_clusters > 1) { // do batching + if (!global.is_full_build()) { // batched: merge this cluster into the global graph raft::copy(this->res, raft::make_device_vector_view(this->batch_distances_d.value().data_handle(), num_data_in_cluster * this->k), raft::make_host_vector_view(refined_distances_h_view.data_handle(), num_data_in_cluster * this->k)); - remap_and_merge_subgraphs( - this->res, - this->inverted_indices_d.value().view(), - inverted_indices.value(), + this->template do_merge( refined_neighbors_h.value().view(), - this->batch_neighbors_h.value().view(), - this->batch_neighbors_d.value().view(), - this->batch_distances_d.value().view(), - global_neighbors.value(), - global_distances.value(), + inverted_indices.value(), + global, num_data_in_cluster, - this->k, cuvs::distance::is_min_close(all_ivf_pq_params.build_params.metric)); - } else { - size_t num_rows = num_data_in_cluster; - // copy resulting indices and distances to device output - raft::copy( - this->res, - raft::make_device_vector_view(this->indices_.value().data_handle(), num_rows * this->k), - raft::make_host_vector_view(refined_neighbors_h_view.data_handle(), - num_rows * this->k)); - if (this->distances_.has_value()) { + } else { // full build: write directly to the device output in the sink + const auto& direct = std::get::direct_t>(global.dest); + size_t num_rows = num_data_in_cluster; + raft::copy(this->res, + raft::make_device_vector_view(direct.first.data_handle(), num_rows * this->k), + raft::make_host_vector_view(refined_neighbors_h_view.data_handle(), + num_rows * this->k)); + if (direct.second.has_value()) { raft::copy( this->res, - raft::make_device_vector_view(this->distances_.value().data_handle(), num_rows * this->k), + raft::make_device_vector_view(direct.second.value().data_handle(), num_rows * this->k), raft::make_host_vector_view(refined_distances_h_view.data_handle(), num_rows * this->k)); } } } - void build_knn( - raft::host_matrix_view dataset, - std::optional> inverted_indices = std::nullopt, - std::optional> global_neighbors = std::nullopt, - std::optional> global_distances = std::nullopt) override + void build_knn(raft::host_matrix_view dataset, + std::optional> inverted_indices, + global_graph_view global) override { // we need data on device for ivfpq build and search. raft::copy(this->res, @@ -265,11 +359,12 @@ struct all_neighbors_builder_ivfpq : public all_neighbors_builder { data_d.value().data_handle(), dataset.extent(0), dataset.extent(1)), dataset, inverted_indices, - global_neighbors, - global_distances); + global); } - void build_knn(raft::device_matrix_view dataset) override + void build_knn(raft::device_matrix_view dataset, + std::optional> /*inverted_indices*/, + global_graph_view global) override { RAFT_EXPECTS(this->n_clusters <= 1, "building all-neighbors knn graph with dataset on device is not supported with " @@ -286,7 +381,9 @@ struct all_neighbors_builder_ivfpq : public all_neighbors_builder { build_knn_common(dataset, raft::make_host_matrix_view( - dataset_h.data_handle(), dataset.extent(0), dataset.extent(1))); + dataset_h.data_handle(), dataset.extent(0), dataset.extent(1)), + std::nullopt, + global); } graph_build_params::ivf_pq_params all_ivf_pq_params; @@ -304,18 +401,14 @@ struct all_neighbors_builder_ivfpq : public all_neighbors_builder { template struct all_neighbors_builder_nn_descent : public all_neighbors_builder { - all_neighbors_builder_nn_descent( - raft::resources const& res, - size_t n_clusters, - size_t min_cluster_size, - size_t max_cluster_size, - size_t k, - graph_build_params::nn_descent_params& params, - std::optional> indices = std::nullopt, - std::optional> distances = std::nullopt, - DistEpilogueT dist_epilogue = DistEpilogueT{}) - : all_neighbors_builder( - res, n_clusters, min_cluster_size, max_cluster_size, k, indices, distances), + all_neighbors_builder_nn_descent(raft::resources const& res, + size_t n_clusters, + size_t min_cluster_size, + size_t max_cluster_size, + size_t k, + graph_build_params::nn_descent_params& params, + DistEpilogueT dist_epilogue = DistEpilogueT{}) + : all_neighbors_builder(res, n_clusters, min_cluster_size, max_cluster_size, k), nnd_params{params}, dist_epilogue{dist_epilogue} { @@ -366,21 +459,16 @@ struct all_neighbors_builder_nn_descent : public all_neighbors_builder } template - void build_knn_common( - mdspan, row_major, Accessor> dataset, - std::optional> inverted_indices = std::nullopt, - std::optional> global_neighbors = std::nullopt, - std::optional> global_distances = std::nullopt) + void build_knn_common(mdspan, row_major, Accessor> dataset, + std::optional> inverted_indices, + global_graph_view global) { - RAFT_EXPECTS( - this->n_clusters <= 1 || (inverted_indices.has_value() && global_neighbors.has_value() && - global_distances.has_value()), - "need valid inverted_indices, global_neighbors, and global_distances for " - "build_knn if doing batching."); + RAFT_EXPECTS(global.is_full_build() || inverted_indices.has_value(), + "need valid inverted_indices for a batched global graph destination"); using ReachabilityPP = cuvs::neighbors::detail::reachability::ReachabilityPostProcess; - if (this->n_clusters > 1) { + if (!global.is_full_build()) { bool return_distances = true; size_t num_data_in_cluster = dataset.extent(0); if constexpr (std::is_same_v) { @@ -415,29 +503,23 @@ struct all_neighbors_builder_nn_descent : public all_neighbors_builder this->batch_distances_d.value().data_handle()); } - remap_and_merge_subgraphs>( - this->res, - this->inverted_indices_d.value().view(), - inverted_indices.value(), + this->template do_merge>( int_graph.value().view(), - this->batch_neighbors_h.value().view(), - this->batch_neighbors_d.value().view(), - this->batch_distances_d.value().view(), - global_neighbors.value(), - global_distances.value(), + inverted_indices.value(), + global, num_data_in_cluster, - this->k, cuvs::distance::is_min_close(nnd_params.metric)); - } else { - size_t num_rows = dataset.extent(0); + } else { // full build: write directly to the device output + const auto& direct = std::get::direct_t>(global.dest); + size_t num_rows = dataset.extent(0); if constexpr (std::is_same_v) { nnd_builder.value().build( dataset.data_handle(), static_cast(num_rows), int_graph.value().data_handle(), - this->distances_.has_value(), - this->distances_.value_or(raft::make_device_matrix(this->res, 0, 0).view()) + direct.second.has_value(), + direct.second.value_or(raft::make_device_matrix(this->res, 0, 0).view()) .data_handle(), cuvs::neighbors::detail::reachability::ReachabilityPostProcess{ dist_epilogue.core_dists, dist_epilogue.alpha, dist_epilogue.n}); @@ -446,8 +528,8 @@ struct all_neighbors_builder_nn_descent : public all_neighbors_builder dataset.data_handle(), static_cast(num_rows), int_graph.value().data_handle(), - this->distances_.has_value(), - this->distances_.value_or(raft::make_device_matrix(this->res, 0, 0).view()) + direct.second.has_value(), + direct.second.value_or(raft::make_device_matrix(this->res, 0, 0).view()) .data_handle(), dist_epilogue); } @@ -463,29 +545,29 @@ struct all_neighbors_builder_nn_descent : public all_neighbors_builder } // copy to final device output - raft::copy(this->res, - raft::make_device_vector_view(this->indices_.value().data_handle(), - tmp_indices.extent(0) * this->k), - raft::make_host_vector_view(tmp_indices.data_handle(), - tmp_indices.extent(0) * this->k)); + raft::copy( + this->res, + raft::make_device_vector_view(direct.first.data_handle(), tmp_indices.extent(0) * this->k), + raft::make_host_vector_view(tmp_indices.data_handle(), + tmp_indices.extent(0) * this->k)); } } - void build_knn( - raft::host_matrix_view dataset, - std::optional> inverted_indices = std::nullopt, - std::optional> global_neighbors = std::nullopt, - std::optional> global_distances = std::nullopt) override + void build_knn(raft::host_matrix_view dataset, + std::optional> inverted_indices, + global_graph_view global) override { - build_knn_common(dataset, inverted_indices, global_neighbors, global_distances); + build_knn_common(dataset, inverted_indices, global); } - void build_knn(raft::device_matrix_view dataset) override + void build_knn(raft::device_matrix_view dataset, + std::optional> inverted_indices, + global_graph_view global) override { RAFT_EXPECTS(this->n_clusters <= 1, "building all-neighbors knn graph with dataset on device is not supported with " "batching (n_clusters > 1)"); - build_knn_common(dataset); + build_knn_common(dataset, inverted_indices, global); } nn_descent::index_params nnd_params; @@ -500,18 +582,14 @@ struct all_neighbors_builder_nn_descent : public all_neighbors_builder template struct all_neighbors_builder_brute_force : public all_neighbors_builder { - all_neighbors_builder_brute_force( - raft::resources const& res, - size_t n_clusters, - size_t min_cluster_size, - size_t max_cluster_size, - size_t k, - graph_build_params::brute_force_params& params, - std::optional> indices = std::nullopt, - std::optional> distances = std::nullopt, - DistEpilogueT dist_epilogue = DistEpilogueT{}) - : all_neighbors_builder( - res, n_clusters, min_cluster_size, max_cluster_size, k, indices, distances), + all_neighbors_builder_brute_force(raft::resources const& res, + size_t n_clusters, + size_t min_cluster_size, + size_t max_cluster_size, + size_t k, + graph_build_params::brute_force_params& params, + DistEpilogueT dist_epilogue = DistEpilogueT{}) + : all_neighbors_builder(res, n_clusters, min_cluster_size, max_cluster_size, k), bf_params{params}, dist_epilogue{dist_epilogue} { @@ -533,19 +611,14 @@ struct all_neighbors_builder_brute_force : public all_neighbors_builder } } - void build_knn_common( - raft::device_matrix_view dataset, - std::optional> inverted_indices = std::nullopt, - std::optional> global_neighbors = std::nullopt, - std::optional> global_distances = std::nullopt) + void build_knn_common(raft::device_matrix_view dataset, + std::optional> inverted_indices, + global_graph_view global) { - RAFT_EXPECTS( - this->n_clusters <= 1 || (inverted_indices.has_value() && global_neighbors.has_value() && - global_distances.has_value()), - "need valid inverted_indices, global_neighbors, and global_distances for " - "build_knn if doing batching."); + RAFT_EXPECTS(global.is_full_build() || inverted_indices.has_value(), + "need valid inverted_indices for a batched global graph destination"); - if (this->n_clusters > 1) { + if (!global.is_full_build()) { size_t num_data_in_cluster = dataset.extent(0); using ReachabilityPP = cuvs::neighbors::detail::reachability::ReachabilityPostProcess; @@ -604,32 +677,23 @@ struct all_neighbors_builder_brute_force : public all_neighbors_builder raft::make_device_vector_view( this->batch_neighbors_d.value().data_handle(), num_data_in_cluster * this->k)); - remap_and_merge_subgraphs>( - this->res, - this->inverted_indices_d.value().view(), - inverted_indices.value(), - this->batch_neighbors_h.value().view(), + this->template do_merge>( this->batch_neighbors_h.value().view(), - this->batch_neighbors_d.value().view(), - this->batch_distances_d.value().view(), - global_neighbors.value(), - global_distances.value(), + inverted_indices.value(), + global, num_data_in_cluster, - this->k, cuvs::distance::is_min_close(bf_params.build_params.metric)); - } else { + } else { // full build: write directly to the device output in the sink + const auto& direct = std::get::direct_t>(global.dest); + auto distances_view = + direct.second.has_value() + ? direct.second.value() + : raft::make_device_matrix(this->res, dataset.extent(0), this->k).view(); if constexpr (std::is_same_v) { auto idx = cuvs::neighbors::brute_force::build(this->res, bf_params.build_params, dataset); cuvs::neighbors::brute_force::search( - this->res, - bf_params.search_params, - idx, - dataset, - this->indices_.value(), - this->distances_.has_value() - ? this->distances_.value() - : raft::make_device_matrix(this->res, dataset.extent(0), this->k).view()); + this->res, bf_params.search_params, idx, dataset, direct.first, distances_view); } else { cuvs::neighbors::detail::tiled_brute_force_knn( this->res, @@ -639,11 +703,8 @@ struct all_neighbors_builder_brute_force : public all_neighbors_builder dataset.extent(0), dataset.extent(1), this->k, - this->distances_.has_value() - ? this->distances_.value().data_handle() - : raft::make_device_matrix(this->res, dataset.extent(0), this->k) - .data_handle(), - this->indices_.value().data_handle(), + distances_view.data_handle(), + direct.first.data_handle(), bf_params.build_params.metric, 2.0, 0, @@ -656,11 +717,9 @@ struct all_neighbors_builder_brute_force : public all_neighbors_builder } } - void build_knn( - raft::host_matrix_view dataset, - std::optional> inverted_indices = std::nullopt, - std::optional> global_neighbors = std::nullopt, - std::optional> global_distances = std::nullopt) override + void build_knn(raft::host_matrix_view dataset, + std::optional> inverted_indices, + global_graph_view global) override { raft::copy(this->res, raft::make_device_vector_view(data_d.value().data_handle(), dataset.size()), @@ -669,16 +728,17 @@ struct all_neighbors_builder_brute_force : public all_neighbors_builder build_knn_common(raft::make_device_matrix_view( data_d.value().data_handle(), dataset.extent(0), dataset.extent(1)), inverted_indices, - global_neighbors, - global_distances); + global); } - void build_knn(raft::device_matrix_view dataset) override + void build_knn(raft::device_matrix_view dataset, + std::optional> inverted_indices, + global_graph_view global) override { RAFT_EXPECTS(this->n_clusters <= 1, "building all-neighbors knn graph with dataset on device is not supported with " "batching (n_clusters > 1)"); - build_knn_common(dataset); + build_knn_common(dataset, inverted_indices, global); } graph_build_params::brute_force_params bf_params; @@ -695,9 +755,7 @@ std::unique_ptr> get_knn_builder( size_t min_cluster_size, size_t max_cluster_size, size_t k, - std::optional> indices = std::nullopt, - std::optional> distances = std::nullopt, - DistEpilogueT dist_epilogue = DistEpilogueT{}) + DistEpilogueT dist_epilogue = DistEpilogueT{}) { if (std::holds_alternative(params.graph_build_params)) { auto brute_force_params = @@ -713,8 +771,6 @@ std::unique_ptr> get_knn_builder( max_cluster_size, k, brute_force_params, - indices, - distances, dist_epilogue); } else if (std::holds_alternative( @@ -732,8 +788,6 @@ std::unique_ptr> get_knn_builder( max_cluster_size, k, nn_descent_params, - indices, - distances, dist_epilogue); } else if (std::holds_alternative(params.graph_build_params)) { auto ivf_pq_params = std::get(params.graph_build_params); @@ -741,14 +795,8 @@ std::unique_ptr> get_knn_builder( RAFT_LOG_WARN("Setting ivfpq_params metric to metric given for batching algorithm"); ivf_pq_params.build_params.metric = params.metric; } - return std::make_unique>(handle, - params.n_clusters, - min_cluster_size, - max_cluster_size, - k, - ivf_pq_params, - indices, - distances); + return std::make_unique>( + handle, params.n_clusters, min_cluster_size, max_cluster_size, k, ivf_pq_params); } else { RAFT_FAIL("Batch KNN build algos only supporting Brute Force, NN Descent, and IVFPQ"); } diff --git a/cpp/src/neighbors/all_neighbors/all_neighbors_merge.cuh b/cpp/src/neighbors/all_neighbors/all_neighbors_merge.cuh index f9088a2730..bf34d345c1 100644 --- a/cpp/src/neighbors/all_neighbors/all_neighbors_merge.cuh +++ b/cpp/src/neighbors/all_neighbors/all_neighbors_merge.cuh @@ -4,6 +4,9 @@ */ #pragma once +#include +#include +#include #include #include #include @@ -11,11 +14,12 @@ #include #include #include -#include #include #include #include #include +#include +#include namespace cuvs::neighbors::all_neighbors::detail { using namespace cuvs::neighbors; @@ -246,7 +250,9 @@ void merge_subgraphs(raft::resources const& res, template + bool SweepAll = false, + typename GNbrView = raft::device_matrix_view, + typename GDistView = raft::device_matrix_view> void remap_and_merge_subgraphs(raft::resources const& res, raft::device_vector_view inverted_indices_d, raft::host_vector_view inverted_indices, @@ -254,8 +260,8 @@ void remap_and_merge_subgraphs(raft::resources const& res, raft::host_matrix_view batch_neighbors_h, raft::device_matrix_view batch_neighbors_d, raft::device_matrix_view batch_distances_d, - raft::managed_matrix_view global_neighbors, - raft::managed_matrix_view global_distances, + GNbrView global_neighbors, + GDistView global_distances, size_t num_data_in_cluster, size_t k, bool select_min) @@ -291,4 +297,75 @@ void remap_and_merge_subgraphs(raft::resources const& res, select_min); } +/** + * Host-side counterpart of merge_subgraphs_kernel + the remap, used when the output lives in host + * memory so the global [num_rows x k] graph never needs to be GPU-resident. + * + * Arguments: + * - [in] inverted_indices: host [num_data_in_cluster] mapping a cluster-local row to its global + * row. + * - [in] indices_for_remap_h: host [num_data_in_cluster x k] cluster-local neighbor ids. + * - [in] batch_distances_h: host [num_data_in_cluster x k] cluster distances. + * - [inout] global_distances / global_neighbors: host [num_rows x k] global knn graph. + * - [in] select_min: whether a smaller distance is closer (sort ascending) or not (descending). + * - [in] row_locks / num_row_locks: optional striped locks. When several GPUs merge into the same + * shared host global graph concurrently (multi-GPU), each row's read-modify-write is guarded by + * row_locks[global_row % num_row_locks] so overlapping rows serialize while distinct rows proceed + * in parallel. nullptr for single-GPU (no contention). + */ +template +void remap_and_merge_subgraphs(raft::resources const& res, + raft::host_vector_view inverted_indices, + raft::host_matrix_view indices_for_remap_h, + raft::host_matrix_view batch_distances_h, + raft::host_matrix_view global_neighbors, + raft::host_matrix_view global_distances, + size_t num_data_in_cluster, + size_t k, + bool select_min, + std::mutex* row_locks = nullptr, + size_t num_row_locks = 0) +{ + // Parallelizing across this cluster's rows. +#pragma omp parallel for + for (size_t i = 0; i < num_data_in_cluster; i++) { + IdxT global_row = inverted_indices(i); + + // Lock this specific row for multi-GPU runs + if (row_locks != nullptr) row_locks[static_cast(global_row) % num_row_locks].lock(); + + // Combine the existing global k neighbors with this cluster's k + std::vector> merged; + merged.reserve(2 * k); + for (size_t j = 0; j < k; j++) { + merged.push_back({global_distances(global_row, j), global_neighbors(global_row, j)}); + } + for (size_t j = 0; j < k; j++) { + merged.push_back({batch_distances_h(i, j), inverted_indices(indices_for_remap_h(i, j))}); + } + + if (select_min) { + std::stable_sort( + merged.begin(), merged.end(), [](const auto& a, const auto& b) { return a.key < b.key; }); + } else { + std::stable_sort( + merged.begin(), merged.end(), [](const auto& a, const auto& b) { return a.key > b.key; }); + } + + // Deduplicate by neighbor id, keeping the first occurrence, and write top-k. + std::unordered_set seen; + seen.reserve(2 * k); + size_t out = 0; + for (size_t m = 0; m < merged.size() && out < k; m++) { + if (seen.insert(merged[m].value).second) { + global_distances(global_row, out) = merged[m].key; + global_neighbors(global_row, out) = merged[m].value; + out++; + } + } + + if (row_locks != nullptr) row_locks[static_cast(global_row) % num_row_locks].unlock(); + } +} + } // namespace cuvs::neighbors::all_neighbors::detail diff --git a/cpp/tests/neighbors/all_neighbors.cuh b/cpp/tests/neighbors/all_neighbors.cuh index 3b6b245a4e..fe04400fa9 100644 --- a/cpp/tests/neighbors/all_neighbors.cuh +++ b/cpp/tests/neighbors/all_neighbors.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ #pragma once @@ -44,6 +44,7 @@ struct AllNeighborsInputs { int k; bool data_on_host; bool mutual_reach; + bool output_on_host; }; inline ::std::ostream& operator<<(::std::ostream& os, const AllNeighborsInputs& p) @@ -51,7 +52,8 @@ inline ::std::ostream& operator<<(::std::ostream& os, const AllNeighborsInputs& os << "dataset shape=" << p.n_rows << "x" << p.dim << ", k=" << p.k << ", metric=" << static_cast(std::get<1>(p.build_algo_metric_recall)) << ", clusters=" << std::get<0>(p.cluster_nearestcluster) - << ", overlap_factor=" << std::get<1>(p.cluster_nearestcluster) << std::endl; + << ", overlap_factor=" << std::get<1>(p.cluster_nearestcluster) + << ", output_on_host=" << p.output_on_host << std::endl; return os; } @@ -149,7 +151,30 @@ void get_graphs(raft::resources& handle, raft::resource::sync_stream(handle); } - { + if (ps.output_on_host) { + // Host-output path: dataset and outputs all live in host memory (no managed memory). + auto database_h = raft::make_host_matrix(ps.n_rows, ps.dim); + raft::copy(database_h.data_handle(), database.data(), ps.n_rows * ps.dim, cuda_stream); + raft::resource::sync_stream(handle); + + auto indices_h = raft::make_host_matrix(ps.n_rows, ps.k); + auto distances_h = raft::make_host_matrix(ps.n_rows, ps.k); + + all_neighbors::build( + handle, + params, + raft::make_const_mdspan(database_h.view()), + indices_h.view(), + std::make_optional(distances_h.view()), + ps.mutual_reach + ? std::make_optional(raft::make_host_vector(ps.n_rows).view()) + : std::nullopt); + + std::copy( + indices_h.data_handle(), indices_h.data_handle() + queries_size, indices_allNN.data()); + std::copy( + distances_h.data_handle(), distances_h.data_handle() + queries_size, distances_allNN.data()); + } else { rmm::device_uvector distances_allNN_dev(queries_size, cuda_stream); rmm::device_uvector indices_allNN_dev(queries_size, cuda_stream); @@ -258,7 +283,8 @@ const std::vector inputsSingle = {64, 137}, // dim {16, 23}, // graph_degree {false, true}, // data on host - {false} // mutual_reach + {false}, // mutual_reach + {false, true} // output on host ); const std::vector inputsBatch = @@ -281,7 +307,8 @@ const std::vector inputsBatch = {64, 137}, // dim {16, 23}, // graph_degree {true}, // data on host - {false} // mutual_reach + {false}, // mutual_reach + {false, true} // output on host ); const std::vector mutualReachSingle = @@ -297,7 +324,8 @@ const std::vector mutualReachSingle = {64, 137}, // dim {16, 23}, // graph_degree {false, true}, // data on host - {true} // mutual_reach + {true}, // mutual_reach + {false, true} // output on host ); const std::vector mutualReachBatch = @@ -317,7 +345,8 @@ const std::vector mutualReachBatch = {64, 137}, // dim {16, 23}, // graph_degree {true}, // data on host - {true} // mutual_reach + {true}, // mutual_reach + {false, true} // output on host ); } // namespace cuvs::neighbors::all_neighbors From 7a7648ceb83a2810628657b89c702f92dbb253eb Mon Sep 17 00:00:00 2001 From: jinsolp Date: Thu, 2 Jul 2026 07:39:59 +0000 Subject: [PATCH 2/4] python wrapper --- c/include/cuvs/neighbors/all_neighbors.h | 17 +-- c/src/neighbors/all_neighbors.cpp | 102 +++++++++++------- .../neighbors/all_neighbors/all_neighbors.pyx | 61 ++++++----- python/cuvs/cuvs/tests/test_all_neighbors.py | 79 ++++++++++++++ 4 files changed, 188 insertions(+), 71 deletions(-) diff --git a/c/include/cuvs/neighbors/all_neighbors.h b/c/include/cuvs/neighbors/all_neighbors.h index 96f0dd4a83..a26919e86c 100644 --- a/c/include/cuvs/neighbors/all_neighbors.h +++ b/c/include/cuvs/neighbors/all_neighbors.h @@ -28,8 +28,10 @@ extern "C" { * provide the dataset on host. * * Notes: - * - Outputs (indices, distances, core_distances) are expected to be on device memory. - * - Host variant accepts host-resident dataset; device variant accepts device-resident dataset. + * - A host-resident dataset accepts either host- or device-resident outputs (indices, distances, + * core_distances); a device-resident dataset requires device-resident outputs. All provided + * outputs must share the same memory space. + * - Host-resident outputs never materialize the full [num_rows x k] graph on the GPU. * - For batching, `overlap_factor < n_clusters` must hold. * - When `core_distances` is provided, mutual-reachability distances are produced (see alpha). */ @@ -94,16 +96,19 @@ CUVS_EXPORT cuvsError_t cuvsAllNeighborsIndexParamsDestroy(cuvsAllNeighborsIndex * resources * @param[in] params Build parameters (see cuvsAllNeighborsIndexParams) * @param[in] dataset 2D tensor [num_rows x dim] on host or device (auto-detected) - * @param[out] indices 2D tensor [num_rows x k] on device (int64) - * @param[out] distances Optional 2D tensor [num_rows x k] on device (float32); can be NULL - * @param[out] core_distances Optional 1D tensor [num_rows] on device (float32); can be NULL + * @param[out] indices 2D tensor [num_rows x k] (int64), host or device + * @param[out] distances Optional 2D tensor [num_rows x k] (float32), host or device; can be + * NULL + * @param[out] core_distances Optional 1D tensor [num_rows] (float32), host or device; can be NULL * @param[in] alpha Mutual-reachability scaling; used only when core_distances is provided * * The function automatically detects whether the dataset is host-resident or device-resident * and calls the appropriate implementation. For host datasets, it partitions data into * `n_clusters` clusters and assigns each row to `overlap_factor` nearest clusters. For device * datasets, `n_clusters` must be 1 (no batching); `overlap_factor` is ignored. - * Outputs always reside in device memory. + * + * Output memory space: a host dataset supports host- or device-resident outputs; a device dataset + * requires device-resident outputs. All provided outputs must share the same memory space. */ CUVS_EXPORT cuvsError_t cuvsAllNeighborsBuild(cuvsResources_t res, cuvsAllNeighborsIndexParams_t params, diff --git a/c/src/neighbors/all_neighbors.cpp b/c/src/neighbors/all_neighbors.cpp index fa110c4662..3b982320a1 100644 --- a/c/src/neighbors/all_neighbors.cpp +++ b/c/src/neighbors/all_neighbors.cpp @@ -5,6 +5,7 @@ #include #include +#include #include #include @@ -80,36 +81,50 @@ static cuvs::neighbors::all_neighbors::all_neighbors_params convert_params( return out; } -static void ensure_indices_dtype_and_device_compatibility(DLManagedTensor* indices) +static void ensure_indices_dtype(DLManagedTensor* indices) { auto dtype = indices->dl_tensor.dtype; RAFT_EXPECTS(dtype.code == kDLInt && dtype.bits == 64, "indices must be int64 output tensor"); - RAFT_EXPECTS(cuvs::core::is_dlpack_device_compatible(indices->dl_tensor), - "indices tensor must be device-compatible"); } -static void ensure_optional_distance_dtype_and_device_compatibility(DLManagedTensor* distances) +static void ensure_optional_distance_dtype(DLManagedTensor* distances) { if (distances == nullptr) { return; } auto dtype = distances->dl_tensor.dtype; RAFT_EXPECTS(dtype.code == kDLFloat && dtype.bits == 32, "distances must be float32 output tensor"); - RAFT_EXPECTS(cuvs::core::is_dlpack_device_compatible(distances->dl_tensor), - "distances tensor must be device-compatible"); } -static void ensure_optional_core_distance_dtype_and_device_compatibility( - DLManagedTensor* core_distances) +static void ensure_optional_core_distance_dtype(DLManagedTensor* core_distances) { if (core_distances == nullptr) { return; } auto dtype = core_distances->dl_tensor.dtype; RAFT_EXPECTS(dtype.code == kDLFloat && dtype.bits == 32, "core_distances must be float32 output tensor"); - RAFT_EXPECTS(cuvs::core::is_dlpack_device_compatible(core_distances->dl_tensor), - "core_distances tensor must be device-compatible"); } -template +// Validate that the outputs (indices/distances/core_distances) all live in the same memory space +static bool validate_output_memory_space(DLManagedTensor* indices, + DLManagedTensor* distances, + DLManagedTensor* core_distances) +{ + const bool host = cuvs::core::is_dlpack_host_compatible(indices->dl_tensor); + RAFT_EXPECTS(host || cuvs::core::is_dlpack_device_compatible(indices->dl_tensor), + "indices tensor must be host- or device-compatible"); + auto same_space = [&](DLManagedTensor* t, const char* name) { + if (t == nullptr) { return; } + RAFT_EXPECTS(cuvs::core::is_dlpack_host_compatible(t->dl_tensor) == host, + "%s tensor must be in the same memory space (host or device) as indices", + name); + }; + same_space(distances, "distances"); + same_space(core_distances, "core_distances"); + return host; +} + +// Build with a host-resident dataset. HostOutput selects whether the outputs live on host or +// device. +template void _build_host(cuvsResources_t res, cuvsAllNeighborsIndexParams_t params, DLManagedTensor* dataset_tensor, @@ -124,24 +139,23 @@ void _build_host(cuvsResources_t res, RAFT_EXPECTS(cuvs::core::is_dlpack_host_compatible(dlt), "Host build expects host-compatible dataset tensor"); - ensure_indices_dtype_and_device_compatibility(indices_tensor); - ensure_optional_distance_dtype_and_device_compatibility(distances_tensor); - ensure_optional_core_distance_dtype_and_device_compatibility(core_distances_tensor); - - // Check dependencies between parameters - if (core_distances_tensor != nullptr && distances_tensor == nullptr) { - RAFT_FAIL("distances tensor must be provided when core_distances tensor is provided"); - } - int64_t n_rows = dlt.shape[0]; int64_t n_cols = dlt.shape[1]; auto cpp_params = convert_params(params, n_rows, n_cols); - using dataset_mdspan_t = raft::host_matrix_view; - using indices_mdspan_t = raft::device_matrix_view; - using distances_mdspan_t = raft::device_matrix_view; - using core_mdspan_t = raft::device_vector_view; + using dataset_mdspan_t = raft::host_matrix_view; + using indices_mdspan_t = + std::conditional_t, + raft::device_matrix_view>; + using distances_mdspan_t = + std::conditional_t, + raft::device_matrix_view>; + using core_mdspan_t = std::conditional_t, + raft::device_vector_view>; auto dataset = cuvs::core::from_dlpack(dataset_tensor); auto indices = cuvs::core::from_dlpack(indices_tensor); @@ -160,6 +174,7 @@ void _build_host(cuvsResources_t res, cpp_res, cpp_params, dataset, indices, distances, core_distances, alpha); } +// Build with a device-resident dataset. Outputs are always device-resident. template void _build_device(cuvsResources_t device_res, cuvsAllNeighborsIndexParams_t params, @@ -175,15 +190,6 @@ void _build_device(cuvsResources_t device_res, RAFT_EXPECTS(cuvs::core::is_dlpack_device_compatible(dlt), "Device build expects device-compatible dataset tensor"); - ensure_indices_dtype_and_device_compatibility(indices_tensor); - ensure_optional_distance_dtype_and_device_compatibility(distances_tensor); - ensure_optional_core_distance_dtype_and_device_compatibility(core_distances_tensor); - - // Check dependencies between parameters - if (core_distances_tensor != nullptr && distances_tensor == nullptr) { - RAFT_FAIL("distances tensor must be provided when core_distances tensor is provided"); - } - int64_t n_rows = dlt.shape[0]; int64_t n_cols = dlt.shape[1]; @@ -250,17 +256,33 @@ extern "C" cuvsError_t cuvsAllNeighborsBuild(cuvsResources_t res, return cuvs::core::translate_exceptions([=] { auto dataset = dataset_tensor->dl_tensor; + ensure_indices_dtype(indices_tensor); + ensure_optional_distance_dtype(distances_tensor); + ensure_optional_core_distance_dtype(core_distances_tensor); + if (core_distances_tensor != nullptr && distances_tensor == nullptr) { + RAFT_FAIL("distances tensor must be provided when core_distances tensor is provided"); + } + + const bool host_output = + validate_output_memory_space(indices_tensor, distances_tensor, core_distances_tensor); + if (dataset.dtype.code == kDLFloat && dataset.dtype.bits == 32) { // Check if dataset is host-compatible or device-compatible if (cuvs::core::is_dlpack_host_compatible(dataset)) { - _build_host(res, - params, - dataset_tensor, - indices_tensor, - distances_tensor, - core_distances_tensor, - alpha); + // Host dataset supports both host- and device-resident outputs. + if (host_output) { + _build_host( + res, params, dataset_tensor, indices_tensor, distances_tensor, core_distances_tensor, + alpha); + } else { + _build_host( + res, params, dataset_tensor, indices_tensor, distances_tensor, core_distances_tensor, + alpha); + } } else if (cuvs::core::is_dlpack_device_compatible(dataset)) { + RAFT_EXPECTS(!host_output, + "A device-resident dataset requires device-resident outputs; put the dataset " + "on host to produce host-resident outputs."); _build_device(res, params, dataset_tensor, diff --git a/python/cuvs/cuvs/neighbors/all_neighbors/all_neighbors.pyx b/python/cuvs/cuvs/neighbors/all_neighbors/all_neighbors.pyx index ce920b47c1..bd420309ce 100644 --- a/python/cuvs/cuvs/neighbors/all_neighbors/all_neighbors.pyx +++ b/python/cuvs/cuvs/neighbors/all_neighbors/all_neighbors.pyx @@ -219,14 +219,17 @@ def build(dataset, k, params, *, Parameters object containing all build settings including algorithm choice and algorithm-specific parameters. indices : array_like, optional - Optional output buffer for indices [num_rows x k] on device - (int64). If not provided, will be allocated automatically. + Optional output buffer for indices [num_rows x k] (int64), on host + or device. If not provided, will be allocated automatically (on + device, unless the other provided outputs are on host). A host + dataset supports host or device outputs; a device dataset requires + device outputs. All provided outputs must share the same memory space. distances : array_like, optional - Optional output buffer for distances [num_rows x k] on device - (float32) + Optional output buffer for distances [num_rows x k] (float32), on + host or device. core_distances : array_like, optional - Optional output buffer for core distances [num_rows] on device - (float32). Requires distances parameter to be provided. + Optional output buffer for core distances [num_rows] (float32), on + host or device. Requires distances parameter to be provided. alpha : float, default=1.0 Mutual-reachability scaling; used only when core_distances is provided @@ -238,9 +241,9 @@ def build(dataset, k, params, *, Returns ------- indices : array_like - k-NN indices for each point [num_rows x k], always on device. - If indices buffer was provided, returns the same array filled - with results. + k-NN indices for each point [num_rows x k], on the same memory space + as the outputs (device by default). If an indices buffer was + provided, returns the same array filled with results. distances : array_like or None k-NN distances if distances buffer was provided, None otherwise core_distances : array_like or None @@ -273,29 +276,37 @@ def build(dataset, k, params, *, "distances must be provided when core_distances is provided" ) - # Validate user-provided outputs (must be device arrays if provided) - if indices is not None and not hasattr( - indices, "__cuda_array_interface__" - ): + # Outputs may live on host (e.g. numpy) or device (CUDA array interface), + # but all provided outputs must be on same memory. A host dataset supports either + # location; a device dataset requires device outputs. When no output is + # given, indices default to device. + def _is_on_device(arr): + return hasattr(arr, "__cuda_array_interface__") + + output_locations = { + _is_on_device(a) + for a in (indices, distances, core_distances) + if a is not None + } + if len(output_locations) > 1: raise ValueError( - "indices must be a device array (CUDA array interface)" + "indices, distances, and core_distances must all be on the same " + "memory space (all host or all device)" ) - if distances is not None and not hasattr( - distances, "__cuda_array_interface__" - ): - raise ValueError( - "distances must be a device array (CUDA array interface)" - ) - if core_distances is not None and not hasattr( - core_distances, "__cuda_array_interface__" - ): + output_on_device = output_locations.pop() if output_locations else True + + if on_device and not output_on_device: raise ValueError( - "core_distances must be a device array (CUDA array interface)" + "A device dataset requires device outputs. Put the dataset on host " + "to produce host outputs." ) # Handle indices array (create if not provided) if indices is None: - indices = device_ndarray.empty((n_rows, k), dtype="int64") + if output_on_device: + indices = device_ndarray.empty((n_rows, k), dtype="int64") + else: + indices = np.empty((n_rows, k), dtype="int64") indices_out = wrap_array(indices) _check_input_array( diff --git a/python/cuvs/cuvs/tests/test_all_neighbors.py b/python/cuvs/cuvs/tests/test_all_neighbors.py index 2de5846054..3d4f947218 100644 --- a/python/cuvs/cuvs/tests/test_all_neighbors.py +++ b/python/cuvs/cuvs/tests/test_all_neighbors.py @@ -250,3 +250,82 @@ def test_all_neighbors_host_build_quality(algo, cluster, snmg): recall = calc_recall(indices_host, bf_indices_host) assert recall > 0.85 + + +@pytest.mark.parametrize("algo", ["nn_descent", "brute_force"]) +@pytest.mark.parametrize("cluster", ["single_cluster", "multi_cluster"]) +@pytest.mark.parametrize("snmg", [False, True]) +def test_all_neighbors_host_output_quality(algo, cluster, snmg): + """Host dataset with host-resident (numpy) outputs: the full graph never + needs to fit on the GPU. Validates recall and that outputs come back on + host. + """ + n_rows, n_cols, k = 7151, 64, 16 + + if cluster == "single_cluster": + n_clusters = 1 + overlap_factor = 0 + else: + n_clusters = 8 + overlap_factor = 3 + + np.random.seed(42) + + X_host, _ = make_blobs( + n_samples=n_rows, + n_features=n_cols, + centers=10, + cluster_std=1.0, + center_box=(-10.0, 10.0), + random_state=42, + ) + X_host = X_host.astype(np.float32) + X_device = device_ndarray(X_host) + + nn_descent_params = None + if algo == "nn_descent": + nn_descent_params = nn_descent.IndexParams( + metric="sqeuclidean", + graph_degree=k, + intermediate_graph_degree=k * 2, + max_iterations=100, + termination_threshold=0.001, + ) + + params = all_neighbors.AllNeighborsParams( + algo=algo, + overlap_factor=overlap_factor, + n_clusters=n_clusters, + metric="sqeuclidean", + nn_descent_params=nn_descent_params, + ) + + res = MultiGpuResources() if snmg else Resources() + + # Host-resident output buffers -> host build, no device-side [n_rows x k]. + indices = np.empty((n_rows, k), dtype=np.int64) + distances = np.empty((n_rows, k), dtype=np.float32) + + indices, distances = all_neighbors.build( + X_host, + k, + params, + indices=indices, + distances=distances, + resources=res, + ) + + assert isinstance(indices, np.ndarray) + assert isinstance(distances, np.ndarray) + assert indices.shape == (n_rows, k) + assert indices.dtype == np.int64 + assert distances.shape == (n_rows, k) + assert distances.dtype == np.float32 + + bf_index = brute_force.build(X_device, metric="sqeuclidean") + _, bf_indices = brute_force.search(bf_index, X_device, k=k) + + recall = calc_recall(indices, cupy.asnumpy(bf_indices)) + print(f"recall: {recall}") + + assert recall > 0.85 From f87bd1580789d1023de61696adc5465c8e78cc4a Mon Sep 17 00:00:00 2001 From: jinsolp Date: Thu, 2 Jul 2026 09:34:43 +0000 Subject: [PATCH 3/4] fix compile error --- cpp/src/neighbors/all_neighbors/all_neighbors_builder.cuh | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/cpp/src/neighbors/all_neighbors/all_neighbors_builder.cuh b/cpp/src/neighbors/all_neighbors/all_neighbors_builder.cuh index 8e56735cc2..27fb9a48c0 100644 --- a/cpp/src/neighbors/all_neighbors/all_neighbors_builder.cuh +++ b/cpp/src/neighbors/all_neighbors/all_neighbors_builder.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once @@ -156,7 +156,7 @@ struct all_neighbors_builder { batch_distances_h.value().data_handle(), num_data_in_cluster, k); raft::copy(res, batch_distances_h_view, - raft::make_device_matrix_view( + raft::make_device_matrix_view( batch_distances_d.value().data_handle(), num_data_in_cluster, k)); raft::resource::sync_stream(res); From 22367527184e3078e9fde41e7963c581b05ffc68 Mon Sep 17 00:00:00 2001 From: jinsolp Date: Fri, 10 Jul 2026 15:40:36 +0000 Subject: [PATCH 4/4] clean --- .../all_neighbors/all_neighbors_merge.cuh | 7 +++---- cpp/tests/neighbors/all_neighbors.cuh | 17 +++++++---------- python/cuvs/cuvs/tests/test_all_neighbors.py | 7 ++----- 3 files changed, 12 insertions(+), 19 deletions(-) diff --git a/cpp/src/neighbors/all_neighbors/all_neighbors_merge.cuh b/cpp/src/neighbors/all_neighbors/all_neighbors_merge.cuh index bf34d345c1..43e891e451 100644 --- a/cpp/src/neighbors/all_neighbors/all_neighbors_merge.cuh +++ b/cpp/src/neighbors/all_neighbors/all_neighbors_merge.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once @@ -299,7 +299,7 @@ void remap_and_merge_subgraphs(raft::resources const& res, /** * Host-side counterpart of merge_subgraphs_kernel + the remap, used when the output lives in host - * memory so the global [num_rows x k] graph never needs to be GPU-resident. + * memory. * * Arguments: * - [in] inverted_indices: host [num_data_in_cluster] mapping a cluster-local row to its global @@ -307,9 +307,8 @@ void remap_and_merge_subgraphs(raft::resources const& res, * - [in] indices_for_remap_h: host [num_data_in_cluster x k] cluster-local neighbor ids. * - [in] batch_distances_h: host [num_data_in_cluster x k] cluster distances. * - [inout] global_distances / global_neighbors: host [num_rows x k] global knn graph. - * - [in] select_min: whether a smaller distance is closer (sort ascending) or not (descending). * - [in] row_locks / num_row_locks: optional striped locks. When several GPUs merge into the same - * shared host global graph concurrently (multi-GPU), each row's read-modify-write is guarded by + * shared host global graph concurrently, each row's read-modify-write is guarded by * row_locks[global_row % num_row_locks] so overlapping rows serialize while distinct rows proceed * in parallel. nullptr for single-GPU (no contention). */ diff --git a/cpp/tests/neighbors/all_neighbors.cuh b/cpp/tests/neighbors/all_neighbors.cuh index fe04400fa9..0b43023eff 100644 --- a/cpp/tests/neighbors/all_neighbors.cuh +++ b/cpp/tests/neighbors/all_neighbors.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once @@ -157,23 +157,20 @@ void get_graphs(raft::resources& handle, raft::copy(database_h.data_handle(), database.data(), ps.n_rows * ps.dim, cuda_stream); raft::resource::sync_stream(handle); - auto indices_h = raft::make_host_matrix(ps.n_rows, ps.k); - auto distances_h = raft::make_host_matrix(ps.n_rows, ps.k); + auto indices_allNN_view = + raft::make_host_matrix_view(indices_allNN.data(), ps.n_rows, ps.k); + auto distances_allNN_view = + raft::make_host_matrix_view(distances_allNN.data(), ps.n_rows, ps.k); all_neighbors::build( handle, params, raft::make_const_mdspan(database_h.view()), - indices_h.view(), - std::make_optional(distances_h.view()), + indices_allNN_view, + std::make_optional(distances_allNN_view), ps.mutual_reach ? std::make_optional(raft::make_host_vector(ps.n_rows).view()) : std::nullopt); - - std::copy( - indices_h.data_handle(), indices_h.data_handle() + queries_size, indices_allNN.data()); - std::copy( - distances_h.data_handle(), distances_h.data_handle() + queries_size, distances_allNN.data()); } else { rmm::device_uvector distances_allNN_dev(queries_size, cuda_stream); rmm::device_uvector indices_allNN_dev(queries_size, cuda_stream); diff --git a/python/cuvs/cuvs/tests/test_all_neighbors.py b/python/cuvs/cuvs/tests/test_all_neighbors.py index 3d4f947218..2f02d22778 100644 --- a/python/cuvs/cuvs/tests/test_all_neighbors.py +++ b/python/cuvs/cuvs/tests/test_all_neighbors.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # @@ -256,10 +256,7 @@ def test_all_neighbors_host_build_quality(algo, cluster, snmg): @pytest.mark.parametrize("cluster", ["single_cluster", "multi_cluster"]) @pytest.mark.parametrize("snmg", [False, True]) def test_all_neighbors_host_output_quality(algo, cluster, snmg): - """Host dataset with host-resident (numpy) outputs: the full graph never - needs to fit on the GPU. Validates recall and that outputs come back on - host. - """ + """Host dataset with host-resident (numpy) outputs""" n_rows, n_cols, k = 7151, 64, 16 if cluster == "single_cluster":