From 5267699c0842b083c92b918b47cf98e1a506f68a Mon Sep 17 00:00:00 2001 From: jamie8johnson Date: Fri, 3 Jul 2026 16:35:38 -0500 Subject: [PATCH] feat(rust): add refine bindings --- rust/cuvs/src/lib.rs | 1 + rust/cuvs/src/refine.rs | 200 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 201 insertions(+) create mode 100644 rust/cuvs/src/refine.rs diff --git a/rust/cuvs/src/lib.rs b/rust/cuvs/src/lib.rs index 52c31392e7..68b1cbc75a 100644 --- a/rust/cuvs/src/lib.rs +++ b/rust/cuvs/src/lib.rs @@ -18,6 +18,7 @@ pub mod dlpack; mod error; pub mod ivf_flat; pub mod ivf_pq; +pub mod refine; mod resources; #[cfg(test)] pub(crate) mod test_utils; diff --git a/rust/cuvs/src/refine.rs b/rust/cuvs/src/refine.rs new file mode 100644 index 0000000000..53b92a5906 --- /dev/null +++ b/rust/cuvs/src/refine.rs @@ -0,0 +1,200 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Refinement of approximate nearest neighbor results + +use crate::distance_type::DistanceType; +use crate::dlpack::{AsDlTensor, AsDlTensorMut}; +use crate::error::{Result, check_cuvs}; +use crate::resources::Resources; + +/// Refine nearest neighbor search results. +/// +/// Refinement is an operation that follows an approximate nearest neighbors +/// search. The approximate search has already selected `n_candidates` neighbor +/// candidates for each query. This narrows the candidate list down to the `k` +/// nearest neighbors by computing the exact distance between each query and its +/// candidates against the original dataset, then selecting the `k` closest. +/// +/// All tensors must reside in the same memory space: either all on the device +/// or all on the host. The dataset and queries may be `f32`, `f16`, `i8`, or +/// `u8` (with matching dtype codes). The candidate and output index tensors +/// must be `i64`, and the output distance tensor must be `f32`. +/// +/// Tensors are borrowed through the [`AsDlTensor`] / +/// [`AsDlTensorMut`] traits: `dataset`, `queries`, and `candidates` are read, +/// while `indices` and `distances` are written in place. See the +/// [`dlpack`](crate::dlpack) module for the tensor model and `examples/cagra.rs` +/// for a device-tensor adapter. +/// +/// # Arguments +/// +/// * `res` - Resources to use +/// * `dataset` - A row-major matrix of the original dataset - shape `(n_rows, dims)` +/// * `queries` - A row-major matrix of the queries - shape `(n_queries, dims)` +/// * `candidates` - A row-major `i64` matrix of candidate indices into `dataset` +/// - shape `(n_queries, n_candidates)`, where `n_candidates >= k` +/// * `metric` - DistanceType used to rank candidates +/// * `indices` - Output `i64` matrix that receives the refined indices - shape +/// `(n_queries, k)`. `k` is inferred from this tensor's shape. +/// * `distances` - Output `f32` matrix that receives the refined distances - +/// shape `(n_queries, k)` +pub fn refine( + res: &Resources, + dataset: &DS, + queries: &Q, + candidates: &C, + metric: DistanceType, + indices: &mut I, + distances: &mut D, +) -> Result<()> +where + DS: AsDlTensor + ?Sized, + Q: AsDlTensor + ?Sized, + C: AsDlTensor + ?Sized, + I: AsDlTensorMut + ?Sized, + D: AsDlTensorMut + ?Sized, +{ + let dataset = dataset.as_dl_tensor()?; + let queries = queries.as_dl_tensor()?; + let candidates = candidates.as_dl_tensor()?; + let indices = indices.as_dl_tensor_mut()?; + let distances = distances.as_dl_tensor_mut()?; + unsafe { + check_cuvs(ffi::cuvsRefine( + res.0, + dataset.to_c().as_mut_ptr(), + queries.to_c().as_mut_ptr(), + candidates.to_c().as_mut_ptr(), + metric, + indices.to_c().as_mut_ptr(), + distances.to_c().as_mut_ptr(), + )) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::test_utils::DeviceTensor; + + /// Refinement must repair a candidate list that contains deliberately + /// wrong entries: after refine, the top-k must equal the exact + /// brute-force top-k. + #[test] + fn test_refine_fixes_wrong_candidates() { + let res = Resources::new().unwrap(); + + // A small, well-separated 2-D dataset. The exact L2 ranking of every + // query is unambiguous, so we can hard-assert the refined output. + // + // index : point + // 0 : (0, 0) + // 1 : (1, 0) + // 2 : (0, 1) + // 3 : (2, 2) + // 4 : (5, 5) + // 5 : (9, 9) + let dataset = ndarray::array![ + [0.0f32, 0.0], + [1.0, 0.0], + [0.0, 1.0], + [2.0, 2.0], + [5.0, 5.0], + [9.0, 9.0], + ]; + + // Two queries near distinct clusters. + // q0 sits next to point 0; true top-3 = [0, 1, 2] + // q1 sits next to point 4; true top-3 = [4, 3, 5] (4 closest, then 3, then 5) + let queries = ndarray::array![[0.1f32, 0.1], [4.9, 4.9]]; + + // Candidate lists are intentionally *wrong order* and include far-away + // points. Each list is a superset of the true top-3 but jumbled, plus a + // planted bad candidate (index 5 for q0, index 0 for q1). Refine must + // re-rank these exactly and select the correct nearest three. + let candidates = ndarray::array![ + [5i64, 2, 0, 1], // q0: true nearest 0 is buried, 5 is far noise + [0i64, 5, 3, 4], // q1: true nearest 4 is last, 0 is far noise + ]; + + let n_queries = 2; + let k = 3; + + let dataset_d = DeviceTensor::from_host(&res, &dataset).unwrap(); + let queries_d = DeviceTensor::from_host(&res, &queries).unwrap(); + let candidates_d = DeviceTensor::from_host(&res, &candidates).unwrap(); + + let mut indices_host = ndarray::Array::::zeros((n_queries, k)); + let mut distances_host = ndarray::Array::::zeros((n_queries, k)); + let mut indices_d = DeviceTensor::::zeros(&res, &[n_queries, k]).unwrap(); + let mut distances_d = DeviceTensor::::zeros(&res, &[n_queries, k]).unwrap(); + + refine( + &res, + &dataset_d, + &queries_d, + &candidates_d, + DistanceType::L2Expanded, + &mut indices_d, + &mut distances_d, + ) + .unwrap(); + + indices_d.copy_to_host(&res, &mut indices_host).unwrap(); + distances_d.copy_to_host(&res, &mut distances_host).unwrap(); + res.sync_stream().unwrap(); + + // Exact brute-force top-3, independent of the candidate ordering. + // q0: distances to (0.1,0.1): 0 -> ~0.14, 1 -> ~0.91, 2 -> ~0.91, ... + // point 0 is strictly nearest; 1 and 2 are tied next. + // q1: distances to (4.9,4.9): 4 -> ~0.14, 3 -> ~4.1, 5 -> ~5.8. + assert_eq!( + indices_host[[0, 0]], + 0, + "q0 nearest must be repaired to index 0, got {:?}", + indices_host.row(0) + ); + assert_eq!( + indices_host[[1, 0]], + 4, + "q1 nearest must be repaired to index 4, got {:?}", + indices_host.row(1) + ); + + // The planted noise candidates (5 for q0, 0 for q1) must be evicted + // from the refined top-k. + let q0: Vec = indices_host.row(0).to_vec(); + let q1: Vec = indices_host.row(1).to_vec(); + assert!(!q0.contains(&5), "q0 must drop far candidate 5, got {:?}", q0); + assert!(!q1.contains(&0), "q1 must drop far candidate 0, got {:?}", q1); + + // The refined top-3 sets must match the exact brute-force top-3 sets. + let mut q0_sorted = q0.clone(); + q0_sorted.sort_unstable(); + assert_eq!(q0_sorted, vec![0, 1, 2], "q0 refined set wrong: {:?}", q0); + + let mut q1_sorted = q1.clone(); + q1_sorted.sort_unstable(); + assert_eq!(q1_sorted, vec![3, 4, 5], "q1 refined set wrong: {:?}", q1); + + // Refined distances must be sorted ascending (nearest first) across + // the full top-k, and the first entry must be the small in-cluster + // distance, not noise. + for q in 0..2 { + for i in 0..2 { + assert!( + distances_host[[q, i]] <= distances_host[[q, i + 1]], + "q{q} distances not ascending at {i}: {:?}", + (distances_host[[q, i]], distances_host[[q, i + 1]]) + ); + } + } + assert!( + distances_host[[0, 0]] < 1.0, + "q0 nearest distance should be small, got {}", + distances_host[[0, 0]] + ); + } +}