Skip to content

[WIP] Add support for returning distances from eps_nn #1815

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
jcrist wants to merge 2 commits into
base: main
Choose a base branch
from
Draft

[WIP] Add support for returning distances from eps_nn #1815

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
2af16c2
Add support for returning distances from `eps_nn`
jcrist Feb 17, 2026
c437774
Merge branch 'main' into eps-nn-dists
cjnolet May 13, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 4 additions & 1 deletion cpp/include/cuvs/neighbors/ball_cover.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Comment on lines +275 to +276

@dantegd dantegd Feb 18, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If I'm not mistaken internally we compute eps2 = eps * eps and compare the raw squared distance dist <= eps2. The dist value written into dists is this squared distance, so we should clarify that in the docstring with something like:

"Distances are squared Euclidean distances which avoid sqrt for performance"

*/
void eps_nn(raft::resources const& handle,
const index<int64_t, float>& index,
Expand All @@ -280,7 +282,8 @@ void eps_nn(raft::resources const& handle,
raft::device_vector_view<int64_t, int64_t> vd,
raft::device_matrix_view<const float, int64_t, raft::row_major> query,
float eps,
std::optional<raft::host_scalar_view<int64_t, int64_t>> max_k = std::nullopt);
std::optional<raft::host_scalar_view<int64_t, int64_t>> max_k = std::nullopt,
std::optional<raft::device_vector_view<float, int64_t>> dists = std::nullopt);

/**
* @ingroup random_ball_cover
Expand Down
7 changes: 4 additions & 3 deletions cpp/src/neighbors/ball_cover.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
*/

Expand Down Expand Up @@ -40,9 +40,10 @@ void eps_nn(raft::resources const& handle,
raft::device_vector_view<int64_t, int64_t> vd,
raft::device_matrix_view<const float, int64_t, raft::row_major> query,
float eps,
std::optional<raft::host_scalar_view<int64_t, int64_t>> max_k)
std::optional<raft::host_scalar_view<int64_t, int64_t>> max_k,
std::optional<raft::device_vector_view<float, int64_t>> dists)
{
detail::eps_nn<int64_t, float>(handle, index, adj_ia, adj_ja, vd, query, eps, max_k);
detail::eps_nn<int64_t, float>(handle, index, adj_ia, adj_ja, vd, query, eps, max_k, dists);
}

void knn_query(raft::resources const& handle,
Expand Down
11 changes: 9 additions & 2 deletions cpp/src/neighbors/ball_cover.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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 <typename idx_t, typename value_t>
void eps_nn(raft::resources const& handle,
Expand All @@ -317,7 +319,8 @@ void eps_nn(raft::resources const& handle,
raft::device_vector_view<idx_t, int64_t> vd,
raft::device_matrix_view<const value_t, int64_t, raft::row_major> query,
value_t eps,
std::optional<raft::host_scalar_view<int64_t, int64_t>> max_k = std::nullopt)
std::optional<raft::host_scalar_view<int64_t, int64_t>> max_k = std::nullopt,
std::optional<raft::device_vector_view<value_t, int64_t>> 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 ||
Expand All @@ -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(); }

@dantegd dantegd Feb 18, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think adding an assert for the unsupported case would be fine here

  ASSERT(!(max_k.has_value() && dists.has_value()),
         "dists output is not supported when max_k is set");


// run query
cuvs::neighbors::ball_cover::detail::rbc_eps_nn_query(
handle,
Expand All @@ -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<value_t, int64_t>());
}

Expand Down
9 changes: 6 additions & 3 deletions cpp/src/neighbors/ball_cover/ball_cover.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<value_idx, value_t>(
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);
}
Expand Down Expand Up @@ -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");
Expand All @@ -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
5 changes: 3 additions & 2 deletions cpp/src/neighbors/ball_cover/detail/ball_cover/registers_00_generate.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -89,7 +89,8 @@
Mdist_func<Mvalue_t, int64_t>& dfunc, \\
Mvalue_idx* adj_ia, \\
Mvalue_idx* adj_ja, \\
Mvalue_idx* vd)
Mvalue_idx* vd, \\
Mvalue_t* dists)

"""

Expand Down
5 changes: 3 additions & 2 deletions cpp/src/neighbors/ball_cover/detail/ball_cover/registers_eps_pass_euclidean.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
*/

Expand Down Expand Up @@ -39,7 +39,8 @@
Mdist_func<Mvalue_t, int64_t>& 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);
Expand Down
42 changes: 34 additions & 8 deletions cpp/src/neighbors/ball_cover/registers.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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; }

@jinsolp jinsolp Feb 20, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

we could cache this condition since the check dists != nullptr is done multiple times in this kernel!

}

const value_t* x_ptr = X + (n_cols * query_id);
Expand Down Expand Up @@ -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);
}
Expand All @@ -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);
}
Expand Down Expand Up @@ -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)
{
Expand All @@ -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);
Expand Down Expand Up @@ -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);
}
Expand All @@ -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);
}
Expand Down Expand Up @@ -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;

Expand All @@ -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;
Expand Down Expand Up @@ -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);
Expand All @@ -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);
Expand Down Expand Up @@ -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) {
Expand All @@ -1340,6 +1361,7 @@ void rbc_eps_pass(raft::resources const& handle,
dfunc,
vd_ptr,
nullptr,
nullptr,
false,
index.n);
} else {
Expand All @@ -1362,6 +1384,7 @@ void rbc_eps_pass(raft::resources const& handle,
dfunc,
vd_ptr,
nullptr,
nullptr,
false);
}

Expand All @@ -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<value_idx, value_t, 64>
<<<raft::ceildiv<int64_t>(n_query_rows, 2),
Expand All @@ -1393,6 +1416,7 @@ void rbc_eps_pass(raft::resources const& handle,
dfunc,
adj_ia,
adj_ja,
dists,
true,
index.n);
} else {
Expand All @@ -1413,6 +1437,7 @@ void rbc_eps_pass(raft::resources const& handle,
dfunc,
adj_ia,
adj_ja,
dists,
true);
}
}
Expand Down Expand Up @@ -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,
Expand Down
5 changes: 3 additions & 2 deletions cpp/src/neighbors/ball_cover/registers.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
*/

Expand Down Expand Up @@ -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
Loading