diff --git a/cpp/include/cuvs/neighbors/ball_cover.hpp b/cpp/include/cuvs/neighbors/ball_cover.hpp index 7471c91c27..d9fd0c211d 100644 --- a/cpp/include/cuvs/neighbors/ball_cover.hpp +++ b/cpp/include/cuvs/neighbors/ball_cover.hpp @@ -272,6 +272,8 @@ void eps_nn(raft::resources const& handle, * that it is not guarantueed to return the nearest neighbors. * Upon return max_k is overwritten with the actual max_k found during * computation. + * @param[out] dists An optional vector (with same length as adj_ja). If provided, will be + * used to store the corresponding distances for the computed neighbors. */ void eps_nn(raft::resources const& handle, const index& index, @@ -280,7 +282,8 @@ void eps_nn(raft::resources const& handle, raft::device_vector_view vd, raft::device_matrix_view query, float eps, - std::optional> max_k = std::nullopt); + std::optional> max_k = std::nullopt, + std::optional> dists = std::nullopt); /** * @ingroup random_ball_cover diff --git a/cpp/src/neighbors/ball_cover.cu b/cpp/src/neighbors/ball_cover.cu index 1ae9a3e3d5..318f4d2d5c 100644 --- a/cpp/src/neighbors/ball_cover.cu +++ b/cpp/src/neighbors/ball_cover.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 */ @@ -40,9 +40,10 @@ void eps_nn(raft::resources const& handle, raft::device_vector_view vd, raft::device_matrix_view query, float eps, - std::optional> max_k) + std::optional> max_k, + std::optional> dists) { - detail::eps_nn(handle, index, adj_ia, adj_ja, vd, query, eps, max_k); + detail::eps_nn(handle, index, adj_ia, adj_ja, vd, query, eps, max_k, dists); } void knn_query(raft::resources const& handle, diff --git a/cpp/src/neighbors/ball_cover.cuh b/cpp/src/neighbors/ball_cover.cuh index 45dbf6054b..0e4ac8a7cf 100644 --- a/cpp/src/neighbors/ball_cover.cuh +++ b/cpp/src/neighbors/ball_cover.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 @@ -308,6 +308,8 @@ void eps_nn(raft::resources const& handle, * that it is not guarantueed to return the nearest neighbors. * Upon return max_k is overwritten with the actual max_k found during * computation. + * @param[out] dists An optional vector (with same length as adj_ja). If provided, will be + * used to store the corresponding distances for the computed neighbors. */ template void eps_nn(raft::resources const& handle, @@ -317,7 +319,8 @@ void eps_nn(raft::resources const& handle, raft::device_vector_view vd, raft::device_matrix_view query, value_t eps, - std::optional> max_k = std::nullopt) + std::optional> max_k = std::nullopt, + std::optional> dists = std::nullopt) { ASSERT(index.n == query.extent(1), "vector dimension needs to be the same for index and queries"); ASSERT(index.metric == cuvs::distance::DistanceType::L2SqrtExpanded || @@ -328,6 +331,9 @@ void eps_nn(raft::resources const& handle, int64_t* max_k_ptr = nullptr; if (max_k.has_value()) { max_k_ptr = max_k.value().data_handle(); } + value_t* dists_ptr = nullptr; + if (dists.has_value()) { dists_ptr = dists.value().data_handle(); } + // run query cuvs::neighbors::ball_cover::detail::rbc_eps_nn_query( handle, @@ -339,6 +345,7 @@ void eps_nn(raft::resources const& handle, adj_ia.data_handle(), adj_ja.data_handle(), vd.data_handle(), + dists_ptr, cuvs::neighbors::ball_cover::detail::EuclideanSqFunc()); } diff --git a/cpp/src/neighbors/ball_cover/ball_cover.cuh b/cpp/src/neighbors/ball_cover/ball_cover.cuh index c39756f7d5..08b1abd00b 100644 --- a/cpp/src/neighbors/ball_cover/ball_cover.cuh +++ b/cpp/src/neighbors/ball_cover/ball_cover.cuh @@ -275,10 +275,11 @@ void perform_rbc_eps_nn_query(raft::resources const& handle, dist_func dfunc, value_idx* adj_ia, value_idx* adj_ja, - value_idx* vd) + value_idx* vd, + value_t* dists) { rbc_eps_pass( - handle, index, query, n_query_pts, eps, max_k, landmarks, dfunc, adj_ia, adj_ja, vd); + handle, index, query, n_query_pts, eps, max_k, landmarks, dfunc, adj_ia, adj_ja, vd, dists); raft::resource::sync_stream(handle); } @@ -515,6 +516,7 @@ void rbc_eps_nn_query(raft::resources const& handle, value_idx* adj_ia, value_idx* adj_ja, value_idx* vd, + value_t* dists, distance_func dfunc) { ASSERT(index.is_index_trained(), "index must be previously trained"); @@ -530,7 +532,8 @@ void rbc_eps_nn_query(raft::resources const& handle, dfunc, adj_ia, adj_ja, - vd); + vd, + dists); } }; // namespace cuvs::neighbors::ball_cover::detail diff --git a/cpp/src/neighbors/ball_cover/detail/ball_cover/registers_00_generate.py b/cpp/src/neighbors/ball_cover/detail/ball_cover/registers_00_generate.py index f5ef49f67f..1da70ee773 100644 --- a/cpp/src/neighbors/ball_cover/detail/ball_cover/registers_00_generate.py +++ b/cpp/src/neighbors/ball_cover/detail/ball_cover/registers_00_generate.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. # SPDX-License-Identifier: Apache-2.0 import datetime @@ -89,7 +89,8 @@ Mdist_func& dfunc, \\ Mvalue_idx* adj_ia, \\ Mvalue_idx* adj_ja, \\ - Mvalue_idx* vd) + Mvalue_idx* vd, \\ + Mvalue_t* dists) """ diff --git a/cpp/src/neighbors/ball_cover/detail/ball_cover/registers_eps_pass_euclidean.cu b/cpp/src/neighbors/ball_cover/detail/ball_cover/registers_eps_pass_euclidean.cu index 5b7f231bf1..d08e7bed16 100644 --- a/cpp/src/neighbors/ball_cover/detail/ball_cover/registers_eps_pass_euclidean.cu +++ b/cpp/src/neighbors/ball_cover/detail/ball_cover/registers_eps_pass_euclidean.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ @@ -39,7 +39,8 @@ Mdist_func& dfunc, \ Mvalue_idx* adj_ia, \ Mvalue_idx* adj_ja, \ - Mvalue_idx* vd) + Mvalue_idx* vd, \ + Mvalue_t* dists) instantiate_cuvs_neighbors_detail_rbc_eps_pass( std::int64_t, float, cuvs::neighbors::ball_cover::detail::EuclideanSqFunc); diff --git a/cpp/src/neighbors/ball_cover/registers.cuh b/cpp/src/neighbors/ball_cover/registers.cuh index 4d381d8ca2..e00f39eb34 100644 --- a/cpp/src/neighbors/ball_cover/registers.cuh +++ b/cpp/src/neighbors/ball_cover/registers.cuh @@ -579,6 +579,7 @@ RAFT_KERNEL block_rbc_kernel_eps_csr_pass(const value_t* X_reordered, distance_func dfunc, value_idx* adj_ia, value_idx* adj_ja, + value_t* dists, bool write_pass) { constexpr int num_warps = tpb / raft::WarpSize; @@ -600,6 +601,7 @@ RAFT_KERNEL block_rbc_kernel_eps_csr_pass(const value_t* X_reordered, // we have no neighbors to fill for this query if (offset == adj_ia[query_id + 1]) return; adj_ja += offset; + if (dists != nullptr) { dists += offset; } } const value_t* x_ptr = X + (n_cols * query_id); @@ -658,8 +660,11 @@ RAFT_KERNEL block_rbc_kernel_eps_csr_pass(const value_t* X_reordered, const uint32_t index = R_1nn_cols[R_start_offset + i]; const uint32_t row_pos = __popc(mask & lid_mask); adj_ja[row_pos] = index; + if (dists != nullptr) { dists[row_pos] = dist; } } - adj_ja += __popc(mask); + const int offset = __popc(mask); + adj_ja += offset; + if (dists != nullptr) { dists += offset; } } else { column_index_offset += (in_range); } @@ -681,8 +686,11 @@ RAFT_KERNEL block_rbc_kernel_eps_csr_pass(const value_t* X_reordered, const uint32_t index = R_1nn_cols[R_start_offset + i0 + lid]; const uint32_t row_pos = __popc(mask & lid_mask); adj_ja[row_pos] = index; + if (dists != nullptr) { dists[row_pos] = dist; } } - adj_ja += __popc(mask); + const int offset = __popc(mask); + adj_ja += offset; + if (dists != nullptr) { dists += offset; } } else { column_index_offset += (in_range); } @@ -718,6 +726,7 @@ RAFT_KERNEL __launch_bounds__(tpb) distance_func dfunc, value_idx* __restrict__ adj_ia, value_idx* __restrict__ adj_ja, + value_t* __restrict__ dists, bool write_pass, int dim) { @@ -740,6 +749,7 @@ RAFT_KERNEL __launch_bounds__(tpb) // we have no neighbors to fill for this query if (offset == adj_ia[query_id + 1]) return; adj_ja += offset; + if (dists != nullptr) { dists += offset; } } const value_t* x_ptr = X + (dim * query_id); @@ -803,8 +813,11 @@ RAFT_KERNEL __launch_bounds__(tpb) const uint32_t index = R_1nn_cols[R_start_offset + i]; const uint32_t row_pos = __popc(mask & lid_mask); adj_ja[row_pos] = index; + if (dists != nullptr) { dists[row_pos] = dist; } } - adj_ja += __popc(mask); + const int offset = __popc(mask); + adj_ja += offset; + if (dists != nullptr) { dists += offset; } } else { column_index_offset += (in_range); } @@ -826,8 +839,11 @@ RAFT_KERNEL __launch_bounds__(tpb) const uint32_t index = R_1nn_cols[R_start_offset + i0 + lid]; const uint32_t row_pos = __popc(mask & lid_mask); adj_ja[row_pos] = index; + if (dists != nullptr) { dists[row_pos] = dist; } } - adj_ja += __popc(mask); + const int offset = __popc(mask); + adj_ja += offset; + if (dists != nullptr) { dists += offset; } } else { column_index_offset += (in_range); } @@ -862,7 +878,8 @@ RAFT_KERNEL block_rbc_kernel_eps_max_k(const value_t* X_reordered, distance_func dfunc, value_idx* vd, int64_t max_k, - value_idx* tmp) + value_idx* tmp, + value_t* dists) { constexpr int num_warps = tpb / raft::WarpSize; @@ -880,6 +897,7 @@ RAFT_KERNEL block_rbc_kernel_eps_max_k(const value_t* X_reordered, const value_t* x_ptr = X + (n_cols * query_id); tmp += query_id * max_k; + if (dists != nullptr) { dists += query_id * max_k; } // we omit the sqrt() in the inner distance compute const value_t eps2 = eps * eps; @@ -936,6 +954,7 @@ RAFT_KERNEL block_rbc_kernel_eps_max_k(const value_t* X_reordered, if (row_pos < max_k) { auto index = R_1nn_cols[R_start_offset + i]; tmp[row_pos] = index; + if (dists != nullptr) { dists[row_pos] = dist; } } } column_count += __popc(mask); @@ -958,6 +977,7 @@ RAFT_KERNEL block_rbc_kernel_eps_max_k(const value_t* X_reordered, if (row_pos < max_k) { auto index = R_1nn_cols[R_start_offset + i0 + lid]; tmp[row_pos] = index; + if (dists != nullptr) { dists[row_pos] = dist; } } } column_count += __popc(mask); @@ -1313,7 +1333,8 @@ void rbc_eps_pass(raft::resources const& handle, dist_func& dfunc, value_idx* adj_ia, value_idx* adj_ja, - value_idx* vd) + value_idx* vd, + value_t* dists) { // if max_k == nullptr we are either pass 1 or pass 2 if (max_k == nullptr) { @@ -1340,6 +1361,7 @@ void rbc_eps_pass(raft::resources const& handle, dfunc, vd_ptr, nullptr, + nullptr, false, index.n); } else { @@ -1362,6 +1384,7 @@ void rbc_eps_pass(raft::resources const& handle, dfunc, vd_ptr, nullptr, + nullptr, false); } @@ -1372,7 +1395,7 @@ void rbc_eps_pass(raft::resources const& handle, (value_idx)0); } else { - // pass 2 -> fill in adj_ja + // pass 2 -> fill in adj_ja and dists (if provided) if (index.n == 2 || index.n == 3) { block_rbc_kernel_eps_csr_pass_xd <<(n_query_rows, 2), @@ -1393,6 +1416,7 @@ void rbc_eps_pass(raft::resources const& handle, dfunc, adj_ia, adj_ja, + dists, true, index.n); } else { @@ -1413,6 +1437,7 @@ void rbc_eps_pass(raft::resources const& handle, dfunc, adj_ia, adj_ja, + dists, true); } } @@ -1440,7 +1465,8 @@ void rbc_eps_pass(raft::resources const& handle, dfunc, vd_ptr, max_k_in, - tmp.data()); + tmp.data(), + dists); int64_t actual_max = thrust::reduce(raft::resource::get_thrust_policy(handle), vd_ptr, diff --git a/cpp/src/neighbors/ball_cover/registers.hpp b/cpp/src/neighbors/ball_cover/registers.hpp index 60069d8cb0..5d67b17e83 100644 --- a/cpp/src/neighbors/ball_cover/registers.hpp +++ b/cpp/src/neighbors/ball_cover/registers.hpp @@ -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 */ @@ -60,6 +60,7 @@ void rbc_eps_pass(raft::resources const& handle, dist_func& dfunc, value_idx* adj_ia, value_idx* adj_ja, - value_idx* vd); + value_idx* vd, + value_t* dists); } // namespace cuvs::neighbors::ball_cover::detail