cuvs/
brute_force.rs

1use std::{marker::PhantomData, ptr, ptr::NonNull};
2
3use crate::{
4    device::{DeviceBuffer, DeviceElement, MatrixLayout, ValidatedDeviceTensor},
5    dlpack::DlpackTensor,
6    error::{Result, check_cuvs, invalid_input, invalid_state, require_nonzero},
7    resources::CuvsResources,
8};
9
10#[derive(Debug, Clone, Copy, PartialEq, Eq)]
11pub enum DistanceMetric {
12    L2Expanded,
13    L2SqrtExpanded,
14    CosineExpanded,
15    InnerProduct,
16}
17
18impl DistanceMetric {
19    pub(crate) fn raw(self) -> cuvs_sys::cuvsDistanceType {
20        match self {
21            DistanceMetric::L2Expanded => cuvs_sys::cuvsDistanceType::L2Expanded,
22            DistanceMetric::L2SqrtExpanded => cuvs_sys::cuvsDistanceType::L2SqrtExpanded,
23            DistanceMetric::CosineExpanded => cuvs_sys::cuvsDistanceType::CosineExpanded,
24            DistanceMetric::InnerProduct => cuvs_sys::cuvsDistanceType::InnerProduct,
25        }
26    }
27}
28
29#[derive(Debug, Clone, Copy, PartialEq, Eq)]
30pub struct KnnSearchOutput {
31    pub n_queries: usize,
32    pub k: usize,
33}
34
35/// A brute-force index that retains the dataset borrow and exact runtime binding.
36pub struct BruteForceIndex<'data> {
37    raw: Option<NonNull<cuvs_sys::cuvsBruteForceIndex>>,
38    resources: CuvsResources,
39    rows: usize,
40    columns: usize,
41    _dataset: PhantomData<&'data f32>,
42}
43
44impl<'data> BruteForceIndex<'data> {
45    /// Builds an index over `dataset`, a contiguous row-major or column-major
46    /// matrix on the same exact runtime binding and CUDA device as `resources`.
47    ///
48    /// The index borrows `dataset` for `'data` and retains `resources`, so the
49    /// dataset memory stays valid until the index is closed or dropped. A
50    /// failed native build destroys the partial index before the error is
51    /// returned.
52    pub fn build(
53        resources: &CuvsResources,
54        dataset: &ValidatedDeviceTensor<'data, f32>,
55        metric: DistanceMetric,
56    ) -> Result<Self> {
57        let operation = "BruteForceIndex::build";
58        validate_matrix_binding(resources, dataset, operation)?;
59        let rows = dataset.rows()?;
60        let columns = dataset.columns()?;
61        matrix_layout(dataset, operation)?;
62
63        resources.with_native(operation, |native| {
64            let mut raw = ptr::null_mut();
65            unsafe { check_cuvs(cuvs_sys::cuvsBruteForceIndexCreate(&mut raw), operation)? };
66            let raw = NonNull::new(raw).ok_or_else(|| invalid_state(operation, "cuVS returned a null index"))?;
67            let mut dataset_tensor = DlpackTensor::from_const(dataset.as_tensor());
68            let status = unsafe {
69                cuvs_sys::cuvsBruteForceBuild(native, dataset_tensor.as_mut_ptr(), metric.raw(), 0.0, raw.as_ptr())
70            };
71            if let Err(error) = check_cuvs(status, operation) {
72                unsafe {
73                    let _ = cuvs_sys::cuvsBruteForceIndexDestroy(raw.as_ptr());
74                }
75                return Err(error);
76            }
77            Ok(Self { raw: Some(raw), resources: resources.retain(), rows, columns, _dataset: PhantomData })
78        })
79    }
80
81    /// Searches `queries` for their `k` nearest indexed rows and writes the
82    /// results into `neighbors` and `distances`.
83    ///
84    /// `k` must be nonzero and at most the indexed row count. `queries` must
85    /// have the index's column count and share its binding and device. Both
86    /// output buffers must contain exactly `n_queries * k` elements, laid out
87    /// as `[n_queries, k]` row-major. Success is returned only after the
88    /// native outputs pass validation on the buffers' execution binding: every
89    /// neighbor index lies in `0..rows()` and every distance is finite.
90    pub fn search_into(
91        &self,
92        queries: &ValidatedDeviceTensor<'_, f32>,
93        k: usize,
94        neighbors: &mut DeviceBuffer<i64>,
95        distances: &mut DeviceBuffer<f32>,
96    ) -> Result<KnnSearchOutput> {
97        let operation = "BruteForceIndex::search_into";
98        require_nonzero(k, operation, "k")?;
99        if k > self.rows {
100            return Err(invalid_input(operation, format!("k ({k}) must be <= indexed rows ({})", self.rows)));
101        }
102        validate_matrix_binding(&self.resources, queries, operation)?;
103        matrix_layout(queries, operation)?;
104        let n_queries = queries.rows()?;
105        if queries.columns()? != self.columns {
106            return Err(invalid_input(
107                operation,
108                format!("query columns {} must match index columns {}", queries.columns()?, self.columns),
109            ));
110        }
111        let output_len = n_queries
112            .checked_mul(k)
113            .ok_or_else(|| invalid_input(operation, "search output element count overflowed usize"))?;
114        if neighbors.len() != output_len || distances.len() != output_len {
115            return Err(invalid_input(
116                operation,
117                format!(
118                    "neighbors and distances must each contain n_queries * k ({output_len}) elements; found {} and {}",
119                    neighbors.len(),
120                    distances.len()
121                ),
122            ));
123        }
124
125        let mut neighbors_view = neighbors.as_tensor_mut(&self.resources, &[n_queries, k], None, 0)?;
126        let mut distances_view = distances.as_tensor_mut(&self.resources, &[n_queries, k], None, 0)?;
127        let raw = self.raw.ok_or_else(|| invalid_state(operation, "the brute-force index has been closed"))?;
128        self.resources.with_native(operation, |native| {
129            let mut queries_tensor = DlpackTensor::from_const(queries.as_tensor());
130            let mut neighbors_tensor = DlpackTensor::from_mut(&mut neighbors_view);
131            let mut distances_tensor = DlpackTensor::from_mut(&mut distances_view);
132            let prefilter = cuvs_sys::cuvsFilter { addr: 0, type_: cuvs_sys::cuvsFilterType::NO_FILTER };
133            unsafe {
134                check_cuvs(
135                    cuvs_sys::cuvsBruteForceSearch(
136                        native,
137                        raw.as_ptr(),
138                        queries_tensor.as_mut_ptr(),
139                        neighbors_tensor.as_mut_ptr(),
140                        distances_tensor.as_mut_ptr(),
141                        prefilter,
142                    ),
143                    operation,
144                )
145            }
146        })?;
147        let indexed_rows = i64::try_from(self.rows)
148            .map_err(|_| invalid_input(operation, "indexed row count exceeds the Int64 neighbor range"))?;
149        neighbors.validate_native_index_range(0, indexed_rows, operation)?;
150        distances.validate_native_input(operation)?;
151        Ok(KnnSearchOutput { n_queries, k })
152    }
153
154    pub fn rows(&self) -> usize {
155        self.rows
156    }
157
158    pub fn columns(&self) -> usize {
159        self.columns
160    }
161
162    /// Reports native destruction errors; failed destruction is retried on drop.
163    pub fn close(mut self) -> Result<()> {
164        self.close_inner("BruteForceIndex::close")
165    }
166
167    fn close_inner(&mut self, operation: &'static str) -> Result<()> {
168        let Some(raw) = self.raw.take() else {
169            return Ok(());
170        };
171        let result = self.resources.with_native(operation, |_| unsafe {
172            check_cuvs(cuvs_sys::cuvsBruteForceIndexDestroy(raw.as_ptr()), operation)
173        });
174        if result.is_err() {
175            self.raw = Some(raw);
176        }
177        result
178    }
179}
180
181impl Drop for BruteForceIndex<'_> {
182    fn drop(&mut self) {
183        let _ = self.close_inner("BruteForceIndex::drop");
184    }
185}
186
187impl std::fmt::Debug for BruteForceIndex<'_> {
188    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
189        formatter
190            .debug_struct("BruteForceIndex")
191            .field("rows", &self.rows)
192            .field("columns", &self.columns)
193            .field("closed", &self.raw.is_none())
194            .finish_non_exhaustive()
195    }
196}
197
198pub(crate) fn validate_matrix_binding<T: DeviceElement>(
199    resources: &CuvsResources,
200    matrix: &ValidatedDeviceTensor<'_, T>,
201    operation: &'static str,
202) -> Result<()> {
203    resources.validate_exact_binding(operation)?;
204    if matrix.binding_id() != resources.binding_id() {
205        return Err(invalid_input(operation, "matrix and cuVS resources must use the same exact runtime binding")
206            .with_context("tensor_binding_id", matrix.binding_id())
207            .with_context("resource_binding_id", resources.binding_id()));
208    }
209    if matrix.device() != resources.device().ordinal() as i32 {
210        return Err(invalid_input(operation, "matrix and cuVS resources must use the same CUDA device")
211            .with_context("tensor_device", matrix.device())
212            .with_context("resource_device", resources.device().ordinal()));
213    }
214    let _ = matrix.rows()?;
215    let _ = matrix.columns()?;
216    Ok(())
217}
218
219pub(crate) fn matrix_layout<T: DeviceElement>(
220    matrix: &ValidatedDeviceTensor<'_, T>,
221    operation: &'static str,
222) -> Result<MatrixLayout> {
223    let rows = matrix.rows()? as i64;
224    let columns = matrix.columns()? as i64;
225    match matrix.strides() {
226        None => Ok(MatrixLayout::RowMajor),
227        Some([row_stride, column_stride]) if *row_stride == columns && *column_stride == 1 => {
228            Ok(MatrixLayout::RowMajor)
229        },
230        Some([row_stride, column_stride]) if *row_stride == 1 && *column_stride == rows => {
231            Ok(MatrixLayout::ColumnMajor)
232        },
233        _ => Err(invalid_input(operation, "cuVS matrices must be contiguous row-major or column-major tensors")),
234    }
235}
cuvs/brute_force.rs

cuvs/
brute_force.rs
