From 2e5b1b9bbf6b62ca6f6315d7fb74805dadf757ed Mon Sep 17 00:00:00 2001 From: Yan Zaretskiy Date: Wed, 1 Jul 2026 13:31:37 -0700 Subject: [PATCH] Rust: builder-based params, per-module errors --- rust/cuvs/Cargo.toml | 1 + rust/cuvs/examples/cagra.rs | 14 +- rust/cuvs/src/cagra/index_params.rs | 196 ------- rust/cuvs/src/cagra/mod.rs | 22 - rust/cuvs/src/cagra/search_params.rs | 154 ------ rust/cuvs/src/cluster/kmeans/mod.rs | 46 +- rust/cuvs/src/cluster/kmeans/params.rs | 179 +++---- rust/cuvs/src/distance/mod.rs | 181 ++++++- rust/cuvs/src/distance_type.rs | 6 - rust/cuvs/src/error.rs | 85 ++- rust/cuvs/src/ivf_flat/index_params.rs | 106 ---- rust/cuvs/src/ivf_flat/mod.rs | 23 - rust/cuvs/src/ivf_flat/search_params.rs | 61 --- rust/cuvs/src/ivf_pq/index_params.rs | 182 ------- rust/cuvs/src/ivf_pq/mod.rs | 23 - rust/cuvs/src/ivf_pq/search_params.rs | 87 --- rust/cuvs/src/lib.rs | 9 +- rust/cuvs/src/{ => neighbors}/brute_force.rs | 96 ++-- rust/cuvs/src/{ => neighbors}/cagra/index.rs | 118 ++--- rust/cuvs/src/neighbors/cagra/mod.rs | 152 ++++++ rust/cuvs/src/neighbors/cagra/params.rs | 497 ++++++++++++++++++ .../src/{ => neighbors}/ivf_flat/index.rs | 84 +-- rust/cuvs/src/neighbors/ivf_flat/mod.rs | 34 ++ rust/cuvs/src/neighbors/ivf_flat/params.rs | 147 ++++++ rust/cuvs/src/{ => neighbors}/ivf_pq/index.rs | 89 ++-- rust/cuvs/src/neighbors/ivf_pq/mod.rs | 36 ++ rust/cuvs/src/neighbors/ivf_pq/params.rs | 184 +++++++ rust/cuvs/src/neighbors/mod.rs | 17 + rust/cuvs/src/{ => neighbors}/vamana/index.rs | 65 +-- rust/cuvs/src/neighbors/vamana/mod.rs | 32 ++ rust/cuvs/src/neighbors/vamana/params.rs | 111 ++++ rust/cuvs/src/resources.rs | 69 ++- rust/cuvs/src/test_utils.rs | 16 +- rust/cuvs/src/vamana/index_params.rs | 129 ----- rust/cuvs/src/vamana/mod.rs | 15 - 35 files changed, 1753 insertions(+), 1513 deletions(-) delete mode 100644 rust/cuvs/src/cagra/index_params.rs delete mode 100644 rust/cuvs/src/cagra/mod.rs delete mode 100644 rust/cuvs/src/cagra/search_params.rs delete mode 100644 rust/cuvs/src/distance_type.rs delete mode 100644 rust/cuvs/src/ivf_flat/index_params.rs delete mode 100644 rust/cuvs/src/ivf_flat/mod.rs delete mode 100644 rust/cuvs/src/ivf_flat/search_params.rs delete mode 100644 rust/cuvs/src/ivf_pq/index_params.rs delete mode 100644 rust/cuvs/src/ivf_pq/mod.rs delete mode 100644 rust/cuvs/src/ivf_pq/search_params.rs rename rust/cuvs/src/{ => neighbors}/brute_force.rs (67%) rename rust/cuvs/src/{ => neighbors}/cagra/index.rs (86%) create mode 100644 rust/cuvs/src/neighbors/cagra/mod.rs create mode 100644 rust/cuvs/src/neighbors/cagra/params.rs rename rust/cuvs/src/{ => neighbors}/ivf_flat/index.rs (69%) create mode 100644 rust/cuvs/src/neighbors/ivf_flat/mod.rs create mode 100644 rust/cuvs/src/neighbors/ivf_flat/params.rs rename rust/cuvs/src/{ => neighbors}/ivf_pq/index.rs (66%) create mode 100644 rust/cuvs/src/neighbors/ivf_pq/mod.rs create mode 100644 rust/cuvs/src/neighbors/ivf_pq/params.rs create mode 100644 rust/cuvs/src/neighbors/mod.rs rename rust/cuvs/src/{ => neighbors}/vamana/index.rs (53%) create mode 100644 rust/cuvs/src/neighbors/vamana/mod.rs create mode 100644 rust/cuvs/src/neighbors/vamana/params.rs delete mode 100644 rust/cuvs/src/vamana/index_params.rs delete mode 100644 rust/cuvs/src/vamana/mod.rs diff --git a/rust/cuvs/Cargo.toml b/rust/cuvs/Cargo.toml index c66fa09993..6cbed00776 100644 --- a/rust/cuvs/Cargo.toml +++ b/rust/cuvs/Cargo.toml @@ -13,6 +13,7 @@ default = [] doc-only = ["cuvs-sys/doc-only"] [dependencies] +bon = "3" cuvs-sys = { workspace = true } thiserror = "2" tinyvec = { version = "1", features = ["alloc"] } diff --git a/rust/cuvs/examples/cagra.rs b/rust/cuvs/examples/cagra.rs index c67397c94c..9adceae17c 100644 --- a/rust/cuvs/examples/cagra.rs +++ b/rust/cuvs/examples/cagra.rs @@ -9,8 +9,8 @@ //! implementing the public [`AsDlTensor`]/[`AsDlTensorMut`] traits. The //! [`CudaTensor`] type manages device memory directly through the CUDA runtime //! (`cudaMalloc`/`cudaFree`) and copies to/from host arrays with `cudaMemcpyAsync` -//! on the cuVS stream, reusing the resources handle's `get_cuda_stream`/ -//! `sync_stream` for stream access and synchronization. +//! on the cuVS stream, reusing the resources handle's `stream`/`sync_stream` +//! for stream access and synchronization. //! //! A real application would likely rely on a helper crate such as `cudarc` //! and its `CudaSlice`. @@ -20,11 +20,11 @@ use std::marker::PhantomData; use std::os::raw::c_int; use cuvs::Resources; -use cuvs::cagra::{Index, IndexParams, SearchParams}; use cuvs::dlpack::{ AsDlTensor, AsDlTensorMut, DLDevice, DLDeviceType, DLPackError, DLTensorView, DLTensorViewMut, DType, }; +use cuvs::neighbors::cagra::{Index, IndexParams, SearchParams}; use ndarray::s; use ndarray_rand::RandomExt; @@ -98,7 +98,7 @@ impl CudaTensor { } let tensor = Self::alloc(host.shape())?; - let stream = res.get_cuda_stream()?; + let stream = res.stream()?; check_cuda(unsafe { cudaMemcpyAsync( tensor.data, @@ -122,7 +122,7 @@ impl CudaTensor { return Err("host array must be contiguous (row-major)".into()); } - let stream = res.get_cuda_stream()?; + let stream = res.stream()?; check_cuda(unsafe { cudaMemcpyAsync( host.as_mut_ptr() as *mut c_void, @@ -188,7 +188,7 @@ fn cagra_example() -> ExampleResult<()> { let dataset = CudaTensor::from_host(&res, &dataset_host)?; // Build the CAGRA index. - let build_params = IndexParams::new()?; + let build_params = IndexParams::builder().build()?; let index = Index::build(&res, &build_params, &dataset)?; println!("Indexed {n_datapoints}x{n_features} datapoints into cagra index"); @@ -201,7 +201,7 @@ fn cagra_example() -> ExampleResult<()> { let mut neighbors = CudaTensor::::alloc(&[n_queries, k])?; let mut distances = CudaTensor::::alloc(&[n_queries, k])?; - let search_params = SearchParams::new()?; + let search_params = SearchParams::builder().build()?; index.search(&res, &search_params, &queries, &mut neighbors, &mut distances)?; // Copy the results back to the host. diff --git a/rust/cuvs/src/cagra/index_params.rs b/rust/cuvs/src/cagra/index_params.rs deleted file mode 100644 index 9425ea060a..0000000000 --- a/rust/cuvs/src/cagra/index_params.rs +++ /dev/null @@ -1,196 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. - * SPDX-License-Identifier: Apache-2.0 - */ - -use crate::error::{Result, check_cuvs}; -use std::fmt; -use std::io::{Write, stderr}; - -pub type BuildAlgo = ffi::cuvsCagraGraphBuildAlgo; - -/// Supplemental parameters to build CAGRA Index -pub struct CompressionParams(pub ffi::cuvsCagraCompressionParams_t); - -impl CompressionParams { - /// Returns a new CompressionParams - pub fn new() -> Result { - unsafe { - let mut params = std::mem::MaybeUninit::::uninit(); - check_cuvs(ffi::cuvsCagraCompressionParamsCreate(params.as_mut_ptr()))?; - Ok(CompressionParams(params.assume_init())) - } - } - - /// The bit length of the vector element after compression by PQ. - pub fn set_pq_bits(self, pq_bits: u32) -> CompressionParams { - unsafe { - (*self.0).pq_bits = pq_bits; - } - self - } - - /// The dimensionality of the vector after compression by PQ. When zero, - /// an optimal value is selected using a heuristic. - pub fn set_pq_dim(self, pq_dim: u32) -> CompressionParams { - unsafe { - (*self.0).pq_dim = pq_dim; - } - self - } - - /// Vector Quantization (VQ) codebook size - number of "coarse cluster - /// centers". When zero, an optimal value is selected using a heuristic. - pub fn set_vq_n_centers(self, vq_n_centers: u32) -> CompressionParams { - unsafe { - (*self.0).vq_n_centers = vq_n_centers; - } - self - } - - /// The number of iterations searching for kmeans centers (both VQ & PQ - /// phases). - pub fn set_kmeans_n_iters(self, kmeans_n_iters: u32) -> CompressionParams { - unsafe { - (*self.0).kmeans_n_iters = kmeans_n_iters; - } - self - } - - /// The fraction of data to use during iterative kmeans building (VQ - /// phase). When zero, an optimal value is selected using a heuristic. - pub fn set_vq_kmeans_trainset_fraction( - self, - vq_kmeans_trainset_fraction: f64, - ) -> CompressionParams { - unsafe { - (*self.0).vq_kmeans_trainset_fraction = vq_kmeans_trainset_fraction; - } - self - } - - /// The fraction of data to use during iterative kmeans building (PQ - /// phase). When zero, an optimal value is selected using a heuristic. - pub fn set_pq_kmeans_trainset_fraction( - self, - pq_kmeans_trainset_fraction: f64, - ) -> CompressionParams { - unsafe { - (*self.0).pq_kmeans_trainset_fraction = pq_kmeans_trainset_fraction; - } - self - } -} - -pub struct IndexParams(pub ffi::cuvsCagraIndexParams_t, Option); - -impl IndexParams { - /// Returns a new IndexParams - pub fn new() -> Result { - unsafe { - let mut params = std::mem::MaybeUninit::::uninit(); - check_cuvs(ffi::cuvsCagraIndexParamsCreate(params.as_mut_ptr()))?; - Ok(IndexParams(params.assume_init(), None)) - } - } - - /// Degree of input graph for pruning - pub fn set_intermediate_graph_degree(self, intermediate_graph_degree: usize) -> IndexParams { - unsafe { - (*self.0).intermediate_graph_degree = intermediate_graph_degree; - } - self - } - - /// Degree of output graph - pub fn set_graph_degree(self, graph_degree: usize) -> IndexParams { - unsafe { - (*self.0).graph_degree = graph_degree; - } - self - } - - /// ANN algorithm to build knn graph - pub fn set_build_algo(self, build_algo: BuildAlgo) -> IndexParams { - unsafe { - (*self.0).build_algo = build_algo; - } - self - } - - /// Number of iterations to run if building with NN_DESCENT - pub fn set_nn_descent_niter(self, nn_descent_niter: usize) -> IndexParams { - unsafe { - (*self.0).nn_descent_niter = nn_descent_niter; - } - self - } - - pub fn set_compression(mut self, compression: CompressionParams) -> IndexParams { - unsafe { - (*self.0).compression = compression.0; - } - // Note: we're moving the ownership of compression here to avoid having it cleaned up - // and leaving a dangling pointer - self.1 = Some(compression); - self - } -} - -impl fmt::Debug for IndexParams { - fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { - // custom debug trait here, default value will show the pointer address - // for the inner params object which isn't that useful. - write!(f, "IndexParams({:?})", unsafe { *self.0 }) - } -} - -impl fmt::Debug for CompressionParams { - fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { - write!(f, "CompressionParams({:?})", unsafe { *self.0 }) - } -} - -impl Drop for IndexParams { - fn drop(&mut self) { - if let Err(e) = check_cuvs(unsafe { ffi::cuvsCagraIndexParamsDestroy(self.0) }) { - write!(stderr(), "failed to call cuvsCagraIndexParamsDestroy {:?}", e) - .expect("failed to write to stderr"); - } - } -} - -impl Drop for CompressionParams { - fn drop(&mut self) { - if let Err(e) = check_cuvs(unsafe { ffi::cuvsCagraCompressionParamsDestroy(self.0) }) { - write!(stderr(), "failed to call cuvsCagraCompressionParamsDestroy {:?}", e) - .expect("failed to write to stderr"); - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_index_params() { - let params = IndexParams::new() - .unwrap() - .set_intermediate_graph_degree(128) - .set_graph_degree(16) - .set_build_algo(BuildAlgo::NN_DESCENT) - .set_nn_descent_niter(10) - .set_compression(CompressionParams::new().unwrap().set_pq_bits(4).set_pq_dim(8)); - - // make sure the setters actually updated internal representation on the c-struct - unsafe { - assert_eq!((*params.0).graph_degree, 16); - assert_eq!((*params.0).intermediate_graph_degree, 128); - assert_eq!((*params.0).build_algo, BuildAlgo::NN_DESCENT); - assert_eq!((*params.0).nn_descent_niter, 10); - assert_eq!((*(*params.0).compression).pq_dim, 8); - assert_eq!((*(*params.0).compression).pq_bits, 4); - } - } -} diff --git a/rust/cuvs/src/cagra/mod.rs b/rust/cuvs/src/cagra/mod.rs deleted file mode 100644 index 050497fc1b..0000000000 --- a/rust/cuvs/src/cagra/mod.rs +++ /dev/null @@ -1,22 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -//! CAGRA: a graph-based approximate nearest neighbors algorithm with -//! state-of-the-art query throughput for both small and large batch sizes. -//! -//! Build an [`Index`] from a dataset, then [`search`](Index::search) it with -//! device-resident queries and output buffers. Tensors are passed through the -//! [`AsDlTensor`](crate::AsDlTensor) / -//! [`AsDlTensorMut`](crate::AsDlTensorMut) traits; see the -//! [`dlpack`](crate::dlpack) module for the tensor model and `examples/cagra.rs` -//! for a complete, runnable example. - -mod index; -mod index_params; -mod search_params; - -pub use index::Index; -pub use index_params::{BuildAlgo, CompressionParams, IndexParams}; -pub use search_params::{HashMode, SearchAlgo, SearchParams}; diff --git a/rust/cuvs/src/cagra/search_params.rs b/rust/cuvs/src/cagra/search_params.rs deleted file mode 100644 index a53fa9de05..0000000000 --- a/rust/cuvs/src/cagra/search_params.rs +++ /dev/null @@ -1,154 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. - * SPDX-License-Identifier: Apache-2.0 - */ - -use crate::error::{Result, check_cuvs}; -use std::fmt; -use std::io::{Write, stderr}; - -pub type SearchAlgo = ffi::cuvsCagraSearchAlgo; -pub type HashMode = ffi::cuvsCagraHashMode; - -/// Supplemental parameters to search CAGRA index -pub struct SearchParams(pub ffi::cuvsCagraSearchParams_t); - -impl SearchParams { - /// Returns a new SearchParams object - pub fn new() -> Result { - unsafe { - let mut params = std::mem::MaybeUninit::::uninit(); - check_cuvs(ffi::cuvsCagraSearchParamsCreate(params.as_mut_ptr()))?; - Ok(SearchParams(params.assume_init())) - } - } - - /// Maximum number of queries to search at the same time (batch size). Auto select when 0 - pub fn set_max_queries(self, max_queries: usize) -> SearchParams { - unsafe { - (*self.0).max_queries = max_queries; - } - self - } - - /// Number of intermediate search results retained during the search. - /// This is the main knob to adjust trade off between accuracy and search speed. - /// Higher values improve the search accuracy - pub fn set_itopk_size(self, itopk_size: usize) -> SearchParams { - unsafe { - (*self.0).itopk_size = itopk_size; - } - self - } - - /// Upper limit of search iterations. Auto select when 0. - pub fn set_max_iterations(self, max_iterations: usize) -> SearchParams { - unsafe { - (*self.0).max_iterations = max_iterations; - } - self - } - - /// Which search implementation to use. - pub fn set_algo(self, algo: SearchAlgo) -> SearchParams { - unsafe { - (*self.0).algo = algo; - } - self - } - - /// Number of threads used to calculate a single distance. 4, 8, 16, or 32. - pub fn set_team_size(self, team_size: usize) -> SearchParams { - unsafe { - (*self.0).team_size = team_size; - } - self - } - - /// Lower limit of search iterations. - pub fn set_min_iterations(self, min_iterations: usize) -> SearchParams { - unsafe { - (*self.0).min_iterations = min_iterations; - } - self - } - - /// Thread block size. 0, 64, 128, 256, 512, 1024. Auto selection when 0. - pub fn set_thread_block_size(self, thread_block_size: usize) -> SearchParams { - unsafe { - (*self.0).thread_block_size = thread_block_size; - } - self - } - - /// Hashmap type. Auto selection when AUTO. - pub fn set_hashmap_mode(self, hashmap_mode: HashMode) -> SearchParams { - unsafe { - (*self.0).hashmap_mode = hashmap_mode; - } - self - } - - /// Lower limit of hashmap bit length. More than 8. - pub fn set_hashmap_min_bitlen(self, hashmap_min_bitlen: usize) -> SearchParams { - unsafe { - (*self.0).hashmap_min_bitlen = hashmap_min_bitlen; - } - self - } - - /// Upper limit of hashmap fill rate. More than 0.1, less than 0.9. - pub fn set_hashmap_max_fill_rate(self, hashmap_max_fill_rate: f32) -> SearchParams { - unsafe { - (*self.0).hashmap_max_fill_rate = hashmap_max_fill_rate; - } - self - } - - /// Number of iterations of initial random seed node selection. 1 or more. - pub fn set_num_random_samplings(self, num_random_samplings: u32) -> SearchParams { - unsafe { - (*self.0).num_random_samplings = num_random_samplings; - } - self - } - - /// Bit mask used for initial random seed node selection. - pub fn set_rand_xor_mask(self, rand_xor_mask: u64) -> SearchParams { - unsafe { - (*self.0).rand_xor_mask = rand_xor_mask; - } - self - } -} - -impl fmt::Debug for SearchParams { - fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { - // custom debug trait here, default value will show the pointer address - // for the inner params object which isn't that useful. - write!(f, "SearchParams {{ params: {:?} }}", unsafe { *self.0 }) - } -} - -impl Drop for SearchParams { - fn drop(&mut self) { - if let Err(e) = check_cuvs(unsafe { ffi::cuvsCagraSearchParamsDestroy(self.0) }) { - write!(stderr(), "failed to call cuvsCagraSearchParamsDestroy {:?}", e) - .expect("failed to write to stderr"); - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_search_params() { - let params = SearchParams::new().unwrap().set_itopk_size(128); - - unsafe { - assert_eq!((*params.0).itopk_size, 128); - } - } -} diff --git a/rust/cuvs/src/cluster/kmeans/mod.rs b/rust/cuvs/src/cluster/kmeans/mod.rs index f5297663db..9212f7911c 100644 --- a/rust/cuvs/src/cluster/kmeans/mod.rs +++ b/rust/cuvs/src/cluster/kmeans/mod.rs @@ -7,26 +7,37 @@ //! //! [`fit`] computes cluster centroids for a dataset, [`predict`] assigns points //! to clusters, and [`cluster_cost`] reports the inertia. All inputs and outputs -//! reside in device memory and are borrowed through the -//! [`AsDlTensor`] / -//! [`AsDlTensorMut`] traits; see the -//! [`dlpack`](crate::dlpack) module for the tensor model. +//! reside in device memory and are borrowed through the `AsDlTensor` / +//! `AsDlTensorMut` traits; see the [`dlpack`](crate::dlpack) module for the +//! tensor model. mod params; pub use params::Params; -use crate::dlpack::{AsDlTensor, AsDlTensorMut}; -use crate::error::{Result, check_cuvs}; +use crate::dlpack::{AsDlTensor, AsDlTensorMut, DLPackError}; +use crate::error::{LibraryError, check_cuvs}; use crate::resources::Resources; +type Result = std::result::Result; + +/// Error type for k-means operations. +#[derive(Debug, thiserror::Error)] +pub enum KMeansError { + /// The cuVS C library reported a failure. + #[error(transparent)] + Library(#[from] LibraryError), + /// Tensor conversion into DLPack metadata failed. + #[error(transparent)] + DLPack(#[from] DLPackError), +} + /// Fits k-means centroids to `x`, returning `(inertia, n_iterations)`. /// /// `x` (shape `m × k`) is the input matrix and `centroids` (shape /// `n_clusters × k`) receives the fitted centroids; `sample_weight` is an /// optional per-sample weight. All reside in device memory and implement -/// [`AsDlTensor`] / -/// [`AsDlTensorMut`]. +/// [`AsDlTensor`] / [`AsDlTensorMut`]. pub fn fit( res: &Resources, params: &Params, @@ -50,8 +61,8 @@ where unsafe { check_cuvs(ffi::cuvsKMeansFit( - res.0, - params.0, + res.handle(), + params.handle(), x.to_c().as_mut_ptr(), sample_weight_ptr, centroids.to_c().as_mut_ptr(), @@ -67,8 +78,7 @@ where /// /// `x` (shape `m × k`), `centroids` (shape `n_clusters × k`), the optional /// `sample_weight`, and `labels` (shape `m × 1`) reside in device memory and -/// implement [`AsDlTensor`] / -/// [`AsDlTensorMut`]. `normalize_weight` selects +/// implement [`AsDlTensor`] / [`AsDlTensorMut`]. `normalize_weight` selects /// whether the sample weights are normalized. pub fn predict( res: &Resources, @@ -96,8 +106,8 @@ where unsafe { check_cuvs(ffi::cuvsKMeansPredict( - res.0, - params.0, + res.handle(), + params.handle(), x.to_c().as_mut_ptr(), sample_weight_ptr, centroids.to_c().as_mut_ptr(), @@ -124,7 +134,7 @@ where unsafe { check_cuvs(ffi::cuvsKMeansClusterCost( - res.0, + res.handle(), x.to_c().as_mut_ptr(), centroids.to_c().as_mut_ptr(), &mut inertia as *mut f64, @@ -146,7 +156,6 @@ mod tests { let n_clusters = 4; - // Create a new random dataset to index let n_datapoints = 256; let n_features = 16; let dataset_host = ndarray::Array::::random( @@ -158,12 +167,10 @@ mod tests { let centroids_host = ndarray::Array::::zeros((n_clusters, n_features)); let mut centroids = DeviceTensor::from_host(&res, ¢roids_host).unwrap(); - let params = Params::new().unwrap().set_n_clusters(n_clusters as i32); + let params = Params::builder().n_clusters(n_clusters as i32).build().unwrap(); - // compute the inertia, before fitting centroids let original_inertia = cluster_cost(&res, &dataset, ¢roids).unwrap(); - // fit the centroids, make sure that inertia has gone down let (inertia, n_iter) = fit(&res, ¶ms, &dataset, None::<&DeviceTensor<'_, f32>>, &mut centroids).unwrap(); @@ -173,7 +180,6 @@ mod tests { let mut labels_host = ndarray::Array::::zeros((n_clusters,)); let mut labels = DeviceTensor::::zeros(&res, &[n_clusters]).unwrap(); - // make sure the prediction for each centroid is the centroid itself predict( &res, ¶ms, diff --git a/rust/cuvs/src/cluster/kmeans/params.rs b/rust/cuvs/src/cluster/kmeans/params.rs index 460b9a7fb6..b5977ff6a1 100644 --- a/rust/cuvs/src/cluster/kmeans/params.rs +++ b/rust/cuvs/src/cluster/kmeans/params.rs @@ -1,124 +1,102 @@ /* - * 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 */ -use crate::distance_type::DistanceType; -use crate::error::{Result, check_cuvs}; -use std::fmt; -use std::io::{Write, stderr}; +//! Builder-pattern parameter type for k-means. +//! +//! All setters are optional; unset values retain the library defaults from the +//! underlying C `cuvsKMeansParamsCreate`. -pub struct Params(pub ffi::cuvsKMeansParams_t); +use std::{fmt, ptr}; -impl Params { - /// Returns a new Params - pub fn new() -> Result { - unsafe { - let mut params = std::mem::MaybeUninit::::uninit(); - check_cuvs(ffi::cuvsKMeansParamsCreate(params.as_mut_ptr()))?; - Ok(Params(params.assume_init())) - } - } - - /// DistanceType to use for fitting kmeans - pub fn set_metric(self, metric: DistanceType) -> Params { - unsafe { - (*self.0).metric = metric; - } - self - } +use bon::bon; - /// The number of clusters to form as well as the number of centroids to generate (default:8). - pub fn set_n_clusters(self, n_clusters: i32) -> Params { - unsafe { - (*self.0).n_clusters = n_clusters; - } - self - } - - /// Maximum number of iterations of the k-means algorithm for a single run. - pub fn set_max_iter(self, max_iter: i32) -> Params { - unsafe { - (*self.0).max_iter = max_iter; - } - self - } - - /// Relative tolerance with regards to inertia to declare convergence. - pub fn set_tol(self, tol: f64) -> Params { - unsafe { - (*self.0).tol = tol; - } - self - } +use crate::distance::DistanceType; +use crate::error::check_cuvs; - /// Number of instance k-means algorithm will be run with different seeds. - pub fn set_n_init(self, n_init: i32) -> Params { - unsafe { - (*self.0).n_init = n_init; - } - self - } +use super::KMeansError; - /// Oversampling factor for use in the k-means|| algorithm - pub fn set_oversampling_factor(self, oversampling_factor: f64) -> Params { - unsafe { - (*self.0).oversampling_factor = oversampling_factor; - } - self - } +/// Parameters for k-means fitting and prediction. +pub struct Params { + handle: ffi::cuvsKMeansParams_t, +} - /** - * batch_samples and batch_centroids are used to tile 1NN computation which is - * useful to optimize/control the memory footprint - * Default tile is [batch_samples x n_clusters] i.e. when batch_centroids is 0 - * then don't tile the centroids. - */ - pub fn set_batch_samples(self, batch_samples: i32) -> Params { - unsafe { - (*self.0).batch_samples = batch_samples; - } - self - } - /// if 0 then batch_centroids = n_clusters - pub fn set_batch_centroids(self, batch_centroids: i32) -> Params { +#[bon] +impl Params { + #[builder] + #[allow(clippy::too_many_arguments)] + pub fn new( + metric: Option, + n_clusters: Option, + max_iter: Option, + tol: Option, + n_init: Option, + oversampling_factor: Option, + batch_samples: Option, + batch_centroids: Option, + hierarchical: Option, + hierarchical_n_iters: Option, + ) -> Result { + let params = Self::try_new()?; unsafe { - (*self.0).batch_centroids = batch_centroids; + if let Some(v) = metric { + (*params.handle).metric = v.into(); + } + if let Some(v) = n_clusters { + (*params.handle).n_clusters = v; + } + if let Some(v) = max_iter { + (*params.handle).max_iter = v; + } + if let Some(v) = tol { + (*params.handle).tol = v; + } + if let Some(v) = n_init { + (*params.handle).n_init = v; + } + if let Some(v) = oversampling_factor { + (*params.handle).oversampling_factor = v; + } + if let Some(v) = batch_samples { + (*params.handle).batch_samples = v; + } + if let Some(v) = batch_centroids { + (*params.handle).batch_centroids = v; + } + if let Some(v) = hierarchical { + (*params.handle).hierarchical = v; + } + if let Some(v) = hierarchical_n_iters { + (*params.handle).hierarchical_n_iters = v; + } } - self + Ok(params) } +} - /// Whether to use hierarchical (balanced) kmeans or not - pub fn set_hierarchical(self, hierarchical: bool) -> Params { - unsafe { - (*self.0).hierarchical = hierarchical; - } - self +impl Params { + /// Allocate parameters populated with the library defaults. + pub fn try_new() -> Result { + let mut handle = ptr::null_mut(); + check_cuvs(unsafe { ffi::cuvsKMeansParamsCreate(&mut handle) })?; + Ok(Self { handle }) } - /// For hierarchical k-means , defines the number of training iterations - pub fn set_hierarchical_n_iters(self, hierarchical_n_iters: i32) -> Params { - unsafe { - (*self.0).hierarchical_n_iters = hierarchical_n_iters; - } - self + pub(super) fn handle(&self) -> ffi::cuvsKMeansParams_t { + self.handle } } impl fmt::Debug for Params { - fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { - // custom debug trait here, default value will show the pointer address - // for the inner params object which isn't that useful. - write!(f, "Params({:?})", unsafe { *self.0 }) + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_tuple("Params").field(unsafe { &*self.handle }).finish() } } impl Drop for Params { fn drop(&mut self) { - if let Err(e) = check_cuvs(unsafe { ffi::cuvsKMeansParamsDestroy(self.0) }) { - write!(stderr(), "failed to call cuvsKMeansParamsDestroy {:?}", e) - .expect("failed to write to stderr"); - } + let _ = unsafe { ffi::cuvsKMeansParamsDestroy(self.handle) }; } } @@ -127,12 +105,11 @@ mod tests { use super::*; #[test] - fn test_params() { - let params = Params::new().unwrap().set_n_clusters(128).set_hierarchical(true); - + fn params_with_values() { + let params = Params::builder().n_clusters(128).hierarchical(true).build().unwrap(); unsafe { - assert_eq!((*params.0).n_clusters, 128); - assert_eq!((*params.0).hierarchical, true); + assert_eq!((*params.handle).n_clusters, 128); + assert!((*params.handle).hierarchical); } } } diff --git a/rust/cuvs/src/distance/mod.rs b/rust/cuvs/src/distance/mod.rs index 5cd250991f..f6bea26aed 100644 --- a/rust/cuvs/src/distance/mod.rs +++ b/rust/cuvs/src/distance/mod.rs @@ -3,34 +3,165 @@ * SPDX-License-Identifier: Apache-2.0 */ -//! Pairwise distance computation. +//! Distance metrics and pairwise distance computation. //! -//! [`pairwise_distance`] computes all pairwise distances between two device -//! matrices. Inputs and output are borrowed through the -//! [`AsDlTensor`] / -//! [`AsDlTensorMut`] traits; see the -//! [`dlpack`](crate::dlpack) module for the tensor model. - -use crate::distance_type::DistanceType; -use crate::dlpack::{AsDlTensor, AsDlTensorMut}; -use crate::error::{Result, check_cuvs}; +//! [`DistanceType`] selects the metric used by every index and by +//! [`pairwise_distance`]. Inputs and output are borrowed through the +//! `AsDlTensor` / `AsDlTensorMut` traits; see the [`dlpack`](crate::dlpack) +//! module for the tensor model. + +use crate::dlpack::{AsDlTensor, AsDlTensorMut, DLPackError}; +use crate::error::{LibraryError, check_cuvs}; use crate::resources::Resources; +const DEFAULT_METRIC_ARG: f32 = 2.0; + +/// Distance metric used for building and searching nearest neighbor indices. +#[derive(Debug, Copy, Clone, PartialEq)] +#[non_exhaustive] +pub enum DistanceType { + /// L2 (squared Euclidean) distance. + L2Expanded, + /// L2 distance with square root. + L2SqrtExpanded, + /// Cosine distance. + CosineExpanded, + /// L1 (Manhattan) distance. + L1, + /// L2 distance (unexpanded form). + L2Unexpanded, + /// L2 distance with square root (unexpanded form). + L2SqrtUnexpanded, + /// Inner product. + InnerProduct, + /// Chebyshev (L-infinity) distance. + Linf, + /// Canberra distance. + Canberra, + /// Generalized Minkowski (Lp) distance with exponent `p`. + LpUnexpanded(f32), + /// Correlation distance. + CorrelationExpanded, + /// Jaccard distance. + JaccardExpanded, + /// Hellinger distance. + HellingerExpanded, + /// Haversine (great-circle) distance. + Haversine, + /// Bray-Curtis distance. + BrayCurtis, + /// Jensen-Shannon divergence. + JensenShannon, + /// Hamming distance. + HammingUnexpanded, + /// Kullback-Leibler divergence. + KLDivergence, + /// Russell-Rao distance. + RusselRaoExpanded, + /// Dice-Sorensen distance. + DiceExpanded, + /// Bitwise Hamming distance. + BitwiseHamming, + /// Precomputed distance matrix. + Precomputed, +} + +impl DistanceType { + /// The `metric_arg` the C API expects: the exponent `p` for Minkowski + /// ([`LpUnexpanded`](DistanceType::LpUnexpanded)), otherwise the default. + pub(crate) fn metric_arg(self) -> f32 { + match self { + Self::LpUnexpanded(p) => p, + _ => DEFAULT_METRIC_ARG, + } + } +} + +impl From for ffi::cuvsDistanceType { + fn from(v: DistanceType) -> Self { + use DistanceType::*; + match v { + L2Expanded => Self::L2Expanded, + L2SqrtExpanded => Self::L2SqrtExpanded, + CosineExpanded => Self::CosineExpanded, + L1 => Self::L1, + L2Unexpanded => Self::L2Unexpanded, + L2SqrtUnexpanded => Self::L2SqrtUnexpanded, + InnerProduct => Self::InnerProduct, + Linf => Self::Linf, + Canberra => Self::Canberra, + LpUnexpanded(_) => Self::LpUnexpanded, + CorrelationExpanded => Self::CorrelationExpanded, + JaccardExpanded => Self::JaccardExpanded, + HellingerExpanded => Self::HellingerExpanded, + Haversine => Self::Haversine, + BrayCurtis => Self::BrayCurtis, + JensenShannon => Self::JensenShannon, + HammingUnexpanded => Self::HammingUnexpanded, + KLDivergence => Self::KLDivergence, + RusselRaoExpanded => Self::RusselRaoExpanded, + DiceExpanded => Self::DiceExpanded, + BitwiseHamming => Self::BitwiseHamming, + Precomputed => Self::Precomputed, + } + } +} + +impl From for DistanceType { + fn from(v: ffi::cuvsDistanceType) -> Self { + use ffi::cuvsDistanceType::*; + match v { + L2Expanded => Self::L2Expanded, + L2SqrtExpanded => Self::L2SqrtExpanded, + CosineExpanded => Self::CosineExpanded, + L1 => Self::L1, + L2Unexpanded => Self::L2Unexpanded, + L2SqrtUnexpanded => Self::L2SqrtUnexpanded, + InnerProduct => Self::InnerProduct, + Linf => Self::Linf, + Canberra => Self::Canberra, + LpUnexpanded => Self::LpUnexpanded(DEFAULT_METRIC_ARG), + CorrelationExpanded => Self::CorrelationExpanded, + JaccardExpanded => Self::JaccardExpanded, + HellingerExpanded => Self::HellingerExpanded, + Haversine => Self::Haversine, + BrayCurtis => Self::BrayCurtis, + JensenShannon => Self::JensenShannon, + HammingUnexpanded => Self::HammingUnexpanded, + KLDivergence => Self::KLDivergence, + RusselRaoExpanded => Self::RusselRaoExpanded, + DiceExpanded => Self::DiceExpanded, + BitwiseHamming => Self::BitwiseHamming, + Precomputed => Self::Precomputed, + } + } +} + +/// Error type for pairwise distance operations. +#[derive(Debug, thiserror::Error)] +pub enum DistanceError { + /// The cuVS C library reported a failure. + #[error(transparent)] + Library(#[from] LibraryError), + /// Tensor conversion into DLPack metadata failed. + #[error(transparent)] + DLPack(#[from] DLPackError), +} + /// Computes all pairwise distances between the rows of `x` (shape `m × k`) and /// `y` (shape `n × k`), writing the `m × n` result into `distances`. /// /// `x`, `y`, and `distances` reside in device memory and implement -/// [`AsDlTensor`] / -/// [`AsDlTensorMut`]. `metric` selects the distance; -/// `metric_arg` is the optional `p` for Minkowski distances (defaults to 2). +/// [`AsDlTensor`] / [`AsDlTensorMut`]. `metric` selects the distance; use +/// [`DistanceType::LpUnexpanded`] to supply the Minkowski exponent `p` (all +/// other metrics use the C API default). pub fn pairwise_distance( res: &Resources, x: &X, y: &Y, distances: &mut D, metric: DistanceType, - metric_arg: Option, -) -> Result<()> +) -> Result<(), DistanceError> where X: AsDlTensor + ?Sized, Y: AsDlTensor + ?Sized, @@ -39,16 +170,17 @@ where let x = x.as_dl_tensor()?; let y = y.as_dl_tensor()?; let distances = distances.as_dl_tensor_mut()?; - unsafe { - check_cuvs(ffi::cuvsPairwiseDistance( - res.0, + check_cuvs(unsafe { + ffi::cuvsPairwiseDistance( + res.handle(), x.to_c().as_mut_ptr(), y.to_c().as_mut_ptr(), distances.to_c().as_mut_ptr(), - metric, - metric_arg.unwrap_or(2.0), - )) - } + metric.into(), + metric.metric_arg(), + ) + })?; + Ok(()) } #[cfg(test)] @@ -62,7 +194,6 @@ mod tests { fn test_pairwise_distance() { let res = Resources::new().unwrap(); - // Create a new random dataset to index let n_datapoints = 256; let n_features = 16; let dataset = ndarray::Array::::random( @@ -81,14 +212,12 @@ mod tests { &dataset_device, &mut distances, DistanceType::L2Expanded, - None, ) .unwrap(); - // Copy back to host memory distances.copy_to_host(&res, &mut distances_host).unwrap(); - // Self distance should be 0 + // Self distance should be 0. assert_eq!(distances_host[[0, 0]], 0.0); } } diff --git a/rust/cuvs/src/distance_type.rs b/rust/cuvs/src/distance_type.rs deleted file mode 100644 index 18b8d14894..0000000000 --- a/rust/cuvs/src/distance_type.rs +++ /dev/null @@ -1,6 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION. - * SPDX-License-Identifier: Apache-2.0 - */ - -pub type DistanceType = ffi::cuvsDistanceType; diff --git a/rust/cuvs/src/error.rs b/rust/cuvs/src/error.rs index f9e58b3d3d..3ff2437946 100644 --- a/rust/cuvs/src/error.rs +++ b/rust/cuvs/src/error.rs @@ -3,64 +3,41 @@ * SPDX-License-Identifier: Apache-2.0 */ -use std::fmt; - -#[derive(Debug, Clone)] -pub struct CuvsError { - code: ffi::cuvsError_t, - text: String, -} - -#[derive(Debug, Clone)] -pub enum Error { - CuvsError(CuvsError), - /// Tensor conversion into DLPack metadata failed. - DLPack(crate::dlpack::DLPackError), - /// The caller passed an argument that could not be forwarded to the C API - /// (e.g. a filename containing an interior NUL byte or invalid UTF-8). - InvalidArgument(String), -} - -impl std::error::Error for Error {} -impl std::error::Error for CuvsError {} - -pub type Result = std::result::Result; - -impl fmt::Display for Error { - fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { - match self { - Error::CuvsError(cuvs_error) => write!(f, "cuvsError={:?}", cuvs_error), - Error::DLPack(err) => write!(f, "DLPack error: {}", err), - Error::InvalidArgument(msg) => write!(f, "invalid argument: {}", msg), - } - } -} - -impl fmt::Display for CuvsError { - fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { - write!(f, "{:?}:{}", self.code, self.text) - } -} - -/// Simple wrapper to convert a cuvsError_t into a Result -pub fn check_cuvs(err: ffi::cuvsError_t) -> Result<()> { - match err { +//! Low-level error handling shared by every cuVS module. +//! +//! [`check_cuvs`] turns a raw `cuvsError_t` status into a [`LibraryError`], which +//! each module's error type wraps via `#[from]`. + +use std::borrow::Cow; + +/// A failure reported by the cuVS C library. +/// +/// Carries the message captured from `cuvsGetLastErrorText` at the point of +/// failure. Every module's error type wraps this via `#[from]`. +#[derive(Debug, Clone, thiserror::Error)] +#[error("{0}")] +pub struct LibraryError(Cow<'static, str>); + +/// Converts a `cuvsError_t` status into a [`LibraryError`]. +/// +/// On failure the thread-local error text is captured immediately, before any +/// subsequent FFI call can overwrite it. +pub(crate) fn check_cuvs(status: ffi::cuvsError_t) -> Result<(), LibraryError> { + match status { ffi::cuvsError_t::CUVS_SUCCESS => Ok(()), _ => { - // get a description of the error from cuvs - let cstr = unsafe { + // SAFETY: `cuvsGetLastErrorText` returns either NULL or a pointer to + // thread-local storage valid until the next FFI call; copy it now. + let text: Cow<'static, str> = unsafe { let text_ptr = ffi::cuvsGetLastErrorText(); - std::ffi::CStr::from_ptr(text_ptr) + if text_ptr.is_null() { + Cow::Borrowed("unknown cuVS error") + } else { + let cstr = std::ffi::CStr::from_ptr(text_ptr); + Cow::Owned(String::from_utf8_lossy(cstr.to_bytes()).into_owned()) + } }; - let text = std::string::String::from_utf8_lossy(cstr.to_bytes()).to_string(); - - Err(Error::CuvsError(CuvsError { code: err, text })) + Err(LibraryError(text)) } } } - -impl From for Error { - fn from(err: crate::dlpack::DLPackError) -> Self { - Self::DLPack(err) - } -} diff --git a/rust/cuvs/src/ivf_flat/index_params.rs b/rust/cuvs/src/ivf_flat/index_params.rs deleted file mode 100644 index aef2487fce..0000000000 --- a/rust/cuvs/src/ivf_flat/index_params.rs +++ /dev/null @@ -1,106 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. - * SPDX-License-Identifier: Apache-2.0 - */ - -use crate::distance_type::DistanceType; -use crate::error::{Result, check_cuvs}; -use std::fmt; -use std::io::{Write, stderr}; - -pub struct IndexParams(pub ffi::cuvsIvfFlatIndexParams_t); - -impl IndexParams { - /// Returns a new IndexParams - pub fn new() -> Result { - unsafe { - let mut params = std::mem::MaybeUninit::::uninit(); - check_cuvs(ffi::cuvsIvfFlatIndexParamsCreate(params.as_mut_ptr()))?; - Ok(IndexParams(params.assume_init())) - } - } - - /// The number of clusters used in the coarse quantizer. - pub fn set_n_lists(self, n_lists: u32) -> IndexParams { - unsafe { - (*self.0).n_lists = n_lists; - } - self - } - - /// DistanceType to use for building the index - pub fn set_metric(self, metric: DistanceType) -> IndexParams { - unsafe { - (*self.0).metric = metric; - } - self - } - - /// The number of iterations searching for kmeans centers during index building. - pub fn set_metric_arg(self, metric_arg: f32) -> IndexParams { - unsafe { - (*self.0).metric_arg = metric_arg; - } - self - } - /// The number of iterations searching for kmeans centers during index building. - pub fn set_kmeans_n_iters(self, kmeans_n_iters: u32) -> IndexParams { - unsafe { - (*self.0).kmeans_n_iters = kmeans_n_iters; - } - self - } - - /// If kmeans_trainset_fraction is less than 1, then the dataset is - /// subsampled, and only n_samples * kmeans_trainset_fraction rows - /// are used for training. - pub fn set_kmeans_trainset_fraction(self, kmeans_trainset_fraction: f64) -> IndexParams { - unsafe { - (*self.0).kmeans_trainset_fraction = kmeans_trainset_fraction; - } - self - } - - /// After training the coarse and fine quantizers, we will populate - /// the index with the dataset if add_data_on_build == true, otherwise - /// the index is left empty, and the extend method can be used - /// to add new vectors to the index. - pub fn set_add_data_on_build(self, add_data_on_build: bool) -> IndexParams { - unsafe { - (*self.0).add_data_on_build = add_data_on_build; - } - self - } -} - -impl fmt::Debug for IndexParams { - fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { - // custom debug trait here, default value will show the pointer address - // for the inner params object which isn't that useful. - write!(f, "IndexParams({:?})", unsafe { *self.0 }) - } -} - -impl Drop for IndexParams { - fn drop(&mut self) { - if let Err(e) = check_cuvs(unsafe { ffi::cuvsIvfFlatIndexParamsDestroy(self.0) }) { - write!(stderr(), "failed to call cuvsIvfFlatIndexParamsDestroy {:?}", e) - .expect("failed to write to stderr"); - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_index_params() { - let params = IndexParams::new().unwrap().set_n_lists(128).set_add_data_on_build(false); - - unsafe { - assert_eq!((*params.0).n_lists, 128); - assert_eq!((*params.0).add_data_on_build, false); - } - } -} diff --git a/rust/cuvs/src/ivf_flat/mod.rs b/rust/cuvs/src/ivf_flat/mod.rs deleted file mode 100644 index f9735a9c0c..0000000000 --- a/rust/cuvs/src/ivf_flat/mod.rs +++ /dev/null @@ -1,23 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -//! IVF-Flat: an inverted-file index over uncompressed ("flat") vectors. It -//! partitions the dataset into `n_lists` clusters and, at query time, scans only -//! the `n_probes` closest clusters — a simple knob to trade recall for speed. -//! -//! Build an [`Index`] from a dataset, then [`search`](Index::search) it with -//! device-resident queries and output buffers. Tensors are passed through the -//! [`AsDlTensor`](crate::AsDlTensor) / -//! [`AsDlTensorMut`](crate::AsDlTensorMut) traits; see the -//! [`dlpack`](crate::dlpack) module for the tensor model and `examples/cagra.rs` -//! for the same build/search workflow. - -mod index; -mod index_params; -mod search_params; - -pub use index::Index; -pub use index_params::IndexParams; -pub use search_params::SearchParams; diff --git a/rust/cuvs/src/ivf_flat/search_params.rs b/rust/cuvs/src/ivf_flat/search_params.rs deleted file mode 100644 index 97da299b24..0000000000 --- a/rust/cuvs/src/ivf_flat/search_params.rs +++ /dev/null @@ -1,61 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. - * SPDX-License-Identifier: Apache-2.0 - */ - -use crate::error::{Result, check_cuvs}; -use std::fmt; -use std::io::{Write, stderr}; - -/// Supplemental parameters to search IvfFlat index -pub struct SearchParams(pub ffi::cuvsIvfFlatSearchParams_t); - -impl SearchParams { - /// Returns a new SearchParams object - pub fn new() -> Result { - unsafe { - let mut params = std::mem::MaybeUninit::::uninit(); - check_cuvs(ffi::cuvsIvfFlatSearchParamsCreate(params.as_mut_ptr()))?; - Ok(SearchParams(params.assume_init())) - } - } - - /// Supplemental parameters to search IVF-Flat index - pub fn set_n_probes(self, n_probes: u32) -> SearchParams { - unsafe { - (*self.0).n_probes = n_probes; - } - self - } -} - -impl fmt::Debug for SearchParams { - fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { - // custom debug trait here, default value will show the pointer address - // for the inner params object which isn't that useful. - write!(f, "SearchParams {{ params: {:?} }}", unsafe { *self.0 }) - } -} - -impl Drop for SearchParams { - fn drop(&mut self) { - if let Err(e) = check_cuvs(unsafe { ffi::cuvsIvfFlatSearchParamsDestroy(self.0) }) { - write!(stderr(), "failed to call cuvsIvfFlatSearchParamsDestroy {:?}", e) - .expect("failed to write to stderr"); - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_search_params() { - let params = SearchParams::new().unwrap().set_n_probes(128); - - unsafe { - assert_eq!((*params.0).n_probes, 128); - } - } -} diff --git a/rust/cuvs/src/ivf_pq/index_params.rs b/rust/cuvs/src/ivf_pq/index_params.rs deleted file mode 100644 index c822aa41c1..0000000000 --- a/rust/cuvs/src/ivf_pq/index_params.rs +++ /dev/null @@ -1,182 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. - * SPDX-License-Identifier: Apache-2.0 - */ - -use crate::distance_type::DistanceType; -use crate::error::{Result, check_cuvs}; -use std::fmt; -use std::io::{Write, stderr}; - -pub use ffi::cuvsIvfPqCodebookGen; -pub use ffi::cuvsIvfPqListLayout; - -pub struct IndexParams(pub ffi::cuvsIvfPqIndexParams_t); - -impl IndexParams { - /// Returns a new IndexParams - pub fn new() -> Result { - unsafe { - let mut params = std::mem::MaybeUninit::::uninit(); - check_cuvs(ffi::cuvsIvfPqIndexParamsCreate(params.as_mut_ptr()))?; - Ok(IndexParams(params.assume_init())) - } - } - - /// The number of clusters used in the coarse quantizer. - pub fn set_n_lists(self, n_lists: u32) -> IndexParams { - unsafe { - (*self.0).n_lists = n_lists; - } - self - } - - /// DistanceType to use for building the index - pub fn set_metric(self, metric: DistanceType) -> IndexParams { - unsafe { - (*self.0).metric = metric; - } - self - } - - /// The number of iterations searching for kmeans centers during index building. - pub fn set_metric_arg(self, metric_arg: f32) -> IndexParams { - unsafe { - (*self.0).metric_arg = metric_arg; - } - self - } - - /// The number of iterations searching for kmeans centers during index building. - pub fn set_kmeans_n_iters(self, kmeans_n_iters: u32) -> IndexParams { - unsafe { - (*self.0).kmeans_n_iters = kmeans_n_iters; - } - self - } - - /// If kmeans_trainset_fraction is less than 1, then the dataset is - /// subsampled, and only n_samples * kmeans_trainset_fraction rows - /// are used for training. - pub fn set_kmeans_trainset_fraction(self, kmeans_trainset_fraction: f64) -> IndexParams { - unsafe { - (*self.0).kmeans_trainset_fraction = kmeans_trainset_fraction; - } - self - } - - /// The bit length of the vector element after quantization. - pub fn set_pq_bits(self, pq_bits: u32) -> IndexParams { - unsafe { - (*self.0).pq_bits = pq_bits; - } - self - } - - /// The dimensionality of a the vector after product quantization. - /// When zero, an optimal value is selected using a heuristic. Note - /// pq_dim * pq_bits must be a multiple of 8. Hint: a smaller 'pq_dim' - /// results in a smaller index size and better search performance, but - /// lower recall. If 'pq_bits' is 8, 'pq_dim' can be set to any number, - /// but multiple of 8 are desirable for good performance. If 'pq_bits' - /// is not 8, 'pq_dim' should be a multiple of 8. For good performance, - /// it is desirable that 'pq_dim' is a multiple of 32. Ideally, - /// 'pq_dim' should be also a divisor of the dataset dim. - pub fn set_pq_dim(self, pq_dim: u32) -> IndexParams { - unsafe { - (*self.0).pq_dim = pq_dim; - } - self - } - - pub fn set_codebook_kind(self, codebook_kind: cuvsIvfPqCodebookGen) -> IndexParams { - unsafe { - (*self.0).codebook_kind = codebook_kind; - } - self - } - - /// Memory layout of the IVF-PQ list data. - /// - FLAT: Codes are stored contiguously, one vector's codes after another. - /// - INTERLEAVED: Codes are interleaved for optimized search performance. - /// This is the default and recommended for search workloads. - pub fn set_codes_layout(self, codes_layout: cuvsIvfPqListLayout) -> IndexParams { - unsafe { - (*self.0).codes_layout = codes_layout; - } - self - } - - /// Apply a random rotation matrix on the input data and queries even - /// if `dim % pq_dim == 0`. Note: if `dim` is not multiple of `pq_dim`, - /// a random rotation is always applied to the input data and queries - /// to transform the working space from `dim` to `rot_dim`, which may - /// be slightly larger than the original space and and is a multiple - /// of `pq_dim` (`rot_dim % pq_dim == 0`). However, this transform is - /// not necessary when `dim` is multiple of `pq_dim` (`dim == rot_dim`, - /// hence no need in adding "extra" data columns / features). By - /// default, if `dim == rot_dim`, the rotation transform is - /// initialized with the identity matrix. When - /// `force_random_rotation == True`, a random orthogonal transform - pub fn set_force_random_rotation(self, force_random_rotation: bool) -> IndexParams { - unsafe { - (*self.0).force_random_rotation = force_random_rotation; - } - self - } - - /// The max number of data points to use per PQ code during PQ codebook training. Using more data - /// points per PQ code may increase the quality of PQ codebook but may also increase the build - /// time. The parameter is applied to both PQ codebook generation methods, i.e., PER_SUBSPACE and - /// PER_CLUSTER. In both cases, we will use `pq_book_size * max_train_points_per_pq_code` training - /// points to train each codebook. - pub fn set_max_train_points_per_pq_code(self, max_pq_points: u32) -> IndexParams { - unsafe { - (*self.0).max_train_points_per_pq_code = max_pq_points; - } - self - } - - /// After training the coarse and fine quantizers, we will populate - /// the index with the dataset if add_data_on_build == true, otherwise - /// the index is left empty, and the extend method can be used - /// to add new vectors to the index. - pub fn set_add_data_on_build(self, add_data_on_build: bool) -> IndexParams { - unsafe { - (*self.0).add_data_on_build = add_data_on_build; - } - self - } -} - -impl fmt::Debug for IndexParams { - fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { - // custom debug trait here, default value will show the pointer address - // for the inner params object which isn't that useful. - write!(f, "IndexParams({:?})", unsafe { *self.0 }) - } -} - -impl Drop for IndexParams { - fn drop(&mut self) { - if let Err(e) = check_cuvs(unsafe { ffi::cuvsIvfPqIndexParamsDestroy(self.0) }) { - write!(stderr(), "failed to call cuvsIvfPqIndexParamsDestroy {:?}", e) - .expect("failed to write to stderr"); - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_index_params() { - let params = IndexParams::new().unwrap().set_n_lists(128).set_add_data_on_build(false); - - unsafe { - assert_eq!((*params.0).n_lists, 128); - assert_eq!((*params.0).add_data_on_build, false); - } - } -} diff --git a/rust/cuvs/src/ivf_pq/mod.rs b/rust/cuvs/src/ivf_pq/mod.rs deleted file mode 100644 index bc796b3691..0000000000 --- a/rust/cuvs/src/ivf_pq/mod.rs +++ /dev/null @@ -1,23 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ -//! IVF-PQ: an inverted-file index that product-quantizes the vectors. Like -//! IVF-Flat it partitions the dataset into `n_lists` clusters and scans the -//! `n_probes` closest at query time, but compresses each vector into `pq_dim` -//! codes of `pq_bits` bits — much smaller, slightly less accurate. -//! -//! Build an [`Index`] from a dataset, then [`search`](Index::search) it with -//! device-resident queries and output buffers. Tensors are passed through the -//! [`AsDlTensor`](crate::AsDlTensor) / -//! [`AsDlTensorMut`](crate::AsDlTensorMut) traits; see the -//! [`dlpack`](crate::dlpack) module for the tensor model and `examples/cagra.rs` -//! for the same build/search workflow. - -mod index; -mod index_params; -mod search_params; - -pub use index::Index; -pub use index_params::IndexParams; -pub use search_params::SearchParams; diff --git a/rust/cuvs/src/ivf_pq/search_params.rs b/rust/cuvs/src/ivf_pq/search_params.rs deleted file mode 100644 index fe880e29b6..0000000000 --- a/rust/cuvs/src/ivf_pq/search_params.rs +++ /dev/null @@ -1,87 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. - * SPDX-License-Identifier: Apache-2.0 - */ - -use crate::error::{Result, check_cuvs}; -use std::fmt; -use std::io::{Write, stderr}; - -pub use ffi::cudaDataType_t; - -/// Supplemental parameters to search IvfPq index -pub struct SearchParams(pub ffi::cuvsIvfPqSearchParams_t); - -impl SearchParams { - /// Returns a new SearchParams object - pub fn new() -> Result { - unsafe { - let mut params = std::mem::MaybeUninit::::uninit(); - check_cuvs(ffi::cuvsIvfPqSearchParamsCreate(params.as_mut_ptr()))?; - Ok(SearchParams(params.assume_init())) - } - } - - /// The number of clusters to search. - pub fn set_n_probes(self, n_probes: u32) -> SearchParams { - unsafe { - (*self.0).n_probes = n_probes; - } - self - } - - /// Data type of look up table to be created dynamically at search - /// time. The use of low-precision types reduces the amount of shared - /// memory required at search time, so fast shared memory kernels can - /// be used even for datasets with large dimansionality. Note that - /// the recall is slightly degraded when low-precision type is - /// selected. - pub fn set_lut_dtype(self, lut_dtype: cudaDataType_t) -> SearchParams { - unsafe { - (*self.0).lut_dtype = lut_dtype; - } - self - } - - /// Storage data type for distance/similarity computation. - pub fn set_internal_distance_dtype( - self, - internal_distance_dtype: cudaDataType_t, - ) -> SearchParams { - unsafe { - (*self.0).internal_distance_dtype = internal_distance_dtype; - } - self - } -} - -impl fmt::Debug for SearchParams { - fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { - // custom debug trait here, default value will show the pointer address - // for the inner params object which isn't that useful. - write!(f, "SearchParams {{ params: {:?} }}", unsafe { *self.0 }) - } -} - -impl Drop for SearchParams { - fn drop(&mut self) { - if let Err(e) = check_cuvs(unsafe { ffi::cuvsIvfPqSearchParamsDestroy(self.0) }) { - write!(stderr(), "failed to call cuvsIvfPqSearchParamsDestroy {:?}", e) - .expect("failed to write to stderr"); - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_search_params() { - let params = SearchParams::new().unwrap().set_n_probes(128); - - unsafe { - assert_eq!((*params.0).n_probes, 128); - } - } -} diff --git a/rust/cuvs/src/lib.rs b/rust/cuvs/src/lib.rs index 52c31392e7..d9b7017e97 100644 --- a/rust/cuvs/src/lib.rs +++ b/rust/cuvs/src/lib.rs @@ -9,20 +9,15 @@ //! approximate nearest neighbors search on the GPU. extern crate cuvs_sys as ffi; -pub mod brute_force; -pub mod cagra; pub mod cluster; pub mod distance; -pub mod distance_type; pub mod dlpack; mod error; -pub mod ivf_flat; -pub mod ivf_pq; +pub mod neighbors; mod resources; #[cfg(test)] pub(crate) mod test_utils; -pub mod vamana; pub use dlpack::{AsDlTensor, AsDlTensorMut, DLPackError, DLTensorView, DLTensorViewMut, DType}; -pub use error::{Error, Result}; +pub use error::LibraryError; pub use resources::Resources; diff --git a/rust/cuvs/src/brute_force.rs b/rust/cuvs/src/neighbors/brute_force.rs similarity index 67% rename from rust/cuvs/src/brute_force.rs rename to rust/cuvs/src/neighbors/brute_force.rs index 318346cc4d..3b054e19bd 100644 --- a/rust/cuvs/src/brute_force.rs +++ b/rust/cuvs/src/neighbors/brute_force.rs @@ -6,20 +6,32 @@ //! //! Build an [`Index`] over a dataset, then [`search`](Index::search) it with //! device-resident queries and output buffers. Tensors are borrowed through the -//! [`AsDlTensor`] / -//! [`AsDlTensorMut`] traits; see the -//! [`dlpack`](crate::dlpack) module for the tensor model and `examples/cagra.rs` -//! for the same build/search workflow. +//! `AsDlTensor` / `AsDlTensorMut` traits; see the [`dlpack`](crate::dlpack) +//! module for the tensor model and `examples/cagra.rs` for the same +//! build/search workflow. use std::io::{Write, stderr}; use std::marker::PhantomData; -use crate::distance_type::DistanceType; -use crate::dlpack::{AsDlTensor, AsDlTensorMut}; -use crate::error::{Result, check_cuvs}; +use crate::distance::DistanceType; +use crate::dlpack::{AsDlTensor, AsDlTensorMut, DLPackError}; +use crate::error::{LibraryError, check_cuvs}; use crate::resources::Resources; -/// Brute Force KNN Index +type Result = std::result::Result; + +/// Error type for brute-force operations. +#[derive(Debug, thiserror::Error)] +pub enum BruteForceError { + /// The cuVS C library reported a failure. + #[error(transparent)] + Library(#[from] LibraryError), + /// Tensor conversion into DLPack metadata failed. + #[error(transparent)] + DLPack(#[from] DLPackError), +} + +/// Brute-force KNN index. #[derive(Debug)] pub struct Index<'d> { inner: ffi::cuvsBruteForceIndex_t, @@ -31,17 +43,12 @@ pub struct Index<'d> { impl<'d> Index<'d> { /// Builds a brute-force index over `dataset` for exact k-NN search. /// - /// `metric` selects the distance and `metric_arg` is the optional `p` for - /// Minkowski distances (defaults to 2). `dataset` is a row-major matrix on - /// the host or device implementing [`AsDlTensor`]; the - /// C++ index keeps a non-owning view of it, so the returned [`Index`] borrows - /// it for `'d` and cannot outlive it. - pub fn build( - res: &Resources, - metric: DistanceType, - metric_arg: Option, - dataset: &'d T, - ) -> Result> + /// `metric` selects the distance (use [`DistanceType::LpUnexpanded`] to set + /// the Minkowski exponent `p`). `dataset` is a row-major matrix on the host + /// or device implementing [`AsDlTensor`]; the C++ index keeps a non-owning + /// view of it, so the returned [`Index`] borrows it for `'d` and cannot + /// outlive it. + pub fn build(res: &Resources, metric: DistanceType, dataset: &'d T) -> Result> where T: AsDlTensor + ?Sized, { @@ -49,17 +56,17 @@ impl<'d> Index<'d> { let index = Index::new()?; unsafe { check_cuvs(ffi::cuvsBruteForceBuild( - res.0, + res.handle(), dataset.to_c().as_mut_ptr(), - metric, - metric_arg.unwrap_or(2.0), + metric.into(), + metric.metric_arg(), index.inner, ))?; } Ok(index) } - /// Creates a new empty index + /// Creates a new empty index. pub fn new() -> Result> { unsafe { let mut index = std::mem::MaybeUninit::::uninit(); @@ -71,8 +78,7 @@ impl<'d> Index<'d> { /// Searches the index for the `k` nearest neighbors of each query. /// /// `queries`, `neighbors`, and `distances` must reside in device memory and - /// implement [`AsDlTensor`] / - /// [`AsDlTensorMut`]. `neighbors` receives the + /// implement [`AsDlTensor`] / [`AsDlTensorMut`]. `neighbors` receives the /// neighbor indices and `distances` their distances; both are written in /// place. pub fn search( @@ -90,18 +96,18 @@ impl<'d> Index<'d> { let queries = queries.as_dl_tensor()?; let neighbors = neighbors.as_dl_tensor_mut()?; let distances = distances.as_dl_tensor_mut()?; - unsafe { - let prefilter = ffi::cuvsFilter { addr: 0, type_: ffi::cuvsFilterType::NO_FILTER }; - - check_cuvs(ffi::cuvsBruteForceSearch( - res.0, + let prefilter = ffi::cuvsFilter { addr: 0, type_: ffi::cuvsFilterType::NO_FILTER }; + check_cuvs(unsafe { + ffi::cuvsBruteForceSearch( + res.handle(), self.inner, queries.to_c().as_mut_ptr(), neighbors.to_c().as_mut_ptr(), distances.to_c().as_mut_ptr(), prefilter, - )) - } + ) + })?; + Ok(()) } } @@ -125,7 +131,7 @@ mod tests { fn test_bfknn(metric: DistanceType) { let res = Resources::new().unwrap(); - // Create a new random dataset to index + // Create a new random dataset to index. let n_datapoints = 16; let n_features = 8; let dataset_host = ndarray::Array::::random( @@ -135,22 +141,17 @@ mod tests { let dataset = DeviceTensor::from_host(&res, &dataset_host).unwrap(); - println!("dataset {:#?}", dataset_host); - - // build the brute force index let index = - Index::build(&res, metric, None, &dataset).expect("failed to create brute force index"); + Index::build(&res, metric, &dataset).expect("failed to create brute force index"); res.sync_stream().unwrap(); - // use the first 4 points from the dataset as queries : will test that we get them back - // as their own nearest neighbor + // Use the first 4 points from the dataset as queries: each should get + // itself back as its own nearest neighbor. let n_queries = 4; let queries = dataset_host.slice(s![0..n_queries, ..]).to_owned(); - let k = 4; - println!("queries! {:#?}", queries); let queries = DeviceTensor::from_host(&res, &queries).unwrap(); let mut neighbors_host = ndarray::Array::::zeros((n_queries, k)); let mut neighbors = DeviceTensor::::zeros(&res, &[n_queries, k]).unwrap(); @@ -160,29 +161,16 @@ mod tests { index.search(&res, &queries, &mut neighbors, &mut distances).unwrap(); - // Copy back to host memory distances.copy_to_host(&res, &mut distances_host).unwrap(); neighbors.copy_to_host(&res, &mut neighbors_host).unwrap(); res.sync_stream().unwrap(); - println!("distances {:#?}", distances_host); - println!("neighbors {:#?}", neighbors_host); - - // nearest neighbors should be themselves, since queries are from the - // dataset assert_eq!(neighbors_host[[0, 0]], 0); assert_eq!(neighbors_host[[1, 0]], 1); assert_eq!(neighbors_host[[2, 0]], 2); assert_eq!(neighbors_host[[3, 0]], 3); } - /* - #[test] - fn test_cosine() { - test_bfknn(DistanceType::CosineExpanded); - } - */ - #[test] fn test_l2() { test_bfknn(DistanceType::L2Expanded); diff --git a/rust/cuvs/src/cagra/index.rs b/rust/cuvs/src/neighbors/cagra/index.rs similarity index 86% rename from rust/cuvs/src/cagra/index.rs rename to rust/cuvs/src/neighbors/cagra/index.rs index a9d39d353f..c9d47db89c 100644 --- a/rust/cuvs/src/cagra/index.rs +++ b/rust/cuvs/src/neighbors/cagra/index.rs @@ -8,11 +8,13 @@ use std::io::{Write, stderr}; use std::marker::PhantomData; use std::path::Path; -use crate::cagra::{IndexParams, SearchParams}; +use super::{CagraError, IndexParams, SearchParams}; use crate::dlpack::{AsDlTensor, AsDlTensorMut}; -use crate::error::{Error, Result, check_cuvs}; +use crate::error::check_cuvs; use crate::resources::Resources; +type Result = std::result::Result; + /// A CAGRA approximate nearest neighbor index. /// /// The lifetime `'d` ties this index to the underlying dataset, @@ -27,14 +29,9 @@ pub struct Index<'d> { } /// Convert a filesystem path into a `CString` suitable for the cuVS C API, -/// returning `Error::InvalidArgument` instead of panicking for paths that are -/// not valid UTF-8 or that contain an interior NUL byte. +/// returning [`CagraError::InvalidPath`] for a path with an interior NUL byte. fn path_to_cstring(path: &Path) -> Result { - let path_str = path - .to_str() - .ok_or_else(|| Error::InvalidArgument(format!("path is not valid UTF-8: {path:?}")))?; - CString::new(path_str) - .map_err(|e| Error::InvalidArgument(format!("path contains an interior NUL byte: {e}"))) + Ok(CString::new(path.as_os_str().as_encoded_bytes())?) } impl<'d> Index<'d> { @@ -52,8 +49,8 @@ impl<'d> Index<'d> { let index = Index::new()?; unsafe { check_cuvs(ffi::cuvsCagraBuild( - res.0, - params.0, + res.handle(), + params.handle(), dataset.to_c().as_mut_ptr(), index.handle, ))?; @@ -93,19 +90,19 @@ impl<'d> Index<'d> { let queries = queries.as_dl_tensor()?; let neighbors = neighbors.as_dl_tensor_mut()?; let distances = distances.as_dl_tensor_mut()?; - unsafe { - let prefilter = ffi::cuvsFilter { addr: 0, type_: ffi::cuvsFilterType::NO_FILTER }; - - check_cuvs(ffi::cuvsCagraSearch( - res.0, - params.0, + let prefilter = ffi::cuvsFilter { addr: 0, type_: ffi::cuvsFilterType::NO_FILTER }; + check_cuvs(unsafe { + ffi::cuvsCagraSearch( + res.handle(), + params.handle(), self.handle, queries.to_c().as_mut_ptr(), neighbors.to_c().as_mut_ptr(), distances.to_c().as_mut_ptr(), prefilter, - )) - } + ) + })?; + Ok(()) } /// Perform a filtered Approximate Nearest Neighbors search on the Index @@ -140,22 +137,22 @@ impl<'d> Index<'d> { // by the search call, so its `ManagedTensorRef` must outlive both. // Hence we keep it bound instead of chaining `to_c().as_mut_ptr()`. let mut bitset_c = bitset.to_c(); - unsafe { - let prefilter = ffi::cuvsFilter { - addr: bitset_c.as_mut_ptr() as usize, - type_: ffi::cuvsFilterType::BITSET, - }; - - check_cuvs(ffi::cuvsCagraSearch( - res.0, - params.0, + let prefilter = ffi::cuvsFilter { + addr: bitset_c.as_mut_ptr() as usize, + type_: ffi::cuvsFilterType::BITSET, + }; + check_cuvs(unsafe { + ffi::cuvsCagraSearch( + res.handle(), + params.handle(), self.handle, queries.to_c().as_mut_ptr(), neighbors.to_c().as_mut_ptr(), distances.to_c().as_mut_ptr(), prefilter, - )) - } + ) + })?; + Ok(()) } /// Save the CAGRA index to file. @@ -170,14 +167,14 @@ impl<'d> Index<'d> { /// /// # Example: /// ```no_run - /// use cuvs::cagra::{Index, IndexParams}; - /// use cuvs::{Resources, Result}; + /// use cuvs::Resources; + /// use cuvs::neighbors::cagra::{Index, IndexParams}; /// - /// fn serialize_example() -> Result<()> { + /// fn serialize_example() -> Result<(), Box> { /// let res = Resources::new()?; /// /// // Build an index (using some dataset) - /// let build_params = IndexParams::new()?; + /// let build_params = IndexParams::builder().build()?; /// // let index = Index::build(&res, &build_params, &dataset)?; /// /// // Save the index to disk (including the dataset) @@ -197,14 +194,10 @@ impl<'d> Index<'d> { include_dataset: bool, ) -> Result<()> { let c_filename = path_to_cstring(filename.as_ref())?; - unsafe { - check_cuvs(ffi::cuvsCagraSerialize( - res.0, - c_filename.as_ptr(), - self.handle, - include_dataset, - )) - } + check_cuvs(unsafe { + ffi::cuvsCagraSerialize(res.handle(), c_filename.as_ptr(), self.handle, include_dataset) + })?; + Ok(()) } /// Save the CAGRA index to file in hnswlib format. @@ -220,9 +213,10 @@ impl<'d> Index<'d> { /// * `filename` - The file path for saving the index pub fn serialize_to_hnswlib>(&self, res: &Resources, filename: P) -> Result<()> { let c_filename = path_to_cstring(filename.as_ref())?; - unsafe { - check_cuvs(ffi::cuvsCagraSerializeToHnswlib(res.0, c_filename.as_ptr(), self.handle)) - } + check_cuvs(unsafe { + ffi::cuvsCagraSerializeToHnswlib(res.handle(), c_filename.as_ptr(), self.handle) + })?; + Ok(()) } /// Load a CAGRA index from file. @@ -237,7 +231,7 @@ impl<'d> Index<'d> { let c_filename = path_to_cstring(filename.as_ref())?; let index = Index::new()?; unsafe { - check_cuvs(ffi::cuvsCagraDeserialize(res.0, c_filename.as_ptr(), index.handle))?; + check_cuvs(ffi::cuvsCagraDeserialize(res.handle(), c_filename.as_ptr(), index.handle))?; } Ok(index) } @@ -282,7 +276,7 @@ mod tests { let mut distances_host = ndarray::Array::::zeros((n_queries, k)); let mut distances = DeviceTensor::::zeros(res, &[n_queries, k]).unwrap(); - let search_params = SearchParams::new().unwrap(); + let search_params = SearchParams::try_new().unwrap(); index .search(res, &search_params, &queries, &mut neighbors, &mut distances) .expect("search failed"); @@ -312,15 +306,17 @@ mod tests { #[test] fn test_cagra_index() { - let build_params = IndexParams::new().unwrap(); + let build_params = IndexParams::try_new().unwrap(); test_cagra(build_params); } #[test] fn test_cagra_compression() { - use crate::cagra::CompressionParams; - let build_params = - IndexParams::new().unwrap().set_compression(CompressionParams::new().unwrap()); + use crate::neighbors::cagra::CompressionParams; + let build_params = IndexParams::builder() + .compression(CompressionParams::builder().build().unwrap()) + .build() + .unwrap(); test_cagra(build_params); } @@ -328,7 +324,7 @@ mod tests { #[test] fn test_cagra_search_with_filter() { let res = Resources::new().unwrap(); - let build_params = IndexParams::new().unwrap(); + let build_params = IndexParams::try_new().unwrap(); let n_datapoints = 256; let n_features = 16; @@ -341,7 +337,7 @@ mod tests { Index::build(&res, &build_params, &*dataset).expect("failed to create cagra index"); // Build a bitset that includes only even-indexed rows - let n_words = (n_datapoints + 31) / 32; + let n_words = n_datapoints.div_ceil(32); let mut bitset_host = ndarray::Array::::zeros(ndarray::Ix1(n_words)); for i in 0..n_datapoints { if i % 2 == 0 { @@ -360,7 +356,7 @@ mod tests { let mut neighbors = DeviceTensor::::zeros(&res, &[n_queries, k]).unwrap(); let mut distances = DeviceTensor::::zeros(&res, &[n_queries, k]).unwrap(); - let search_params = SearchParams::new().unwrap(); + let search_params = SearchParams::try_new().unwrap(); index .search_with_filter( @@ -396,7 +392,7 @@ mod tests { #[test] fn test_cagra_multiple_searches() { let res = Resources::new().unwrap(); - let build_params = IndexParams::new().unwrap(); + let build_params = IndexParams::try_new().unwrap(); let dataset = ndarray::Array::::random( (N_DATAPOINTS, N_FEATURES), Uniform::new(0., 1.0).unwrap(), @@ -412,7 +408,7 @@ mod tests { #[test] fn test_cagra_serialize_deserialize() { let res = Resources::new().unwrap(); - let build_params = IndexParams::new().unwrap(); + let build_params = IndexParams::try_new().unwrap(); let dataset = ndarray::Array::::random( (N_DATAPOINTS, N_FEATURES), Uniform::new(0., 1.0).unwrap(), @@ -442,7 +438,7 @@ mod tests { #[test] fn test_cagra_serialize_without_dataset() { let res = Resources::new().unwrap(); - let build_params = IndexParams::new().unwrap(); + let build_params = IndexParams::try_new().unwrap(); let dataset = ndarray::Array::::random( (N_DATAPOINTS, N_FEATURES), Uniform::new(0., 1.0).unwrap(), @@ -463,7 +459,7 @@ mod tests { #[test] fn test_cagra_serialize_to_hnswlib() { let res = Resources::new().unwrap(); - let build_params = IndexParams::new().unwrap(); + let build_params = IndexParams::try_new().unwrap(); let dataset = ndarray::Array::::random( (N_DATAPOINTS, N_FEATURES), Uniform::new(0., 1.0).unwrap(), @@ -486,11 +482,11 @@ mod tests { } /// Passing a filename containing an interior NUL byte must surface as an - /// `InvalidArgument` error rather than panicking inside the serializer. + /// `InvalidPath` error rather than panicking inside the serializer. #[test] fn test_cagra_serialize_rejects_interior_nul() { let res = Resources::new().unwrap(); - let build_params = IndexParams::new().unwrap(); + let build_params = IndexParams::try_new().unwrap(); let dataset = ndarray::Array::::random( (N_DATAPOINTS, N_FEATURES), Uniform::new(0., 1.0).unwrap(), @@ -504,6 +500,6 @@ mod tests { let err = index .serialize(&res, &bad_path, true) .expect_err("serialize should reject paths with interior NUL"); - assert!(matches!(err, Error::InvalidArgument(_)), "expected InvalidArgument, got {err:?}"); + assert!(matches!(err, CagraError::InvalidPath(_)), "expected InvalidPath, got {err:?}"); } } diff --git a/rust/cuvs/src/neighbors/cagra/mod.rs b/rust/cuvs/src/neighbors/cagra/mod.rs new file mode 100644 index 0000000000..0e445be90f --- /dev/null +++ b/rust/cuvs/src/neighbors/cagra/mod.rs @@ -0,0 +1,152 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! CAGRA: a graph-based approximate nearest neighbors algorithm with +//! state-of-the-art query throughput for both small and large batch sizes. +//! +//! Build an [`Index`] from a dataset, then [`search`](Index::search) it with +//! device-resident queries and output buffers. Tensors are passed through the +//! `AsDlTensor` / `AsDlTensorMut` traits; see the [`dlpack`](crate::dlpack) +//! module for the tensor model and `examples/cagra.rs` for a complete, runnable +//! example. +//! +//! Parameter types ([`IndexParams`], [`SearchParams`], ...) use the [`bon`] +//! builder pattern: every setter is optional and unset values keep the cuVS C +//! library defaults. Values are validated when the builder's `build()` runs, +//! returning [`CagraError::Validation`] for out-of-range inputs. + +mod index; +mod params; + +pub use index::Index; +pub use params::{CompressionParams, IndexParams, SearchParams}; + +use crate::dlpack::DLPackError; +use crate::error::LibraryError; + +/// Algorithm for building the internal k-NN graph. +#[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)] +#[non_exhaustive] +pub enum GraphBuildAlgo { + /// Automatically select the best algorithm. + Auto, + /// Build using IVF-PQ. + IvfPq, + /// Build using NN-Descent. + NnDescent, + /// Build using iterative CAGRA search. + IterativeCagraSearch, + /// Build using ACE (Augmented Core Extraction) for large datasets. + Ace, +} + +impl From for ffi::cuvsCagraGraphBuildAlgo { + fn from(v: GraphBuildAlgo) -> Self { + match v { + GraphBuildAlgo::Auto => Self::AUTO_SELECT, + GraphBuildAlgo::IvfPq => Self::IVF_PQ, + GraphBuildAlgo::NnDescent => Self::NN_DESCENT, + GraphBuildAlgo::IterativeCagraSearch => Self::ITERATIVE_CAGRA_SEARCH, + GraphBuildAlgo::Ace => Self::ACE, + } + } +} + +impl From for GraphBuildAlgo { + fn from(v: ffi::cuvsCagraGraphBuildAlgo) -> Self { + match v { + ffi::cuvsCagraGraphBuildAlgo::AUTO_SELECT => Self::Auto, + ffi::cuvsCagraGraphBuildAlgo::IVF_PQ => Self::IvfPq, + ffi::cuvsCagraGraphBuildAlgo::NN_DESCENT => Self::NnDescent, + ffi::cuvsCagraGraphBuildAlgo::ITERATIVE_CAGRA_SEARCH => Self::IterativeCagraSearch, + ffi::cuvsCagraGraphBuildAlgo::ACE => Self::Ace, + } + } +} + +/// Search kernel implementation. +#[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)] +#[non_exhaustive] +pub enum SearchAlgo { + /// Single CTA -- best for large batch sizes. + SingleCta, + /// Multi CTA -- best for small batch sizes. + MultiCta, + /// Multi kernel -- best for small batch sizes. + MultiKernel, + /// Automatically select the best kernel. + Auto, +} + +impl From for ffi::cuvsCagraSearchAlgo { + fn from(v: SearchAlgo) -> Self { + match v { + SearchAlgo::SingleCta => Self::SINGLE_CTA, + SearchAlgo::MultiCta => Self::MULTI_CTA, + SearchAlgo::MultiKernel => Self::MULTI_KERNEL, + SearchAlgo::Auto => Self::AUTO, + } + } +} + +impl From for SearchAlgo { + fn from(v: ffi::cuvsCagraSearchAlgo) -> Self { + match v { + ffi::cuvsCagraSearchAlgo::SINGLE_CTA => Self::SingleCta, + ffi::cuvsCagraSearchAlgo::MULTI_CTA => Self::MultiCta, + ffi::cuvsCagraSearchAlgo::MULTI_KERNEL => Self::MultiKernel, + ffi::cuvsCagraSearchAlgo::AUTO => Self::Auto, + } + } +} + +/// Hash-table mode used during search. +#[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)] +#[non_exhaustive] +pub enum HashMode { + /// Standard hash table. + Hash, + /// Small hash table optimised for low memory. + Small, + /// Automatically select the best mode. + Auto, +} + +impl From for ffi::cuvsCagraHashMode { + fn from(v: HashMode) -> Self { + match v { + HashMode::Hash => Self::HASH, + HashMode::Small => Self::SMALL, + HashMode::Auto => Self::AUTO_HASH, + } + } +} + +impl From for HashMode { + fn from(v: ffi::cuvsCagraHashMode) -> Self { + match v { + ffi::cuvsCagraHashMode::HASH => Self::Hash, + ffi::cuvsCagraHashMode::SMALL => Self::Small, + ffi::cuvsCagraHashMode::AUTO_HASH => Self::Auto, + } + } +} + +/// Error type for CAGRA operations. +#[derive(Debug, thiserror::Error)] +pub enum CagraError { + /// The cuVS C library reported a failure. + #[error(transparent)] + Library(#[from] LibraryError), + /// Tensor conversion into DLPack metadata failed. + #[error(transparent)] + DLPack(#[from] DLPackError), + /// A file path contained an interior NUL byte. + #[error("path contains an interior NUL byte")] + InvalidPath(#[from] std::ffi::NulError), + /// A parameter value failed validation. + #[error("invalid parameter: {0}")] + Validation(String), +} diff --git a/rust/cuvs/src/neighbors/cagra/params.rs b/rust/cuvs/src/neighbors/cagra/params.rs new file mode 100644 index 0000000000..03df19aab1 --- /dev/null +++ b/rust/cuvs/src/neighbors/cagra/params.rs @@ -0,0 +1,497 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Builder-pattern parameter types for CAGRA index build and search. +//! +//! Each parameter type owns its C params handle directly. The generated `bon` +//! builder configures that handle in the constructor, so there is no duplicate +//! Rust field-bag to keep in sync with the FFI state. All setters are optional; +//! unset values retain the library defaults from the underlying C +//! `*ParamsCreate` functions. Out-of-range values are rejected by `build()` with +//! [`CagraError::Validation`]. + +use std::{fmt, ptr}; + +use bon::bon; + +use crate::distance::DistanceType; +use crate::error::check_cuvs; + +use super::{CagraError, GraphBuildAlgo, HashMode, SearchAlgo}; + +// --------------------------------------------------------------------------- +// CompressionParams +// --------------------------------------------------------------------------- + +/// VPQ (Vector-Product Quantization) compression parameters. +/// +/// Attach to [`IndexParams`] to enable compressed dataset storage. +pub struct CompressionParams { + handle: ffi::cuvsCagraCompressionParams_t, +} + +#[bon] +impl CompressionParams { + #[builder] + pub fn new( + pq_bits: Option, + pq_dim: Option, + vq_n_centers: Option, + kmeans_n_iters: Option, + vq_kmeans_trainset_fraction: Option, + pq_kmeans_trainset_fraction: Option, + ) -> Result { + if let Some(bits) = pq_bits + && !(4..=16).contains(&bits) + { + return Err(CagraError::Validation(format!( + "pq_bits must be within [4, 16], got {bits}" + ))); + } + + let params = Self::try_new()?; + unsafe { + if let Some(v) = pq_bits { + (*params.handle).pq_bits = v; + } + if let Some(v) = pq_dim { + (*params.handle).pq_dim = v; + } + if let Some(v) = vq_n_centers { + (*params.handle).vq_n_centers = v; + } + if let Some(v) = kmeans_n_iters { + (*params.handle).kmeans_n_iters = v; + } + if let Some(v) = vq_kmeans_trainset_fraction { + (*params.handle).vq_kmeans_trainset_fraction = v; + } + if let Some(v) = pq_kmeans_trainset_fraction { + (*params.handle).pq_kmeans_trainset_fraction = v; + } + } + + Ok(params) + } +} + +impl CompressionParams { + /// Allocate parameters populated with the library defaults. + pub fn try_new() -> Result { + let mut handle = ptr::null_mut(); + check_cuvs(unsafe { ffi::cuvsCagraCompressionParamsCreate(&mut handle) })?; + Ok(Self { handle }) + } + + fn handle(&self) -> ffi::cuvsCagraCompressionParams_t { + self.handle + } +} + +impl fmt::Debug for CompressionParams { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_tuple("CompressionParams").field(unsafe { &*self.handle }).finish() + } +} + +impl Drop for CompressionParams { + fn drop(&mut self) { + let _ = unsafe { ffi::cuvsCagraCompressionParamsDestroy(self.handle) }; + } +} + +// --------------------------------------------------------------------------- +// IndexParams +// --------------------------------------------------------------------------- + +/// Parameters for building a CAGRA index. +/// +/// ```ignore +/// use cuvs::neighbors::cagra::IndexParams; +/// use cuvs::distance::DistanceType; +/// +/// let params = IndexParams::builder() +/// .metric(DistanceType::InnerProduct) +/// .graph_degree(64) +/// .build()?; +/// ``` +pub struct IndexParams { + handle: ffi::cuvsCagraIndexParams_t, + // Keep the compression params alive for as long as the index params + // reference them through `(*handle).compression`. + _compression: Option, +} + +#[bon] +impl IndexParams { + #[builder] + pub fn new( + metric: Option, + intermediate_graph_degree: Option, + graph_degree: Option, + build_algo: Option, + nn_descent_niter: Option, + compression: Option, + ) -> Result { + if let Some(d) = graph_degree + && d == 0 + { + return Err(CagraError::Validation("graph_degree must be > 0".into())); + } + + if let (Some(inter), Some(graph)) = (intermediate_graph_degree, graph_degree) + && inter < graph + { + return Err(CagraError::Validation(format!( + "intermediate_graph_degree ({inter}) must be >= graph_degree ({graph})" + ))); + } + + if let Some(0) = nn_descent_niter { + return Err(CagraError::Validation("nn_descent_niter must be > 0".into())); + } + + let metric_supports_compression = metric.is_none_or(|v| v == DistanceType::L2Expanded); + if compression.is_some() && !metric_supports_compression { + return Err(CagraError::Validation( + "VPQ compression is only supported with L2Expanded distance metric".into(), + )); + } + + let mut params = Self::try_new()?; + + unsafe { + if let Some(v) = metric { + (*params.handle).metric = v.into(); + } + if let Some(v) = intermediate_graph_degree { + (*params.handle).intermediate_graph_degree = v; + } + if let Some(v) = graph_degree { + (*params.handle).graph_degree = v; + } + if let Some(v) = build_algo { + (*params.handle).build_algo = v.into(); + } + if let Some(v) = nn_descent_niter { + (*params.handle).nn_descent_niter = v; + } + } + + if let Some(compression) = compression { + unsafe { (*params.handle).compression = compression.handle() }; + params._compression = Some(compression); + } + + Ok(params) + } +} + +impl IndexParams { + /// Allocate parameters populated with the library defaults. + pub fn try_new() -> Result { + let mut handle = ptr::null_mut(); + check_cuvs(unsafe { ffi::cuvsCagraIndexParamsCreate(&mut handle) })?; + Ok(Self { handle, _compression: None }) + } + + pub(super) fn handle(&self) -> ffi::cuvsCagraIndexParams_t { + self.handle + } +} + +impl fmt::Debug for IndexParams { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_tuple("IndexParams").field(unsafe { &*self.handle }).finish() + } +} + +impl Drop for IndexParams { + fn drop(&mut self) { + let _ = unsafe { ffi::cuvsCagraIndexParamsDestroy(self.handle) }; + } +} + +// --------------------------------------------------------------------------- +// SearchParams +// --------------------------------------------------------------------------- + +/// Parameters for searching a CAGRA index. +/// +/// ```ignore +/// use cuvs::neighbors::cagra::SearchParams; +/// +/// let params = SearchParams::builder().itopk_size(128).build()?; +/// ``` +pub struct SearchParams { + handle: ffi::cuvsCagraSearchParams_t, +} + +#[bon] +impl SearchParams { + #[builder] + #[allow(clippy::too_many_arguments)] + pub fn new( + max_queries: Option, + itopk_size: Option, + max_iterations: Option, + algo: Option, + team_size: Option, + min_iterations: Option, + thread_block_size: Option, + hashmap_mode: Option, + hashmap_min_bitlen: Option, + hashmap_max_fill_rate: Option, + num_random_samplings: Option, + rand_xor_mask: Option, + ) -> Result { + let params = Self::try_new()?; + + let effective_algo = algo.unwrap_or(unsafe { (*params.handle).algo.into() }); + let effective_hashmap_mode = + hashmap_mode.unwrap_or(unsafe { (*params.handle).hashmap_mode.into() }); + + if let Some(n) = itopk_size + && effective_algo == SearchAlgo::SingleCta + && n > 512 + { + return Err(CagraError::Validation(format!( + "itopk_size cannot be larger than 512 for SingleCta, got {n}" + ))); + } + + if let Some(n) = team_size + && !matches!(n, 0 | 8 | 16 | 32) + { + return Err(CagraError::Validation(format!( + "team_size must be 0 (auto), 8, 16, or 32, got {n}" + ))); + } + + if let Some(n) = thread_block_size + && !matches!(n, 0 | 64 | 128 | 256 | 512 | 1024) + { + return Err(CagraError::Validation(format!( + "thread_block_size must be 0, 64, 128, 256, 512, or 1024, got {n}" + ))); + } + + if let Some(bitlen) = hashmap_min_bitlen + && bitlen > 20 + { + return Err(CagraError::Validation(format!( + "hashmap_min_bitlen must be <= 20, got {bitlen}" + ))); + } + + if let Some(rate) = hashmap_max_fill_rate + && !(0.1..0.9).contains(&rate) + { + return Err(CagraError::Validation(format!( + "hashmap_max_fill_rate must be in [0.1, 0.9), got {rate}" + ))); + } + + if effective_algo == SearchAlgo::MultiCta && effective_hashmap_mode == HashMode::Small { + return Err(CagraError::Validation( + "`small_hash` is not available when 'search_mode' is \"multi-cta\"".into(), + )); + } + + unsafe { + if let Some(v) = max_queries { + (*params.handle).max_queries = v; + } + if let Some(v) = itopk_size { + (*params.handle).itopk_size = v; + } + if let Some(v) = max_iterations { + (*params.handle).max_iterations = v; + } + if let Some(v) = algo { + (*params.handle).algo = v.into(); + } + if let Some(v) = team_size { + (*params.handle).team_size = v; + } + if let Some(v) = min_iterations { + (*params.handle).min_iterations = v; + } + if let Some(v) = thread_block_size { + (*params.handle).thread_block_size = v; + } + if let Some(v) = hashmap_mode { + (*params.handle).hashmap_mode = v.into(); + } + if let Some(v) = hashmap_min_bitlen { + (*params.handle).hashmap_min_bitlen = v; + } + if let Some(v) = hashmap_max_fill_rate { + (*params.handle).hashmap_max_fill_rate = v; + } + if let Some(v) = num_random_samplings { + (*params.handle).num_random_samplings = v; + } + if let Some(v) = rand_xor_mask { + (*params.handle).rand_xor_mask = v; + } + } + + Ok(params) + } +} + +impl SearchParams { + /// Allocate parameters populated with the library defaults. + pub fn try_new() -> Result { + let mut handle = ptr::null_mut(); + check_cuvs(unsafe { ffi::cuvsCagraSearchParamsCreate(&mut handle) })?; + Ok(Self { handle }) + } + + pub(super) fn handle(&self) -> ffi::cuvsCagraSearchParams_t { + self.handle + } +} + +impl fmt::Debug for SearchParams { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_tuple("SearchParams").field(unsafe { &*self.handle }).finish() + } +} + +impl Drop for SearchParams { + fn drop(&mut self) { + let _ = unsafe { ffi::cuvsCagraSearchParamsDestroy(self.handle) }; + } +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn index_params_all_defaults() { + let params = IndexParams::try_new().unwrap(); + unsafe { + assert_eq!((*params.handle).metric, ffi::cuvsDistanceType::L2Expanded); + assert_eq!((*params.handle).graph_degree, 64); + } + } + + #[test] + fn index_params_with_values() { + let params = IndexParams::builder() + .metric(DistanceType::InnerProduct) + .graph_degree(64) + .intermediate_graph_degree(128) + .build_algo(GraphBuildAlgo::NnDescent) + .nn_descent_niter(10) + .build() + .unwrap(); + + unsafe { + assert_eq!((*params.handle).metric, ffi::cuvsDistanceType::InnerProduct); + assert_eq!((*params.handle).graph_degree, 64); + assert_eq!((*params.handle).intermediate_graph_degree, 128); + assert_eq!((*params.handle).build_algo, ffi::cuvsCagraGraphBuildAlgo::NN_DESCENT); + assert_eq!((*params.handle).nn_descent_niter, 10); + } + } + + #[test] + fn index_params_rejects_zero_graph_degree() { + let err = IndexParams::builder().graph_degree(0).build().unwrap_err(); + assert!(err.to_string().contains("graph_degree must be > 0")); + } + + #[test] + fn index_params_rejects_invalid_intermediate_degree() { + let err = IndexParams::builder() + .graph_degree(64) + .intermediate_graph_degree(32) + .build() + .unwrap_err(); + assert!( + err.to_string().contains("intermediate_graph_degree (32) must be >= graph_degree (64)") + ); + } + + #[test] + fn index_params_rejects_zero_niter() { + let err = IndexParams::builder().nn_descent_niter(0).build().unwrap_err(); + assert!(err.to_string().contains("nn_descent_niter must be > 0")); + } + + #[test] + fn index_params_rejects_non_l2_metric_with_compression() { + let compression = CompressionParams::builder().pq_bits(8).build().unwrap(); + let err = IndexParams::builder() + .metric(DistanceType::InnerProduct) + .compression(compression) + .build() + .unwrap_err(); + assert!(err.to_string().contains("VPQ compression is only supported with L2Expanded")); + } + + #[test] + fn index_params_with_compression() { + let params = IndexParams::builder() + .compression(CompressionParams::builder().pq_bits(4).pq_dim(8).build().unwrap()) + .build() + .unwrap(); + unsafe { + let c = (*params.handle).compression; + assert!(!c.is_null()); + assert_eq!((*c).pq_bits, 4); + assert_eq!((*c).pq_dim, 8); + } + } + + #[test] + fn compression_params_rejects_pq_bits_below_range() { + let err = CompressionParams::builder().pq_bits(3).build().unwrap_err(); + assert!(err.to_string().contains("pq_bits")); + } + + #[test] + fn search_params_all_defaults() { + let params = SearchParams::try_new().unwrap(); + unsafe { + assert_eq!((*params.handle).itopk_size, 64); + assert_eq!((*params.handle).algo, ffi::cuvsCagraSearchAlgo::SINGLE_CTA); + } + } + + #[test] + fn search_params_rejects_invalid_team_size() { + let err = SearchParams::builder().team_size(4).build().unwrap_err(); + assert!(err.to_string().contains("team_size must be")); + } + + #[test] + fn search_params_rejects_single_cta_itopk_above_limit() { + let err = SearchParams::builder() + .algo(SearchAlgo::SingleCta) + .itopk_size(513) + .build() + .unwrap_err(); + assert!(err.to_string().contains("512")); + } + + #[test] + fn search_params_rejects_small_hash_with_multi_cta() { + let err = SearchParams::builder() + .algo(SearchAlgo::MultiCta) + .hashmap_mode(HashMode::Small) + .build() + .unwrap_err(); + assert!(err.to_string().contains("small_hash")); + } +} diff --git a/rust/cuvs/src/ivf_flat/index.rs b/rust/cuvs/src/neighbors/ivf_flat/index.rs similarity index 69% rename from rust/cuvs/src/ivf_flat/index.rs rename to rust/cuvs/src/neighbors/ivf_flat/index.rs index 2aedf243f0..e03502c600 100644 --- a/rust/cuvs/src/ivf_flat/index.rs +++ b/rust/cuvs/src/neighbors/ivf_flat/index.rs @@ -5,12 +5,14 @@ use std::io::{Write, stderr}; +use super::{IndexParams, IvfFlatError, SearchParams}; use crate::dlpack::{AsDlTensor, AsDlTensorMut}; -use crate::error::{Result, check_cuvs}; -use crate::ivf_flat::{IndexParams, SearchParams}; +use crate::error::check_cuvs; use crate::resources::Resources; -/// Ivf-Flat ANN Index +type Result = std::result::Result; + +/// IVF-Flat ANN index. #[derive(Debug)] pub struct Index(ffi::cuvsIvfFlatIndex_t); @@ -18,9 +20,8 @@ impl Index { /// Builds an IVF-Flat index over `dataset` for efficient search. /// /// `dataset` is a row-major matrix on the host or device implementing - /// [`AsDlTensor`]. It is copied into the index, so the - /// caller may free it once this call returns (hence `Index` carries no - /// lifetime). + /// [`AsDlTensor`]. It is copied into the index, so the caller may free it + /// once this call returns (hence `Index` carries no lifetime). /// /// Supported dataset/query dtypes in the current C-backed implementation are /// `f32`, `f16`, `i8`, and `u8`. @@ -32,8 +33,8 @@ impl Index { let index = Index::new()?; unsafe { check_cuvs(ffi::cuvsIvfFlatBuild( - res.0, - params.0, + res.handle(), + params.handle(), dataset.to_c().as_mut_ptr(), index.0, ))?; @@ -41,7 +42,7 @@ impl Index { Ok(index) } - /// Creates a new empty index + /// Creates a new empty index. pub fn new() -> Result { unsafe { let mut index = std::mem::MaybeUninit::::uninit(); @@ -53,8 +54,7 @@ impl Index { /// Searches the index for the `k` nearest neighbors of each query. /// /// `queries`, `neighbors`, and `distances` must reside in device memory and - /// implement [`AsDlTensor`] / - /// [`AsDlTensorMut`]. `neighbors` receives the + /// implement [`AsDlTensor`] / [`AsDlTensorMut`]. `neighbors` receives the /// neighbor indices and `distances` their distances; both are written in /// place. pub fn search( @@ -73,19 +73,19 @@ impl Index { let queries = queries.as_dl_tensor()?; let neighbors = neighbors.as_dl_tensor_mut()?; let distances = distances.as_dl_tensor_mut()?; - unsafe { - let prefilter = ffi::cuvsFilter { addr: 0, type_: ffi::cuvsFilterType::NO_FILTER }; - - check_cuvs(ffi::cuvsIvfFlatSearch( - res.0, - params.0, + let prefilter = ffi::cuvsFilter { addr: 0, type_: ffi::cuvsFilterType::NO_FILTER }; + check_cuvs(unsafe { + ffi::cuvsIvfFlatSearch( + res.handle(), + params.handle(), self.0, queries.to_c().as_mut_ptr(), neighbors.to_c().as_mut_ptr(), distances.to_c().as_mut_ptr(), prefilter, - )) - } + ) + })?; + Ok(()) } } @@ -108,11 +108,9 @@ mod tests { #[test] fn test_ivf_flat() { - let build_params = IndexParams::new().unwrap().set_n_lists(64); - + let build_params = IndexParams::builder().n_lists(64).build().unwrap(); let res = Resources::new().unwrap(); - // Create a new random dataset to index let n_datapoints = 1024; let n_features = 16; let dataset = ndarray::Array::::random( @@ -122,20 +120,13 @@ mod tests { let dataset_device = DeviceTensor::from_host(&res, &dataset).unwrap(); - // build the ivf-flat index let index = Index::build(&res, &build_params, &dataset_device) .expect("failed to create ivf-flat index"); - // use the first 4 points from the dataset as queries : will test that we get them back - // as their own nearest neighbor let n_queries = 4; let queries = dataset.slice(s![0..n_queries, ..]).to_owned(); - let k = 10; - // IvfFlat search API requires queries and outputs to be on device memory - // copy query data over, and allocate new device memory for the distances/ neighbors - // outputs let queries = DeviceTensor::from_host(&res, &queries).unwrap(); let mut neighbors_host = ndarray::Array::::zeros((n_queries, k)); let mut neighbors = DeviceTensor::::zeros(&res, &[n_queries, k]).unwrap(); @@ -143,30 +134,26 @@ mod tests { let mut distances_host = ndarray::Array::::zeros((n_queries, k)); let mut distances = DeviceTensor::::zeros(&res, &[n_queries, k]).unwrap(); - let search_params = SearchParams::new().unwrap(); + let search_params = SearchParams::builder().build().unwrap(); index.search(&res, &search_params, &queries, &mut neighbors, &mut distances).unwrap(); - // Copy back to host memory distances.copy_to_host(&res, &mut distances_host).unwrap(); neighbors.copy_to_host(&res, &mut neighbors_host).unwrap(); - // nearest neighbors should be themselves, since queries are from the - // dataset assert_eq!(neighbors_host[[0, 0]], 0); assert_eq!(neighbors_host[[1, 0]], 1); assert_eq!(neighbors_host[[2, 0]], 2); assert_eq!(neighbors_host[[3, 0]], 3); } - /// Test that an index can be searched multiple times without rebuilding. - /// This validates that search() takes &self instead of self. + /// Searching the same index multiple times validates that `search` takes + /// `&self` rather than consuming the index. #[test] fn test_ivf_flat_multiple_searches() { - let build_params = IndexParams::new().unwrap().set_n_lists(64); + let build_params = IndexParams::builder().n_lists(64).build().unwrap(); let res = Resources::new().unwrap(); - // Create a random dataset let n_datapoints = 1024; let n_features = 16; let dataset = ndarray::Array::::random( @@ -175,42 +162,27 @@ mod tests { ); let dataset_device = DeviceTensor::from_host(&res, &dataset).unwrap(); - - // Build the index once let index = Index::build(&res, &build_params, &dataset_device) .expect("failed to create ivf-flat index"); - let search_params = SearchParams::new().unwrap(); + let search_params = SearchParams::builder().build().unwrap(); let k = 5; - // Perform multiple searches on the same index - for search_iter in 0..3 { + for _ in 0..3 { let n_queries = 4; let queries = dataset.slice(s![0..n_queries, ..]).to_owned(); let queries = DeviceTensor::from_host(&res, &queries).unwrap(); let mut neighbors_host = ndarray::Array::::zeros((n_queries, k)); let mut neighbors = DeviceTensor::::zeros(&res, &[n_queries, k]).unwrap(); - - let mut distances_host = ndarray::Array::::zeros((n_queries, k)); let mut distances = DeviceTensor::::zeros(&res, &[n_queries, k]).unwrap(); - // This should work on every iteration because search() takes &self index .search(&res, &search_params, &queries, &mut neighbors, &mut distances) - .expect(&format!("search iteration {} failed", search_iter)); + .expect("search failed"); - // Copy back to host memory - distances.copy_to_host(&res, &mut distances_host).unwrap(); neighbors.copy_to_host(&res, &mut neighbors_host).unwrap(); - - // Verify results are consistent - assert_eq!( - neighbors_host[[0, 0]], - 0, - "iteration {}: first query should find itself", - search_iter - ); + assert_eq!(neighbors_host[[0, 0]], 0, "first query should find itself"); } } } diff --git a/rust/cuvs/src/neighbors/ivf_flat/mod.rs b/rust/cuvs/src/neighbors/ivf_flat/mod.rs new file mode 100644 index 0000000000..9f467afb7b --- /dev/null +++ b/rust/cuvs/src/neighbors/ivf_flat/mod.rs @@ -0,0 +1,34 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! IVF-Flat: an inverted-file index over uncompressed ("flat") vectors. It +//! partitions the dataset into `n_lists` clusters and, at query time, scans only +//! the `n_probes` closest clusters — a simple knob to trade recall for speed. +//! +//! Build an [`Index`] from a dataset, then [`search`](Index::search) it with +//! device-resident queries and output buffers. Tensors are borrowed through the +//! `AsDlTensor` / `AsDlTensorMut` traits; see the [`dlpack`](crate::dlpack) +//! module for the tensor model and `examples/cagra.rs` for the same build/search +//! workflow. + +mod index; +mod params; + +pub use index::Index; +pub use params::{IndexParams, SearchParams}; + +use crate::dlpack::DLPackError; +use crate::error::LibraryError; + +/// Error type for IVF-Flat operations. +#[derive(Debug, thiserror::Error)] +pub enum IvfFlatError { + /// The cuVS C library reported a failure. + #[error(transparent)] + Library(#[from] LibraryError), + /// Tensor conversion into DLPack metadata failed. + #[error(transparent)] + DLPack(#[from] DLPackError), +} diff --git a/rust/cuvs/src/neighbors/ivf_flat/params.rs b/rust/cuvs/src/neighbors/ivf_flat/params.rs new file mode 100644 index 0000000000..27f805996e --- /dev/null +++ b/rust/cuvs/src/neighbors/ivf_flat/params.rs @@ -0,0 +1,147 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Builder-pattern parameter types for IVF-Flat build and search. +//! +//! All setters are optional; unset values retain the library defaults from the +//! underlying C `*ParamsCreate` functions. + +use std::{fmt, ptr}; + +use bon::bon; + +use crate::distance::DistanceType; +use crate::error::check_cuvs; + +use super::IvfFlatError; + +/// Parameters for building an IVF-Flat index. +pub struct IndexParams { + handle: ffi::cuvsIvfFlatIndexParams_t, +} + +#[bon] +impl IndexParams { + #[builder] + pub fn new( + n_lists: Option, + metric: Option, + kmeans_n_iters: Option, + kmeans_trainset_fraction: Option, + add_data_on_build: Option, + ) -> Result { + let params = Self::try_new()?; + unsafe { + if let Some(v) = n_lists { + (*params.handle).n_lists = v; + } + if let Some(v) = metric { + (*params.handle).metric = v.into(); + (*params.handle).metric_arg = v.metric_arg(); + } + if let Some(v) = kmeans_n_iters { + (*params.handle).kmeans_n_iters = v; + } + if let Some(v) = kmeans_trainset_fraction { + (*params.handle).kmeans_trainset_fraction = v; + } + if let Some(v) = add_data_on_build { + (*params.handle).add_data_on_build = v; + } + } + Ok(params) + } +} + +impl IndexParams { + /// Allocate parameters populated with the library defaults. + pub fn try_new() -> Result { + let mut handle = ptr::null_mut(); + check_cuvs(unsafe { ffi::cuvsIvfFlatIndexParamsCreate(&mut handle) })?; + Ok(Self { handle }) + } + + pub(super) fn handle(&self) -> ffi::cuvsIvfFlatIndexParams_t { + self.handle + } +} + +impl fmt::Debug for IndexParams { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_tuple("IndexParams").field(unsafe { &*self.handle }).finish() + } +} + +impl Drop for IndexParams { + fn drop(&mut self) { + let _ = unsafe { ffi::cuvsIvfFlatIndexParamsDestroy(self.handle) }; + } +} + +/// Parameters for searching an IVF-Flat index. +pub struct SearchParams { + handle: ffi::cuvsIvfFlatSearchParams_t, +} + +#[bon] +impl SearchParams { + #[builder] + pub fn new(n_probes: Option) -> Result { + let params = Self::try_new()?; + unsafe { + if let Some(v) = n_probes { + (*params.handle).n_probes = v; + } + } + Ok(params) + } +} + +impl SearchParams { + /// Allocate parameters populated with the library defaults. + pub fn try_new() -> Result { + let mut handle = ptr::null_mut(); + check_cuvs(unsafe { ffi::cuvsIvfFlatSearchParamsCreate(&mut handle) })?; + Ok(Self { handle }) + } + + pub(super) fn handle(&self) -> ffi::cuvsIvfFlatSearchParams_t { + self.handle + } +} + +impl fmt::Debug for SearchParams { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_tuple("SearchParams").field(unsafe { &*self.handle }).finish() + } +} + +impl Drop for SearchParams { + fn drop(&mut self) { + let _ = unsafe { ffi::cuvsIvfFlatSearchParamsDestroy(self.handle) }; + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn index_params_with_values() { + let params = IndexParams::builder().n_lists(128).add_data_on_build(false).build().unwrap(); + unsafe { + assert_eq!((*params.handle).n_lists, 128); + assert!(!(*params.handle).add_data_on_build); + } + } + + #[test] + fn search_params_with_values() { + let params = SearchParams::builder().n_probes(128).build().unwrap(); + unsafe { + assert_eq!((*params.handle).n_probes, 128); + } + } +} diff --git a/rust/cuvs/src/ivf_pq/index.rs b/rust/cuvs/src/neighbors/ivf_pq/index.rs similarity index 66% rename from rust/cuvs/src/ivf_pq/index.rs rename to rust/cuvs/src/neighbors/ivf_pq/index.rs index abd4e0d00d..889c266efd 100644 --- a/rust/cuvs/src/ivf_pq/index.rs +++ b/rust/cuvs/src/neighbors/ivf_pq/index.rs @@ -5,25 +5,23 @@ use std::io::{Write, stderr}; +use super::{IndexParams, IvfPqError, SearchParams}; use crate::dlpack::{AsDlTensor, AsDlTensorMut}; -use crate::error::{Result, check_cuvs}; -use crate::ivf_pq::{IndexParams, SearchParams}; +use crate::error::check_cuvs; use crate::resources::Resources; -/// Ivf-Pq ANN Index +type Result = std::result::Result; + +/// IVF-PQ ANN index. #[derive(Debug)] pub struct Index(ffi::cuvsIvfPqIndex_t); impl Index { - /// Builds an IVF-PQ index over `dataset` for efficient search. + /// Builds an IVF-PQ index over `dataset` for compressed, efficient search. /// /// `dataset` is a row-major matrix on the host or device implementing - /// [`AsDlTensor`]. It is copied (and quantized) into - /// the index, so the caller may free it once this call returns (hence - /// `Index` carries no lifetime). - /// - /// Supported dataset/query dtypes in the current C-backed implementation are - /// `f32`, `f16`, `i8`, and `u8`. + /// [`AsDlTensor`]. It is copied into the index, so the caller may free it + /// once this call returns (hence `Index` carries no lifetime). pub fn build(res: &Resources, params: &IndexParams, dataset: &T) -> Result where T: AsDlTensor + ?Sized, @@ -31,12 +29,17 @@ impl Index { let dataset = dataset.as_dl_tensor()?; let index = Index::new()?; unsafe { - check_cuvs(ffi::cuvsIvfPqBuild(res.0, params.0, dataset.to_c().as_mut_ptr(), index.0))?; + check_cuvs(ffi::cuvsIvfPqBuild( + res.handle(), + params.handle(), + dataset.to_c().as_mut_ptr(), + index.0, + ))?; } Ok(index) } - /// Creates a new empty index + /// Creates a new empty index. pub fn new() -> Result { unsafe { let mut index = std::mem::MaybeUninit::::uninit(); @@ -48,8 +51,7 @@ impl Index { /// Searches the index for the `k` nearest neighbors of each query. /// /// `queries`, `neighbors`, and `distances` must reside in device memory and - /// implement [`AsDlTensor`] / - /// [`AsDlTensorMut`]. `neighbors` receives the + /// implement [`AsDlTensor`] / [`AsDlTensorMut`]. `neighbors` receives the /// neighbor indices and `distances` their distances; both are written in /// place. pub fn search( @@ -68,16 +70,17 @@ impl Index { let queries = queries.as_dl_tensor()?; let neighbors = neighbors.as_dl_tensor_mut()?; let distances = distances.as_dl_tensor_mut()?; - unsafe { - check_cuvs(ffi::cuvsIvfPqSearch( - res.0, - params.0, + check_cuvs(unsafe { + ffi::cuvsIvfPqSearch( + res.handle(), + params.handle(), self.0, queries.to_c().as_mut_ptr(), neighbors.to_c().as_mut_ptr(), distances.to_c().as_mut_ptr(), - )) - } + ) + })?; + Ok(()) } } @@ -100,11 +103,9 @@ mod tests { #[test] fn test_ivf_pq() { - let build_params = IndexParams::new().unwrap().set_n_lists(64); - + let build_params = IndexParams::builder().n_lists(64).build().unwrap(); let res = Resources::new().unwrap(); - // Create a new random dataset to index let n_datapoints = 1024; let n_features = 16; let dataset = ndarray::Array::::random( @@ -114,20 +115,13 @@ mod tests { let dataset_device = DeviceTensor::from_host(&res, &dataset).unwrap(); - // build the ivf-pq index let index = Index::build(&res, &build_params, &dataset_device) .expect("failed to create ivf-pq index"); - // use the first 4 points from the dataset as queries : will test that we get them back - // as their own nearest neighbor let n_queries = 4; let queries = dataset.slice(s![0..n_queries, ..]).to_owned(); - let k = 10; - // Ivf-Pq search API requires queries and outputs to be on device memory - // copy query data over, and allocate new device memory for the distances/ neighbors - // outputs let queries = DeviceTensor::from_host(&res, &queries).unwrap(); let mut neighbors_host = ndarray::Array::::zeros((n_queries, k)); let mut neighbors = DeviceTensor::::zeros(&res, &[n_queries, k]).unwrap(); @@ -135,30 +129,26 @@ mod tests { let mut distances_host = ndarray::Array::::zeros((n_queries, k)); let mut distances = DeviceTensor::::zeros(&res, &[n_queries, k]).unwrap(); - let search_params = SearchParams::new().unwrap(); + let search_params = SearchParams::builder().build().unwrap(); index.search(&res, &search_params, &queries, &mut neighbors, &mut distances).unwrap(); - // Copy back to host memory distances.copy_to_host(&res, &mut distances_host).unwrap(); neighbors.copy_to_host(&res, &mut neighbors_host).unwrap(); - // nearest neighbors should be themselves, since queries are from the - // dataset assert_eq!(neighbors_host[[0, 0]], 0); assert_eq!(neighbors_host[[1, 0]], 1); assert_eq!(neighbors_host[[2, 0]], 2); assert_eq!(neighbors_host[[3, 0]], 3); } - /// Test that an index can be searched multiple times without rebuilding. - /// This validates that search() takes &self instead of self. + /// Searching the same index multiple times validates that `search` takes + /// `&self` rather than consuming the index. #[test] fn test_ivf_pq_multiple_searches() { - let build_params = IndexParams::new().unwrap().set_n_lists(64); + let build_params = IndexParams::builder().n_lists(64).build().unwrap(); let res = Resources::new().unwrap(); - // Create a random dataset let n_datapoints = 1024; let n_features = 16; let dataset = ndarray::Array::::random( @@ -167,42 +157,27 @@ mod tests { ); let dataset_device = DeviceTensor::from_host(&res, &dataset).unwrap(); - - // Build the index once let index = Index::build(&res, &build_params, &dataset_device) .expect("failed to create ivf-pq index"); - let search_params = SearchParams::new().unwrap(); + let search_params = SearchParams::builder().build().unwrap(); let k = 5; - // Perform multiple searches on the same index - for search_iter in 0..3 { + for _ in 0..3 { let n_queries = 4; let queries = dataset.slice(s![0..n_queries, ..]).to_owned(); let queries = DeviceTensor::from_host(&res, &queries).unwrap(); let mut neighbors_host = ndarray::Array::::zeros((n_queries, k)); let mut neighbors = DeviceTensor::::zeros(&res, &[n_queries, k]).unwrap(); - - let mut distances_host = ndarray::Array::::zeros((n_queries, k)); let mut distances = DeviceTensor::::zeros(&res, &[n_queries, k]).unwrap(); - // This should work on every iteration because search() takes &self index .search(&res, &search_params, &queries, &mut neighbors, &mut distances) - .expect(&format!("search iteration {} failed", search_iter)); + .expect("search failed"); - // Copy back to host memory - distances.copy_to_host(&res, &mut distances_host).unwrap(); neighbors.copy_to_host(&res, &mut neighbors_host).unwrap(); - - // Verify results are consistent - assert_eq!( - neighbors_host[[0, 0]], - 0, - "iteration {}: first query should find itself", - search_iter - ); + assert_eq!(neighbors_host[[0, 0]], 0, "first query should find itself"); } } } diff --git a/rust/cuvs/src/neighbors/ivf_pq/mod.rs b/rust/cuvs/src/neighbors/ivf_pq/mod.rs new file mode 100644 index 0000000000..abdf146bf3 --- /dev/null +++ b/rust/cuvs/src/neighbors/ivf_pq/mod.rs @@ -0,0 +1,36 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! IVF-PQ: an inverted-file index that product-quantizes the vectors. Like +//! IVF-Flat it partitions the dataset into `n_lists` clusters and scans the +//! `n_probes` closest at query time, but compresses each vector into `pq_dim` +//! codes of `pq_bits` bits — much smaller, slightly less accurate. +//! +//! Build an [`Index`] from a dataset, then [`search`](Index::search) it with +//! device-resident queries and output buffers. Tensors are borrowed through the +//! `AsDlTensor` / `AsDlTensorMut` traits; see the [`dlpack`](crate::dlpack) +//! module for the tensor model and `examples/cagra.rs` for the same build/search +//! workflow. + +mod index; +mod params; + +pub use index::Index; +pub use params::{ + IndexParams, SearchParams, cudaDataType_t, cuvsIvfPqCodebookGen, cuvsIvfPqListLayout, +}; + +use crate::dlpack::DLPackError; +use crate::error::LibraryError; + +/// Error type for IVF-PQ operations. +#[derive(Debug, thiserror::Error)] +pub enum IvfPqError { + /// The cuVS C library reported a failure. + #[error(transparent)] + Library(#[from] LibraryError), + /// Tensor conversion into DLPack metadata failed. + #[error(transparent)] + DLPack(#[from] DLPackError), +} diff --git a/rust/cuvs/src/neighbors/ivf_pq/params.rs b/rust/cuvs/src/neighbors/ivf_pq/params.rs new file mode 100644 index 0000000000..66a678a93f --- /dev/null +++ b/rust/cuvs/src/neighbors/ivf_pq/params.rs @@ -0,0 +1,184 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Builder-pattern parameter types for IVF-PQ build and search. +//! +//! All setters are optional; unset values retain the library defaults from the +//! underlying C `*ParamsCreate` functions. + +use std::{fmt, ptr}; + +use bon::bon; + +use crate::distance::DistanceType; +use crate::error::check_cuvs; + +use super::IvfPqError; + +pub use ffi::{cudaDataType_t, cuvsIvfPqCodebookGen, cuvsIvfPqListLayout}; + +/// Parameters for building an IVF-PQ index. +pub struct IndexParams { + handle: ffi::cuvsIvfPqIndexParams_t, +} + +#[bon] +impl IndexParams { + #[builder] + #[allow(clippy::too_many_arguments)] + pub fn new( + n_lists: Option, + metric: Option, + kmeans_n_iters: Option, + kmeans_trainset_fraction: Option, + pq_bits: Option, + pq_dim: Option, + codebook_kind: Option, + codes_layout: Option, + force_random_rotation: Option, + max_train_points_per_pq_code: Option, + add_data_on_build: Option, + ) -> Result { + let params = Self::try_new()?; + unsafe { + if let Some(v) = n_lists { + (*params.handle).n_lists = v; + } + if let Some(v) = metric { + (*params.handle).metric = v.into(); + (*params.handle).metric_arg = v.metric_arg(); + } + if let Some(v) = kmeans_n_iters { + (*params.handle).kmeans_n_iters = v; + } + if let Some(v) = kmeans_trainset_fraction { + (*params.handle).kmeans_trainset_fraction = v; + } + if let Some(v) = pq_bits { + (*params.handle).pq_bits = v; + } + if let Some(v) = pq_dim { + (*params.handle).pq_dim = v; + } + if let Some(v) = codebook_kind { + (*params.handle).codebook_kind = v; + } + if let Some(v) = codes_layout { + (*params.handle).codes_layout = v; + } + if let Some(v) = force_random_rotation { + (*params.handle).force_random_rotation = v; + } + if let Some(v) = max_train_points_per_pq_code { + (*params.handle).max_train_points_per_pq_code = v; + } + if let Some(v) = add_data_on_build { + (*params.handle).add_data_on_build = v; + } + } + Ok(params) + } +} + +impl IndexParams { + /// Allocate parameters populated with the library defaults. + pub fn try_new() -> Result { + let mut handle = ptr::null_mut(); + check_cuvs(unsafe { ffi::cuvsIvfPqIndexParamsCreate(&mut handle) })?; + Ok(Self { handle }) + } + + pub(super) fn handle(&self) -> ffi::cuvsIvfPqIndexParams_t { + self.handle + } +} + +impl fmt::Debug for IndexParams { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_tuple("IndexParams").field(unsafe { &*self.handle }).finish() + } +} + +impl Drop for IndexParams { + fn drop(&mut self) { + let _ = unsafe { ffi::cuvsIvfPqIndexParamsDestroy(self.handle) }; + } +} + +/// Parameters for searching an IVF-PQ index. +pub struct SearchParams { + handle: ffi::cuvsIvfPqSearchParams_t, +} + +#[bon] +impl SearchParams { + #[builder] + pub fn new( + n_probes: Option, + lut_dtype: Option, + internal_distance_dtype: Option, + ) -> Result { + let params = Self::try_new()?; + unsafe { + if let Some(v) = n_probes { + (*params.handle).n_probes = v; + } + if let Some(v) = lut_dtype { + (*params.handle).lut_dtype = v; + } + if let Some(v) = internal_distance_dtype { + (*params.handle).internal_distance_dtype = v; + } + } + Ok(params) + } +} + +impl SearchParams { + /// Allocate parameters populated with the library defaults. + pub fn try_new() -> Result { + let mut handle = ptr::null_mut(); + check_cuvs(unsafe { ffi::cuvsIvfPqSearchParamsCreate(&mut handle) })?; + Ok(Self { handle }) + } + + pub(super) fn handle(&self) -> ffi::cuvsIvfPqSearchParams_t { + self.handle + } +} + +impl fmt::Debug for SearchParams { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_tuple("SearchParams").field(unsafe { &*self.handle }).finish() + } +} + +impl Drop for SearchParams { + fn drop(&mut self) { + let _ = unsafe { ffi::cuvsIvfPqSearchParamsDestroy(self.handle) }; + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn index_params_with_values() { + let params = IndexParams::builder().n_lists(128).add_data_on_build(false).build().unwrap(); + unsafe { + assert_eq!((*params.handle).n_lists, 128); + assert!(!(*params.handle).add_data_on_build); + } + } + + #[test] + fn search_params_with_values() { + let params = SearchParams::builder().n_probes(128).build().unwrap(); + unsafe { + assert_eq!((*params.handle).n_probes, 128); + } + } +} diff --git a/rust/cuvs/src/neighbors/mod.rs b/rust/cuvs/src/neighbors/mod.rs new file mode 100644 index 0000000000..78fe534ebb --- /dev/null +++ b/rust/cuvs/src/neighbors/mod.rs @@ -0,0 +1,17 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Nearest neighbor search algorithms. +//! +//! Mirrors the C++ `cuvs::neighbors` namespace: each submodule wraps one index +//! type. Build an [`Index`](cagra::Index) from a dataset, then search it with +//! device-resident queries and output buffers; see the [`dlpack`](crate::dlpack) +//! module for the tensor model. + +pub mod brute_force; +pub mod cagra; +pub mod ivf_flat; +pub mod ivf_pq; +pub mod vamana; diff --git a/rust/cuvs/src/vamana/index.rs b/rust/cuvs/src/neighbors/vamana/index.rs similarity index 53% rename from rust/cuvs/src/vamana/index.rs rename to rust/cuvs/src/neighbors/vamana/index.rs index 41a18ef1b0..4fe4949e98 100644 --- a/rust/cuvs/src/vamana/index.rs +++ b/rust/cuvs/src/neighbors/vamana/index.rs @@ -6,39 +6,39 @@ use std::ffi::CString; use std::io::{Write, stderr}; +use super::{IndexParams, VamanaError}; use crate::dlpack::AsDlTensor; -use crate::error::{Result, check_cuvs}; +use crate::error::check_cuvs; use crate::resources::Resources; -use crate::vamana::IndexParams; -/// Vamana ANN Index +type Result = std::result::Result; + +/// Vamana ANN index. #[derive(Debug)] pub struct Index(ffi::cuvsVamanaIndex_t); impl Index { - /// Builds Vamana Index for efficient DiskANN search + /// Builds a Vamana index for efficient DiskANN search. /// - /// The build uses the Vamana insertion-based algorithm to create the graph. The algorithm - /// starts with an empty graph and iteratively inserts batches of nodes. Each batch involves - /// performing a greedy search for each vector to be inserted, and inserting it with edges to - /// all nodes traversed during the search. Reverse edges are also inserted and robustPrune is applied - /// to improve graph quality. The index_params struct controls the degree of the final graph. + /// The build uses the Vamana insertion-based algorithm: starting from an + /// empty graph it iteratively inserts batches of nodes, performing a greedy + /// search for each inserted vector and connecting it to all nodes traversed; + /// reverse edges are added and `robustPrune` is applied to improve quality. + /// [`IndexParams`] controls the degree of the final graph. /// /// `dataset` is a row-major matrix on the host or device implementing - /// [`AsDlTensor`]; it is copied into the index. + /// [`AsDlTensor`]; it is copied into the index, so `Index` carries no + /// lifetime. pub fn build(res: &Resources, params: &IndexParams, dataset: &T) -> Result where T: AsDlTensor + ?Sized, { let dataset = dataset.as_dl_tensor()?; let index = Index::new()?; - // `cuvsVamanaBuild` copies the dataset into the index, so the index does not - // retain a view into `dataset`; the borrow only needs to be valid for the - // duration of this call. That is why `Index` carries no lifetime. unsafe { check_cuvs(ffi::cuvsVamanaBuild( - res.0, - params.0, + res.handle(), + params.handle(), dataset.to_c().as_mut_ptr(), index.0, ))?; @@ -46,7 +46,7 @@ impl Index { Ok(index) } - /// Creates a new empty index + /// Creates a new empty index. pub fn new() -> Result { unsafe { let mut index = std::mem::MaybeUninit::::uninit(); @@ -55,27 +55,19 @@ impl Index { } } - /// Save Vamana index to file - /// - /// Matches the file format used by the DiskANN open-source repository, allowing cross-compatibility. - /// - /// Serialized Index is to be used by the DiskANN open-source repository for graph search. + /// Saves the Vamana index to a file. /// - /// # Arguments + /// Matches the on-disk format used by the DiskANN open-source repository, + /// so the serialized index can be consumed there for graph search. /// - /// * `res` - Resources to use - /// * `filename` - The file prefix for where the index is sazved - /// * `include_dataset` - whether to include the dataset in the serialized index + /// `filename` is the file prefix under which the index is saved; + /// `include_dataset` controls whether the dataset is embedded. pub fn serialize(self, res: &Resources, filename: &str, include_dataset: bool) -> Result<()> { - let c_filename = CString::new(filename).unwrap(); - unsafe { - check_cuvs(ffi::cuvsVamanaSerialize( - res.0, - c_filename.as_ptr(), - self.0, - include_dataset, - )) - } + let c_filename = CString::new(filename)?; + check_cuvs(unsafe { + ffi::cuvsVamanaSerialize(res.handle(), c_filename.as_ptr(), self.0, include_dataset) + })?; + Ok(()) } } @@ -97,11 +89,9 @@ mod tests { #[test] fn test_vamana() { - let build_params = IndexParams::new().unwrap(); - + let build_params = IndexParams::builder().build().unwrap(); let res = Resources::new().unwrap(); - // Create a new random dataset to index let n_datapoints = 1024; let n_features = 16; let dataset = ndarray::Array::::random( @@ -111,7 +101,6 @@ mod tests { let dataset_device = DeviceTensor::from_host(&res, &dataset).unwrap(); - // build the vamana index let _index = Index::build(&res, &build_params, &dataset_device) .expect("failed to create vamana index"); } diff --git a/rust/cuvs/src/neighbors/vamana/mod.rs b/rust/cuvs/src/neighbors/vamana/mod.rs new file mode 100644 index 0000000000..f7bb396978 --- /dev/null +++ b/rust/cuvs/src/neighbors/vamana/mod.rs @@ -0,0 +1,32 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +//! Vamana: builds a DiskANN-style Vamana graph over a dataset. +//! +//! Build an [`Index`] from a dataset (then typically serialize it). The dataset +//! is borrowed through the `AsDlTensor` trait; see the [`dlpack`](crate::dlpack) +//! module for the tensor model. + +mod index; +mod params; + +pub use index::Index; +pub use params::IndexParams; + +use crate::dlpack::DLPackError; +use crate::error::LibraryError; + +/// Error type for Vamana operations. +#[derive(Debug, thiserror::Error)] +pub enum VamanaError { + /// The cuVS C library reported a failure. + #[error(transparent)] + Library(#[from] LibraryError), + /// Tensor conversion into DLPack metadata failed. + #[error(transparent)] + DLPack(#[from] DLPackError), + /// A file path contained an interior NUL byte. + #[error("path contains an interior NUL byte")] + InvalidPath(#[from] std::ffi::NulError), +} diff --git a/rust/cuvs/src/neighbors/vamana/params.rs b/rust/cuvs/src/neighbors/vamana/params.rs new file mode 100644 index 0000000000..e2e7a57dbe --- /dev/null +++ b/rust/cuvs/src/neighbors/vamana/params.rs @@ -0,0 +1,111 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Builder-pattern parameter type for Vamana index build. +//! +//! All setters are optional; unset values retain the library defaults from the +//! underlying C `cuvsVamanaIndexParamsCreate`. + +use std::{fmt, ptr}; + +use bon::bon; + +use crate::distance::DistanceType; +use crate::error::check_cuvs; + +use super::VamanaError; + +/// Parameters for building a Vamana index. +pub struct IndexParams { + handle: ffi::cuvsVamanaIndexParams_t, +} + +#[bon] +impl IndexParams { + #[builder] + #[allow(clippy::too_many_arguments)] + pub fn new( + metric: Option, + graph_degree: Option, + visited_size: Option, + vamana_iters: Option, + alpha: Option, + max_fraction: Option, + batch_base: Option, + queue_size: Option, + reverse_batchsize: Option, + ) -> Result { + let params = Self::try_new()?; + unsafe { + if let Some(v) = metric { + (*params.handle).metric = v.into(); + } + if let Some(v) = graph_degree { + (*params.handle).graph_degree = v; + } + if let Some(v) = visited_size { + (*params.handle).visited_size = v; + } + if let Some(v) = vamana_iters { + (*params.handle).vamana_iters = v; + } + if let Some(v) = alpha { + (*params.handle).alpha = v; + } + if let Some(v) = max_fraction { + (*params.handle).max_fraction = v; + } + if let Some(v) = batch_base { + (*params.handle).batch_base = v; + } + if let Some(v) = queue_size { + (*params.handle).queue_size = v; + } + if let Some(v) = reverse_batchsize { + (*params.handle).reverse_batchsize = v; + } + } + Ok(params) + } +} + +impl IndexParams { + /// Allocate parameters populated with the library defaults. + pub fn try_new() -> Result { + let mut handle = ptr::null_mut(); + check_cuvs(unsafe { ffi::cuvsVamanaIndexParamsCreate(&mut handle) })?; + Ok(Self { handle }) + } + + pub(super) fn handle(&self) -> ffi::cuvsVamanaIndexParams_t { + self.handle + } +} + +impl fmt::Debug for IndexParams { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_tuple("IndexParams").field(unsafe { &*self.handle }).finish() + } +} + +impl Drop for IndexParams { + fn drop(&mut self) { + let _ = unsafe { ffi::cuvsVamanaIndexParamsDestroy(self.handle) }; + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn index_params_with_values() { + let params = IndexParams::builder().alpha(1.0).visited_size(128).build().unwrap(); + unsafe { + assert_eq!((*params.handle).alpha, 1.0); + assert_eq!((*params.handle).visited_size, 128); + } + } +} diff --git a/rust/cuvs/src/resources.rs b/rust/cuvs/src/resources.rs index 5e0c9ee7af..cacae4b42c 100644 --- a/rust/cuvs/src/resources.rs +++ b/rust/cuvs/src/resources.rs @@ -3,57 +3,80 @@ * SPDX-License-Identifier: Apache-2.0 */ -use crate::error::{Result, check_cuvs}; +//! GPU resource management with RAII semantics. + +use crate::error::{LibraryError, check_cuvs}; use std::io::{Write, stderr}; +type Result = std::result::Result; + +/// Error type for resource operations. +#[derive(Debug, thiserror::Error)] +pub enum ResourcesError { + /// The cuVS C library reported a failure. + #[error(transparent)] + Library(#[from] LibraryError), +} + /// Resources are objects that are shared between function calls, /// and includes things like CUDA streams, cuBLAS handles and other /// resources that are expensive to create. #[derive(Debug)] -pub struct Resources(pub ffi::cuvsResources_t); +pub struct Resources { + handle: ffi::cuvsResources_t, +} impl Resources { - /// Returns a new Resources object + /// Creates a new resources handle bound to the current CUDA device. pub fn new() -> Result { - let mut res: ffi::cuvsResources_t = 0; - unsafe { - check_cuvs(ffi::cuvsResourcesCreate(&mut res))?; - } - Ok(Resources(res)) + let mut handle: ffi::cuvsResources_t = 0; + check_cuvs(unsafe { ffi::cuvsResourcesCreate(&mut handle) })?; + Ok(Resources { handle }) } - /// Sets the current cuda stream + /// Creates a resources handle that enqueues work on `stream` instead of the + /// default internal stream. + /// + /// The stream is bound once, at construction. /// /// # Safety /// - /// `stream` must be a valid CUDA stream that remains valid for as long as it - /// is used by this resources handle. - pub unsafe fn set_cuda_stream(&self, stream: ffi::cudaStream_t) -> Result<()> { - unsafe { check_cuvs(ffi::cuvsStreamSet(self.0, stream)) } + /// `stream` must be a valid CUDA stream for the current device and must + /// remain valid for as long as this handle uses it. + pub unsafe fn with_stream(stream: ffi::cudaStream_t) -> Result { + let res = Resources::new()?; + // SAFETY: the caller guarantees `stream` is valid for this device and + // outlives the handle. + check_cuvs(unsafe { ffi::cuvsStreamSet(res.handle, stream) })?; + Ok(res) } - /// Gets the current cuda stream - pub fn get_cuda_stream(&self) -> Result { + /// Returns the current CUDA stream associated with this handle. + pub fn stream(&self) -> Result { unsafe { let mut stream = std::mem::MaybeUninit::::uninit(); - check_cuvs(ffi::cuvsStreamGet(self.0, stream.as_mut_ptr()))?; + check_cuvs(ffi::cuvsStreamGet(self.handle, stream.as_mut_ptr()))?; Ok(stream.assume_init()) } } - /// Syncs the current cuda stream + /// Blocks until all operations on the current CUDA stream have completed. pub fn sync_stream(&self) -> Result<()> { - unsafe { check_cuvs(ffi::cuvsStreamSync(self.0)) } + check_cuvs(unsafe { ffi::cuvsStreamSync(self.handle) })?; + Ok(()) + } + + /// Raw handle for FFI calls in other modules. + pub(crate) fn handle(&self) -> ffi::cuvsResources_t { + self.handle } } impl Drop for Resources { fn drop(&mut self) { - unsafe { - if let Err(e) = check_cuvs(ffi::cuvsResourcesDestroy(self.0)) { - write!(stderr(), "failed to call cuvsResourcesDestroy {:?}", e) - .expect("failed to write to stderr"); - } + if let Err(e) = check_cuvs(unsafe { ffi::cuvsResourcesDestroy(self.handle) }) { + write!(stderr(), "failed to call cuvsResourcesDestroy {:?}", e) + .expect("failed to write to stderr"); } } } diff --git a/rust/cuvs/src/test_utils.rs b/rust/cuvs/src/test_utils.rs index 30d6b6accf..65af3cb557 100644 --- a/rust/cuvs/src/test_utils.rs +++ b/rust/cuvs/src/test_utils.rs @@ -12,10 +12,14 @@ use std::marker::PhantomData; use crate::dlpack::{AsDlTensor, AsDlTensorMut, DLPackError, DLTensorView, DLTensorViewMut, DType}; -use crate::error::{Result, check_cuvs}; +use crate::error::check_cuvs; use crate::ffi; use crate::resources::Resources; +// Test helpers can fail with either a `LibraryError` or a `DLPackError`; a boxed +// error keeps the (test-only) surface simple. +type Result = std::result::Result>; + pub(crate) struct DeviceTensor<'res, T: DType> { data: *mut std::ffi::c_void, shape: Vec, @@ -29,7 +33,7 @@ impl<'res, T: DType> DeviceTensor<'res, T> { let capacity_bytes = shape.iter().product::() * std::mem::size_of::(); let mut data: *mut std::ffi::c_void = std::ptr::null_mut(); unsafe { - check_cuvs(ffi::cuvsRMMAlloc(res.0, &mut data, capacity_bytes))?; + check_cuvs(ffi::cuvsRMMAlloc(res.handle(), &mut data, capacity_bytes))?; } Ok(Self { @@ -66,7 +70,7 @@ impl<'res, T: DType> DeviceTensor<'res, T> { let device_view = device.as_dl_tensor_mut()?; unsafe { check_cuvs(ffi::cuvsMatrixCopy( - res.0, + res.handle(), host.to_c().as_mut_ptr(), device_view.to_c().as_mut_ptr(), ))?; @@ -101,7 +105,7 @@ impl<'res, T: DType> DeviceTensor<'res, T> { let device = self.as_dl_tensor()?; unsafe { check_cuvs(ffi::cuvsMatrixCopy( - res.0, + res.handle(), device.to_c().as_mut_ptr(), host.to_c().as_mut_ptr(), ))?; @@ -114,7 +118,9 @@ impl<'res, T: DType> DeviceTensor<'res, T> { impl Drop for DeviceTensor<'_, T> { fn drop(&mut self) { if !self.data.is_null() { - let _ = unsafe { ffi::cuvsRMMFree(self.resources.0, self.data, self.capacity_bytes) }; + let _ = unsafe { + ffi::cuvsRMMFree(self.resources.handle(), self.data, self.capacity_bytes) + }; } } } diff --git a/rust/cuvs/src/vamana/index_params.rs b/rust/cuvs/src/vamana/index_params.rs deleted file mode 100644 index 40a1d21e52..0000000000 --- a/rust/cuvs/src/vamana/index_params.rs +++ /dev/null @@ -1,129 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. - * SPDX-License-Identifier: Apache-2.0 - */ - -use crate::distance_type::DistanceType; -use crate::error::{Result, check_cuvs}; -use std::fmt; -use std::io::{Write, stderr}; - -pub struct IndexParams(pub ffi::cuvsVamanaIndexParams_t); - -impl IndexParams { - /// Returns a new IndexParams - pub fn new() -> Result { - unsafe { - let mut params = std::mem::MaybeUninit::::uninit(); - check_cuvs(ffi::cuvsVamanaIndexParamsCreate(params.as_mut_ptr()))?; - Ok(IndexParams(params.assume_init())) - } - } - - /// DistanceType to use for building the index - pub fn set_metric(self, metric: DistanceType) -> IndexParams { - unsafe { - (*self.0).metric = metric; - } - self - } - - /// Maximum degree of output graph corresponds to the R parameter in the original Vamana - /// literature. - pub fn set_graph_degree(self, graph_degree: u32) -> IndexParams { - unsafe { - (*self.0).graph_degree = graph_degree; - } - self - } - - /// Maximum number of visited nodes per search corresponds to the L parameter in the Vamana - /// literature - pub fn set_visited_size(self, visited_size: u32) -> IndexParams { - unsafe { - (*self.0).visited_size = visited_size; - } - self - } - - /// Number of Vamana vector insertion iterations (each iteration inserts all vectors). - pub fn set_vamana_iters(self, vamana_iters: f32) -> IndexParams { - unsafe { - (*self.0).vamana_iters = vamana_iters; - } - self - } - - /// Alpha for pruning parameter - pub fn set_alpha(self, alpha: f32) -> IndexParams { - unsafe { - (*self.0).alpha = alpha; - } - self - } - - /// Maximum fraction of dataset inserted per batch. - /// Larger max batch decreases graph quality, but improves speed - pub fn set_max_fraction(self, max_fraction: f32) -> IndexParams { - unsafe { - (*self.0).max_fraction = max_fraction; - } - self - } - - /// Base of growth rate of batch sizes - pub fn set_batch_base(self, batch_base: f32) -> IndexParams { - unsafe { - (*self.0).batch_base = batch_base; - } - self - } - - /// Size of candidate queue structure - should be (2^x)-1 - pub fn set_queue_size(self, queue_size: u32) -> IndexParams { - unsafe { - (*self.0).queue_size = queue_size; - } - self - } - - /// Max batchsize of reverse edge processing (reduces memory footprint) - pub fn set_reverse_batchsize(self, reverse_batchsize: u32) -> IndexParams { - unsafe { - (*self.0).reverse_batchsize = reverse_batchsize; - } - self - } -} - -impl fmt::Debug for IndexParams { - fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { - // custom debug trait here, default value will show the pointer address - // for the inner params object which isn't that useful. - write!(f, "IndexParams({:?})", unsafe { *self.0 }) - } -} - -impl Drop for IndexParams { - fn drop(&mut self) { - if let Err(e) = check_cuvs(unsafe { ffi::cuvsVamanaIndexParamsDestroy(self.0) }) { - write!(stderr(), "failed to call cuvsVamanaIndexParamsDestroy {:?}", e) - .expect("failed to write to stderr"); - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_index_params() { - let params = IndexParams::new().unwrap().set_alpha(1.0).set_visited_size(128); - - unsafe { - assert_eq!((*params.0).alpha, 1.0); - assert_eq!((*params.0).visited_size, 128); - } - } -} diff --git a/rust/cuvs/src/vamana/mod.rs b/rust/cuvs/src/vamana/mod.rs deleted file mode 100644 index 84105d37e0..0000000000 --- a/rust/cuvs/src/vamana/mod.rs +++ /dev/null @@ -1,15 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ -//! Vamana: builds a DiskANN-style Vamana graph over a dataset. -//! -//! Build an [`Index`] from a dataset (then typically serialize it). The dataset -//! is borrowed through the [`AsDlTensor`](crate::AsDlTensor) trait; see the -//! [`dlpack`](crate::dlpack) module for the tensor model. - -mod index; -mod index_params; - -pub use index::Index; -pub use index_params::IndexParams;