From 2e5b1b9bbf6b62ca6f6315d7fb74805dadf757ed Mon Sep 17 00:00:00 2001 From: Yan Zaretskiy Date: Wed, 1 Jul 2026 13:31:37 -0700 Subject: [PATCH 1/3] 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; From 77f1f98a3e994c3038c3a6d6c6cc443c5fae9138 Mon Sep 17 00:00:00 2001 From: Yan Zaretskiy Date: Thu, 2 Jul 2026 09:51:09 -0700 Subject: [PATCH 2/3] Update Fern docs for the Rust bindings --- fern/docs.yml | 62 ++--- fern/pages/neighbors/bruteforce.md | 2 +- fern/pages/neighbors/cagra.md | 2 +- fern/pages/neighbors/ivfflat.md | 2 +- fern/pages/neighbors/ivfpq.md | 2 +- fern/pages/neighbors/vamana.md | 2 +- fern/pages/rust_api/index.md | 31 +-- .../rust_api/rust-api-cuvs-brute-force.md | 87 ------ .../rust-api-cuvs-cagra-index-params.md | 201 -------------- .../rust_api/rust-api-cuvs-cagra-index.md | 181 ------------ .../rust-api-cuvs-cagra-search-params.md | 187 ------------- fern/pages/rust_api/rust-api-cuvs-cagra.md | 119 -------- .../rust-api-cuvs-cluster-kmeans-params.md | 139 ++-------- .../rust_api/rust-api-cuvs-cluster-kmeans.md | 123 ++++----- .../rust_api/rust-api-cuvs-distance-type.md | 17 -- fern/pages/rust_api/rust-api-cuvs-distance.md | 66 +++-- fern/pages/rust_api/rust-api-cuvs-dlpack.md | 257 +++++++++++++++--- fern/pages/rust_api/rust-api-cuvs-error.md | 42 +-- .../rust-api-cuvs-ivf-flat-index-params.md | 106 -------- .../rust_api/rust-api-cuvs-ivf-flat-index.md | 85 ------ .../rust-api-cuvs-ivf-flat-search-params.md | 48 ---- fern/pages/rust_api/rust-api-cuvs-ivf-flat.md | 97 ------- .../rust-api-cuvs-ivf-pq-index-params.md | 211 -------------- .../rust_api/rust-api-cuvs-ivf-pq-index.md | 85 ------ .../rust-api-cuvs-ivf-pq-search-params.md | 86 ------ fern/pages/rust_api/rust-api-cuvs-ivf-pq.md | 95 ------- .../rust-api-cuvs-neighbors-brute-force.md | 104 +++++++ .../rust-api-cuvs-neighbors-cagra-index.md | 208 ++++++++++++++ .../rust-api-cuvs-neighbors-cagra-params.md | 178 ++++++++++++ .../rust_api/rust-api-cuvs-neighbors-cagra.md | 94 +++++++ .../rust-api-cuvs-neighbors-ivf-flat-index.md | 85 ++++++ ...rust-api-cuvs-neighbors-ivf-flat-params.md | 96 +++++++ .../rust-api-cuvs-neighbors-ivf-flat.md | 48 ++++ .../rust-api-cuvs-neighbors-ivf-pq-index.md | 82 ++++++ .../rust-api-cuvs-neighbors-ivf-pq-params.md | 115 ++++++++ .../rust-api-cuvs-neighbors-ivf-pq.md | 51 ++++ .../rust-api-cuvs-neighbors-vamana-index.md | 78 ++++++ .../rust-api-cuvs-neighbors-vamana-params.md | 63 +++++ .../rust-api-cuvs-neighbors-vamana.md | 44 +++ .../pages/rust_api/rust-api-cuvs-neighbors.md | 56 ++++ .../pages/rust_api/rust-api-cuvs-resources.md | 59 ++-- .../rust-api-cuvs-vamana-index-params.md | 137 ---------- .../rust_api/rust-api-cuvs-vamana-index.md | 87 ------ fern/pages/rust_api/rust-api-cuvs-vamana.md | 27 -- fern/pages/rust_api/rust-api-cuvs.md | 60 +--- fern/scripts/generate_api_reference.py | 4 +- 46 files changed, 1779 insertions(+), 2232 deletions(-) delete mode 100644 fern/pages/rust_api/rust-api-cuvs-brute-force.md delete mode 100644 fern/pages/rust_api/rust-api-cuvs-cagra-index-params.md delete mode 100644 fern/pages/rust_api/rust-api-cuvs-cagra-index.md delete mode 100644 fern/pages/rust_api/rust-api-cuvs-cagra-search-params.md delete mode 100644 fern/pages/rust_api/rust-api-cuvs-cagra.md delete mode 100644 fern/pages/rust_api/rust-api-cuvs-distance-type.md delete mode 100644 fern/pages/rust_api/rust-api-cuvs-ivf-flat-index-params.md delete mode 100644 fern/pages/rust_api/rust-api-cuvs-ivf-flat-index.md delete mode 100644 fern/pages/rust_api/rust-api-cuvs-ivf-flat-search-params.md delete mode 100644 fern/pages/rust_api/rust-api-cuvs-ivf-flat.md delete mode 100644 fern/pages/rust_api/rust-api-cuvs-ivf-pq-index-params.md delete mode 100644 fern/pages/rust_api/rust-api-cuvs-ivf-pq-index.md delete mode 100644 fern/pages/rust_api/rust-api-cuvs-ivf-pq-search-params.md delete mode 100644 fern/pages/rust_api/rust-api-cuvs-ivf-pq.md create mode 100644 fern/pages/rust_api/rust-api-cuvs-neighbors-brute-force.md create mode 100644 fern/pages/rust_api/rust-api-cuvs-neighbors-cagra-index.md create mode 100644 fern/pages/rust_api/rust-api-cuvs-neighbors-cagra-params.md create mode 100644 fern/pages/rust_api/rust-api-cuvs-neighbors-cagra.md create mode 100644 fern/pages/rust_api/rust-api-cuvs-neighbors-ivf-flat-index.md create mode 100644 fern/pages/rust_api/rust-api-cuvs-neighbors-ivf-flat-params.md create mode 100644 fern/pages/rust_api/rust-api-cuvs-neighbors-ivf-flat.md create mode 100644 fern/pages/rust_api/rust-api-cuvs-neighbors-ivf-pq-index.md create mode 100644 fern/pages/rust_api/rust-api-cuvs-neighbors-ivf-pq-params.md create mode 100644 fern/pages/rust_api/rust-api-cuvs-neighbors-ivf-pq.md create mode 100644 fern/pages/rust_api/rust-api-cuvs-neighbors-vamana-index.md create mode 100644 fern/pages/rust_api/rust-api-cuvs-neighbors-vamana-params.md create mode 100644 fern/pages/rust_api/rust-api-cuvs-neighbors-vamana.md create mode 100644 fern/pages/rust_api/rust-api-cuvs-neighbors.md delete mode 100644 fern/pages/rust_api/rust-api-cuvs-vamana-index-params.md delete mode 100644 fern/pages/rust_api/rust-api-cuvs-vamana-index.md delete mode 100644 fern/pages/rust_api/rust-api-cuvs-vamana.md diff --git a/fern/docs.yml b/fern/docs.yml index 882dd665aa..16b95f4745 100644 --- a/fern/docs.yml +++ b/fern/docs.yml @@ -536,44 +536,38 @@ navigation: path: "./pages/rust_api/rust-api-cuvs-dlpack.md" - page: "cuVS Error" path: "./pages/rust_api/rust-api-cuvs-error.md" + - page: "cuVS Neighbors" + path: "./pages/rust_api/rust-api-cuvs-neighbors.md" - page: "cuVS Resources" path: "./pages/rust_api/rust-api-cuvs-resources.md" - page: "cuVS Distance" path: "./pages/rust_api/rust-api-cuvs-distance.md" - - page: "cuVS Distance Type" - path: "./pages/rust_api/rust-api-cuvs-distance-type.md" - - page: "cuVS Brute Force" - path: "./pages/rust_api/rust-api-cuvs-brute-force.md" - - page: "cuVS Cagra" - path: "./pages/rust_api/rust-api-cuvs-cagra.md" - - page: "cuVS Cagra Index" - path: "./pages/rust_api/rust-api-cuvs-cagra-index.md" - - page: "cuVS Cagra Index Params" - path: "./pages/rust_api/rust-api-cuvs-cagra-index-params.md" - - page: "cuVS Cagra Search Params" - path: "./pages/rust_api/rust-api-cuvs-cagra-search-params.md" - - page: "cuVS IVF Flat" - path: "./pages/rust_api/rust-api-cuvs-ivf-flat.md" - - page: "cuVS IVF Flat Index" - path: "./pages/rust_api/rust-api-cuvs-ivf-flat-index.md" - - page: "cuVS IVF Flat Index Params" - path: "./pages/rust_api/rust-api-cuvs-ivf-flat-index-params.md" - - page: "cuVS IVF Flat Search Params" - path: "./pages/rust_api/rust-api-cuvs-ivf-flat-search-params.md" - - page: "cuVS IVF PQ" - path: "./pages/rust_api/rust-api-cuvs-ivf-pq.md" - - page: "cuVS IVF PQ Index" - path: "./pages/rust_api/rust-api-cuvs-ivf-pq-index.md" - - page: "cuVS IVF PQ Index Params" - path: "./pages/rust_api/rust-api-cuvs-ivf-pq-index-params.md" - - page: "cuVS IVF PQ Search Params" - path: "./pages/rust_api/rust-api-cuvs-ivf-pq-search-params.md" - - page: "cuVS Vamana" - path: "./pages/rust_api/rust-api-cuvs-vamana.md" - - page: "cuVS Vamana Index" - path: "./pages/rust_api/rust-api-cuvs-vamana-index.md" - - page: "cuVS Vamana Index Params" - path: "./pages/rust_api/rust-api-cuvs-vamana-index-params.md" + - page: "cuVS Neighbors Brute Force" + path: "./pages/rust_api/rust-api-cuvs-neighbors-brute-force.md" + - page: "cuVS Neighbors Cagra" + path: "./pages/rust_api/rust-api-cuvs-neighbors-cagra.md" + - page: "cuVS Neighbors Cagra Index" + path: "./pages/rust_api/rust-api-cuvs-neighbors-cagra-index.md" + - page: "cuVS Neighbors Cagra Params" + path: "./pages/rust_api/rust-api-cuvs-neighbors-cagra-params.md" + - page: "cuVS Neighbors IVF Flat" + path: "./pages/rust_api/rust-api-cuvs-neighbors-ivf-flat.md" + - page: "cuVS Neighbors IVF Flat Index" + path: "./pages/rust_api/rust-api-cuvs-neighbors-ivf-flat-index.md" + - page: "cuVS Neighbors IVF Flat Params" + path: "./pages/rust_api/rust-api-cuvs-neighbors-ivf-flat-params.md" + - page: "cuVS Neighbors IVF PQ" + path: "./pages/rust_api/rust-api-cuvs-neighbors-ivf-pq.md" + - page: "cuVS Neighbors IVF PQ Index" + path: "./pages/rust_api/rust-api-cuvs-neighbors-ivf-pq-index.md" + - page: "cuVS Neighbors IVF PQ Params" + path: "./pages/rust_api/rust-api-cuvs-neighbors-ivf-pq-params.md" + - page: "cuVS Neighbors Vamana" + path: "./pages/rust_api/rust-api-cuvs-neighbors-vamana.md" + - page: "cuVS Neighbors Vamana Index" + path: "./pages/rust_api/rust-api-cuvs-neighbors-vamana-index.md" + - page: "cuVS Neighbors Vamana Params" + path: "./pages/rust_api/rust-api-cuvs-neighbors-vamana-params.md" - section: "Go API Documentation" path: "./pages/go_api/index.md" contents: diff --git a/fern/pages/neighbors/bruteforce.md b/fern/pages/neighbors/bruteforce.md index 70c22eceb5..281bb050db 100644 --- a/fern/pages/neighbors/bruteforce.md +++ b/fern/pages/neighbors/bruteforce.md @@ -6,7 +6,7 @@ Brute-force works well when exact results matter, the dataset is small enough to ## Example API Usage -[C API](/api-reference/c-api-neighbors-brute-force) | [C++ API](/api-reference/cpp-api-neighbors-brute-force) | [Python API](/api-reference/python-api-neighbors-brute-force) | [Java API](/api-reference/java-api-com-nvidia-cuvs-bruteforceindex) | [Rust API](/api-reference/rust-api-cuvs-brute-force) | [Go API](/api-reference/go-api-brute-force) +[C API](/api-reference/c-api-neighbors-brute-force) | [C++ API](/api-reference/cpp-api-neighbors-brute-force) | [Python API](/api-reference/python-api-neighbors-brute-force) | [Java API](/api-reference/java-api-com-nvidia-cuvs-bruteforceindex) | [Rust API](/api-reference/rust-api-cuvs-neighbors-brute-force) | [Go API](/api-reference/go-api-brute-force) ### Building an index diff --git a/fern/pages/neighbors/cagra.md b/fern/pages/neighbors/cagra.md index c1eb80faea..9724332b20 100644 --- a/fern/pages/neighbors/cagra.md +++ b/fern/pages/neighbors/cagra.md @@ -6,7 +6,7 @@ CAGRA works well when you want strong recall, high GPU throughput, and fast grap ## Example API Usage -[C API](/api-reference/c-api-neighbors-cagra) | [C++ API](/api-reference/cpp-api-neighbors-cagra) | [Python API](/api-reference/python-api-neighbors-cagra) | [Java API](/api-reference/java-api-com-nvidia-cuvs-cagraindex) | [Rust API](/api-reference/rust-api-cuvs-cagra) | [Go API](/api-reference/go-api-cagra) +[C API](/api-reference/c-api-neighbors-cagra) | [C++ API](/api-reference/cpp-api-neighbors-cagra) | [Python API](/api-reference/python-api-neighbors-cagra) | [Java API](/api-reference/java-api-com-nvidia-cuvs-cagraindex) | [Rust API](/api-reference/rust-api-cuvs-neighbors-cagra) | [Go API](/api-reference/go-api-cagra) ### Building an index diff --git a/fern/pages/neighbors/ivfflat.md b/fern/pages/neighbors/ivfflat.md index ea4127b6ed..6001d28388 100644 --- a/fern/pages/neighbors/ivfflat.md +++ b/fern/pages/neighbors/ivfflat.md @@ -8,7 +8,7 @@ IVF-Flat works well when the index fits in GPU memory, exact recall is not requi ## Example API Usage -[C API](/api-reference/c-api-neighbors-ivf-flat) | [C++ API](/api-reference/cpp-api-neighbors-ivf-flat) | [Python API](/api-reference/python-api-neighbors-ivf-flat) | [Rust API](/api-reference/rust-api-cuvs-ivf-flat) | [Go API](/api-reference/go-api-ivf-flat) +[C API](/api-reference/c-api-neighbors-ivf-flat) | [C++ API](/api-reference/cpp-api-neighbors-ivf-flat) | [Python API](/api-reference/python-api-neighbors-ivf-flat) | [Rust API](/api-reference/rust-api-cuvs-neighbors-ivf-flat) | [Go API](/api-reference/go-api-ivf-flat) ### Building an index diff --git a/fern/pages/neighbors/ivfpq.md b/fern/pages/neighbors/ivfpq.md index 8a2db90afc..31de0d162e 100644 --- a/fern/pages/neighbors/ivfpq.md +++ b/fern/pages/neighbors/ivfpq.md @@ -8,7 +8,7 @@ IVF-PQ works well when the dataset is too large for full-precision storage on th ## Example API Usage -[C API](/api-reference/c-api-neighbors-ivf-pq) | [C++ API](/api-reference/cpp-api-neighbors-ivf-pq) | [Python API](/api-reference/python-api-neighbors-ivf-pq) | [Java IVF-PQ Params](/api-reference/java-api-com-nvidia-cuvs-cuvsivfpqparams) | [Rust API](/api-reference/rust-api-cuvs-ivf-pq) | [Go API](/api-reference/go-api-ivf-pq) +[C API](/api-reference/c-api-neighbors-ivf-pq) | [C++ API](/api-reference/cpp-api-neighbors-ivf-pq) | [Python API](/api-reference/python-api-neighbors-ivf-pq) | [Java IVF-PQ Params](/api-reference/java-api-com-nvidia-cuvs-cuvsivfpqparams) | [Rust API](/api-reference/rust-api-cuvs-neighbors-ivf-pq) | [Go API](/api-reference/go-api-ivf-pq) Java currently exposes IVF-PQ parameter classes for CAGRA graph construction, not a standalone IVF-PQ index/search binding. The runnable standalone examples below cover C, C++, Python, Rust, and Go. diff --git a/fern/pages/neighbors/vamana.md b/fern/pages/neighbors/vamana.md index 11ce69ba98..6a8a3b25a5 100644 --- a/fern/pages/neighbors/vamana.md +++ b/fern/pages/neighbors/vamana.md @@ -8,7 +8,7 @@ Vamana works well when you want to build large DiskANN-compatible graph indexes ## Example API Usage -[C API](/api-reference/c-api-neighbors-vamana) | [C++ API](/api-reference/cpp-api-neighbors-vamana) | [Python API](/api-reference/python-api-neighbors-vamana) | [Rust API](/api-reference/rust-api-cuvs-vamana) +[C API](/api-reference/c-api-neighbors-vamana) | [C++ API](/api-reference/cpp-api-neighbors-vamana) | [Python API](/api-reference/python-api-neighbors-vamana) | [Rust API](/api-reference/rust-api-cuvs-neighbors-vamana) Vamana currently supports build and serialize operations in NVIDIA cuVS. Search is performed by loading the serialized index with DiskANN. Java and Go do not currently expose standalone Vamana bindings. diff --git a/fern/pages/rust_api/index.md b/fern/pages/rust_api/index.md index 92d2151338..75878288fb 100644 --- a/fern/pages/rust_api/index.md +++ b/fern/pages/rust_api/index.md @@ -13,28 +13,25 @@ These pages are generated from the Rust crate sources under `rust/cuvs/src`. - [`cuvs`](/api-reference/rust-api-cuvs) - [`cuvs::dlpack`](/api-reference/rust-api-cuvs-dlpack) - [`cuvs::error`](/api-reference/rust-api-cuvs-error) +- [`cuvs::neighbors`](/api-reference/rust-api-cuvs-neighbors) - [`cuvs::resources`](/api-reference/rust-api-cuvs-resources) ## Distance - [`cuvs::distance`](/api-reference/rust-api-cuvs-distance) -- [`cuvs::distance_type`](/api-reference/rust-api-cuvs-distance-type) ## Nearest Neighbors -- [`cuvs::brute_force`](/api-reference/rust-api-cuvs-brute-force) -- [`cuvs::cagra`](/api-reference/rust-api-cuvs-cagra) -- [`cuvs::cagra::index`](/api-reference/rust-api-cuvs-cagra-index) -- [`cuvs::cagra::index_params`](/api-reference/rust-api-cuvs-cagra-index-params) -- [`cuvs::cagra::search_params`](/api-reference/rust-api-cuvs-cagra-search-params) -- [`cuvs::ivf_flat`](/api-reference/rust-api-cuvs-ivf-flat) -- [`cuvs::ivf_flat::index`](/api-reference/rust-api-cuvs-ivf-flat-index) -- [`cuvs::ivf_flat::index_params`](/api-reference/rust-api-cuvs-ivf-flat-index-params) -- [`cuvs::ivf_flat::search_params`](/api-reference/rust-api-cuvs-ivf-flat-search-params) -- [`cuvs::ivf_pq`](/api-reference/rust-api-cuvs-ivf-pq) -- [`cuvs::ivf_pq::index`](/api-reference/rust-api-cuvs-ivf-pq-index) -- [`cuvs::ivf_pq::index_params`](/api-reference/rust-api-cuvs-ivf-pq-index-params) -- [`cuvs::ivf_pq::search_params`](/api-reference/rust-api-cuvs-ivf-pq-search-params) -- [`cuvs::vamana`](/api-reference/rust-api-cuvs-vamana) -- [`cuvs::vamana::index`](/api-reference/rust-api-cuvs-vamana-index) -- [`cuvs::vamana::index_params`](/api-reference/rust-api-cuvs-vamana-index-params) +- [`cuvs::neighbors::brute_force`](/api-reference/rust-api-cuvs-neighbors-brute-force) +- [`cuvs::neighbors::cagra`](/api-reference/rust-api-cuvs-neighbors-cagra) +- [`cuvs::neighbors::cagra::index`](/api-reference/rust-api-cuvs-neighbors-cagra-index) +- [`cuvs::neighbors::cagra::params`](/api-reference/rust-api-cuvs-neighbors-cagra-params) +- [`cuvs::neighbors::ivf_flat`](/api-reference/rust-api-cuvs-neighbors-ivf-flat) +- [`cuvs::neighbors::ivf_flat::index`](/api-reference/rust-api-cuvs-neighbors-ivf-flat-index) +- [`cuvs::neighbors::ivf_flat::params`](/api-reference/rust-api-cuvs-neighbors-ivf-flat-params) +- [`cuvs::neighbors::ivf_pq`](/api-reference/rust-api-cuvs-neighbors-ivf-pq) +- [`cuvs::neighbors::ivf_pq::index`](/api-reference/rust-api-cuvs-neighbors-ivf-pq-index) +- [`cuvs::neighbors::ivf_pq::params`](/api-reference/rust-api-cuvs-neighbors-ivf-pq-params) +- [`cuvs::neighbors::vamana`](/api-reference/rust-api-cuvs-neighbors-vamana) +- [`cuvs::neighbors::vamana::index`](/api-reference/rust-api-cuvs-neighbors-vamana-index) +- [`cuvs::neighbors::vamana::params`](/api-reference/rust-api-cuvs-neighbors-vamana-params) diff --git a/fern/pages/rust_api/rust-api-cuvs-brute-force.md b/fern/pages/rust_api/rust-api-cuvs-brute-force.md deleted file mode 100644 index a0b27685c5..0000000000 --- a/fern/pages/rust_api/rust-api-cuvs-brute-force.md +++ /dev/null @@ -1,87 +0,0 @@ ---- -slug: api-reference/rust-api-cuvs-brute-force ---- - -# Brute Force Module - -_Rust module: `cuvs::brute_force`_ - -_Source: `rust/cuvs/src/brute_force.rs`_ - -Brute Force KNN - -## Index - -```rust -#[derive(Debug)] -pub struct Index { - /* private fields */ -} -``` - -Brute Force KNN Index - -**Methods** - -| Name | Source | -| --- | --- | -| `build` | `rust/cuvs/src/brute_force.rs:32` | -| `new` | `rust/cuvs/src/brute_force.rs:54` | -| `search` | `rust/cuvs/src/brute_force.rs:70` | - -### build - -```rust -pub fn build>( -res: &Resources, -metric: DistanceType, -metric_arg: Option, -dataset: T, -) -> Result -``` - -Builds a new Brute Force KNN Index from the dataset for efficient search. - -#### Arguments - -* `res` - Resources to use -* `metric` - DistanceType to use for building the index -* `metric_arg` - Optional value of `p` for Minkowski distances -* `dataset` - A row-major matrix on either the host or device to index - -_Source: `rust/cuvs/src/brute_force.rs:32`_ - -### new - -```rust -pub fn new() -> Result -``` - -Creates a new empty index - -_Source: `rust/cuvs/src/brute_force.rs:54`_ - -### search - -```rust -pub fn search( -&self, -res: &Resources, -queries: &ManagedTensor, -neighbors: &ManagedTensor, -distances: &ManagedTensor, -) -> Result<()> -``` - -Perform a Nearest Neighbors search on the Index - -#### Arguments - -* `res` - Resources to use -* `queries` - A matrix in device memory to query for -* `neighbors` - Matrix in device memory that receives the indices of the nearest neighbors -* `distances` - Matrix in device memory that receives the distances of the nearest neighbors - -_Source: `rust/cuvs/src/brute_force.rs:70`_ - -_Source: `rust/cuvs/src/brute_force.rs:16`_ diff --git a/fern/pages/rust_api/rust-api-cuvs-cagra-index-params.md b/fern/pages/rust_api/rust-api-cuvs-cagra-index-params.md deleted file mode 100644 index 66c8b6e2c4..0000000000 --- a/fern/pages/rust_api/rust-api-cuvs-cagra-index-params.md +++ /dev/null @@ -1,201 +0,0 @@ ---- -slug: api-reference/rust-api-cuvs-cagra-index-params ---- - -# Cagra Index Params Module - -_Rust module: `cuvs::cagra::index_params`_ - -_Source: `rust/cuvs/src/cagra/index_params.rs`_ - -## BuildAlgo - -```rust -pub type BuildAlgo = ffi::cuvsCagraGraphBuildAlgo; -``` - -_Source: `rust/cuvs/src/cagra/index_params.rs:10`_ - -## CompressionParams - -```rust -pub struct CompressionParams(pub ffi::cuvsCagraCompressionParams_t); { - /* private fields */ -} -``` - -Supplemental parameters to build CAGRA Index - -**Methods** - -| Name | Source | -| --- | --- | -| `new` | `rust/cuvs/src/cagra/index_params.rs:17` | -| `set_pq_bits` | `rust/cuvs/src/cagra/index_params.rs:26` | -| `set_pq_dim` | `rust/cuvs/src/cagra/index_params.rs:35` | -| `set_vq_n_centers` | `rust/cuvs/src/cagra/index_params.rs:44` | -| `set_kmeans_n_iters` | `rust/cuvs/src/cagra/index_params.rs:53` | -| `set_vq_kmeans_trainset_fraction` | `rust/cuvs/src/cagra/index_params.rs:62` | -| `set_pq_kmeans_trainset_fraction` | `rust/cuvs/src/cagra/index_params.rs:74` | - -### new - -```rust -pub fn new() -> Result -``` - -Returns a new CompressionParams - -_Source: `rust/cuvs/src/cagra/index_params.rs:17`_ - -### set_pq_bits - -```rust -pub fn set_pq_bits(self, pq_bits: u32) -> CompressionParams -``` - -The bit length of the vector element after compression by PQ. - -_Source: `rust/cuvs/src/cagra/index_params.rs:26`_ - -### set_pq_dim - -```rust -pub fn set_pq_dim(self, pq_dim: u32) -> CompressionParams -``` - -The dimensionality of the vector after compression by PQ. When zero, -an optimal value is selected using a heuristic. - -_Source: `rust/cuvs/src/cagra/index_params.rs:35`_ - -### set_vq_n_centers - -```rust -pub fn set_vq_n_centers(self, vq_n_centers: u32) -> CompressionParams -``` - -Vector Quantization (VQ) codebook size - number of "coarse cluster -centers". When zero, an optimal value is selected using a heuristic. - -_Source: `rust/cuvs/src/cagra/index_params.rs:44`_ - -### set_kmeans_n_iters - -```rust -pub fn set_kmeans_n_iters(self, kmeans_n_iters: u32) -> CompressionParams -``` - -The number of iterations searching for kmeans centers (both VQ & PQ -phases). - -_Source: `rust/cuvs/src/cagra/index_params.rs:53`_ - -### set_vq_kmeans_trainset_fraction - -```rust -pub fn set_vq_kmeans_trainset_fraction( -self, -vq_kmeans_trainset_fraction: f64, -) -> CompressionParams -``` - -The fraction of data to use during iterative kmeans building (VQ -phase). When zero, an optimal value is selected using a heuristic. - -_Source: `rust/cuvs/src/cagra/index_params.rs:62`_ - -### set_pq_kmeans_trainset_fraction - -```rust -pub fn set_pq_kmeans_trainset_fraction( -self, -pq_kmeans_trainset_fraction: f64, -) -> CompressionParams -``` - -The fraction of data to use during iterative kmeans building (PQ -phase). When zero, an optimal value is selected using a heuristic. - -_Source: `rust/cuvs/src/cagra/index_params.rs:74`_ - -_Source: `rust/cuvs/src/cagra/index_params.rs:13`_ - -## IndexParams - -```rust -pub struct IndexParams(pub ffi::cuvsCagraIndexParams_t, Option); { - /* private fields */ -} -``` - -**Methods** - -| Name | Source | -| --- | --- | -| `new` | `rust/cuvs/src/cagra/index_params.rs:89` | -| `set_intermediate_graph_degree` | `rust/cuvs/src/cagra/index_params.rs:98` | -| `set_graph_degree` | `rust/cuvs/src/cagra/index_params.rs:106` | -| `set_build_algo` | `rust/cuvs/src/cagra/index_params.rs:114` | -| `set_nn_descent_niter` | `rust/cuvs/src/cagra/index_params.rs:122` | -| `set_compression` | `rust/cuvs/src/cagra/index_params.rs:129` | - -### new - -```rust -pub fn new() -> Result -``` - -Returns a new IndexParams - -_Source: `rust/cuvs/src/cagra/index_params.rs:89`_ - -### set_intermediate_graph_degree - -```rust -pub fn set_intermediate_graph_degree(self, intermediate_graph_degree: usize) -> IndexParams -``` - -Degree of input graph for pruning - -_Source: `rust/cuvs/src/cagra/index_params.rs:98`_ - -### set_graph_degree - -```rust -pub fn set_graph_degree(self, graph_degree: usize) -> IndexParams -``` - -Degree of output graph - -_Source: `rust/cuvs/src/cagra/index_params.rs:106`_ - -### set_build_algo - -```rust -pub fn set_build_algo(self, build_algo: BuildAlgo) -> IndexParams -``` - -ANN algorithm to build knn graph - -_Source: `rust/cuvs/src/cagra/index_params.rs:114`_ - -### set_nn_descent_niter - -```rust -pub fn set_nn_descent_niter(self, nn_descent_niter: usize) -> IndexParams -``` - -Number of iterations to run if building with NN_DESCENT - -_Source: `rust/cuvs/src/cagra/index_params.rs:122`_ - -### set_compression - -```rust -pub fn set_compression(mut self, compression: CompressionParams) -> IndexParams -``` - -_Source: `rust/cuvs/src/cagra/index_params.rs:129`_ - -_Source: `rust/cuvs/src/cagra/index_params.rs:85`_ diff --git a/fern/pages/rust_api/rust-api-cuvs-cagra-index.md b/fern/pages/rust_api/rust-api-cuvs-cagra-index.md deleted file mode 100644 index 9d98ceb273..0000000000 --- a/fern/pages/rust_api/rust-api-cuvs-cagra-index.md +++ /dev/null @@ -1,181 +0,0 @@ ---- -slug: api-reference/rust-api-cuvs-cagra-index ---- - -# Cagra Index Module - -_Rust module: `cuvs::cagra::index`_ - -_Source: `rust/cuvs/src/cagra/index.rs`_ - -## Index - -```rust -#[derive(Debug)] -pub struct Index(ffi::cuvsCagraIndex_t); { - /* private fields */ -} -``` - -CAGRA ANN Index - -**Methods** - -| Name | Source | -| --- | --- | -| `build` | `rust/cuvs/src/cagra/index.rs:38` | -| `new` | `rust/cuvs/src/cagra/index.rs:52` | -| `search` | `rust/cuvs/src/cagra/index.rs:69` | -| `search_with_filter` | `rust/cuvs/src/cagra/index.rs:107` | -| `serialize` | `rust/cuvs/src/cagra/index.rs:143` | -| `serialize_to_hnswlib` | `rust/cuvs/src/cagra/index.rs:166` | -| `deserialize` | `rust/cuvs/src/cagra/index.rs:179` | - -### build - -```rust -pub fn build>( -res: &Resources, -params: &IndexParams, -dataset: T, -) -> Result -``` - -Builds a new Index from the dataset for efficient search. - -#### Arguments - -* `res` - Resources to use -* `params` - Parameters for building the index -* `dataset` - A row-major matrix on either the host or device to index - -_Source: `rust/cuvs/src/cagra/index.rs:38`_ - -### new - -```rust -pub fn new() -> Result -``` - -Creates a new empty index - -_Source: `rust/cuvs/src/cagra/index.rs:52`_ - -### search - -```rust -pub fn search( -&self, -res: &Resources, -params: &SearchParams, -queries: &ManagedTensor, -neighbors: &ManagedTensor, -distances: &ManagedTensor, -) -> Result<()> -``` - -Perform a Approximate Nearest Neighbors search on the Index - -#### Arguments - -* `res` - Resources to use -* `params` - Parameters to use in searching the index -* `queries` - A matrix in device memory to query for -* `neighbors` - Matrix in device memory that receives the indices of the nearest neighbors -* `distances` - Matrix in device memory that receives the distances of the nearest neighbors - -_Source: `rust/cuvs/src/cagra/index.rs:69`_ - -### search_with_filter - -```rust -pub fn search_with_filter( -&self, -res: &Resources, -params: &SearchParams, -queries: &ManagedTensor, -neighbors: &ManagedTensor, -distances: &ManagedTensor, -bitset: &ManagedTensor, -) -> Result<()> -``` - -Perform a filtered Approximate Nearest Neighbors search on the Index - -Like [`search`](Self::search), but accepts a bitset filter to exclude -vectors during graph traversal. Filtered vectors are never visited, -giving better recall than post-filtering. - -#### Arguments - -* `res` - Resources to use -* `params` - Parameters to use in searching the index -* `queries` - A matrix in device memory to query for -* `neighbors` - Matrix in device memory that receives the indices of the nearest neighbors -* `distances` - Matrix in device memory that receives the distances of the nearest neighbors -* `bitset` - A 1-D `uint32` device tensor with `ceil(n_rows / 32)` elements. -Each bit corresponds to a dataset row: bit 1 = include, bit 0 = exclude. - -_Source: `rust/cuvs/src/cagra/index.rs:107`_ - -### serialize - -```rust -pub fn serialize>( -&self, -res: &Resources, -filename: P, -include_dataset: bool, -) -> Result<()> -``` - -Save the CAGRA index to file. - -Experimental, both the API and the serialization format are subject to change. - -#### Arguments - -* `res` - Resources to use -* `filename` - The file path for saving the index -* `include_dataset` - Whether to write out the dataset to the file - -_Source: `rust/cuvs/src/cagra/index.rs:143`_ - -### serialize_to_hnswlib - -```rust -pub fn serialize_to_hnswlib>(&self, res: &Resources, filename: P) -> Result<()> -``` - -Save the CAGRA index to file in hnswlib format. - -NOTE: The saved index can only be read by the hnswlib wrapper in cuVS, -as the serialization format is not compatible with the original hnswlib. - -Experimental, both the API and the serialization format are subject to change. - -#### Arguments - -* `res` - Resources to use -* `filename` - The file path for saving the index - -_Source: `rust/cuvs/src/cagra/index.rs:166`_ - -### deserialize - -```rust -pub fn deserialize>(res: &Resources, filename: P) -> Result -``` - -Load a CAGRA index from file. - -Experimental, both the API and the serialization format are subject to change. - -#### Arguments - -* `res` - Resources to use -* `filename` - The path of the file that stores the index - -_Source: `rust/cuvs/src/cagra/index.rs:179`_ - -_Source: `rust/cuvs/src/cagra/index.rs:17`_ diff --git a/fern/pages/rust_api/rust-api-cuvs-cagra-search-params.md b/fern/pages/rust_api/rust-api-cuvs-cagra-search-params.md deleted file mode 100644 index e5bb684b43..0000000000 --- a/fern/pages/rust_api/rust-api-cuvs-cagra-search-params.md +++ /dev/null @@ -1,187 +0,0 @@ ---- -slug: api-reference/rust-api-cuvs-cagra-search-params ---- - -# Cagra Search Params Module - -_Rust module: `cuvs::cagra::search_params`_ - -_Source: `rust/cuvs/src/cagra/search_params.rs`_ - -## SearchAlgo - -```rust -pub type SearchAlgo = ffi::cuvsCagraSearchAlgo; -``` - -_Source: `rust/cuvs/src/cagra/search_params.rs:10`_ - -## HashMode - -```rust -pub type HashMode = ffi::cuvsCagraHashMode; -``` - -_Source: `rust/cuvs/src/cagra/search_params.rs:11`_ - -## SearchParams - -```rust -pub struct SearchParams(pub ffi::cuvsCagraSearchParams_t); { - /* private fields */ -} -``` - -Supplemental parameters to search CAGRA index - -**Methods** - -| Name | Source | -| --- | --- | -| `new` | `rust/cuvs/src/cagra/search_params.rs:18` | -| `set_max_queries` | `rust/cuvs/src/cagra/search_params.rs:27` | -| `set_itopk_size` | `rust/cuvs/src/cagra/search_params.rs:37` | -| `set_max_iterations` | `rust/cuvs/src/cagra/search_params.rs:45` | -| `set_algo` | `rust/cuvs/src/cagra/search_params.rs:53` | -| `set_team_size` | `rust/cuvs/src/cagra/search_params.rs:61` | -| `set_min_iterations` | `rust/cuvs/src/cagra/search_params.rs:69` | -| `set_thread_block_size` | `rust/cuvs/src/cagra/search_params.rs:77` | -| `set_hashmap_mode` | `rust/cuvs/src/cagra/search_params.rs:85` | -| `set_hashmap_min_bitlen` | `rust/cuvs/src/cagra/search_params.rs:93` | -| `set_hashmap_max_fill_rate` | `rust/cuvs/src/cagra/search_params.rs:101` | -| `set_num_random_samplings` | `rust/cuvs/src/cagra/search_params.rs:109` | -| `set_rand_xor_mask` | `rust/cuvs/src/cagra/search_params.rs:117` | - -### new - -```rust -pub fn new() -> Result -``` - -Returns a new SearchParams object - -_Source: `rust/cuvs/src/cagra/search_params.rs:18`_ - -### set_max_queries - -```rust -pub fn set_max_queries(self, max_queries: usize) -> SearchParams -``` - -Maximum number of queries to search at the same time (batch size). Auto select when 0 - -_Source: `rust/cuvs/src/cagra/search_params.rs:27`_ - -### set_itopk_size - -```rust -pub fn set_itopk_size(self, itopk_size: usize) -> SearchParams -``` - -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 - -_Source: `rust/cuvs/src/cagra/search_params.rs:37`_ - -### set_max_iterations - -```rust -pub fn set_max_iterations(self, max_iterations: usize) -> SearchParams -``` - -Upper limit of search iterations. Auto select when 0. - -_Source: `rust/cuvs/src/cagra/search_params.rs:45`_ - -### set_algo - -```rust -pub fn set_algo(self, algo: SearchAlgo) -> SearchParams -``` - -Which search implementation to use. - -_Source: `rust/cuvs/src/cagra/search_params.rs:53`_ - -### set_team_size - -```rust -pub fn set_team_size(self, team_size: usize) -> SearchParams -``` - -Number of threads used to calculate a single distance. 4, 8, 16, or 32. - -_Source: `rust/cuvs/src/cagra/search_params.rs:61`_ - -### set_min_iterations - -```rust -pub fn set_min_iterations(self, min_iterations: usize) -> SearchParams -``` - -Lower limit of search iterations. - -_Source: `rust/cuvs/src/cagra/search_params.rs:69`_ - -### set_thread_block_size - -```rust -pub fn set_thread_block_size(self, thread_block_size: usize) -> SearchParams -``` - -Thread block size. 0, 64, 128, 256, 512, 1024. Auto selection when 0. - -_Source: `rust/cuvs/src/cagra/search_params.rs:77`_ - -### set_hashmap_mode - -```rust -pub fn set_hashmap_mode(self, hashmap_mode: HashMode) -> SearchParams -``` - -Hashmap type. Auto selection when AUTO. - -_Source: `rust/cuvs/src/cagra/search_params.rs:85`_ - -### set_hashmap_min_bitlen - -```rust -pub fn set_hashmap_min_bitlen(self, hashmap_min_bitlen: usize) -> SearchParams -``` - -Lower limit of hashmap bit length. More than 8. - -_Source: `rust/cuvs/src/cagra/search_params.rs:93`_ - -### set_hashmap_max_fill_rate - -```rust -pub fn set_hashmap_max_fill_rate(self, hashmap_max_fill_rate: f32) -> SearchParams -``` - -Upper limit of hashmap fill rate. More than 0.1, less than 0.9. - -_Source: `rust/cuvs/src/cagra/search_params.rs:101`_ - -### set_num_random_samplings - -```rust -pub fn set_num_random_samplings(self, num_random_samplings: u32) -> SearchParams -``` - -Number of iterations of initial random seed node selection. 1 or more. - -_Source: `rust/cuvs/src/cagra/search_params.rs:109`_ - -### set_rand_xor_mask - -```rust -pub fn set_rand_xor_mask(self, rand_xor_mask: u64) -> SearchParams -``` - -Bit mask used for initial random seed node selection. - -_Source: `rust/cuvs/src/cagra/search_params.rs:117`_ - -_Source: `rust/cuvs/src/cagra/search_params.rs:14`_ diff --git a/fern/pages/rust_api/rust-api-cuvs-cagra.md b/fern/pages/rust_api/rust-api-cuvs-cagra.md deleted file mode 100644 index 3fc7753ef7..0000000000 --- a/fern/pages/rust_api/rust-api-cuvs-cagra.md +++ /dev/null @@ -1,119 +0,0 @@ ---- -slug: api-reference/rust-api-cuvs-cagra ---- - -# Cagra Module - -_Rust module: `cuvs::cagra`_ - -_Source: `rust/cuvs/src/cagra/mod.rs`_ - -CAGRA is a graph-based nearest neighbors implementation with state-of-the art -query performance for both small- and large-batch sized search. - -Example: -``` - -use cuvs::cagra::{Index, IndexParams, SearchParams}; -use cuvs::{ManagedTensor, Resources, Result}; - -use ndarray::s; -use ndarray_rand::rand_distr::Uniform; -use ndarray_rand::RandomExt; - -fn cagra_example() -> Result<()> { -let res = Resources::new()?; - -// Create a new random dataset to index -let n_datapoints = 65536; -let n_features = 512; -let dataset = -ndarray::Array::::random((n_datapoints, n_features), Uniform::new(0., 1.0)); - -// build the cagra index -let build_params = IndexParams::new()?; -let index = Index::build(&res, &build_params, &dataset)?; -println!( -"Indexed {}x{} datapoints into cagra index", -n_datapoints, n_features -); - -// 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, ..]); - -let k = 10; - -// CAGRA 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 = ManagedTensor::from(&queries).to_device(&res)?; -let mut neighbors_host = ndarray::Array::::zeros((n_queries, k)); -let neighbors = ManagedTensor::from(&neighbors_host).to_device(&res)?; - -let mut distances_host = ndarray::Array::::zeros((n_queries, k)); -let distances = ManagedTensor::from(&distances_host).to_device(&res)?; - -let search_params = SearchParams::new()?; - -index.search(&res, &search_params, &queries, &neighbors, &distances)?; - -// Copy back to host memory -distances.to_host(&res, &mut distances_host)?; -neighbors.to_host(&res, &mut neighbors_host)?; - -// nearest neighbors should be themselves, since queries are from the -// dataset -println!("Neighbors {:?}", neighbors_host); -println!("Distances {:?}", distances_host); -Ok(()) -} -``` - -Serialization example: -```no_run -use cuvs::cagra::{Index, IndexParams}; -use cuvs::{Resources, Result}; - -fn serialize_example() -> Result<()> { -let res = Resources::new()?; - -// Build an index (using some dataset) -let build_params = IndexParams::new()?; -// let index = Index::build(&res, &build_params, &dataset)?; - -// Save the index to disk (including the dataset) -// index.serialize(&res, "/path/to/index.bin", true)?; - -// Later, load the index from disk -let loaded_index = Index::deserialize(&res, "/path/to/index.bin")?; - -// The loaded index can be used for search just like the original -Ok(()) -} -``` - -## index::Index - -```rust -pub use index::Index; -``` - -_Source: `rust/cuvs/src/cagra/mod.rs:96`_ - -## index_params::\{BuildAlgo, CompressionParams, IndexParams\} - -```rust -pub use index_params::{BuildAlgo, CompressionParams, IndexParams}; -``` - -_Source: `rust/cuvs/src/cagra/mod.rs:97`_ - -## search_params::\{HashMode, SearchAlgo, SearchParams\} - -```rust -pub use search_params::{HashMode, SearchAlgo, SearchParams}; -``` - -_Source: `rust/cuvs/src/cagra/mod.rs:98`_ diff --git a/fern/pages/rust_api/rust-api-cuvs-cluster-kmeans-params.md b/fern/pages/rust_api/rust-api-cuvs-cluster-kmeans-params.md index ee31aca4ae..74230aad87 100644 --- a/fern/pages/rust_api/rust-api-cuvs-cluster-kmeans-params.md +++ b/fern/pages/rust_api/rust-api-cuvs-cluster-kmeans-params.md @@ -8,136 +8,57 @@ _Rust module: `cuvs::cluster::kmeans::params`_ _Source: `rust/cuvs/src/cluster/kmeans/params.rs`_ +Builder-pattern parameter type for k-means. + +All setters are optional; unset values retain the library defaults from the +underlying C `cuvsKMeansParamsCreate`. + ## Params ```rust -pub struct Params(pub ffi::cuvsKMeansParams_t); { +pub struct Params { /* private fields */ } ``` +Parameters for k-means fitting and prediction. + **Methods** | Name | Source | | --- | --- | -| `new` | `rust/cuvs/src/cluster/kmeans/params.rs:15` | -| `set_metric` | `rust/cuvs/src/cluster/kmeans/params.rs:24` | -| `set_n_clusters` | `rust/cuvs/src/cluster/kmeans/params.rs:32` | -| `set_max_iter` | `rust/cuvs/src/cluster/kmeans/params.rs:40` | -| `set_tol` | `rust/cuvs/src/cluster/kmeans/params.rs:48` | -| `set_n_init` | `rust/cuvs/src/cluster/kmeans/params.rs:56` | -| `set_oversampling_factor` | `rust/cuvs/src/cluster/kmeans/params.rs:64` | -| `set_batch_samples` | `rust/cuvs/src/cluster/kmeans/params.rs:77` | -| `set_batch_centroids` | `rust/cuvs/src/cluster/kmeans/params.rs:84` | -| `set_hierarchical` | `rust/cuvs/src/cluster/kmeans/params.rs:92` | -| `set_hierarchical_n_iters` | `rust/cuvs/src/cluster/kmeans/params.rs:100` | +| `new` | `rust/cuvs/src/cluster/kmeans/params.rs:29` | +| `try_new` | `rust/cuvs/src/cluster/kmeans/params.rs:80` | ### new ```rust -pub fn new() -> Result -``` - -Returns a new Params - -_Source: `rust/cuvs/src/cluster/kmeans/params.rs:15`_ - -### set_metric - -```rust -pub fn set_metric(self, metric: DistanceType) -> Params -``` - -DistanceType to use for fitting kmeans - -_Source: `rust/cuvs/src/cluster/kmeans/params.rs:24`_ - -### set_n_clusters - -```rust -pub fn set_n_clusters(self, n_clusters: i32) -> 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 ``` -The number of clusters to form as well as the number of centroids to generate (default:8). - -_Source: `rust/cuvs/src/cluster/kmeans/params.rs:32`_ - -### set_max_iter - -```rust -pub fn set_max_iter(self, max_iter: i32) -> Params -``` - -Maximum number of iterations of the k-means algorithm for a single run. - -_Source: `rust/cuvs/src/cluster/kmeans/params.rs:40`_ - -### set_tol - -```rust -pub fn set_tol(self, tol: f64) -> Params -``` - -Relative tolerance with regards to inertia to declare convergence. - -_Source: `rust/cuvs/src/cluster/kmeans/params.rs:48`_ - -### set_n_init - -```rust -pub fn set_n_init(self, n_init: i32) -> Params -``` - -Number of instance k-means algorithm will be run with different seeds. - -_Source: `rust/cuvs/src/cluster/kmeans/params.rs:56`_ - -### set_oversampling_factor - -```rust -pub fn set_oversampling_factor(self, oversampling_factor: f64) -> Params -``` - -Oversampling factor for use in the k-means\|\| algorithm - -_Source: `rust/cuvs/src/cluster/kmeans/params.rs:64`_ - -### set_batch_samples - -```rust -pub fn set_batch_samples(self, batch_samples: i32) -> Params -``` - -_Source: `rust/cuvs/src/cluster/kmeans/params.rs:77`_ - -### set_batch_centroids - -```rust -pub fn set_batch_centroids(self, batch_centroids: i32) -> Params -``` - -if 0 then batch_centroids = n_clusters - -_Source: `rust/cuvs/src/cluster/kmeans/params.rs:84`_ - -### set_hierarchical - -```rust -pub fn set_hierarchical(self, hierarchical: bool) -> Params -``` - -Whether to use hierarchical (balanced) kmeans or not - -_Source: `rust/cuvs/src/cluster/kmeans/params.rs:92`_ +_Source: `rust/cuvs/src/cluster/kmeans/params.rs:29`_ -### set_hierarchical_n_iters +### try_new ```rust -pub fn set_hierarchical_n_iters(self, hierarchical_n_iters: i32) -> Params +pub fn try_new() -> Result ``` -For hierarchical k-means , defines the number of training iterations +Allocate parameters populated with the library defaults. -_Source: `rust/cuvs/src/cluster/kmeans/params.rs:100`_ +_Source: `rust/cuvs/src/cluster/kmeans/params.rs:80`_ -_Source: `rust/cuvs/src/cluster/kmeans/params.rs:11`_ +_Source: `rust/cuvs/src/cluster/kmeans/params.rs:21`_ diff --git a/fern/pages/rust_api/rust-api-cuvs-cluster-kmeans.md b/fern/pages/rust_api/rust-api-cuvs-cluster-kmeans.md index 372a192ac9..742843515f 100644 --- a/fern/pages/rust_api/rust-api-cuvs-cluster-kmeans.md +++ b/fern/pages/rust_api/rust-api-cuvs-cluster-kmeans.md @@ -8,110 +8,101 @@ _Rust module: `cuvs::cluster::kmeans`_ _Source: `rust/cuvs/src/cluster/kmeans/mod.rs`_ -Kmeans clustering API's +K-means clustering. -Example: -``` - -use cuvs::cluster::kmeans; -use cuvs::{ManagedTensor, Resources, Result}; +[`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. -use ndarray_rand::rand_distr::Uniform; -use ndarray_rand::RandomExt; +## params::Params -fn kmeans_example() -> Result<()> { -let res = Resources::new()?; +```rust +pub use params::Params; +``` -// Create a new random dataset to index -let n_datapoints = 65536; -let n_features = 512; -let n_clusters = 8; -let dataset = -ndarray::Array::::random((n_datapoints, n_features), Uniform::new(0., 1.0)); -let dataset = ManagedTensor::from(&dataset).to_device(&res)?; +_Source: `rust/cuvs/src/cluster/kmeans/mod.rs:16`_ -let centroids_host = ndarray::Array::::zeros((n_clusters, n_features)); -let mut centroids = ManagedTensor::from(¢roids_host).to_device(&res)?; +## KMeansError -// find the centroids with the kmeans index -let kmeans_params = kmeans::Params::new()?.set_n_clusters(n_clusters as i32); -let (inertia, n_iter) = kmeans::fit(&res, &kmeans_params, &dataset, &None, &mut centroids)?; -Ok(()) +```rust +#[derive(Debug, thiserror::Error)] +pub enum KMeansError { + /* variants omitted */ } ``` -## params::Params +Error type for k-means operations. -```rust -pub use params::Params; -``` - -_Source: `rust/cuvs/src/cluster/kmeans/mod.rs:40`_ +_Source: `rust/cuvs/src/cluster/kmeans/mod.rs:26`_ ## fit ```rust -pub fn fit( +pub fn fit( res: &Resources, params: &Params, -x: &ManagedTensor, -sample_weight: &Option, -centroids: &mut ManagedTensor, +x: &X, +sample_weight: Option<&W>, +centroids: &mut C, ) -> Result<(f64, i32)> +where +X: AsDlTensor + ?Sized, +W: AsDlTensor + ?Sized, +C: AsDlTensorMut + ?Sized, ``` -Find clusters with the k-means algorithm - -#### Arguments +Fits k-means centroids to `x`, returning `(inertia, n_iterations)`. -* `res` - Resources to use -* `params` - Parameters to use to fit KMeans model -* `x` - A matrix in device memory - shape (m, k) -* `sample_weight` - Optional device matrix shape (n_clusters, 1) -* `centroids` - Output device matrix, that has the centroids for each cluster -shape (n_clusters, k) +`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`]. -_Source: `rust/cuvs/src/cluster/kmeans/mod.rs:56`_ +_Source: `rust/cuvs/src/cluster/kmeans/mod.rs:41`_ ## predict ```rust -pub fn predict( +pub fn predict( res: &Resources, params: &Params, -x: &ManagedTensor, -sample_weight: &Option, -centroids: &ManagedTensor, -labels: &mut ManagedTensor, +x: &X, +sample_weight: Option<&W>, +centroids: &C, +labels: &mut L, normalize_weight: bool, ) -> Result +where +X: AsDlTensor + ?Sized, +W: AsDlTensor + ?Sized, +C: AsDlTensor + ?Sized, +L: AsDlTensorMut + ?Sized, ``` -Predict clusters with the k-means algorithm - -#### Arguments +Assigns each row of `x` to its nearest centroid, writing cluster labels into +`labels` and returning the inertia. -* `res` - Resources to use -* `params` - Parameters to use to fit KMeans model -* `x` - Input matrix in device memory - shape (m, k) -* `sample_weight` - Optional device matrix shape (n_clusters, 1) -* `centroids` - Centroids calculated by fit in device memory, shape (n_clusters, k) -* `labels` - preallocated CUDA array interface matrix shape (m, 1) to hold the output labels -* `normalize_weight` - whether or not to normalize the weights +`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 +whether the sample weights are normalized. -_Source: `rust/cuvs/src/cluster/kmeans/mod.rs:95`_ +_Source: `rust/cuvs/src/cluster/kmeans/mod.rs:83`_ ## cluster_cost ```rust -pub fn cluster_cost(res: &Resources, x: &ManagedTensor, centroids: &ManagedTensor) -> Result +pub fn cluster_cost(res: &Resources, x: &X, centroids: &C) -> Result +where +X: AsDlTensor + ?Sized, +C: AsDlTensor + ?Sized, ``` -Compute cluster cost given an input matrix and existing centroids -#### Arguments +Computes the k-means cost (inertia) of `x` against existing `centroids`. -* `res` - Resources to use -* `x` - Input matrix in device memory - shape (m, k) -* `centroids` - Centroids calculated by fit in device memory, shape (n_clusters, k) +`x` (shape `m × k`) and `centroids` (shape `n_clusters × k`) reside in device +memory and implement [`AsDlTensor`]. -_Source: `rust/cuvs/src/cluster/kmeans/mod.rs:131`_ +_Source: `rust/cuvs/src/cluster/kmeans/mod.rs:126`_ diff --git a/fern/pages/rust_api/rust-api-cuvs-distance-type.md b/fern/pages/rust_api/rust-api-cuvs-distance-type.md deleted file mode 100644 index e6750a0d02..0000000000 --- a/fern/pages/rust_api/rust-api-cuvs-distance-type.md +++ /dev/null @@ -1,17 +0,0 @@ ---- -slug: api-reference/rust-api-cuvs-distance-type ---- - -# Distance Type Module - -_Rust module: `cuvs::distance_type`_ - -_Source: `rust/cuvs/src/distance_type.rs`_ - -## DistanceType - -```rust -pub type DistanceType = ffi::cuvsDistanceType; -``` - -_Source: `rust/cuvs/src/distance_type.rs:6`_ diff --git a/fern/pages/rust_api/rust-api-cuvs-distance.md b/fern/pages/rust_api/rust-api-cuvs-distance.md index 89c3ab726f..a513143269 100644 --- a/fern/pages/rust_api/rust-api-cuvs-distance.md +++ b/fern/pages/rust_api/rust-api-cuvs-distance.md @@ -8,28 +8,62 @@ _Rust module: `cuvs::distance`_ _Source: `rust/cuvs/src/distance/mod.rs`_ +Distance metrics and pairwise distance computation. + +[`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. + +## DistanceType + +```rust +#[derive(Debug, Copy, Clone, PartialEq)] +#[non_exhaustive] +pub enum DistanceType { + /* variants omitted */ +} +``` + +Distance metric used for building and searching nearest neighbor indices. + +_Source: `rust/cuvs/src/distance/mod.rs:22`_ + +## DistanceError + +```rust +#[derive(Debug, thiserror::Error)] +pub enum DistanceError { + /* variants omitted */ +} +``` + +Error type for pairwise distance operations. + +_Source: `rust/cuvs/src/distance/mod.rs:142`_ + ## pairwise_distance ```rust -pub fn pairwise_distance( +pub fn pairwise_distance( res: &Resources, -x: &ManagedTensor, -y: &ManagedTensor, -distances: &ManagedTensor, +x: &X, +y: &Y, +distances: &mut D, metric: DistanceType, -metric_arg: Option, -) -> Result<()> +) -> Result<(), DistanceError> +where +X: AsDlTensor + ?Sized, +Y: AsDlTensor + ?Sized, +D: AsDlTensorMut + ?Sized, ``` -Compute pairwise distances between X and Y - -#### Arguments +Computes all pairwise distances between the rows of `x` (shape `m × k`) and +`y` (shape `n × k`), writing the `m × n` result into `distances`. -* `res` - Resources to use -* `x` - A matrix in device memory - shape (m, k) -* `y` - A matrix in device memory - shape (n, k) -* `distances` - A matrix in device memory that receives the output distances - shape (m, n) -* `metric` - DistanceType to use for building the index -* `metric_arg` - Optional value of `p` for Minkowski distances +`x`, `y`, and `distances` reside in device memory and implement +[`AsDlTensor`] / [`AsDlTensorMut`]. `metric` selects the distance; use +[`DistanceType::LpUnexpanded`] to supply the Minkowski exponent `p` (all +other metrics use the C API default). -_Source: `rust/cuvs/src/distance/mod.rs:21`_ +_Source: `rust/cuvs/src/distance/mod.rs:158`_ diff --git a/fern/pages/rust_api/rust-api-cuvs-dlpack.md b/fern/pages/rust_api/rust-api-cuvs-dlpack.md index 6cddcb2596..1847adb60a 100644 --- a/fern/pages/rust_api/rust-api-cuvs-dlpack.md +++ b/fern/pages/rust_api/rust-api-cuvs-dlpack.md @@ -8,71 +8,262 @@ _Rust module: `cuvs::dlpack`_ _Source: `rust/cuvs/src/dlpack.rs`_ -## ManagedTensor +DLPack tensor interop. + +cuVS exchanges tensors with the C library through the DLPack ABI. This crate +never owns tensor storage: every entry point borrows a value that exposes a +view through the [`AsDlTensor`] / [`AsDlTensorMut`] traits: + +* [`DLTensorView`] — a read-only view, for inputs the C API only reads +(datasets, queries). +* [`DLTensorViewMut`] — a writable view, for outputs the C API writes +(neighbors, distances). + +A view is non-owning. The traits take `&self` / `&mut self`, so the compiler +ties the view's lifetime to the borrow of the value that owns the underlying +buffer. The view materializes a stack-local [`DLManagedTensor`] only for the +duration of each FFI call. + +The crate ships no public tensor type. To hand your own GPU (or host) buffer +to cuVS, implement [`AsDlTensor`] / [`AsDlTensorMut`] for it on top of +[`DLTensorView::from_raw_parts`]. Most algorithms require search inputs and +outputs to live in device memory. See `examples/cagra.rs` for a complete, +runnable adapter built on the raw CUDA runtime. + +## ffi::\{DLDataType, DLDataTypeCode, DLDevice, DLDeviceType, DLManagedTensor, DLTensor\} + +```rust +pub use ffi::{DLDataType, DLDataTypeCode, DLDevice, DLDeviceType, DLManagedTensor, DLTensor}; +``` + +_Source: `rust/cuvs/src/dlpack.rs:34`_ + +## AsDlTensor + +```rust +pub trait AsDlTensor { + /* required methods omitted */ +} +``` + +Borrows a tensor as a read-only [`DLTensorView`] for tensor inputs. + +Implement this for your own tensor type by calling +[`DLTensorView::from_raw_parts`] inside a small `unsafe` block, upholding its +safety contract. + +#### Examples + +A minimal adapter for a row-major matrix in device memory: + +``` +use cuvs::dlpack::{AsDlTensor, DLDevice, DLDeviceType, DLPackError, DLTensorView, DType}; + +struct GpuMatrix { +ptr: *mut T, +rows: usize, +cols: usize, +} + +impl AsDlTensor for GpuMatrix { +fn as_dl_tensor(&self) -> Result, DLPackError> { +let shape = [self.rows as i64, self.cols as i64]; +// SAFETY: `ptr` points to `rows * cols` initialized elements of `T` +// in device 0's memory, valid while `self` is borrowed, and is +// row-major contiguous. +unsafe { +DLTensorView::from_raw_parts( +self.ptr.cast(), +DLDevice { device_type: DLDeviceType::kDLCUDA, device_id: 0 }, +&shape, +None, +T::dl_dtype(), +) +} +} +} +``` + +_Source: `rust/cuvs/src/dlpack.rs:79`_ + +## AsDlTensorMut ```rust -#[derive(Debug)] -pub struct ManagedTensor(ffi::DLManagedTensor); { +pub trait AsDlTensorMut { + /* required methods omitted */ +} +``` + +Borrows a tensor as a writable [`DLTensorViewMut`] for tensor outputs. + +In addition to the [`DLTensorView::from_raw_parts`] invariants, writable +adapters must guarantee exclusive access to the data region. The `&mut self` +receiver makes the compiler enforce that exclusivity for the borrow. + +_Source: `rust/cuvs/src/dlpack.rs:88`_ + +## DType + +```rust +pub trait DType { + /* required methods omitted */ +} +``` + +Maps a Rust element type to a DLPack [`DLDataType`]. + +_Source: `rust/cuvs/src/dlpack.rs:93`_ + +## DLPackError + +```rust +#[derive(Debug, Clone, thiserror::Error)] +#[non_exhaustive] +pub enum DLPackError { + /* variants omitted */ +} +``` + +Error when converting an external tensor to a DLPack view. + +_Source: `rust/cuvs/src/dlpack.rs:121`_ + +## DLTensorView + +```rust +#[must_use] +pub struct DLTensorView<'a> { /* private fields */ } ``` -ManagedTensor is a wrapper around a dlpack DLManagedTensor object. -This lets you pass matrices in device or host memory into cuvs. +A non-owning, read-only DLPack tensor view. **Methods** | Name | Source | | --- | --- | -| `as_ptr` | `rust/cuvs/src/dlpack.rs:21` | -| `to_device` | `rust/cuvs/src/dlpack.rs:27` | -| `to_host` | `rust/cuvs/src/dlpack.rs:47` | +| `from_raw_parts` | `rust/cuvs/src/dlpack.rs:176` | +| `ndim` | `rust/cuvs/src/dlpack.rs:226` | +| `shape` | `rust/cuvs/src/dlpack.rs:231` | +| `strides` | `rust/cuvs/src/dlpack.rs:236` | +| `dtype` | `rust/cuvs/src/dlpack.rs:241` | +| `device` | `rust/cuvs/src/dlpack.rs:246` | -### as_ptr +### from_raw_parts ```rust -pub fn as_ptr(&self) -> *mut ffi::DLManagedTensor +pub unsafe fn from_raw_parts( +data: *mut std::ffi::c_void, +device: ffi::DLDevice, +shape: &[i64], +strides: Option<&[i64]>, +dtype: ffi::DLDataType, +) -> std::result::Result ``` -_Source: `rust/cuvs/src/dlpack.rs:21`_ +Construct a DLPack view from raw tensor metadata. + +#### Safety -### to_device +The caller must guarantee that: +- `data` points to initialized storage matching `shape`, `strides`, and +`dtype`, residing on the device described by `device`; +- that storage remains valid for the lifetime `'a`; +- the C API consumes the resulting [`DLManagedTensor`] (including its +`shape`/`strides` pointers) only for the duration of the FFI call and +does not retain it afterward — cuVS upholds this. + +_Source: `rust/cuvs/src/dlpack.rs:176`_ + +### ndim ```rust -pub fn to_device(&self, res: &Resources) -> Result +pub fn ndim(&self) -> usize ``` -Creates a new ManagedTensor on the current GPU device, and copies -the data into it. +Number of dimensions. -_Source: `rust/cuvs/src/dlpack.rs:27`_ +_Source: `rust/cuvs/src/dlpack.rs:226`_ -### to_host +### shape ```rust -pub fn to_host< -T: IntoDtype, -S: ndarray::RawData + ndarray::RawDataMut, -D: ndarray::Dimension, ->( -&self, -res: &Resources, -arr: &mut ndarray::ArrayBase, -) -> Result<()> +pub fn shape(&self) -> &[i64] ``` -Copies data from device memory into host memory +Shape of the tensor. -_Source: `rust/cuvs/src/dlpack.rs:47`_ +_Source: `rust/cuvs/src/dlpack.rs:231`_ -_Source: `rust/cuvs/src/dlpack.rs:14`_ +### strides -## IntoDtype +```rust +pub fn strides(&self) -> Option<&[i64]> +``` + +Strides, if non-contiguous. `None` means row-major contiguous. + +_Source: `rust/cuvs/src/dlpack.rs:236`_ + +### dtype ```rust -pub trait IntoDtype { - /* required methods omitted */ +pub fn dtype(&self) -> ffi::DLDataType +``` + +Element data type. + +_Source: `rust/cuvs/src/dlpack.rs:241`_ + +### device + +```rust +pub fn device(&self) -> ffi::DLDevice +``` + +Device where the data resides. + +_Source: `rust/cuvs/src/dlpack.rs:246`_ + +_Source: `rust/cuvs/src/dlpack.rs:155`_ + +## DLTensorViewMut + +```rust +#[must_use] +pub struct DLTensorViewMut<'a> { + /* private fields */ } ``` -_Source: `rust/cuvs/src/dlpack.rs:16`_ +A non-owning, writable DLPack tensor view. + +**Methods** + +| Name | Source | +| --- | --- | +| `from_raw_parts` | `rust/cuvs/src/dlpack.rs:274` | + +### from_raw_parts + +```rust +pub unsafe fn from_raw_parts( +data: *mut std::ffi::c_void, +device: ffi::DLDevice, +shape: &[i64], +strides: Option<&[i64]>, +dtype: ffi::DLDataType, +) -> std::result::Result +``` + +Construct a writable DLPack view from raw tensor metadata. + +#### Safety + +In addition to the [`DLTensorView::from_raw_parts`] invariants, the +caller must guarantee the storage is exclusively writable for `'a`. + +_Source: `rust/cuvs/src/dlpack.rs:274`_ + +_Source: `rust/cuvs/src/dlpack.rs:262`_ diff --git a/fern/pages/rust_api/rust-api-cuvs-error.md b/fern/pages/rust_api/rust-api-cuvs-error.md index d8dcd92541..297bd35036 100644 --- a/fern/pages/rust_api/rust-api-cuvs-error.md +++ b/fern/pages/rust_api/rust-api-cuvs-error.md @@ -8,42 +8,24 @@ _Rust module: `cuvs::error`_ _Source: `rust/cuvs/src/error.rs`_ -## CuvsError +Low-level error handling shared by every cuVS module. -```rust -#[derive(Debug, Clone)] -pub struct CuvsError { - /* private fields */ -} -``` +[`check_cuvs`] turns a raw `cuvsError_t` status into a [`LibraryError`], which +each module's error type wraps via `#[from]`. -_Source: `rust/cuvs/src/error.rs:9`_ - -## Error +## LibraryError ```rust -#[derive(Debug, Clone)] -pub enum Error { - /* variants omitted */ +#[derive(Debug, Clone, thiserror::Error)] +#[error("{0}")] +pub struct LibraryError(Cow<'static, str>); { + /* private fields */ } ``` -_Source: `rust/cuvs/src/error.rs:15`_ - -## Result - -```rust -pub type Result = std::result::Result; -``` - -_Source: `rust/cuvs/src/error.rs:25`_ - -## check_cuvs - -```rust -pub fn check_cuvs(err: ffi::cuvsError_t) -> Result<()> -``` +A failure reported by the cuVS C library. -Simple wrapper to convert a cuvsError_t into a Result +Carries the message captured from `cuvsGetLastErrorText` at the point of +failure. Every module's error type wraps this via `#[from]`. -_Source: `rust/cuvs/src/error.rs:43`_ +_Source: `rust/cuvs/src/error.rs:19`_ diff --git a/fern/pages/rust_api/rust-api-cuvs-ivf-flat-index-params.md b/fern/pages/rust_api/rust-api-cuvs-ivf-flat-index-params.md deleted file mode 100644 index 35e0ca92a7..0000000000 --- a/fern/pages/rust_api/rust-api-cuvs-ivf-flat-index-params.md +++ /dev/null @@ -1,106 +0,0 @@ ---- -slug: api-reference/rust-api-cuvs-ivf-flat-index-params ---- - -# Ivf Flat Index Params Module - -_Rust module: `cuvs::ivf_flat::index_params`_ - -_Source: `rust/cuvs/src/ivf_flat/index_params.rs`_ - -## IndexParams - -```rust -pub struct IndexParams(pub ffi::cuvsIvfFlatIndexParams_t); { - /* private fields */ -} -``` - -**Methods** - -| Name | Source | -| --- | --- | -| `new` | `rust/cuvs/src/ivf_flat/index_params.rs:15` | -| `set_n_lists` | `rust/cuvs/src/ivf_flat/index_params.rs:24` | -| `set_metric` | `rust/cuvs/src/ivf_flat/index_params.rs:32` | -| `set_metric_arg` | `rust/cuvs/src/ivf_flat/index_params.rs:40` | -| `set_kmeans_n_iters` | `rust/cuvs/src/ivf_flat/index_params.rs:47` | -| `set_kmeans_trainset_fraction` | `rust/cuvs/src/ivf_flat/index_params.rs:57` | -| `set_add_data_on_build` | `rust/cuvs/src/ivf_flat/index_params.rs:68` | - -### new - -```rust -pub fn new() -> Result -``` - -Returns a new IndexParams - -_Source: `rust/cuvs/src/ivf_flat/index_params.rs:15`_ - -### set_n_lists - -```rust -pub fn set_n_lists(self, n_lists: u32) -> IndexParams -``` - -The number of clusters used in the coarse quantizer. - -_Source: `rust/cuvs/src/ivf_flat/index_params.rs:24`_ - -### set_metric - -```rust -pub fn set_metric(self, metric: DistanceType) -> IndexParams -``` - -DistanceType to use for building the index - -_Source: `rust/cuvs/src/ivf_flat/index_params.rs:32`_ - -### set_metric_arg - -```rust -pub fn set_metric_arg(self, metric_arg: f32) -> IndexParams -``` - -The number of iterations searching for kmeans centers during index building. - -_Source: `rust/cuvs/src/ivf_flat/index_params.rs:40`_ - -### set_kmeans_n_iters - -```rust -pub fn set_kmeans_n_iters(self, kmeans_n_iters: u32) -> IndexParams -``` - -The number of iterations searching for kmeans centers during index building. - -_Source: `rust/cuvs/src/ivf_flat/index_params.rs:47`_ - -### set_kmeans_trainset_fraction - -```rust -pub fn set_kmeans_trainset_fraction(self, kmeans_trainset_fraction: f64) -> IndexParams -``` - -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. - -_Source: `rust/cuvs/src/ivf_flat/index_params.rs:57`_ - -### set_add_data_on_build - -```rust -pub fn set_add_data_on_build(self, add_data_on_build: bool) -> IndexParams -``` - -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. - -_Source: `rust/cuvs/src/ivf_flat/index_params.rs:68`_ - -_Source: `rust/cuvs/src/ivf_flat/index_params.rs:11`_ diff --git a/fern/pages/rust_api/rust-api-cuvs-ivf-flat-index.md b/fern/pages/rust_api/rust-api-cuvs-ivf-flat-index.md deleted file mode 100644 index 5f985d092a..0000000000 --- a/fern/pages/rust_api/rust-api-cuvs-ivf-flat-index.md +++ /dev/null @@ -1,85 +0,0 @@ ---- -slug: api-reference/rust-api-cuvs-ivf-flat-index ---- - -# Ivf Flat Index Module - -_Rust module: `cuvs::ivf_flat::index`_ - -_Source: `rust/cuvs/src/ivf_flat/index.rs`_ - -## Index - -```rust -#[derive(Debug)] -pub struct Index(ffi::cuvsIvfFlatIndex_t); { - /* private fields */ -} -``` - -Ivf-Flat ANN Index - -**Methods** - -| Name | Source | -| --- | --- | -| `build` | `rust/cuvs/src/ivf_flat/index.rs:25` | -| `new` | `rust/cuvs/src/ivf_flat/index.rs:39` | -| `search` | `rust/cuvs/src/ivf_flat/index.rs:56` | - -### build - -```rust -pub fn build>( -res: &Resources, -params: &IndexParams, -dataset: T, -) -> Result -``` - -Builds a new Index from the dataset for efficient search. - -#### Arguments - -* `res` - Resources to use -* `params` - Parameters for building the index -* `dataset` - A row-major matrix on either the host or device to index - -_Source: `rust/cuvs/src/ivf_flat/index.rs:25`_ - -### new - -```rust -pub fn new() -> Result -``` - -Creates a new empty index - -_Source: `rust/cuvs/src/ivf_flat/index.rs:39`_ - -### search - -```rust -pub fn search( -&self, -res: &Resources, -params: &SearchParams, -queries: &ManagedTensor, -neighbors: &ManagedTensor, -distances: &ManagedTensor, -) -> Result<()> -``` - -Perform a Approximate Nearest Neighbors search on the Index - -#### Arguments - -* `res` - Resources to use -* `params` - Parameters to use in searching the index -* `queries` - A matrix in device memory to query for -* `neighbors` - Matrix in device memory that receives the indices of the nearest neighbors -* `distances` - Matrix in device memory that receives the distances of the nearest neighbors - -_Source: `rust/cuvs/src/ivf_flat/index.rs:56`_ - -_Source: `rust/cuvs/src/ivf_flat/index.rs:15`_ diff --git a/fern/pages/rust_api/rust-api-cuvs-ivf-flat-search-params.md b/fern/pages/rust_api/rust-api-cuvs-ivf-flat-search-params.md deleted file mode 100644 index 7c6ace29b8..0000000000 --- a/fern/pages/rust_api/rust-api-cuvs-ivf-flat-search-params.md +++ /dev/null @@ -1,48 +0,0 @@ ---- -slug: api-reference/rust-api-cuvs-ivf-flat-search-params ---- - -# Ivf Flat Search Params Module - -_Rust module: `cuvs::ivf_flat::search_params`_ - -_Source: `rust/cuvs/src/ivf_flat/search_params.rs`_ - -## SearchParams - -```rust -pub struct SearchParams(pub ffi::cuvsIvfFlatSearchParams_t); { - /* private fields */ -} -``` - -Supplemental parameters to search IvfFlat index - -**Methods** - -| Name | Source | -| --- | --- | -| `new` | `rust/cuvs/src/ivf_flat/search_params.rs:15` | -| `set_n_probes` | `rust/cuvs/src/ivf_flat/search_params.rs:24` | - -### new - -```rust -pub fn new() -> Result -``` - -Returns a new SearchParams object - -_Source: `rust/cuvs/src/ivf_flat/search_params.rs:15`_ - -### set_n_probes - -```rust -pub fn set_n_probes(self, n_probes: u32) -> SearchParams -``` - -Supplemental parameters to search IVF-Flat index - -_Source: `rust/cuvs/src/ivf_flat/search_params.rs:24`_ - -_Source: `rust/cuvs/src/ivf_flat/search_params.rs:11`_ diff --git a/fern/pages/rust_api/rust-api-cuvs-ivf-flat.md b/fern/pages/rust_api/rust-api-cuvs-ivf-flat.md deleted file mode 100644 index b34e3ba369..0000000000 --- a/fern/pages/rust_api/rust-api-cuvs-ivf-flat.md +++ /dev/null @@ -1,97 +0,0 @@ ---- -slug: api-reference/rust-api-cuvs-ivf-flat ---- - -# Ivf Flat Module - -_Rust module: `cuvs::ivf_flat`_ - -_Source: `rust/cuvs/src/ivf_flat/mod.rs`_ - -The IVF-Flat method is an ANN algorithm. It uses an inverted file index (IVF) with -unmodified (that is, flat) vectors. This algorithm provides simple knobs to reduce -the overall search space and to trade-off accuracy for speed. - -Example: -``` - -use cuvs::ivf_flat::{Index, IndexParams, SearchParams}; -use cuvs::{ManagedTensor, Resources, Result}; - -use ndarray::s; -use ndarray_rand::rand_distr::Uniform; -use ndarray_rand::RandomExt; - -fn ivf_flat_example() -> Result<()> { -let res = Resources::new()?; - -// Create a new random dataset to index -let n_datapoints = 65536; -let n_features = 512; -let dataset = -ndarray::Array::::random((n_datapoints, n_features), Uniform::new(0., 1.0)); - -// build the ivf-flat index -let build_params = IndexParams::new()?; -let index = Index::build(&res, &build_params, &dataset)?; -println!( -"Indexed {}x{} datapoints into ivf-flat index", -n_datapoints, n_features -); - -// 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, ..]); - -let k = 10; - -// Ivf-Flat 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 = ManagedTensor::from(&queries).to_device(&res)?; -let mut neighbors_host = ndarray::Array::::zeros((n_queries, k)); -let neighbors = ManagedTensor::from(&neighbors_host).to_device(&res)?; - -let mut distances_host = ndarray::Array::::zeros((n_queries, k)); -let distances = ManagedTensor::from(&distances_host).to_device(&res)?; - -let search_params = SearchParams::new()?; - -index.search(&res, &search_params, &queries, &neighbors, &distances)?; - -// Copy back to host memory -distances.to_host(&res, &mut distances_host)?; -neighbors.to_host(&res, &mut neighbors_host)?; - -// nearest neighbors should be themselves, since queries are from the -// dataset -println!("Neighbors {:?}", neighbors_host); -println!("Distances {:?}", distances_host); -Ok(()) -} -``` - -## index::Index - -```rust -pub use index::Index; -``` - -_Source: `rust/cuvs/src/ivf_flat/mod.rs:74`_ - -## index_params::IndexParams - -```rust -pub use index_params::IndexParams; -``` - -_Source: `rust/cuvs/src/ivf_flat/mod.rs:75`_ - -## search_params::SearchParams - -```rust -pub use search_params::SearchParams; -``` - -_Source: `rust/cuvs/src/ivf_flat/mod.rs:76`_ diff --git a/fern/pages/rust_api/rust-api-cuvs-ivf-pq-index-params.md b/fern/pages/rust_api/rust-api-cuvs-ivf-pq-index-params.md deleted file mode 100644 index ba32c97ddb..0000000000 --- a/fern/pages/rust_api/rust-api-cuvs-ivf-pq-index-params.md +++ /dev/null @@ -1,211 +0,0 @@ ---- -slug: api-reference/rust-api-cuvs-ivf-pq-index-params ---- - -# Ivf Pq Index Params Module - -_Rust module: `cuvs::ivf_pq::index_params`_ - -_Source: `rust/cuvs/src/ivf_pq/index_params.rs`_ - -## ffi::cuvsIvfPqCodebookGen - -```rust -pub use ffi::cuvsIvfPqCodebookGen; -``` - -_Source: `rust/cuvs/src/ivf_pq/index_params.rs:11`_ - -## ffi::cuvsIvfPqListLayout - -```rust -pub use ffi::cuvsIvfPqListLayout; -``` - -_Source: `rust/cuvs/src/ivf_pq/index_params.rs:12`_ - -## IndexParams - -```rust -pub struct IndexParams(pub ffi::cuvsIvfPqIndexParams_t); { - /* private fields */ -} -``` - -**Methods** - -| Name | Source | -| --- | --- | -| `new` | `rust/cuvs/src/ivf_pq/index_params.rs:18` | -| `set_n_lists` | `rust/cuvs/src/ivf_pq/index_params.rs:27` | -| `set_metric` | `rust/cuvs/src/ivf_pq/index_params.rs:35` | -| `set_metric_arg` | `rust/cuvs/src/ivf_pq/index_params.rs:43` | -| `set_kmeans_n_iters` | `rust/cuvs/src/ivf_pq/index_params.rs:51` | -| `set_kmeans_trainset_fraction` | `rust/cuvs/src/ivf_pq/index_params.rs:61` | -| `set_pq_bits` | `rust/cuvs/src/ivf_pq/index_params.rs:69` | -| `set_pq_dim` | `rust/cuvs/src/ivf_pq/index_params.rs:85` | -| `set_codebook_kind` | `rust/cuvs/src/ivf_pq/index_params.rs:92` | -| `set_codes_layout` | `rust/cuvs/src/ivf_pq/index_params.rs:103` | -| `set_force_random_rotation` | `rust/cuvs/src/ivf_pq/index_params.rs:121` | -| `set_max_train_points_per_pq_code` | `rust/cuvs/src/ivf_pq/index_params.rs:133` | -| `set_add_data_on_build` | `rust/cuvs/src/ivf_pq/index_params.rs:144` | - -### new - -```rust -pub fn new() -> Result -``` - -Returns a new IndexParams - -_Source: `rust/cuvs/src/ivf_pq/index_params.rs:18`_ - -### set_n_lists - -```rust -pub fn set_n_lists(self, n_lists: u32) -> IndexParams -``` - -The number of clusters used in the coarse quantizer. - -_Source: `rust/cuvs/src/ivf_pq/index_params.rs:27`_ - -### set_metric - -```rust -pub fn set_metric(self, metric: DistanceType) -> IndexParams -``` - -DistanceType to use for building the index - -_Source: `rust/cuvs/src/ivf_pq/index_params.rs:35`_ - -### set_metric_arg - -```rust -pub fn set_metric_arg(self, metric_arg: f32) -> IndexParams -``` - -The number of iterations searching for kmeans centers during index building. - -_Source: `rust/cuvs/src/ivf_pq/index_params.rs:43`_ - -### set_kmeans_n_iters - -```rust -pub fn set_kmeans_n_iters(self, kmeans_n_iters: u32) -> IndexParams -``` - -The number of iterations searching for kmeans centers during index building. - -_Source: `rust/cuvs/src/ivf_pq/index_params.rs:51`_ - -### set_kmeans_trainset_fraction - -```rust -pub fn set_kmeans_trainset_fraction(self, kmeans_trainset_fraction: f64) -> IndexParams -``` - -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. - -_Source: `rust/cuvs/src/ivf_pq/index_params.rs:61`_ - -### set_pq_bits - -```rust -pub fn set_pq_bits(self, pq_bits: u32) -> IndexParams -``` - -The bit length of the vector element after quantization. - -_Source: `rust/cuvs/src/ivf_pq/index_params.rs:69`_ - -### set_pq_dim - -```rust -pub fn set_pq_dim(self, pq_dim: u32) -> IndexParams -``` - -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. - -_Source: `rust/cuvs/src/ivf_pq/index_params.rs:85`_ - -### set_codebook_kind - -```rust -pub fn set_codebook_kind(self, codebook_kind: cuvsIvfPqCodebookGen) -> IndexParams -``` - -_Source: `rust/cuvs/src/ivf_pq/index_params.rs:92`_ - -### set_codes_layout - -```rust -pub fn set_codes_layout(self, codes_layout: cuvsIvfPqListLayout) -> IndexParams -``` - -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. - -_Source: `rust/cuvs/src/ivf_pq/index_params.rs:103`_ - -### set_force_random_rotation - -```rust -pub fn set_force_random_rotation(self, force_random_rotation: bool) -> IndexParams -``` - -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 - -_Source: `rust/cuvs/src/ivf_pq/index_params.rs:121`_ - -### set_max_train_points_per_pq_code - -```rust -pub fn set_max_train_points_per_pq_code(self, max_pq_points: u32) -> IndexParams -``` - -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. - -_Source: `rust/cuvs/src/ivf_pq/index_params.rs:133`_ - -### set_add_data_on_build - -```rust -pub fn set_add_data_on_build(self, add_data_on_build: bool) -> IndexParams -``` - -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. - -_Source: `rust/cuvs/src/ivf_pq/index_params.rs:144`_ - -_Source: `rust/cuvs/src/ivf_pq/index_params.rs:14`_ diff --git a/fern/pages/rust_api/rust-api-cuvs-ivf-pq-index.md b/fern/pages/rust_api/rust-api-cuvs-ivf-pq-index.md deleted file mode 100644 index b72bada9da..0000000000 --- a/fern/pages/rust_api/rust-api-cuvs-ivf-pq-index.md +++ /dev/null @@ -1,85 +0,0 @@ ---- -slug: api-reference/rust-api-cuvs-ivf-pq-index ---- - -# Ivf Pq Index Module - -_Rust module: `cuvs::ivf_pq::index`_ - -_Source: `rust/cuvs/src/ivf_pq/index.rs`_ - -## Index - -```rust -#[derive(Debug)] -pub struct Index(ffi::cuvsIvfPqIndex_t); { - /* private fields */ -} -``` - -Ivf-Pq ANN Index - -**Methods** - -| Name | Source | -| --- | --- | -| `build` | `rust/cuvs/src/ivf_pq/index.rs:25` | -| `new` | `rust/cuvs/src/ivf_pq/index.rs:39` | -| `search` | `rust/cuvs/src/ivf_pq/index.rs:56` | - -### build - -```rust -pub fn build>( -res: &Resources, -params: &IndexParams, -dataset: T, -) -> Result -``` - -Builds a new Index from the dataset for efficient search. - -#### Arguments - -* `res` - Resources to use -* `params` - Parameters for building the index -* `dataset` - A row-major matrix on either the host or device to index - -_Source: `rust/cuvs/src/ivf_pq/index.rs:25`_ - -### new - -```rust -pub fn new() -> Result -``` - -Creates a new empty index - -_Source: `rust/cuvs/src/ivf_pq/index.rs:39`_ - -### search - -```rust -pub fn search( -&self, -res: &Resources, -params: &SearchParams, -queries: &ManagedTensor, -neighbors: &ManagedTensor, -distances: &ManagedTensor, -) -> Result<()> -``` - -Perform a Approximate Nearest Neighbors search on the Index - -#### Arguments - -* `res` - Resources to use -* `params` - Parameters to use in searching the index -* `queries` - A matrix in device memory to query for -* `neighbors` - Matrix in device memory that receives the indices of the nearest neighbors -* `distances` - Matrix in device memory that receives the distances of the nearest neighbors - -_Source: `rust/cuvs/src/ivf_pq/index.rs:56`_ - -_Source: `rust/cuvs/src/ivf_pq/index.rs:15`_ diff --git a/fern/pages/rust_api/rust-api-cuvs-ivf-pq-search-params.md b/fern/pages/rust_api/rust-api-cuvs-ivf-pq-search-params.md deleted file mode 100644 index dbc2ce1b41..0000000000 --- a/fern/pages/rust_api/rust-api-cuvs-ivf-pq-search-params.md +++ /dev/null @@ -1,86 +0,0 @@ ---- -slug: api-reference/rust-api-cuvs-ivf-pq-search-params ---- - -# Ivf Pq Search Params Module - -_Rust module: `cuvs::ivf_pq::search_params`_ - -_Source: `rust/cuvs/src/ivf_pq/search_params.rs`_ - -## ffi::cudaDataType_t - -```rust -pub use ffi::cudaDataType_t; -``` - -_Source: `rust/cuvs/src/ivf_pq/search_params.rs:10`_ - -## SearchParams - -```rust -pub struct SearchParams(pub ffi::cuvsIvfPqSearchParams_t); { - /* private fields */ -} -``` - -Supplemental parameters to search IvfPq index - -**Methods** - -| Name | Source | -| --- | --- | -| `new` | `rust/cuvs/src/ivf_pq/search_params.rs:17` | -| `set_n_probes` | `rust/cuvs/src/ivf_pq/search_params.rs:26` | -| `set_lut_dtype` | `rust/cuvs/src/ivf_pq/search_params.rs:39` | -| `set_internal_distance_dtype` | `rust/cuvs/src/ivf_pq/search_params.rs:47` | - -### new - -```rust -pub fn new() -> Result -``` - -Returns a new SearchParams object - -_Source: `rust/cuvs/src/ivf_pq/search_params.rs:17`_ - -### set_n_probes - -```rust -pub fn set_n_probes(self, n_probes: u32) -> SearchParams -``` - -The number of clusters to search. - -_Source: `rust/cuvs/src/ivf_pq/search_params.rs:26`_ - -### set_lut_dtype - -```rust -pub fn set_lut_dtype(self, lut_dtype: cudaDataType_t) -> SearchParams -``` - -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. - -_Source: `rust/cuvs/src/ivf_pq/search_params.rs:39`_ - -### set_internal_distance_dtype - -```rust -pub fn set_internal_distance_dtype( -self, -internal_distance_dtype: cudaDataType_t, -) -> SearchParams -``` - -Storage data type for distance/similarity computation. - -_Source: `rust/cuvs/src/ivf_pq/search_params.rs:47`_ - -_Source: `rust/cuvs/src/ivf_pq/search_params.rs:13`_ diff --git a/fern/pages/rust_api/rust-api-cuvs-ivf-pq.md b/fern/pages/rust_api/rust-api-cuvs-ivf-pq.md deleted file mode 100644 index 08c4d2195e..0000000000 --- a/fern/pages/rust_api/rust-api-cuvs-ivf-pq.md +++ /dev/null @@ -1,95 +0,0 @@ ---- -slug: api-reference/rust-api-cuvs-ivf-pq ---- - -# Ivf Pq Module - -_Rust module: `cuvs::ivf_pq`_ - -_Source: `rust/cuvs/src/ivf_pq/mod.rs`_ - -Inverted File Product Quantization - -Example: -``` - -use cuvs::ivf_pq::{Index, IndexParams, SearchParams}; -use cuvs::{ManagedTensor, Resources, Result}; - -use ndarray::s; -use ndarray_rand::rand_distr::Uniform; -use ndarray_rand::RandomExt; - -fn ivf_pq_example() -> Result<()> { -let res = Resources::new()?; - -// Create a new random dataset to index -let n_datapoints = 65536; -let n_features = 512; -let dataset = -ndarray::Array::::random((n_datapoints, n_features), Uniform::new(0., 1.0)); - -// build the ivf-pq index -let build_params = IndexParams::new()?; -let index = Index::build(&res, &build_params, &dataset)?; -println!( -"Indexed {}x{} datapoints into ivf-pq index", -n_datapoints, n_features -); - -// 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, ..]); - -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 = ManagedTensor::from(&queries).to_device(&res)?; -let mut neighbors_host = ndarray::Array::::zeros((n_queries, k)); -let neighbors = ManagedTensor::from(&neighbors_host).to_device(&res)?; - -let mut distances_host = ndarray::Array::::zeros((n_queries, k)); -let distances = ManagedTensor::from(&distances_host).to_device(&res)?; - -let search_params = SearchParams::new()?; - -index.search(&res, &search_params, &queries, &neighbors, &distances)?; - -// Copy back to host memory -distances.to_host(&res, &mut distances_host)?; -neighbors.to_host(&res, &mut neighbors_host)?; - -// nearest neighbors should be themselves, since queries are from the -// dataset -println!("Neighbors {:?}", neighbors_host); -println!("Distances {:?}", distances_host); -Ok(()) -} -``` - -## index::Index - -```rust -pub use index::Index; -``` - -_Source: `rust/cuvs/src/ivf_pq/mod.rs:71`_ - -## index_params::IndexParams - -```rust -pub use index_params::IndexParams; -``` - -_Source: `rust/cuvs/src/ivf_pq/mod.rs:72`_ - -## search_params::SearchParams - -```rust -pub use search_params::SearchParams; -``` - -_Source: `rust/cuvs/src/ivf_pq/mod.rs:73`_ diff --git a/fern/pages/rust_api/rust-api-cuvs-neighbors-brute-force.md b/fern/pages/rust_api/rust-api-cuvs-neighbors-brute-force.md new file mode 100644 index 0000000000..d0e2e8e5ac --- /dev/null +++ b/fern/pages/rust_api/rust-api-cuvs-neighbors-brute-force.md @@ -0,0 +1,104 @@ +--- +slug: api-reference/rust-api-cuvs-neighbors-brute-force +--- + +# Neighbors Brute Force Module + +_Rust module: `cuvs::neighbors::brute_force`_ + +_Source: `rust/cuvs/src/neighbors/brute_force.rs`_ + +Brute-force (exact) k-NN. + +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. + +## BruteForceError + +```rust +#[derive(Debug, thiserror::Error)] +pub enum BruteForceError { + /* variants omitted */ +} +``` + +Error type for brute-force operations. + +_Source: `rust/cuvs/src/neighbors/brute_force.rs:25`_ + +## Index + +```rust +#[derive(Debug)] +pub struct Index<'d> { + /* private fields */ +} +``` + +Brute-force KNN index. + +**Methods** + +| Name | Source | +| --- | --- | +| `build` | `rust/cuvs/src/neighbors/brute_force.rs:51` | +| `new` | `rust/cuvs/src/neighbors/brute_force.rs:70` | +| `search` | `rust/cuvs/src/neighbors/brute_force.rs:84` | + +### build + +```rust +pub fn build(res: &Resources, metric: DistanceType, dataset: &'d T) -> Result> +where +T: AsDlTensor + ?Sized, +``` + +Builds a brute-force index over `dataset` for exact k-NN search. + +`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. + +_Source: `rust/cuvs/src/neighbors/brute_force.rs:51`_ + +### new + +```rust +pub fn new() -> Result> +``` + +Creates a new empty index. + +_Source: `rust/cuvs/src/neighbors/brute_force.rs:70`_ + +### search + +```rust +pub fn search( +&self, +res: &Resources, +queries: &Q, +neighbors: &mut N, +distances: &mut D, +) -> Result<()> +where +Q: AsDlTensor + ?Sized, +N: AsDlTensorMut + ?Sized, +D: AsDlTensorMut + ?Sized, +``` + +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 +neighbor indices and `distances` their distances; both are written in +place. + +_Source: `rust/cuvs/src/neighbors/brute_force.rs:84`_ + +_Source: `rust/cuvs/src/neighbors/brute_force.rs:36`_ diff --git a/fern/pages/rust_api/rust-api-cuvs-neighbors-cagra-index.md b/fern/pages/rust_api/rust-api-cuvs-neighbors-cagra-index.md new file mode 100644 index 0000000000..488b931210 --- /dev/null +++ b/fern/pages/rust_api/rust-api-cuvs-neighbors-cagra-index.md @@ -0,0 +1,208 @@ +--- +slug: api-reference/rust-api-cuvs-neighbors-cagra-index +--- + +# Neighbors Cagra Index Module + +_Rust module: `cuvs::neighbors::cagra::index`_ + +_Source: `rust/cuvs/src/neighbors/cagra/index.rs`_ + +## Index + +```rust +#[derive(Debug)] +pub struct Index<'d> { + /* private fields */ +} +``` + +A CAGRA approximate nearest neighbor index. + +The lifetime `'d` ties this index to the underlying dataset, +passed at construction time. The C library may store a non-owning view +of properly aligned device-resident data, so the dataset must outlive +the index. When an index is deserialized from disk, the data is +self-contained and its lifetime is `'static`. + +**Methods** + +| Name | Source | +| --- | --- | +| `build` | `rust/cuvs/src/neighbors/cagra/index.rs:44` | +| `new` | `rust/cuvs/src/neighbors/cagra/index.rs:62` | +| `search` | `rust/cuvs/src/neighbors/cagra/index.rs:77` | +| `search_with_filter` | `rust/cuvs/src/neighbors/cagra/index.rs:117` | +| `serialize` | `rust/cuvs/src/neighbors/cagra/index.rs:190` | +| `serialize_to_hnswlib` | `rust/cuvs/src/neighbors/cagra/index.rs:214` | +| `deserialize` | `rust/cuvs/src/neighbors/cagra/index.rs:230` | + +### build + +```rust +pub fn build(res: &Resources, params: &IndexParams, dataset: &'d T) -> Result> +where +T: AsDlTensor + ?Sized, +``` + +Builds a CAGRA index over `dataset` for efficient search. + +`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 `dataset` for `'d` and +cannot outlive it. + +_Source: `rust/cuvs/src/neighbors/cagra/index.rs:44`_ + +### new + +```rust +pub fn new() -> Result> +``` + +Creates a new empty index + +_Source: `rust/cuvs/src/neighbors/cagra/index.rs:62`_ + +### search + +```rust +pub fn search( +&self, +res: &Resources, +params: &SearchParams, +queries: &Q, +neighbors: &mut N, +distances: &mut D, +) -> Result<()> +where +Q: AsDlTensor + ?Sized, +N: AsDlTensorMut + ?Sized, +D: AsDlTensorMut + ?Sized, +``` + +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` (shape +`n_queries × k`) receives the neighbor indices and `distances` their +distances; both are written in place. + +_Source: `rust/cuvs/src/neighbors/cagra/index.rs:77`_ + +### search_with_filter + +```rust +pub fn search_with_filter( +&self, +res: &Resources, +params: &SearchParams, +queries: &Q, +neighbors: &mut N, +distances: &mut D, +bitset: &B, +) -> Result<()> +where +Q: AsDlTensor + ?Sized, +N: AsDlTensorMut + ?Sized, +D: AsDlTensorMut + ?Sized, +B: AsDlTensor + ?Sized, +``` + +Perform a filtered Approximate Nearest Neighbors search on the Index + +Like [`search`](Self::search), but applies a bitset filter to exclude +vectors during graph traversal. Filtered vectors are never visited, which +gives better recall than post-filtering. + +`queries`, `neighbors`, and `distances` are as in [`search`](Self::search). +`bitset` is a 1-D `uint32` device tensor of `ceil(n_rows / 32)` elements, +where each bit maps to a dataset row (1 = include, 0 = exclude). + +_Source: `rust/cuvs/src/neighbors/cagra/index.rs:117`_ + +### serialize + +```rust +pub fn serialize>( +&self, +res: &Resources, +filename: P, +include_dataset: bool, +) -> Result<()> +``` + +Save the CAGRA index to file. + +Experimental, both the API and the serialization format are subject to change. + +#### Arguments + +* `res` - Resources to use +* `filename` - The file path for saving the index +* `include_dataset` - Whether to write out the dataset to the file + +#### Example: +```no_run +use cuvs::Resources; +use cuvs::neighbors::cagra::{Index, IndexParams}; + +fn serialize_example() -> Result<(), Box> { +let res = Resources::new()?; + +// Build an index (using some dataset) +let build_params = IndexParams::builder().build()?; +// let index = Index::build(&res, &build_params, &dataset)?; + +// Save the index to disk (including the dataset) +// index.serialize(&res, "/path/to/index.bin", true)?; + +// Later, load the index from disk +let loaded_index = Index::deserialize(&res, "/path/to/index.bin")?; + +// The loaded index can be used for search just like the original +Ok(()) +} +``` + +_Source: `rust/cuvs/src/neighbors/cagra/index.rs:190`_ + +### serialize_to_hnswlib + +```rust +pub fn serialize_to_hnswlib>(&self, res: &Resources, filename: P) -> Result<()> +``` + +Save the CAGRA index to file in hnswlib format. + +NOTE: The saved index can only be read by the hnswlib wrapper in cuVS, +as the serialization format is not compatible with the original hnswlib. + +Experimental, both the API and the serialization format are subject to change. + +#### Arguments + +* `res` - Resources to use +* `filename` - The file path for saving the index + +_Source: `rust/cuvs/src/neighbors/cagra/index.rs:214`_ + +### deserialize + +```rust +pub fn deserialize>(res: &Resources, filename: P) -> Result> +``` + +Load a CAGRA index from file. + +Experimental, both the API and the serialization format are subject to change. + +#### Arguments + +* `res` - Resources to use +* `filename` - The path of the file that stores the index + +_Source: `rust/cuvs/src/neighbors/cagra/index.rs:230`_ + +_Source: `rust/cuvs/src/neighbors/cagra/index.rs:26`_ diff --git a/fern/pages/rust_api/rust-api-cuvs-neighbors-cagra-params.md b/fern/pages/rust_api/rust-api-cuvs-neighbors-cagra-params.md new file mode 100644 index 0000000000..a8e96f1fd3 --- /dev/null +++ b/fern/pages/rust_api/rust-api-cuvs-neighbors-cagra-params.md @@ -0,0 +1,178 @@ +--- +slug: api-reference/rust-api-cuvs-neighbors-cagra-params +--- + +# Neighbors Cagra Params Module + +_Rust module: `cuvs::neighbors::cagra::params`_ + +_Source: `rust/cuvs/src/neighbors/cagra/params.rs`_ + +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`]. + +## CompressionParams + +```rust +pub struct CompressionParams { + /* private fields */ +} +``` + +VPQ (Vector-Product Quantization) compression parameters. + +Attach to [`IndexParams`] to enable compressed dataset storage. + +**Methods** + +| Name | Source | +| --- | --- | +| `new` | `rust/cuvs/src/neighbors/cagra/params.rs:38` | +| `try_new` | `rust/cuvs/src/neighbors/cagra/params.rs:82` | + +### new + +```rust +#[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 +``` + +_Source: `rust/cuvs/src/neighbors/cagra/params.rs:38`_ + +### try_new + +```rust +pub fn try_new() -> Result +``` + +Allocate parameters populated with the library defaults. + +_Source: `rust/cuvs/src/neighbors/cagra/params.rs:82`_ + +_Source: `rust/cuvs/src/neighbors/cagra/params.rs:31`_ + +## IndexParams + +```rust +pub struct IndexParams { + /* private fields */ +} +``` + +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()?; +``` + +**Methods** + +| Name | Source | +| --- | --- | +| `new` | `rust/cuvs/src/neighbors/cagra/params.rs:130` | +| `try_new` | `rust/cuvs/src/neighbors/cagra/params.rs:194` | + +### new + +```rust +#[builder] +pub fn new( +metric: Option, +intermediate_graph_degree: Option, +graph_degree: Option, +build_algo: Option, +nn_descent_niter: Option, +compression: Option, +) -> Result +``` + +_Source: `rust/cuvs/src/neighbors/cagra/params.rs:130`_ + +### try_new + +```rust +pub fn try_new() -> Result +``` + +Allocate parameters populated with the library defaults. + +_Source: `rust/cuvs/src/neighbors/cagra/params.rs:194`_ + +_Source: `rust/cuvs/src/neighbors/cagra/params.rs:120`_ + +## SearchParams + +```rust +pub struct SearchParams { + /* private fields */ +} +``` + +Parameters for searching a CAGRA index. + +```ignore +use cuvs::neighbors::cagra::SearchParams; + +let params = SearchParams::builder().itopk_size(128).build()?; +``` + +**Methods** + +| Name | Source | +| --- | --- | +| `new` | `rust/cuvs/src/neighbors/cagra/params.rs:236` | +| `try_new` | `rust/cuvs/src/neighbors/cagra/params.rs:348` | + +### new + +```rust +#[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 +``` + +_Source: `rust/cuvs/src/neighbors/cagra/params.rs:236`_ + +### try_new + +```rust +pub fn try_new() -> Result +``` + +Allocate parameters populated with the library defaults. + +_Source: `rust/cuvs/src/neighbors/cagra/params.rs:348`_ + +_Source: `rust/cuvs/src/neighbors/cagra/params.rs:228`_ diff --git a/fern/pages/rust_api/rust-api-cuvs-neighbors-cagra.md b/fern/pages/rust_api/rust-api-cuvs-neighbors-cagra.md new file mode 100644 index 0000000000..c08b919b36 --- /dev/null +++ b/fern/pages/rust_api/rust-api-cuvs-neighbors-cagra.md @@ -0,0 +1,94 @@ +--- +slug: api-reference/rust-api-cuvs-neighbors-cagra +--- + +# Neighbors Cagra Module + +_Rust module: `cuvs::neighbors::cagra`_ + +_Source: `rust/cuvs/src/neighbors/cagra/mod.rs`_ + +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. + +## index::Index + +```rust +pub use index::Index; +``` + +_Source: `rust/cuvs/src/neighbors/cagra/mod.rs:23`_ + +## params::\{CompressionParams, IndexParams, SearchParams\} + +```rust +pub use params::{CompressionParams, IndexParams, SearchParams}; +``` + +_Source: `rust/cuvs/src/neighbors/cagra/mod.rs:24`_ + +## GraphBuildAlgo + +```rust +#[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)] +#[non_exhaustive] +pub enum GraphBuildAlgo { + /* variants omitted */ +} +``` + +Algorithm for building the internal k-NN graph. + +_Source: `rust/cuvs/src/neighbors/cagra/mod.rs:32`_ + +## SearchAlgo + +```rust +#[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)] +#[non_exhaustive] +pub enum SearchAlgo { + /* variants omitted */ +} +``` + +Search kernel implementation. + +_Source: `rust/cuvs/src/neighbors/cagra/mod.rs:72`_ + +## HashMode + +```rust +#[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)] +#[non_exhaustive] +pub enum HashMode { + /* variants omitted */ +} +``` + +Hash-table mode used during search. + +_Source: `rust/cuvs/src/neighbors/cagra/mod.rs:108`_ + +## CagraError + +```rust +#[derive(Debug, thiserror::Error)] +pub enum CagraError { + /* variants omitted */ +} +``` + +Error type for CAGRA operations. + +_Source: `rust/cuvs/src/neighbors/cagra/mod.rs:139`_ diff --git a/fern/pages/rust_api/rust-api-cuvs-neighbors-ivf-flat-index.md b/fern/pages/rust_api/rust-api-cuvs-neighbors-ivf-flat-index.md new file mode 100644 index 0000000000..ef19b98b01 --- /dev/null +++ b/fern/pages/rust_api/rust-api-cuvs-neighbors-ivf-flat-index.md @@ -0,0 +1,85 @@ +--- +slug: api-reference/rust-api-cuvs-neighbors-ivf-flat-index +--- + +# Neighbors Ivf Flat Index Module + +_Rust module: `cuvs::neighbors::ivf_flat::index`_ + +_Source: `rust/cuvs/src/neighbors/ivf_flat/index.rs`_ + +## Index + +```rust +#[derive(Debug)] +pub struct Index(ffi::cuvsIvfFlatIndex_t); { + /* private fields */ +} +``` + +IVF-Flat ANN index. + +**Methods** + +| Name | Source | +| --- | --- | +| `build` | `rust/cuvs/src/neighbors/ivf_flat/index.rs:28` | +| `new` | `rust/cuvs/src/neighbors/ivf_flat/index.rs:46` | +| `search` | `rust/cuvs/src/neighbors/ivf_flat/index.rs:60` | + +### build + +```rust +pub fn build(res: &Resources, params: &IndexParams, dataset: &T) -> Result +where +T: AsDlTensor + ?Sized, +``` + +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). + +Supported dataset/query dtypes in the current C-backed implementation are +`f32`, `f16`, `i8`, and `u8`. + +_Source: `rust/cuvs/src/neighbors/ivf_flat/index.rs:28`_ + +### new + +```rust +pub fn new() -> Result +``` + +Creates a new empty index. + +_Source: `rust/cuvs/src/neighbors/ivf_flat/index.rs:46`_ + +### search + +```rust +pub fn search( +&self, +res: &Resources, +params: &SearchParams, +queries: &Q, +neighbors: &mut N, +distances: &mut D, +) -> Result<()> +where +Q: AsDlTensor + ?Sized, +N: AsDlTensorMut + ?Sized, +D: AsDlTensorMut + ?Sized, +``` + +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 +neighbor indices and `distances` their distances; both are written in +place. + +_Source: `rust/cuvs/src/neighbors/ivf_flat/index.rs:60`_ + +_Source: `rust/cuvs/src/neighbors/ivf_flat/index.rs:17`_ diff --git a/fern/pages/rust_api/rust-api-cuvs-neighbors-ivf-flat-params.md b/fern/pages/rust_api/rust-api-cuvs-neighbors-ivf-flat-params.md new file mode 100644 index 0000000000..d460562a9c --- /dev/null +++ b/fern/pages/rust_api/rust-api-cuvs-neighbors-ivf-flat-params.md @@ -0,0 +1,96 @@ +--- +slug: api-reference/rust-api-cuvs-neighbors-ivf-flat-params +--- + +# Neighbors Ivf Flat Params Module + +_Rust module: `cuvs::neighbors::ivf_flat::params`_ + +_Source: `rust/cuvs/src/neighbors/ivf_flat/params.rs`_ + +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. + +## IndexParams + +```rust +pub struct IndexParams { + /* private fields */ +} +``` + +Parameters for building an IVF-Flat index. + +**Methods** + +| Name | Source | +| --- | --- | +| `new` | `rust/cuvs/src/neighbors/ivf_flat/params.rs:28` | +| `try_new` | `rust/cuvs/src/neighbors/ivf_flat/params.rs:60` | + +### new + +```rust +#[builder] +pub fn new( +n_lists: Option, +metric: Option, +kmeans_n_iters: Option, +kmeans_trainset_fraction: Option, +add_data_on_build: Option, +) -> Result +``` + +_Source: `rust/cuvs/src/neighbors/ivf_flat/params.rs:28`_ + +### try_new + +```rust +pub fn try_new() -> Result +``` + +Allocate parameters populated with the library defaults. + +_Source: `rust/cuvs/src/neighbors/ivf_flat/params.rs:60`_ + +_Source: `rust/cuvs/src/neighbors/ivf_flat/params.rs:21`_ + +## SearchParams + +```rust +pub struct SearchParams { + /* private fields */ +} +``` + +Parameters for searching an IVF-Flat index. + +**Methods** + +| Name | Source | +| --- | --- | +| `new` | `rust/cuvs/src/neighbors/ivf_flat/params.rs:91` | +| `try_new` | `rust/cuvs/src/neighbors/ivf_flat/params.rs:104` | + +### new + +```rust +#[builder] +pub fn new(n_probes: Option) -> Result +``` + +_Source: `rust/cuvs/src/neighbors/ivf_flat/params.rs:91`_ + +### try_new + +```rust +pub fn try_new() -> Result +``` + +Allocate parameters populated with the library defaults. + +_Source: `rust/cuvs/src/neighbors/ivf_flat/params.rs:104`_ + +_Source: `rust/cuvs/src/neighbors/ivf_flat/params.rs:84`_ diff --git a/fern/pages/rust_api/rust-api-cuvs-neighbors-ivf-flat.md b/fern/pages/rust_api/rust-api-cuvs-neighbors-ivf-flat.md new file mode 100644 index 0000000000..2a48116624 --- /dev/null +++ b/fern/pages/rust_api/rust-api-cuvs-neighbors-ivf-flat.md @@ -0,0 +1,48 @@ +--- +slug: api-reference/rust-api-cuvs-neighbors-ivf-flat +--- + +# Neighbors Ivf Flat Module + +_Rust module: `cuvs::neighbors::ivf_flat`_ + +_Source: `rust/cuvs/src/neighbors/ivf_flat/mod.rs`_ + +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. + +## index::Index + +```rust +pub use index::Index; +``` + +_Source: `rust/cuvs/src/neighbors/ivf_flat/mod.rs:19`_ + +## params::\{IndexParams, SearchParams\} + +```rust +pub use params::{IndexParams, SearchParams}; +``` + +_Source: `rust/cuvs/src/neighbors/ivf_flat/mod.rs:20`_ + +## IvfFlatError + +```rust +#[derive(Debug, thiserror::Error)] +pub enum IvfFlatError { + /* variants omitted */ +} +``` + +Error type for IVF-Flat operations. + +_Source: `rust/cuvs/src/neighbors/ivf_flat/mod.rs:27`_ diff --git a/fern/pages/rust_api/rust-api-cuvs-neighbors-ivf-pq-index.md b/fern/pages/rust_api/rust-api-cuvs-neighbors-ivf-pq-index.md new file mode 100644 index 0000000000..a8d66531fc --- /dev/null +++ b/fern/pages/rust_api/rust-api-cuvs-neighbors-ivf-pq-index.md @@ -0,0 +1,82 @@ +--- +slug: api-reference/rust-api-cuvs-neighbors-ivf-pq-index +--- + +# Neighbors Ivf Pq Index Module + +_Rust module: `cuvs::neighbors::ivf_pq::index`_ + +_Source: `rust/cuvs/src/neighbors/ivf_pq/index.rs`_ + +## Index + +```rust +#[derive(Debug)] +pub struct Index(ffi::cuvsIvfPqIndex_t); { + /* private fields */ +} +``` + +IVF-PQ ANN index. + +**Methods** + +| Name | Source | +| --- | --- | +| `build` | `rust/cuvs/src/neighbors/ivf_pq/index.rs:25` | +| `new` | `rust/cuvs/src/neighbors/ivf_pq/index.rs:43` | +| `search` | `rust/cuvs/src/neighbors/ivf_pq/index.rs:57` | + +### build + +```rust +pub fn build(res: &Resources, params: &IndexParams, dataset: &T) -> Result +where +T: AsDlTensor + ?Sized, +``` + +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 into the index, so the caller may free it +once this call returns (hence `Index` carries no lifetime). + +_Source: `rust/cuvs/src/neighbors/ivf_pq/index.rs:25`_ + +### new + +```rust +pub fn new() -> Result +``` + +Creates a new empty index. + +_Source: `rust/cuvs/src/neighbors/ivf_pq/index.rs:43`_ + +### search + +```rust +pub fn search( +&self, +res: &Resources, +params: &SearchParams, +queries: &Q, +neighbors: &mut N, +distances: &mut D, +) -> Result<()> +where +Q: AsDlTensor + ?Sized, +N: AsDlTensorMut + ?Sized, +D: AsDlTensorMut + ?Sized, +``` + +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 +neighbor indices and `distances` their distances; both are written in +place. + +_Source: `rust/cuvs/src/neighbors/ivf_pq/index.rs:57`_ + +_Source: `rust/cuvs/src/neighbors/ivf_pq/index.rs:17`_ diff --git a/fern/pages/rust_api/rust-api-cuvs-neighbors-ivf-pq-params.md b/fern/pages/rust_api/rust-api-cuvs-neighbors-ivf-pq-params.md new file mode 100644 index 0000000000..bb742fa122 --- /dev/null +++ b/fern/pages/rust_api/rust-api-cuvs-neighbors-ivf-pq-params.md @@ -0,0 +1,115 @@ +--- +slug: api-reference/rust-api-cuvs-neighbors-ivf-pq-params +--- + +# Neighbors Ivf Pq Params Module + +_Rust module: `cuvs::neighbors::ivf_pq::params`_ + +_Source: `rust/cuvs/src/neighbors/ivf_pq/params.rs`_ + +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. + +## ffi::\{cudaDataType_t, cuvsIvfPqCodebookGen, cuvsIvfPqListLayout\} + +```rust +pub use ffi::{cudaDataType_t, cuvsIvfPqCodebookGen, cuvsIvfPqListLayout}; +``` + +_Source: `rust/cuvs/src/neighbors/ivf_pq/params.rs:20`_ + +## IndexParams + +```rust +pub struct IndexParams { + /* private fields */ +} +``` + +Parameters for building an IVF-PQ index. + +**Methods** + +| Name | Source | +| --- | --- | +| `new` | `rust/cuvs/src/neighbors/ivf_pq/params.rs:31` | +| `try_new` | `rust/cuvs/src/neighbors/ivf_pq/params.rs:87` | + +### new + +```rust +#[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 +``` + +_Source: `rust/cuvs/src/neighbors/ivf_pq/params.rs:31`_ + +### try_new + +```rust +pub fn try_new() -> Result +``` + +Allocate parameters populated with the library defaults. + +_Source: `rust/cuvs/src/neighbors/ivf_pq/params.rs:87`_ + +_Source: `rust/cuvs/src/neighbors/ivf_pq/params.rs:23`_ + +## SearchParams + +```rust +pub struct SearchParams { + /* private fields */ +} +``` + +Parameters for searching an IVF-PQ index. + +**Methods** + +| Name | Source | +| --- | --- | +| `new` | `rust/cuvs/src/neighbors/ivf_pq/params.rs:118` | +| `try_new` | `rust/cuvs/src/neighbors/ivf_pq/params.rs:141` | + +### new + +```rust +#[builder] +pub fn new( +n_probes: Option, +lut_dtype: Option, +internal_distance_dtype: Option, +) -> Result +``` + +_Source: `rust/cuvs/src/neighbors/ivf_pq/params.rs:118`_ + +### try_new + +```rust +pub fn try_new() -> Result +``` + +Allocate parameters populated with the library defaults. + +_Source: `rust/cuvs/src/neighbors/ivf_pq/params.rs:141`_ + +_Source: `rust/cuvs/src/neighbors/ivf_pq/params.rs:111`_ diff --git a/fern/pages/rust_api/rust-api-cuvs-neighbors-ivf-pq.md b/fern/pages/rust_api/rust-api-cuvs-neighbors-ivf-pq.md new file mode 100644 index 0000000000..d018d31541 --- /dev/null +++ b/fern/pages/rust_api/rust-api-cuvs-neighbors-ivf-pq.md @@ -0,0 +1,51 @@ +--- +slug: api-reference/rust-api-cuvs-neighbors-ivf-pq +--- + +# Neighbors Ivf Pq Module + +_Rust module: `cuvs::neighbors::ivf_pq`_ + +_Source: `rust/cuvs/src/neighbors/ivf_pq/mod.rs`_ + +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. + +## index::Index + +```rust +pub use index::Index; +``` + +_Source: `rust/cuvs/src/neighbors/ivf_pq/mod.rs:19`_ + +## params::\{ IndexParams, SearchParams, cudaDataType_t, cuvsIvfPqCodebookGen, cuvsIvfPqListLayout, \} + +```rust +pub use params::{ +IndexParams, SearchParams, cudaDataType_t, cuvsIvfPqCodebookGen, cuvsIvfPqListLayout, +}; +``` + +_Source: `rust/cuvs/src/neighbors/ivf_pq/mod.rs:20`_ + +## IvfPqError + +```rust +#[derive(Debug, thiserror::Error)] +pub enum IvfPqError { + /* variants omitted */ +} +``` + +Error type for IVF-PQ operations. + +_Source: `rust/cuvs/src/neighbors/ivf_pq/mod.rs:29`_ diff --git a/fern/pages/rust_api/rust-api-cuvs-neighbors-vamana-index.md b/fern/pages/rust_api/rust-api-cuvs-neighbors-vamana-index.md new file mode 100644 index 0000000000..9f0aa8499b --- /dev/null +++ b/fern/pages/rust_api/rust-api-cuvs-neighbors-vamana-index.md @@ -0,0 +1,78 @@ +--- +slug: api-reference/rust-api-cuvs-neighbors-vamana-index +--- + +# Neighbors Vamana Index Module + +_Rust module: `cuvs::neighbors::vamana::index`_ + +_Source: `rust/cuvs/src/neighbors/vamana/index.rs`_ + +## Index + +```rust +#[derive(Debug)] +pub struct Index(ffi::cuvsVamanaIndex_t); { + /* private fields */ +} +``` + +Vamana ANN index. + +**Methods** + +| Name | Source | +| --- | --- | +| `build` | `rust/cuvs/src/neighbors/vamana/index.rs:32` | +| `new` | `rust/cuvs/src/neighbors/vamana/index.rs:50` | +| `serialize` | `rust/cuvs/src/neighbors/vamana/index.rs:65` | + +### build + +```rust +pub fn build(res: &Resources, params: &IndexParams, dataset: &T) -> Result +where +T: AsDlTensor + ?Sized, +``` + +Builds a Vamana index for efficient DiskANN search. + +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, so `Index` carries no +lifetime. + +_Source: `rust/cuvs/src/neighbors/vamana/index.rs:32`_ + +### new + +```rust +pub fn new() -> Result +``` + +Creates a new empty index. + +_Source: `rust/cuvs/src/neighbors/vamana/index.rs:50`_ + +### serialize + +```rust +pub fn serialize(self, res: &Resources, filename: &str, include_dataset: bool) -> Result<()> +``` + +Saves the Vamana index to a file. + +Matches the on-disk format used by the DiskANN open-source repository, +so the serialized index can be consumed there for graph search. + +`filename` is the file prefix under which the index is saved; +`include_dataset` controls whether the dataset is embedded. + +_Source: `rust/cuvs/src/neighbors/vamana/index.rs:65`_ + +_Source: `rust/cuvs/src/neighbors/vamana/index.rs:18`_ diff --git a/fern/pages/rust_api/rust-api-cuvs-neighbors-vamana-params.md b/fern/pages/rust_api/rust-api-cuvs-neighbors-vamana-params.md new file mode 100644 index 0000000000..613cd6cdbd --- /dev/null +++ b/fern/pages/rust_api/rust-api-cuvs-neighbors-vamana-params.md @@ -0,0 +1,63 @@ +--- +slug: api-reference/rust-api-cuvs-neighbors-vamana-params +--- + +# Neighbors Vamana Params Module + +_Rust module: `cuvs::neighbors::vamana::params`_ + +_Source: `rust/cuvs/src/neighbors/vamana/params.rs`_ + +Builder-pattern parameter type for Vamana index build. + +All setters are optional; unset values retain the library defaults from the +underlying C `cuvsVamanaIndexParamsCreate`. + +## IndexParams + +```rust +pub struct IndexParams { + /* private fields */ +} +``` + +Parameters for building a Vamana index. + +**Methods** + +| Name | Source | +| --- | --- | +| `new` | `rust/cuvs/src/neighbors/vamana/params.rs:29` | +| `try_new` | `rust/cuvs/src/neighbors/vamana/params.rs:76` | + +### new + +```rust +#[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 +``` + +_Source: `rust/cuvs/src/neighbors/vamana/params.rs:29`_ + +### try_new + +```rust +pub fn try_new() -> Result +``` + +Allocate parameters populated with the library defaults. + +_Source: `rust/cuvs/src/neighbors/vamana/params.rs:76`_ + +_Source: `rust/cuvs/src/neighbors/vamana/params.rs:21`_ diff --git a/fern/pages/rust_api/rust-api-cuvs-neighbors-vamana.md b/fern/pages/rust_api/rust-api-cuvs-neighbors-vamana.md new file mode 100644 index 0000000000..985b5a3e81 --- /dev/null +++ b/fern/pages/rust_api/rust-api-cuvs-neighbors-vamana.md @@ -0,0 +1,44 @@ +--- +slug: api-reference/rust-api-cuvs-neighbors-vamana +--- + +# Neighbors Vamana Module + +_Rust module: `cuvs::neighbors::vamana`_ + +_Source: `rust/cuvs/src/neighbors/vamana/mod.rs`_ + +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. + +## index::Index + +```rust +pub use index::Index; +``` + +_Source: `rust/cuvs/src/neighbors/vamana/mod.rs:14`_ + +## params::IndexParams + +```rust +pub use params::IndexParams; +``` + +_Source: `rust/cuvs/src/neighbors/vamana/mod.rs:15`_ + +## VamanaError + +```rust +#[derive(Debug, thiserror::Error)] +pub enum VamanaError { + /* variants omitted */ +} +``` + +Error type for Vamana operations. + +_Source: `rust/cuvs/src/neighbors/vamana/mod.rs:22`_ diff --git a/fern/pages/rust_api/rust-api-cuvs-neighbors.md b/fern/pages/rust_api/rust-api-cuvs-neighbors.md new file mode 100644 index 0000000000..b81316dfc3 --- /dev/null +++ b/fern/pages/rust_api/rust-api-cuvs-neighbors.md @@ -0,0 +1,56 @@ +--- +slug: api-reference/rust-api-cuvs-neighbors +--- + +# Neighbors Module + +_Rust module: `cuvs::neighbors`_ + +_Source: `rust/cuvs/src/neighbors/mod.rs`_ + +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. + +## brute_force + +```rust +pub mod brute_force; +``` + +_Source: `rust/cuvs/src/neighbors/mod.rs:13`_ + +## cagra + +```rust +pub mod cagra; +``` + +_Source: `rust/cuvs/src/neighbors/mod.rs:14`_ + +## ivf_flat + +```rust +pub mod ivf_flat; +``` + +_Source: `rust/cuvs/src/neighbors/mod.rs:15`_ + +## ivf_pq + +```rust +pub mod ivf_pq; +``` + +_Source: `rust/cuvs/src/neighbors/mod.rs:16`_ + +## vamana + +```rust +pub mod vamana; +``` + +_Source: `rust/cuvs/src/neighbors/mod.rs:17`_ diff --git a/fern/pages/rust_api/rust-api-cuvs-resources.md b/fern/pages/rust_api/rust-api-cuvs-resources.md index 1aa375c324..c83378a43a 100644 --- a/fern/pages/rust_api/rust-api-cuvs-resources.md +++ b/fern/pages/rust_api/rust-api-cuvs-resources.md @@ -8,11 +8,26 @@ _Rust module: `cuvs::resources`_ _Source: `rust/cuvs/src/resources.rs`_ +GPU resource management with RAII semantics. + +## ResourcesError + +```rust +#[derive(Debug, thiserror::Error)] +pub enum ResourcesError { + /* variants omitted */ +} +``` + +Error type for resource operations. + +_Source: `rust/cuvs/src/resources.rs:15`_ + ## Resources ```rust #[derive(Debug)] -pub struct Resources(pub ffi::cuvsResources_t); { +pub struct Resources { /* private fields */ } ``` @@ -25,10 +40,10 @@ resources that are expensive to create. | Name | Source | | --- | --- | -| `new` | `rust/cuvs/src/resources.rs:17` | -| `set_cuda_stream` | `rust/cuvs/src/resources.rs:26` | -| `get_cuda_stream` | `rust/cuvs/src/resources.rs:31` | -| `sync_stream` | `rust/cuvs/src/resources.rs:40` | +| `new` | `rust/cuvs/src/resources.rs:31` | +| `with_stream` | `rust/cuvs/src/resources.rs:46` | +| `stream` | `rust/cuvs/src/resources.rs:55` | +| `sync_stream` | `rust/cuvs/src/resources.rs:64` | ### new @@ -36,29 +51,37 @@ resources that are expensive to create. pub fn new() -> Result ``` -Returns a new Resources object +Creates a new resources handle bound to the current CUDA device. -_Source: `rust/cuvs/src/resources.rs:17`_ +_Source: `rust/cuvs/src/resources.rs:31`_ -### set_cuda_stream +### with_stream ```rust -pub fn set_cuda_stream(&self, stream: ffi::cudaStream_t) -> Result<()> +pub unsafe fn with_stream(stream: ffi::cudaStream_t) -> Result ``` -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. -_Source: `rust/cuvs/src/resources.rs:26`_ +#### Safety -### get_cuda_stream +`stream` must be a valid CUDA stream for the current device and must +remain valid for as long as this handle uses it. + +_Source: `rust/cuvs/src/resources.rs:46`_ + +### stream ```rust -pub fn get_cuda_stream(&self) -> Result +pub fn stream(&self) -> Result ``` -Gets the current cuda stream +Returns the current CUDA stream associated with this handle. -_Source: `rust/cuvs/src/resources.rs:31`_ +_Source: `rust/cuvs/src/resources.rs:55`_ ### sync_stream @@ -66,8 +89,8 @@ _Source: `rust/cuvs/src/resources.rs:31`_ pub fn sync_stream(&self) -> Result<()> ``` -Syncs the current cuda stream +Blocks until all operations on the current CUDA stream have completed. -_Source: `rust/cuvs/src/resources.rs:40`_ +_Source: `rust/cuvs/src/resources.rs:64`_ -_Source: `rust/cuvs/src/resources.rs:13`_ +_Source: `rust/cuvs/src/resources.rs:25`_ diff --git a/fern/pages/rust_api/rust-api-cuvs-vamana-index-params.md b/fern/pages/rust_api/rust-api-cuvs-vamana-index-params.md deleted file mode 100644 index 33334febcf..0000000000 --- a/fern/pages/rust_api/rust-api-cuvs-vamana-index-params.md +++ /dev/null @@ -1,137 +0,0 @@ ---- -slug: api-reference/rust-api-cuvs-vamana-index-params ---- - -# Vamana Index Params Module - -_Rust module: `cuvs::vamana::index_params`_ - -_Source: `rust/cuvs/src/vamana/index_params.rs`_ - -## IndexParams - -```rust -pub struct IndexParams(pub ffi::cuvsVamanaIndexParams_t); { - /* private fields */ -} -``` - -**Methods** - -| Name | Source | -| --- | --- | -| `new` | `rust/cuvs/src/vamana/index_params.rs:15` | -| `set_metric` | `rust/cuvs/src/vamana/index_params.rs:24` | -| `set_graph_degree` | `rust/cuvs/src/vamana/index_params.rs:33` | -| `set_visited_size` | `rust/cuvs/src/vamana/index_params.rs:42` | -| `set_vamana_iters` | `rust/cuvs/src/vamana/index_params.rs:50` | -| `set_alpha` | `rust/cuvs/src/vamana/index_params.rs:58` | -| `set_max_fraction` | `rust/cuvs/src/vamana/index_params.rs:67` | -| `set_batch_base` | `rust/cuvs/src/vamana/index_params.rs:75` | -| `set_queue_size` | `rust/cuvs/src/vamana/index_params.rs:83` | -| `set_reverse_batchsize` | `rust/cuvs/src/vamana/index_params.rs:91` | - -### new - -```rust -pub fn new() -> Result -``` - -Returns a new IndexParams - -_Source: `rust/cuvs/src/vamana/index_params.rs:15`_ - -### set_metric - -```rust -pub fn set_metric(self, metric: DistanceType) -> IndexParams -``` - -DistanceType to use for building the index - -_Source: `rust/cuvs/src/vamana/index_params.rs:24`_ - -### set_graph_degree - -```rust -pub fn set_graph_degree(self, graph_degree: u32) -> IndexParams -``` - -Maximum degree of output graph corresponds to the R parameter in the original Vamana -literature. - -_Source: `rust/cuvs/src/vamana/index_params.rs:33`_ - -### set_visited_size - -```rust -pub fn set_visited_size(self, visited_size: u32) -> IndexParams -``` - -Maximum number of visited nodes per search corresponds to the L parameter in the Vamana -literature - -_Source: `rust/cuvs/src/vamana/index_params.rs:42`_ - -### set_vamana_iters - -```rust -pub fn set_vamana_iters(self, vamana_iters: f32) -> IndexParams -``` - -Number of Vamana vector insertion iterations (each iteration inserts all vectors). - -_Source: `rust/cuvs/src/vamana/index_params.rs:50`_ - -### set_alpha - -```rust -pub fn set_alpha(self, alpha: f32) -> IndexParams -``` - -Alpha for pruning parameter - -_Source: `rust/cuvs/src/vamana/index_params.rs:58`_ - -### set_max_fraction - -```rust -pub fn set_max_fraction(self, max_fraction: f32) -> IndexParams -``` - -Maximum fraction of dataset inserted per batch. -Larger max batch decreases graph quality, but improves speed - -_Source: `rust/cuvs/src/vamana/index_params.rs:67`_ - -### set_batch_base - -```rust -pub fn set_batch_base(self, batch_base: f32) -> IndexParams -``` - -Base of growth rate of batch sizes - -_Source: `rust/cuvs/src/vamana/index_params.rs:75`_ - -### set_queue_size - -```rust -pub fn set_queue_size(self, queue_size: u32) -> IndexParams -``` - -Size of candidate queue structure - should be (2^x)-1 - -_Source: `rust/cuvs/src/vamana/index_params.rs:83`_ - -### set_reverse_batchsize - -```rust -pub fn set_reverse_batchsize(self, reverse_batchsize: u32) -> IndexParams -``` - -Max batchsize of reverse edge processing (reduces memory footprint) - -_Source: `rust/cuvs/src/vamana/index_params.rs:91`_ - -_Source: `rust/cuvs/src/vamana/index_params.rs:11`_ diff --git a/fern/pages/rust_api/rust-api-cuvs-vamana-index.md b/fern/pages/rust_api/rust-api-cuvs-vamana-index.md deleted file mode 100644 index 4e6050c9fc..0000000000 --- a/fern/pages/rust_api/rust-api-cuvs-vamana-index.md +++ /dev/null @@ -1,87 +0,0 @@ ---- -slug: api-reference/rust-api-cuvs-vamana-index ---- - -# Vamana Index Module - -_Rust module: `cuvs::vamana::index`_ - -_Source: `rust/cuvs/src/vamana/index.rs`_ - -## Index - -```rust -#[derive(Debug)] -pub struct Index(ffi::cuvsVamanaIndex_t); { - /* private fields */ -} -``` - -Vamana ANN Index - -**Methods** - -| Name | Source | -| --- | --- | -| `build` | `rust/cuvs/src/vamana/index.rs:33` | -| `new` | `rust/cuvs/src/vamana/index.rs:47` | -| `serialize` | `rust/cuvs/src/vamana/index.rs:66` | - -### build - -```rust -pub fn build>( -res: &Resources, -params: &IndexParams, -dataset: T, -) -> Result -``` - -Builds 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. - - -#### Arguments - -* `res` - Resources to use -* `params` - Parameters for building the index -* `dataset` - A row-major matrix on either the host or device to index - -_Source: `rust/cuvs/src/vamana/index.rs:33`_ - -### new - -```rust -pub fn new() -> Result -``` - -Creates a new empty index - -_Source: `rust/cuvs/src/vamana/index.rs:47`_ - -### serialize - -```rust -pub fn serialize(self, res: &Resources, filename: &str, include_dataset: bool) -> Result<()> -``` - -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. - -#### Arguments - -* `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 - -_Source: `rust/cuvs/src/vamana/index.rs:66`_ - -_Source: `rust/cuvs/src/vamana/index.rs:16`_ diff --git a/fern/pages/rust_api/rust-api-cuvs-vamana.md b/fern/pages/rust_api/rust-api-cuvs-vamana.md deleted file mode 100644 index 4bff7c0c79..0000000000 --- a/fern/pages/rust_api/rust-api-cuvs-vamana.md +++ /dev/null @@ -1,27 +0,0 @@ ---- -slug: api-reference/rust-api-cuvs-vamana ---- - -# Vamana Module - -_Rust module: `cuvs::vamana`_ - -_Source: `rust/cuvs/src/vamana/mod.rs`_ - -Vamana - -## index::Index - -```rust -pub use index::Index; -``` - -_Source: `rust/cuvs/src/vamana/mod.rs:10`_ - -## index_params::IndexParams - -```rust -pub use index_params::IndexParams; -``` - -_Source: `rust/cuvs/src/vamana/mod.rs:11`_ diff --git a/fern/pages/rust_api/rust-api-cuvs.md b/fern/pages/rust_api/rust-api-cuvs.md index 547ed2d7ae..546da6c913 100644 --- a/fern/pages/rust_api/rust-api-cuvs.md +++ b/fern/pages/rust_api/rust-api-cuvs.md @@ -13,29 +13,13 @@ cuVS: Rust bindings for Vector Search on the GPU This crate provides Rust bindings for cuVS, allowing you to run approximate nearest neighbors search on the GPU. -## brute_force - -```rust -pub mod brute_force; -``` - -_Source: `rust/cuvs/src/lib.rs:12`_ - -## cagra - -```rust -pub mod cagra; -``` - -_Source: `rust/cuvs/src/lib.rs:13`_ - ## cluster ```rust pub mod cluster; ``` -_Source: `rust/cuvs/src/lib.rs:14`_ +_Source: `rust/cuvs/src/lib.rs:12`_ ## distance @@ -43,60 +27,44 @@ _Source: `rust/cuvs/src/lib.rs:14`_ pub mod distance; ``` -_Source: `rust/cuvs/src/lib.rs:15`_ +_Source: `rust/cuvs/src/lib.rs:13`_ -## distance_type +## dlpack ```rust -pub mod distance_type; +pub mod dlpack; ``` -_Source: `rust/cuvs/src/lib.rs:16`_ +_Source: `rust/cuvs/src/lib.rs:14`_ -## ivf_flat +## neighbors ```rust -pub mod ivf_flat; +pub mod neighbors; ``` -_Source: `rust/cuvs/src/lib.rs:19`_ +_Source: `rust/cuvs/src/lib.rs:16`_ -## ivf_pq +## dlpack::\{AsDlTensor, AsDlTensorMut, DLPackError, DLTensorView, DLTensorViewMut, DType\} ```rust -pub mod ivf_pq; +pub use dlpack::{AsDlTensor, AsDlTensorMut, DLPackError, DLTensorView, DLTensorViewMut, DType}; ``` -_Source: `rust/cuvs/src/lib.rs:20`_ +_Source: `rust/cuvs/src/lib.rs:21`_ -## vamana +## error::LibraryError ```rust -pub mod vamana; +pub use error::LibraryError; ``` _Source: `rust/cuvs/src/lib.rs:22`_ -## dlpack::ManagedTensor - -```rust -pub use dlpack::ManagedTensor; -``` - -_Source: `rust/cuvs/src/lib.rs:24`_ - -## error::\{Error, Result\} - -```rust -pub use error::{Error, Result}; -``` - -_Source: `rust/cuvs/src/lib.rs:25`_ - ## resources::Resources ```rust pub use resources::Resources; ``` -_Source: `rust/cuvs/src/lib.rs:26`_ +_Source: `rust/cuvs/src/lib.rs:23`_ diff --git a/fern/scripts/generate_api_reference.py b/fern/scripts/generate_api_reference.py index 7541b11c8b..dd7f23442e 100755 --- a/fern/scripts/generate_api_reference.py +++ b/fern/scripts/generate_api_reference.py @@ -1,5 +1,5 @@ #!/usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """Generate Fern API reference pages from cuVS source code. @@ -4773,7 +4773,7 @@ def generate_rust_api_pages() -> None: def collect_rust_pages() -> list[RustPage]: pages: list[RustPage] = [] for path in sorted(RUST_SOURCE_DIR.rglob("*.rs")): - if path.name.endswith("_test.rs"): + if path.name.endswith("_test.rs") or path.name == "test_utils.rs": continue page = parse_rust_page(path) if page.module_doc or page.items: From a0690719afbb5b3187e5ab4b26b3872d03999e26 Mon Sep 17 00:00:00 2001 From: Yan Zaretskiy Date: Thu, 2 Jul 2026 15:14:21 -0700 Subject: [PATCH 3/3] PR feedback --- .github/CODEOWNERS | 1 - README.md | 89 ++++++++------------ examples/rust/Cargo.toml | 9 -- examples/rust/README.md | 33 -------- examples/rust/src/main.rs | 67 --------------- rust/cuvs/src/lib.rs | 6 ++ rust/cuvs/src/neighbors/cagra/params.rs | 4 +- rust/cuvs/src/neighbors/ivf_pq/mod.rs | 101 ++++++++++++++++++++++- rust/cuvs/src/neighbors/ivf_pq/params.rs | 20 ++--- 9 files changed, 151 insertions(+), 179 deletions(-) delete mode 100644 examples/rust/Cargo.toml delete mode 100644 examples/rust/README.md delete mode 100644 examples/rust/src/main.rs diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index a684459392..b897d49c50 100755 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -16,7 +16,6 @@ python/ @NVIDIA/cuvs-python-codeowners #rust code owners rust/ @NVIDIA/cuvs-rust-codeowners -examples/rust/ @NVIDIA/cuvs-rust-codeowners #docs code owners docs/ @NVIDIA/cuvs-docs-codeowners diff --git a/README.md b/README.md index fbed9b5481..32f8290f74 100755 --- a/README.md +++ b/README.md @@ -143,66 +143,49 @@ For more code examples of the C APIs, including drop-in Cmake project templates, ### Rust API -```rust -use cuvs::cagra::{Index, IndexParams, SearchParams}; -use cuvs::{ManagedTensor, Resources, Result}; - -use ndarray::s; -use ndarray_rand::rand_distr::Uniform; -use ndarray_rand::RandomExt; +```rust,no_run +use cuvs::distance::DistanceType; +use cuvs::neighbors::cagra::{Index, IndexParams, SearchParams}; +use cuvs::{AsDlTensor, AsDlTensorMut, DLPackError, DLTensorView, DLTensorViewMut, Resources}; + +// cuVS is agnostic about where your vectors live: `build` and `search` accept +// any type implementing `AsDlTensor` (inputs) / `AsDlTensorMut` (outputs). Wrap +// your own GPU buffer by implementing these traits. See `rust/cuvs/examples` +// for a complete, runnable CUDA-backed implementation. +struct GpuTensor; +impl AsDlTensor for GpuTensor { + fn as_dl_tensor(&self) -> Result, DLPackError> { + unimplemented!("wrap your device buffer; see rust/cuvs/examples") + } +} +impl AsDlTensorMut for GpuTensor { + fn as_dl_tensor_mut(&mut self) -> Result, DLPackError> { + unimplemented!("wrap your device buffer; see rust/cuvs/examples") + } +} -/// Example showing how to index and search data with CAGRA -fn cagra_example() -> Result<()> { +fn main() -> Result<(), Box> { let res = Resources::new()?; - // Create a new random dataset to index - let n_datapoints = 65536; - let n_features = 512; - let dataset = - ndarray::Array::::random((n_datapoints, n_features), Uniform::new(0., 1.0)); - - // build the cagra index - let build_params = IndexParams::new()?; - let index = Index::build(&res, &build_params, &dataset)?; - println!( - "Indexed {}x{} datapoints into cagra index", - n_datapoints, n_features - ); - - // 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, ..]); - - let k = 10; - - // CAGRA 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 = ManagedTensor::from(&queries).to_device(&res)?; - let mut neighbors_host = ndarray::Array::::zeros((n_queries, k)); - let neighbors = ManagedTensor::from(&neighbors_host).to_device(&res)?; - - let mut distances_host = ndarray::Array::::zeros((n_queries, k)); - let distances = ManagedTensor::from(&distances_host).to_device(&res)?; - - let search_params = SearchParams::new()?; - - index.search(&res, &search_params, &queries, &neighbors, &distances)?; - - // Copy back to host memory - distances.to_host(&res, &mut distances_host)?; - neighbors.to_host(&res, &mut neighbors_host)?; - - // nearest neighbors should be themselves, since queries are from the - // dataset - println!("Neighbors {:?}", neighbors_host); - println!("Distances {:?}", distances_host); + // Build a CAGRA index over your dataset. + let dataset = GpuTensor; + let index_params = IndexParams::builder() + .metric(DistanceType::L2Expanded) + .graph_degree(64) + .build()?; + let index = Index::build(&res, &index_params, &dataset)?; + + // Search for the k nearest neighbors of each query, writing the results into + // the neighbor and distance device buffers. + let queries = GpuTensor; + let (mut neighbors, mut distances) = (GpuTensor, GpuTensor); + let search_params = SearchParams::builder().itopk_size(64).build()?; + index.search(&res, &search_params, &queries, &mut neighbors, &mut distances)?; + Ok(()) } ``` -For more code examples of the Rust APIs, including a drop-in project templates, please refer to the [Rust examples](https://github.com/rapidsai/cuvs/tree/main/examples/rust). ## Contributing diff --git a/examples/rust/Cargo.toml b/examples/rust/Cargo.toml deleted file mode 100644 index 7955e08e69..0000000000 --- a/examples/rust/Cargo.toml +++ /dev/null @@ -1,9 +0,0 @@ -[package] -name = "cuvs-rust-example" -version = "0.1.0" -edition = "2021" - -[dependencies] -cuvs = ">=24.6.0" -ndarray = "0.15" -ndarray-rand = "0.14" diff --git a/examples/rust/README.md b/examples/rust/README.md deleted file mode 100644 index 00058de5e8..0000000000 --- a/examples/rust/README.md +++ /dev/null @@ -1,33 +0,0 @@ -# cuVS Rust Example - -This template project provides a drop-in sample to for using cuVS in a rust project. - -First, please refer to our [installation docs](https://docs.rapids.ai/api/cuvs/stable/build.html#cuda-gpu-requirements) for the minimum requirements to use cuVS. Note that you will have to have the libcuvs.so and libcuvs_c.so binaries installed to compile this example project. - -Once the minimum requirements are satisfied, this example template application can be built using 'cargo build', and this directory can be copied directly in order to build a new application with cuVS. - -You may follow these steps to quickly get set up: - -```bash -conda env create --name rust -f conda/environments/rust_cuda-133_arch-$(uname -m).yaml -conda activate rust -``` -You may prefer to use `mamba`, as it provides significant speedup over `conda`. - -1. Set up the required environment variables: -```bash -LIBCLANG_PATH=$(dirname "$(find "$CONDA_PREFIX" -name libclang.so | head -n 1)") -export LIBCLANG_PATH -echo "LIBCLANG_PATH=$LIBCLANG_PATH" -``` - -2. Add the cuvs dependency in your project: -```TOML -# Cargo.toml -[dependencies] -cuvs = ">=24.6.0" -#...rest of your dependencies -``` -Then you can run your project with `cargo run`. - -See [main.rs](./src/main.rs) for an example implementation. diff --git a/examples/rust/src/main.rs b/examples/rust/src/main.rs deleted file mode 100644 index b118c3abc5..0000000000 --- a/examples/rust/src/main.rs +++ /dev/null @@ -1,67 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION. - * SPDX-License-Identifier: Apache-2.0 - */ - -use cuvs::cagra::{Index, IndexParams, SearchParams}; -use cuvs::{ManagedTensor, Resources, Result}; - -use ndarray::s; -use ndarray_rand::rand_distr::Uniform; -use ndarray_rand::RandomExt; - -/// Example showing how to index and search data with CAGRA -fn cagra_example() -> Result<()> { - let res = Resources::new()?; - - // Create a new random dataset to index - let n_datapoints = 65536; - let n_features = 512; - let dataset = - ndarray::Array::::random((n_datapoints, n_features), Uniform::new(0., 1.0)); - - // build the cagra index - let build_params = IndexParams::new()?; - let index = Index::build(&res, &build_params, &dataset)?; - println!( - "Indexed {}x{} datapoints into cagra index", - n_datapoints, n_features - ); - - // 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, ..]); - - let k = 10; - - // CAGRA 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 = ManagedTensor::from(&queries).to_device(&res)?; - let mut neighbors_host = ndarray::Array::::zeros((n_queries, k)); - let neighbors = ManagedTensor::from(&neighbors_host).to_device(&res)?; - - let mut distances_host = ndarray::Array::::zeros((n_queries, k)); - let distances = ManagedTensor::from(&distances_host).to_device(&res)?; - - let search_params = SearchParams::new()?; - - index.search(&res, &search_params, &queries, &neighbors, &distances)?; - - // Copy back to host memory - distances.to_host(&res, &mut distances_host)?; - neighbors.to_host(&res, &mut neighbors_host)?; - - // nearest neighbors should be themselves, since queries are from the - // dataset - println!("Neighbors {:?}", neighbors_host); - println!("Distances {:?}", distances_host); - Ok(()) -} - -fn main() { - if let Err(e) = cagra_example() { - println!("Failed to run CAGRA: {:?}", e); - } -} diff --git a/rust/cuvs/src/lib.rs b/rust/cuvs/src/lib.rs index d9b7017e97..71338b38da 100644 --- a/rust/cuvs/src/lib.rs +++ b/rust/cuvs/src/lib.rs @@ -21,3 +21,9 @@ pub(crate) mod test_utils; pub use dlpack::{AsDlTensor, AsDlTensorMut, DLPackError, DLTensorView, DLTensorViewMut, DType}; pub use error::LibraryError; pub use resources::Resources; + +// Compile the Rust code blocks in the top-level README as doctests so the +// documented examples can't drift from the API. +#[cfg(doctest)] +#[doc = include_str!("../../../README.md")] +pub struct ReadmeDocTests; diff --git a/rust/cuvs/src/neighbors/cagra/params.rs b/rust/cuvs/src/neighbors/cagra/params.rs index 03df19aab1..c0204a08f8 100644 --- a/rust/cuvs/src/neighbors/cagra/params.rs +++ b/rust/cuvs/src/neighbors/cagra/params.rs @@ -249,9 +249,9 @@ impl SearchParams { ) -> Result { let params = Self::try_new()?; - let effective_algo = algo.unwrap_or(unsafe { (*params.handle).algo.into() }); + let effective_algo = algo.unwrap_or_else(|| unsafe { (*params.handle).algo.into() }); let effective_hashmap_mode = - hashmap_mode.unwrap_or(unsafe { (*params.handle).hashmap_mode.into() }); + hashmap_mode.unwrap_or_else(|| unsafe { (*params.handle).hashmap_mode.into() }); if let Some(n) = itopk_size && effective_algo == SearchAlgo::SingleCta diff --git a/rust/cuvs/src/neighbors/ivf_pq/mod.rs b/rust/cuvs/src/neighbors/ivf_pq/mod.rs index abdf146bf3..bce0b54cc0 100644 --- a/rust/cuvs/src/neighbors/ivf_pq/mod.rs +++ b/rust/cuvs/src/neighbors/ivf_pq/mod.rs @@ -17,13 +17,108 @@ mod index; mod params; pub use index::Index; -pub use params::{ - IndexParams, SearchParams, cudaDataType_t, cuvsIvfPqCodebookGen, cuvsIvfPqListLayout, -}; +pub use params::{IndexParams, SearchParams}; use crate::dlpack::DLPackError; use crate::error::LibraryError; +/// Strategy for creating PQ codebooks. +#[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)] +#[non_exhaustive] +pub enum CodebookGen { + /// One codebook per PQ subspace. + PerSubspace, + /// One codebook per IVF cluster. + PerCluster, +} + +impl From for ffi::cuvsIvfPqCodebookGen { + fn from(v: CodebookGen) -> Self { + match v { + CodebookGen::PerSubspace => Self::CUVS_IVF_PQ_CODEBOOK_GEN_PER_SUBSPACE, + CodebookGen::PerCluster => Self::CUVS_IVF_PQ_CODEBOOK_GEN_PER_CLUSTER, + } + } +} + +impl From for CodebookGen { + fn from(v: ffi::cuvsIvfPqCodebookGen) -> Self { + match v { + ffi::cuvsIvfPqCodebookGen::CUVS_IVF_PQ_CODEBOOK_GEN_PER_SUBSPACE => Self::PerSubspace, + ffi::cuvsIvfPqCodebookGen::CUVS_IVF_PQ_CODEBOOK_GEN_PER_CLUSTER => Self::PerCluster, + } + } +} + +/// Memory layout of the IVF-PQ list data. +#[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)] +#[non_exhaustive] +pub enum ListLayout { + /// Codes stored contiguously, one vector's codes after another. + Flat, + /// Codes interleaved for optimized search performance (the default). + Interleaved, +} + +impl From for ffi::cuvsIvfPqListLayout { + fn from(v: ListLayout) -> Self { + match v { + ListLayout::Flat => Self::CUVS_IVF_PQ_LIST_LAYOUT_FLAT, + ListLayout::Interleaved => Self::CUVS_IVF_PQ_LIST_LAYOUT_INTERLEAVED, + } + } +} + +impl From for ListLayout { + fn from(v: ffi::cuvsIvfPqListLayout) -> Self { + match v { + ffi::cuvsIvfPqListLayout::CUVS_IVF_PQ_LIST_LAYOUT_FLAT => Self::Flat, + ffi::cuvsIvfPqListLayout::CUVS_IVF_PQ_LIST_LAYOUT_INTERLEAVED => Self::Interleaved, + } + } +} + +/// Lookup-table dtype used during IVF-PQ search. +#[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)] +#[non_exhaustive] +pub enum LutDType { + /// 32-bit floating-point lookup tables. + F32, + /// 16-bit floating-point lookup tables. + F16, + /// 8-bit unsigned lookup tables. + U8, +} + +impl From for ffi::cudaDataType_t { + fn from(v: LutDType) -> Self { + match v { + LutDType::F32 => Self::CUDA_R_32F, + LutDType::F16 => Self::CUDA_R_16F, + LutDType::U8 => Self::CUDA_R_8U, + } + } +} + +/// Accumulator dtype used for internal IVF-PQ distance computation. +#[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)] +#[non_exhaustive] +pub enum InternalDistanceDType { + /// 32-bit floating-point accumulators. + F32, + /// 16-bit floating-point accumulators. + F16, +} + +impl From for ffi::cudaDataType_t { + fn from(v: InternalDistanceDType) -> Self { + match v { + InternalDistanceDType::F32 => Self::CUDA_R_32F, + InternalDistanceDType::F16 => Self::CUDA_R_16F, + } + } +} + /// Error type for IVF-PQ operations. #[derive(Debug, thiserror::Error)] pub enum IvfPqError { diff --git a/rust/cuvs/src/neighbors/ivf_pq/params.rs b/rust/cuvs/src/neighbors/ivf_pq/params.rs index 66a678a93f..818880173d 100644 --- a/rust/cuvs/src/neighbors/ivf_pq/params.rs +++ b/rust/cuvs/src/neighbors/ivf_pq/params.rs @@ -15,9 +15,7 @@ use bon::bon; use crate::distance::DistanceType; use crate::error::check_cuvs; -use super::IvfPqError; - -pub use ffi::{cudaDataType_t, cuvsIvfPqCodebookGen, cuvsIvfPqListLayout}; +use super::{CodebookGen, InternalDistanceDType, IvfPqError, ListLayout, LutDType}; /// Parameters for building an IVF-PQ index. pub struct IndexParams { @@ -35,8 +33,8 @@ impl IndexParams { kmeans_trainset_fraction: Option, pq_bits: Option, pq_dim: Option, - codebook_kind: Option, - codes_layout: Option, + codebook_kind: Option, + codes_layout: Option, force_random_rotation: Option, max_train_points_per_pq_code: Option, add_data_on_build: Option, @@ -63,10 +61,10 @@ impl IndexParams { (*params.handle).pq_dim = v; } if let Some(v) = codebook_kind { - (*params.handle).codebook_kind = v; + (*params.handle).codebook_kind = v.into(); } if let Some(v) = codes_layout { - (*params.handle).codes_layout = v; + (*params.handle).codes_layout = v.into(); } if let Some(v) = force_random_rotation { (*params.handle).force_random_rotation = v; @@ -117,8 +115,8 @@ impl SearchParams { #[builder] pub fn new( n_probes: Option, - lut_dtype: Option, - internal_distance_dtype: Option, + lut_dtype: Option, + internal_distance_dtype: Option, ) -> Result { let params = Self::try_new()?; unsafe { @@ -126,10 +124,10 @@ impl SearchParams { (*params.handle).n_probes = v; } if let Some(v) = lut_dtype { - (*params.handle).lut_dtype = v; + (*params.handle).lut_dtype = v.into(); } if let Some(v) = internal_distance_dtype { - (*params.handle).internal_distance_dtype = v; + (*params.handle).internal_distance_dtype = v.into(); } } Ok(params)