cuvs 26.8.1

RAPIDS vector search library
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135

/*
 * 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};

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<CodebookGen> 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<ffi::cuvsIvfPqCodebookGen> 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<ListLayout> 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<ffi::cuvsIvfPqListLayout> 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<LutDType> 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<InternalDistanceDType> 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)]
#[non_exhaustive]
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),
    /// A parameter value failed validation.
    #[error("invalid parameter: {0}")]
    Validation(String),
}