cuvs/
ivf_pq.rs

1use std::{marker::PhantomData, ptr, ptr::NonNull};
2
3use cudf::TypeId;
4
5use crate::{
6    brute_force::{DistanceMetric, KnnSearchOutput, matrix_layout, validate_matrix_binding},
7    device::{DeviceBuffer, MatrixLayout, ValidatedDeviceTensor},
8    dlpack::DlpackTensor,
9    elements::IvfElement,
10    error::{Result, check_cuvs, from_cudf, invalid_input, invalid_state, require_nonzero, unsupported},
11    params_handle::ParamsHandle,
12    resources::CuvsResources,
13};
14
15#[derive(Debug, Clone, Copy, PartialEq, Eq)]
16pub enum IvfPqCodebookKind {
17    PerSubspace,
18    PerCluster,
19}
20
21#[derive(Debug, Clone, Copy)]
22pub struct IvfPqIndexParams {
23    pub metric: DistanceMetric,
24    pub n_lists: u32,
25    pub kmeans_n_iters: u32,
26    pub kmeans_trainset_fraction: f64,
27    pub pq_dim: u32,
28    pub pq_bits: u32,
29    pub codebook_kind: IvfPqCodebookKind,
30}
31
32impl IvfPqIndexParams {
33    pub fn new(metric: DistanceMetric, n_lists: u32) -> Self {
34        Self {
35            metric,
36            n_lists,
37            kmeans_n_iters: 20,
38            kmeans_trainset_fraction: 0.5,
39            pq_dim: 0,
40            pq_bits: 8,
41            codebook_kind: IvfPqCodebookKind::PerSubspace,
42        }
43    }
44}
45
46#[derive(Debug, Clone, Copy)]
47pub struct IvfPqSearchParams {
48    pub n_probes: u32,
49}
50
51impl Default for IvfPqSearchParams {
52    fn default() -> Self {
53        Self { n_probes: 20 }
54    }
55}
56
57/// Query-local owned IVF-PQ index. Native build copies the indexed vectors.
58pub struct IvfPqIndex<T: IvfElement> {
59    raw: Option<NonNull<cuvs_sys::cuvsIvfPqIndex>>,
60    resources: CuvsResources,
61    rows: usize,
62    columns: usize,
63    n_lists: u32,
64    metric: DistanceMetric,
65    _element: PhantomData<T>,
66}
67
68impl<T: IvfElement> IvfPqIndex<T> {
69    pub fn build(
70        resources: &CuvsResources,
71        dataset: &ValidatedDeviceTensor<'_, T>,
72        params: IvfPqIndexParams,
73    ) -> Result<Self> {
74        let operation = "IvfPqIndex::build";
75        validate_matrix_binding(resources, dataset, operation)?;
76        if matrix_layout(dataset, operation)? != MatrixLayout::RowMajor {
77            return Err(invalid_input(operation, "IVF-PQ dataset must be contiguous row-major"));
78        }
79        let rows = dataset.rows()?;
80        let columns = dataset.columns()?;
81        require_nonzero(rows, operation, "dataset rows")?;
82        if params.n_lists == 0 || params.n_lists as usize > rows {
83            return Err(invalid_input(
84                operation,
85                format!("n_lists ({}) must be in 1..=dataset rows ({rows})", params.n_lists),
86            ));
87        }
88        if params.kmeans_n_iters == 0
89            || !params.kmeans_trainset_fraction.is_finite()
90            || !(0.0 < params.kmeans_trainset_fraction && params.kmeans_trainset_fraction <= 1.0)
91        {
92            return Err(invalid_input(
93                operation,
94                "kmeans_n_iters must be positive and kmeans_trainset_fraction must be in (0, 1]",
95            ));
96        }
97        if !(4..=8).contains(&params.pq_bits)
98            || (params.pq_dim != 0 && !(params.pq_dim as u64 * params.pq_bits as u64).is_multiple_of(8))
99        {
100            return Err(invalid_input(
101                operation,
102                "pq_bits must be in 4..=8 and pq_dim * pq_bits must be a multiple of 8",
103            ));
104        }
105        if params.metric == DistanceMetric::CosineExpanded && columns < 2 {
106            return Err(invalid_input(operation, "cosine IVF-PQ requires at least two dimensions"));
107        }
108        // ivf_pq_search.cuh:658-665 skips integer cosine score rescaling.
109        if params.metric == DistanceMetric::CosineExpanded && T::TYPE_ID != TypeId::Float32 {
110            return Err(unsupported(
111                operation,
112                "integer IVF-PQ cosine returns distances on the native integer score scale",
113            )
114            .with_context("element", format!("{:?}", T::TYPE_ID))
115            .with_context("metric", "CosineExpanded"));
116        }
117        resources.with_native(operation, |native| {
118            let mut config = ParamsHandle::new(
119                cuvs_sys::cuvsIvfPqIndexParamsCreate,
120                cuvs_sys::cuvsIvfPqIndexParamsDestroy,
121                operation,
122                "IVF-PQ params",
123            )?;
124            unsafe {
125                let config_ptr = config.as_ptr();
126                (*config_ptr).metric = params.metric.raw();
127                (*config_ptr).n_lists = params.n_lists;
128                (*config_ptr).kmeans_n_iters = params.kmeans_n_iters;
129                (*config_ptr).kmeans_trainset_fraction = params.kmeans_trainset_fraction;
130                (*config_ptr).pq_dim = params.pq_dim;
131                (*config_ptr).pq_bits = params.pq_bits;
132                (*config_ptr).codebook_kind = match params.codebook_kind {
133                    IvfPqCodebookKind::PerSubspace => {
134                        cuvs_sys::cuvsIvfPqCodebookGen::CUVS_IVF_PQ_CODEBOOK_GEN_PER_SUBSPACE
135                    },
136                    IvfPqCodebookKind::PerCluster => {
137                        cuvs_sys::cuvsIvfPqCodebookGen::CUVS_IVF_PQ_CODEBOOK_GEN_PER_CLUSTER
138                    },
139                };
140            }
141            let mut raw = ptr::null_mut();
142            unsafe { check_cuvs(cuvs_sys::cuvsIvfPqIndexCreate(&mut raw), operation)? };
143            let raw = NonNull::new(raw).ok_or_else(|| invalid_state(operation, "cuVS returned a null IVF-PQ index"))?;
144            let mut tensor = DlpackTensor::from_const(dataset.as_tensor());
145            let built = unsafe {
146                check_cuvs(
147                    cuvs_sys::cuvsIvfPqBuild(native, config.as_ptr(), tensor.as_mut_ptr(), raw.as_ptr()),
148                    operation,
149                )
150            };
151            let closed = config.close(operation);
152            if let Err(error) = built.and(closed) {
153                unsafe {
154                    let _ = cuvs_sys::cuvsIvfPqIndexDestroy(raw.as_ptr());
155                }
156                return Err(error);
157            }
158            Ok(Self {
159                raw: Some(raw),
160                resources: resources.retain(),
161                rows,
162                columns,
163                n_lists: params.n_lists,
164                metric: params.metric,
165                _element: PhantomData,
166            })
167        })
168    }
169
170    /// Outputs are `[queries.rows(), k]` row-major. Missing neighbors have ordinal `i64::MAX`.
171    pub fn search_into(
172        &self,
173        queries: &ValidatedDeviceTensor<'_, T>,
174        k: usize,
175        params: IvfPqSearchParams,
176        neighbors: &mut DeviceBuffer<i64>,
177        distances: &mut DeviceBuffer<f32>,
178    ) -> Result<KnnSearchOutput> {
179        let operation = "IvfPqIndex::search_into";
180        require_nonzero(k, operation, "k")?;
181        if k > self.rows || k > u32::MAX as usize {
182            return Err(invalid_input(
183                operation,
184                format!("k ({k}) must be <= indexed rows ({}) and fit u32", self.rows),
185            ));
186        }
187        if params.n_probes == 0 || params.n_probes > self.n_lists {
188            return Err(invalid_input(
189                operation,
190                format!("n_probes ({}) must be in 1..=n_lists ({})", params.n_probes, self.n_lists),
191            ));
192        }
193        validate_matrix_binding(&self.resources, queries, operation)?;
194        if matrix_layout(queries, operation)? != MatrixLayout::RowMajor || queries.columns()? != self.columns {
195            return Err(invalid_input(
196                operation,
197                "IVF-PQ queries must be contiguous row-major with the indexed dimension",
198            ));
199        }
200        let n_queries = queries.rows()?;
201        let output_len = n_queries
202            .checked_mul(k)
203            .ok_or_else(|| invalid_input(operation, "search output element count overflowed usize"))?;
204        if neighbors.len() != output_len || distances.len() != output_len {
205            return Err(invalid_input(
206                operation,
207                format!("neighbors and distances must each contain n_queries * k ({output_len}) elements"),
208            ));
209        }
210        let mut neighbors_view = neighbors.as_tensor_mut(&self.resources, &[n_queries, k], None, 0)?;
211        let mut distances_view = distances.as_tensor_mut(&self.resources, &[n_queries, k], None, 0)?;
212        let raw = self.raw.ok_or_else(|| invalid_state(operation, "the IVF-PQ index has been closed"))?;
213        self.resources.with_native(operation, |native| {
214            let mut config = ParamsHandle::new(
215                cuvs_sys::cuvsIvfPqSearchParamsCreate,
216                cuvs_sys::cuvsIvfPqSearchParamsDestroy,
217                operation,
218                "IVF-PQ search params",
219            )?;
220            unsafe {
221                (*config.as_ptr()).n_probes = params.n_probes;
222            }
223            let mut query = DlpackTensor::from_const(queries.as_tensor());
224            let mut neighbor = DlpackTensor::from_mut(&mut neighbors_view);
225            let mut distance = DlpackTensor::from_mut(&mut distances_view);
226            let result = unsafe {
227                check_cuvs(
228                    cuvs_sys::cuvsIvfPqSearch(
229                        native,
230                        config.as_ptr(),
231                        raw.as_ptr(),
232                        query.as_mut_ptr(),
233                        neighbor.as_mut_ptr(),
234                        distance.as_mut_ptr(),
235                    ),
236                    operation,
237                )
238            };
239            result.and(config.close(operation))
240        })?;
241        let indexed_rows =
242            i64::try_from(self.rows).map_err(|_| invalid_input(operation, "indexed row count exceeds Int64"))?;
243        if self.metric == DistanceMetric::CosineExpanded {
244            // ivf_pq_search.cuh:658-665 leaves negative cosine scores unconverted; 1 + score is cosine distance.
245            distances.inner_mut().add_scalar_to_float32(1.0).map_err(|error| from_cudf(operation, error))?;
246        }
247        // cuVS may return an in-range placeholder ID with an infinite distance for a missing slot.
248        neighbors
249            .inner_mut()
250            .fill_int64_where_float32_infinite(distances.inner(), i64::MAX)
251            .map_err(|error| from_cudf(operation, error))?;
252        neighbors.validate_native_index_range_or_sentinel(0, indexed_rows, i64::MAX, operation)?;
253        distances.validate_native_distances_for_valid_neighbors(neighbors, i64::MAX, operation)?;
254        Ok(KnnSearchOutput { n_queries, k })
255    }
256
257    pub fn rows(&self) -> usize {
258        self.rows
259    }
260    pub fn columns(&self) -> usize {
261        self.columns
262    }
263
264    /// Report native destruction errors; a failed destruction is retried on drop.
265    pub fn close(mut self) -> Result<()> {
266        self.close_inner("IvfPqIndex::close")
267    }
268
269    fn close_inner(&mut self, operation: &'static str) -> Result<()> {
270        let Some(raw) = self.raw.take() else { return Ok(()) };
271        let result = self.resources.with_native(operation, |_| unsafe {
272            check_cuvs(cuvs_sys::cuvsIvfPqIndexDestroy(raw.as_ptr()), operation)
273        });
274        if result.is_err() {
275            self.raw = Some(raw);
276        }
277        result
278    }
279}
280
281impl<T: IvfElement> Drop for IvfPqIndex<T> {
282    fn drop(&mut self) {
283        let _ = self.close_inner("IvfPqIndex::drop");
284    }
285}
286
287impl<T: IvfElement> std::fmt::Debug for IvfPqIndex<T> {
288    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
289        formatter
290            .debug_struct("IvfPqIndex")
291            .field("rows", &self.rows)
292            .field("columns", &self.columns)
293            .finish_non_exhaustive()
294    }
295}
cuvs/ivf_pq.rs

cuvs/
ivf_pq.rs
