cuvs/
cagra.rs

1use std::{marker::PhantomData, mem, ptr, ptr::NonNull};
2
3use crate::{
4    brute_force::{DistanceMetric, KnnSearchOutput, matrix_layout, validate_matrix_binding},
5    device::{DeviceBuffer, MatrixLayout, ValidatedDeviceTensor},
6    dlpack::DlpackTensor,
7    elements::CagraElement,
8    error::{Result, check_cuvs, from_cudf, invalid_input, invalid_state, require_nonzero, unsupported},
9    params_handle::ParamsHandle,
10    resources::CuvsResources,
11};
12
13#[derive(Debug, Clone, Copy, PartialEq, Eq)]
14pub enum CagraBuildAlgo {
15    IvfPq,
16    NnDescent,
17}
18
19#[derive(Debug, Clone, Copy, PartialEq, Eq)]
20pub enum CagraSearchAlgo {
21    Auto,
22    SingleCta,
23    MultiCta,
24    MultiKernel,
25}
26
27#[derive(Debug, Clone, Copy)]
28pub struct CagraIndexParams {
29    pub metric: DistanceMetric,
30    pub graph_degree: usize,
31    pub intermediate_graph_degree: usize,
32    pub build_algo: CagraBuildAlgo,
33}
34
35impl CagraIndexParams {
36    pub fn new(metric: DistanceMetric) -> Self {
37        Self { metric, graph_degree: 64, intermediate_graph_degree: 128, build_algo: CagraBuildAlgo::IvfPq }
38    }
39}
40
41#[derive(Debug, Clone, Copy)]
42pub struct CagraSearchParams {
43    pub itopk_size: usize,
44    pub search_width: usize,
45    pub max_iterations: usize,
46    pub search_algo: CagraSearchAlgo,
47}
48
49impl Default for CagraSearchParams {
50    fn default() -> Self {
51        Self { itopk_size: 64, search_width: 1, max_iterations: 0, search_algo: CagraSearchAlgo::Auto }
52    }
53}
54
55// cagra.cpp:1959-1994 allocates IVF-PQ graph params on create but destroys them
56// according to the current build_algo. NN-descent must restore IVF_PQ on every exit.
57unsafe extern "C" fn destroy_index_params(params: *mut cuvs_sys::cuvsCagraIndexParams) -> cuvs_sys::cuvsError_t {
58    unsafe {
59        (*params).build_algo = cuvs_sys::cuvsCagraGraphBuildAlgo::IVF_PQ;
60        cuvs_sys::cuvsCagraIndexParamsDestroy(params)
61    }
62}
63
64struct DatasetHandle(NonNull<cuvs_sys::cuvsDataset>);
65
66impl Drop for DatasetHandle {
67    fn drop(&mut self) {
68        unsafe {
69            let _ = cuvs_sys::cuvsDatasetDestroy(self.0.as_ptr());
70        }
71    }
72}
73
74/// Query-local CAGRA index with a device-padded dataset. Rows already aligned
75/// to 16 bytes use a borrowed padded view; other rows use an owning padded copy.
76/// Native search requires a device-padded index (components/cuvs/cuvs/c/src/neighbors/cagra.cpp:1843).
77pub struct CagraIndex<'data, T: CagraElement> {
78    raw: Option<NonNull<cuvs_sys::cuvsCagraIndex>>,
79    dataset: DatasetHandle,
80    resources: CuvsResources,
81    rows: usize,
82    columns: usize,
83    _dataset_borrow: PhantomData<&'data T>,
84}
85
86impl<'data, T: CagraElement> CagraIndex<'data, T> {
87    pub fn build(
88        resources: &CuvsResources,
89        dataset: &ValidatedDeviceTensor<'data, T>,
90        params: CagraIndexParams,
91    ) -> Result<Self> {
92        let operation = "CagraIndex::build";
93        validate_matrix_binding(resources, dataset, operation)?;
94        if matrix_layout(dataset, operation)? != MatrixLayout::RowMajor {
95            return Err(invalid_input(operation, "CAGRA dataset must be contiguous row-major"));
96        }
97        let rows = dataset.rows()?;
98        let columns = dataset.columns()?;
99        require_nonzero(rows, operation, "dataset rows")?;
100        if rows > u32::MAX as usize {
101            return Err(invalid_input(operation, "CAGRA graph indices require at most UInt32::MAX dataset rows"));
102        }
103        if params.graph_degree == 0
104            || params.graph_degree >= rows
105            || params.intermediate_graph_degree <= params.graph_degree
106            || params.intermediate_graph_degree >= rows
107        {
108            return Err(invalid_input(
109                operation,
110                format!(
111                    "graph_degree must be in 1..dataset rows ({rows}) and intermediate_graph_degree must be greater than graph_degree and less than dataset rows"
112                ),
113            ));
114        }
115        if params.metric == DistanceMetric::L2SqrtExpanded
116            || (params.metric == DistanceMetric::CosineExpanded && params.build_algo == CagraBuildAlgo::NnDescent)
117        {
118            return Err(invalid_input(operation, "unsupported CAGRA metric and build algorithm combination"));
119        }
120        let parts = dataset.as_tensor().parts();
121        let row_pitch_aligned = columns.is_multiple_of(16 / mem::size_of::<T>());
122        let effective_address = parts.data.as_ptr().wrapping_byte_add(parts.byte_offset).addr();
123        if row_pitch_aligned && !effective_address.is_multiple_of(16) {
124            // The pinned owning-copy path rejects device input with an already padded row pitch.
125            return Err(unsupported(
126                operation,
127                "CAGRA requires a 16-byte-aligned dataset address when row pitch is padded",
128            ));
129        }
130        resources.with_native(operation, |native| {
131            let mut config = ParamsHandle::new(
132                cuvs_sys::cuvsCagraIndexParamsCreate,
133                destroy_index_params,
134                operation,
135                "CAGRA params",
136            )?;
137            unsafe {
138                let ptr = config.as_ptr();
139                (*ptr).metric = params.metric.raw();
140                (*ptr).graph_degree = params.graph_degree;
141                (*ptr).intermediate_graph_degree = params.intermediate_graph_degree;
142                (*ptr).build_algo = match params.build_algo {
143                    CagraBuildAlgo::IvfPq => cuvs_sys::cuvsCagraGraphBuildAlgo::IVF_PQ,
144                    CagraBuildAlgo::NnDescent => cuvs_sys::cuvsCagraGraphBuildAlgo::NN_DESCENT,
145                };
146            }
147            let mut tensor = DlpackTensor::from_const(dataset.as_tensor());
148            let mut padded = ptr::null_mut();
149            // common.hpp:1068-1078,1157-1165 rejects an owning copy when the
150            // device row width already has the required 16-byte alignment.
151            unsafe {
152                if row_pitch_aligned {
153                    check_cuvs(
154                        cuvs_sys::cuvsDatasetMakePaddedView(native, tensor.as_mut_ptr(), &mut padded),
155                        operation,
156                    )?;
157                } else {
158                    check_cuvs(
159                        cuvs_sys::cuvsDatasetMakePadded(
160                            native,
161                            tensor.as_mut_ptr(),
162                            cuvs_sys::cuvsDatasetMemType_t::CUVS_DATASET_MEM_TYPE_DEVICE,
163                            &mut padded,
164                        ),
165                        operation,
166                    )?;
167                }
168            };
169            let padded = DatasetHandle(
170                NonNull::new(padded).ok_or_else(|| invalid_state(operation, "cuVS returned a null padded dataset"))?,
171            );
172            let mut raw = ptr::null_mut();
173            unsafe { check_cuvs(cuvs_sys::cuvsCagraIndexCreate(&mut raw), operation)? };
174            let raw = NonNull::new(raw).ok_or_else(|| invalid_state(operation, "cuVS returned a null CAGRA index"))?;
175            let built = unsafe {
176                check_cuvs(
177                    cuvs_sys::cuvsCagraBuild(native, config.as_ptr(), padded.0.as_ptr(), raw.as_ptr()),
178                    operation,
179                )
180            };
181            let closed = config.close(operation);
182            if let Err(error) = built.and(closed) {
183                unsafe {
184                    let _ = cuvs_sys::cuvsCagraIndexDestroy(raw.as_ptr());
185                }
186                return Err(error);
187            }
188            Ok(Self {
189                raw: Some(raw),
190                dataset: padded,
191                resources: resources.retain(),
192                rows,
193                columns,
194                _dataset_borrow: PhantomData,
195            })
196        })
197    }
198
199    /// Search outputs are `[queries.rows(), k]`. Missing neighbors have ordinal `i64::MAX`.
200    pub fn search_into(
201        &self,
202        queries: &ValidatedDeviceTensor<'_, T>,
203        k: usize,
204        params: CagraSearchParams,
205        neighbors: &mut DeviceBuffer<i64>,
206        distances: &mut DeviceBuffer<f32>,
207    ) -> Result<KnnSearchOutput> {
208        let operation = "CagraIndex::search_into";
209        require_nonzero(k, operation, "k")?;
210        if k > self.rows
211            || k > u32::MAX as usize
212            || params.itopk_size < k
213            || params.search_width == 0
214            || (params.search_algo == CagraSearchAlgo::SingleCta && params.itopk_size > 1024)
215        {
216            return Err(invalid_input(
217                operation,
218                "k must not exceed dataset rows or itopk_size; search_width must be positive and single_cta itopk_size must not exceed 1024",
219            ));
220        }
221        validate_matrix_binding(&self.resources, queries, operation)?;
222        if matrix_layout(queries, operation)? != MatrixLayout::RowMajor || queries.columns()? != self.columns {
223            return Err(invalid_input(
224                operation,
225                "CAGRA queries must be contiguous row-major with the indexed dimension",
226            ));
227        }
228        let n_queries = queries.rows()?;
229        let output_len = n_queries
230            .checked_mul(k)
231            .ok_or_else(|| invalid_input(operation, "search output element count overflowed usize"))?;
232        if neighbors.len() != output_len || distances.len() != output_len {
233            return Err(invalid_input(
234                operation,
235                format!("neighbors and distances must each contain n_queries * k ({output_len}) elements"),
236            ));
237        }
238        let mut neighbors_view = neighbors.as_tensor_mut(&self.resources, &[n_queries, k], None, 0)?;
239        let mut distances_view = distances.as_tensor_mut(&self.resources, &[n_queries, k], None, 0)?;
240        let raw = self.raw.ok_or_else(|| invalid_state(operation, "the CAGRA index has been closed"))?;
241        self.resources.with_native(operation, |native| {
242            let mut config = ParamsHandle::new(
243                cuvs_sys::cuvsCagraSearchParamsCreate,
244                cuvs_sys::cuvsCagraSearchParamsDestroy,
245                operation,
246                "CAGRA search params",
247            )?;
248            unsafe {
249                let ptr = config.as_ptr();
250                (*ptr).itopk_size = params.itopk_size;
251                (*ptr).search_width = params.search_width;
252                (*ptr).max_iterations = params.max_iterations;
253                (*ptr).algo = match params.search_algo {
254                    CagraSearchAlgo::Auto => cuvs_sys::cuvsCagraSearchAlgo::AUTO,
255                    CagraSearchAlgo::SingleCta => cuvs_sys::cuvsCagraSearchAlgo::SINGLE_CTA,
256                    CagraSearchAlgo::MultiCta => cuvs_sys::cuvsCagraSearchAlgo::MULTI_CTA,
257                    CagraSearchAlgo::MultiKernel => cuvs_sys::cuvsCagraSearchAlgo::MULTI_KERNEL,
258                };
259            }
260            let mut query = DlpackTensor::from_const(queries.as_tensor());
261            let mut neighbor = DlpackTensor::from_mut(&mut neighbors_view);
262            let mut distance = DlpackTensor::from_mut(&mut distances_view);
263            let filter = cuvs_sys::cuvsFilter { addr: 0, type_: cuvs_sys::cuvsFilterType::NO_FILTER };
264            let result = unsafe {
265                check_cuvs(
266                    cuvs_sys::cuvsCagraSearch(
267                        native,
268                        config.as_ptr(),
269                        raw.as_ptr(),
270                        query.as_mut_ptr(),
271                        neighbor.as_mut_ptr(),
272                        distance.as_mut_ptr(),
273                        filter,
274                    ),
275                    operation,
276                )
277            };
278            result.and(config.close(operation))
279        })?;
280        // Native CAGRA writes UINT32_MAX and FLT_MAX for exhausted graph slots,
281        // including when the requested output index type is Int64.
282        neighbors
283            .inner_mut()
284            .replace_int64_value_and_set_paired_float32_infinite(distances.inner_mut(), u32::MAX as i64, i64::MAX)
285            .map_err(|error| from_cudf(operation, error))?;
286        neighbors
287            .inner_mut()
288            .fill_int64_where_float32_infinite(distances.inner(), i64::MAX)
289            .map_err(|error| from_cudf(operation, error))?;
290        neighbors.validate_native_index_range_or_sentinel(
291            0,
292            i64::try_from(self.rows).map_err(|_| invalid_input(operation, "indexed row count exceeds Int64"))?,
293            i64::MAX,
294            operation,
295        )?;
296        distances.validate_native_distances_for_valid_neighbors(neighbors, i64::MAX, operation)?;
297        Ok(KnnSearchOutput { n_queries, k })
298    }
299
300    pub fn rows(&self) -> usize {
301        self.rows
302    }
303    pub fn columns(&self) -> usize {
304        self.columns
305    }
306
307    /// Reports native destruction errors; failed destruction is retried on drop.
308    pub fn close(mut self) -> Result<()> {
309        self.close_inner("CagraIndex::close")
310    }
311
312    fn close_inner(&mut self, operation: &'static str) -> Result<()> {
313        let Some(raw) = self.raw.take() else { return Ok(()) };
314        let result = self.resources.with_native(operation, |_| unsafe {
315            check_cuvs(cuvs_sys::cuvsCagraIndexDestroy(raw.as_ptr()), operation)
316        });
317        if result.is_err() {
318            self.raw = Some(raw);
319        }
320        result
321    }
322}
323
324impl<T: CagraElement> Drop for CagraIndex<'_, T> {
325    fn drop(&mut self) {
326        let _ = self.close_inner("CagraIndex::drop");
327    }
328}
329
330impl<T: CagraElement> std::fmt::Debug for CagraIndex<'_, T> {
331    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
332        formatter
333            .debug_struct("CagraIndex")
334            .field("rows", &self.rows)
335            .field("columns", &self.columns)
336            .field("dataset", &self.dataset.0)
337            .finish_non_exhaustive()
338    }
339}
cuvs/cagra.rs

cuvs/
cagra.rs
