IVF PQ

Source header: cuvs/neighbors/ivf_pq.hpp

IVF-PQ index build parameters

neighbors::ivf_pq::codebook_gen

A type for specifying how PQ codebooks are created.

1
enum class codebook_gen { ... };

neighbors::ivf_pq::list_layout

A type for specifying the memory layout of PQ codes in IVF lists.

1
enum class list_layout { ... };

neighbors::ivf_pq::index_params::from_dataset

Creates index_params based on shape of the input dataset.

1
static index_params from_dataset(
2
raft::matrix_extent<int64_t> dataset,
3
cuvs::distance::DistanceType metric = cuvs::distance::DistanceType::L2Expanded);

Usage example:

Parameters

Returns

static index_params

IVF-PQ index search parameters

neighbors::ivf_pq::search_params

IVF-PQ index search parameters

1
struct search_params : cuvs::neighbors::search_params { ... };

Fields

Types

list_extents

PQ-encoded data stored in the interleaved format:

1
using list_extents = raft::
2
extents<SizeT, raft::dynamic_extent, raft::dynamic_extent, kIndexGroupSize, kIndexGroupVecLen>;

neighbors::ivf_pq::list_spec_flat

Flat (non-interleaved) storage specification for PQ-encoded data.

This stores each vector’s PQ codes contiguously: [n_rows, bytes_per_vector] where bytes_per_vector = ceildiv(pq_dim * pq_bits, 8)

1
template <typename SizeT, typename IdxT>
2
struct list_spec_flat { ... };

Fields

neighbors::graph_build_params::ivf_pq_params

Specialized parameters utilizing IVF-PQ to build knn graph

1
struct ivf_pq_params { ... };

Fields

IVF-PQ index

neighbors::ivf_pq::index

IVF-PQ index.

In the IVF-PQ index, a database vector y is approximated with two level quantization:

y = Q_1(y) + Q_2(y - Q_1(y))

The first level quantizer (Q_1), maps the vector y to the nearest cluster center. The number of clusters is n_lists.

The second quantizer encodes the residual, and it is defined as a product quantizer [1].

A product quantizer encodes a dim dimensional vector with a pq_dim dimensional vector. First we split the input vector into pq_dim subvectors (denoted by u), where each u vector contains pq_len distinct components of y

y_1, y_2, … y_{pq_len}, y_{pq_len+1}, … y_{2*pq_len}, … y_{dim-pq_len+1} … y_{dim} u_1 u_2 u_{pq_dim}

Then each subvector encoded with a separate quantizer q_i, end the results are concatenated

Q_2(y) = q_1(u_1),q_2(u_2),…,q_{pq_dim}(u_pq_dim})

Each quantizer q_i outputs a code with pq_bit bits. The second level quantizers are also defined by k-means clustering in the corresponding sub-space: the reproduction values are the centroids, and the set of reproduction values is the codebook.

When the data dimensionality dim is not multiple of pq_dim, the feature space is transformed using a random orthogonal matrix to have rot_dim = pq_dim * pq_len dimensions (rot_dim >= dim).

The second-level quantizers are trained either for each subspace or for each cluster: (a) codebook_gen::PER_SUBSPACE: creates pq_dim second-level quantizers - one for each slice of the data along features; (b) codebook_gen::PER_CLUSTER: creates n_lists second-level quantizers - one for each first-level cluster. In either case, the centroids are again found using k-means clustering interpreting the data as having pq_len dimensions.

[1] Product quantization for nearest neighbor search Herve Jegou, Matthijs Douze, Cordelia Schmid

1
template <typename IdxT>
2
class index : public index_iface<IdxT>, cuvs::neighbors::index { ... };

neighbors::ivf_pq::index::index

Construct an empty index.

1
index(raft::resources const& handle);

Constructs an empty index. This index will either need to be trained with build or loaded from a saved copy with deserialize

Parameters

Returns

void

Additional overload: neighbors::ivf_pq::index::index

Construct an index with specified parameters.

1
index(raft::resources const& handle,
2
cuvs::distance::DistanceType metric,
3
codebook_gen codebook_kind,
4
uint32_t n_lists,
5
uint32_t dim,
6
uint32_t pq_bits                    = 8,
7
uint32_t pq_dim                     = 0,
8
bool conservative_memory_allocation = false);

This constructor creates an owning index with the given parameters.

Parameters

Returns

void

Additional overload: neighbors::ivf_pq::index::index

Construct an index from index parameters.

1
index(raft::resources const& handle, const index_params& params, uint32_t dim);

Parameters

Returns

void

neighbors::ivf_pq::index::size

Total length of the index.

1
IdxT size() const noexcept override;

Returns

IdxT

neighbors::ivf_pq::index::dim

Dimensionality of the input data.

1
uint32_t dim() const noexcept override;

Returns

uint32_t

neighbors::ivf_pq::index::dim_ext

Dimensionality of the cluster centers:

1
uint32_t dim_ext() const noexcept;

input data dim extended with vector norms and padded to 8 elems.

Returns

uint32_t

neighbors::ivf_pq::index::rot_dim

Dimensionality of the data after transforming it for PQ processing

1
uint32_t rot_dim() const noexcept;

(rotated and augmented to be muplitple of pq_dim).

Returns

uint32_t

neighbors::ivf_pq::index::pq_bits

The bit length of an encoded vector element after compression by PQ.

1
uint32_t pq_bits() const noexcept override;

Returns

uint32_t

neighbors::ivf_pq::index::pq_dim

The dimensionality of an encoded vector after compression by PQ.

1
uint32_t pq_dim() const noexcept override;

Returns

uint32_t

neighbors::ivf_pq::index::pq_len

Dimensionality of a subspace, i.e. the number of vector components mapped to a subspace

1
uint32_t pq_len() const noexcept;

Returns

uint32_t

neighbors::ivf_pq::index::pq_book_size

The number of vectors in a PQ codebook (1 << pq_bits).

1
uint32_t pq_book_size() const noexcept;

Returns

uint32_t

neighbors::ivf_pq::index::metric

Distance metric used for clustering.

1
cuvs::distance::DistanceType metric() const noexcept override;

Returns

cuvs::distance::DistanceType

neighbors::ivf_pq::index::codebook_kind

How PQ codebooks are created.

1
codebook_gen codebook_kind() const noexcept override;

Returns

codebook_gen

neighbors::ivf_pq::index::codes_layout

Memory layout of PQ codes in IVF lists.

1
list_layout codes_layout() const noexcept override;

Returns

list_layout

neighbors::ivf_pq::index::n_lists

Number of clusters/inverted lists (first level quantization).

1
uint32_t n_lists() const noexcept;

Returns

uint32_t

neighbors::ivf_pq::index::conservative_memory_allocation

Whether to use conservative memory allocation when extending the list (cluster) data

1
bool conservative_memory_allocation() const noexcept override;

(see index_params.conservative_memory_allocation).

Returns

bool

neighbors::ivf_pq::index::pq_centers

PQ cluster centers

1
raft::device_mdspan<const float, pq_centers_extents, raft::row_major> pq_centers()
2
const noexcept override;

• codebook_gen::PER_SUBSPACE: [pq_dim , pq_len, pq_book_size]
• codebook_gen::PER_CLUSTER: [n_lists, pq_len, pq_book_size]

Returns

raft::device_mdspan<const float, pq_centers_extents, raft::row_major>

neighbors::ivf_pq::index::lists

Lists’ data and indices (polymorphic, works for both FLAT and INTERLEAVED layouts).

1
std::vector<std::shared_ptr<list_data_base<IdxT>>>& lists() noexcept override;

Returns

std::vector<std::shared_ptr<list_data_base<IdxT>>>&

neighbors::ivf_pq::index::data_ptrs

Pointers to the inverted lists (clusters) data [n_lists].

1
raft::device_vector_view<uint8_t*, uint32_t, raft::row_major> data_ptrs() noexcept override;

Returns

raft::device_vector_view<uint8_t*, uint32_t, raft::row_major>

neighbors::ivf_pq::index::inds_ptrs

Pointers to the inverted lists (clusters) indices [n_lists].

1
raft::device_vector_view<IdxT*, uint32_t, raft::row_major> inds_ptrs() noexcept override;

Returns

raft::device_vector_view<IdxT*, uint32_t, raft::row_major>

neighbors::ivf_pq::index::rotation_matrix

The transform matrix (original space -> rotated padded space) [rot_dim, dim]

1
raft::device_matrix_view<const float, uint32_t, raft::row_major> rotation_matrix()
2
const noexcept override;

Returns

raft::device_matrix_view<const float, uint32_t, raft::row_major>

neighbors::ivf_pq::index::accum_sorted_sizes

Accumulated list sizes, sorted in descending order [n_lists + 1].

1
raft::host_vector_view<IdxT, uint32_t, raft::row_major> accum_sorted_sizes() noexcept override;

The last value contains the total length of the index. The value at index zero is always zero.

That is, the content of this span is as if the list_sizes was sorted and then accumulated.

This span is used during search to estimate the maximum size of the workspace.

Returns

raft::host_vector_view<IdxT, uint32_t, raft::row_major>

neighbors::ivf_pq::index::list_sizes

Sizes of the lists [n_lists].

1
raft::device_vector_view<uint32_t, uint32_t, raft::row_major> list_sizes() noexcept override;

Returns

raft::device_vector_view<uint32_t, uint32_t, raft::row_major>

neighbors::ivf_pq::index::centers

Cluster centers corresponding to the lists in the original space [n_lists, dim_ext]

1
raft::device_matrix_view<const float, uint32_t, raft::row_major> centers()
2
const noexcept override;

Returns

raft::device_matrix_view<const float, uint32_t, raft::row_major>

neighbors::ivf_pq::index::centers_rot

Cluster centers corresponding to the lists in the rotated space [n_lists, rot_dim]

1
raft::device_matrix_view<const float, uint32_t, raft::row_major> centers_rot()
2
const noexcept override;

Returns

raft::device_matrix_view<const float, uint32_t, raft::row_major>

neighbors::ivf_pq::index::get_list_size_in_bytes

fetch size of a particular IVF list in bytes using the list extents.

1
uint32_t get_list_size_in_bytes(uint32_t label) const override;

Usage example:

Parameters

Returns

uint32_t

Additional overload: neighbors::ivf_pq::index::index

Construct index from implementation pointer.

1
explicit index(std::unique_ptr<index_iface<IdxT>> impl);

This constructor is used internally by build/extend/deserialize functions.

Parameters

Returns

explicit

IVF-PQ index build

neighbors::ivf_pq::build

Build the index from the dataset for efficient search.

1
auto build(raft::resources const& handle,
2
const cuvs::neighbors::ivf_pq::index_params& index_params,
3
raft::device_matrix_view<const float, int64_t, raft::row_major> dataset)
4
-> cuvs::neighbors::ivf_pq::index<int64_t>;

Usage example:

Parameters

Returns

cuvs::neighbors::ivf_pq::index<int64_t>

Additional overload: neighbors::ivf_pq::build

Build the index from the dataset for efficient search.

1
void build(raft::resources const& handle,
2
const cuvs::neighbors::ivf_pq::index_params& index_params,
3
raft::device_matrix_view<const float, int64_t, raft::row_major> dataset,
4
cuvs::neighbors::ivf_pq::index<int64_t>* idx);

NB: Currently, the following distance metrics are supported:

• L2Expanded
• L2Unexpanded
• InnerProduct
• CosineExpanded

Usage example:

Parameters

Returns

void

Additional overload: neighbors::ivf_pq::build

Build the index from the dataset for efficient search.

1
auto build(raft::resources const& handle,
2
const cuvs::neighbors::ivf_pq::index_params& index_params,
3
raft::device_matrix_view<const half, int64_t, raft::row_major> dataset)
4
-> cuvs::neighbors::ivf_pq::index<int64_t>;

Usage example:

Parameters

Returns

cuvs::neighbors::ivf_pq::index<int64_t>

Additional overload: neighbors::ivf_pq::build

Build the index from the dataset for efficient search.

1
void build(raft::resources const& handle,
2
const cuvs::neighbors::ivf_pq::index_params& index_params,
3
raft::device_matrix_view<const half, int64_t, raft::row_major> dataset,
4
cuvs::neighbors::ivf_pq::index<int64_t>* idx);

NB: Currently, the following distance metrics are supported:

• L2Expanded
• L2Unexpanded
• InnerProduct
• CosineExpanded

Usage example:

Parameters

Returns

void

Additional overload: neighbors::ivf_pq::build

Build the index from the dataset for efficient search.

1
auto build(raft::resources const& handle,
2
const cuvs::neighbors::ivf_pq::index_params& index_params,
3
raft::device_matrix_view<const int8_t, int64_t, raft::row_major> dataset)
4
-> cuvs::neighbors::ivf_pq::index<int64_t>;

Usage example:

Parameters

Returns

cuvs::neighbors::ivf_pq::index<int64_t>

Additional overload: neighbors::ivf_pq::build

Build the index from the dataset for efficient search.

1
void build(raft::resources const& handle,
2
const cuvs::neighbors::ivf_pq::index_params& index_params,
3
raft::device_matrix_view<const int8_t, int64_t, raft::row_major> dataset,
4
cuvs::neighbors::ivf_pq::index<int64_t>* idx);

NB: Currently, the following distance metrics are supported:

• L2Expanded
• L2Unexpanded
• InnerProduct
• CosineExpanded

Usage example:

Parameters

Returns

void

Additional overload: neighbors::ivf_pq::build

Build the index from the dataset for efficient search.

1
auto build(raft::resources const& handle,
2
const cuvs::neighbors::ivf_pq::index_params& index_params,
3
raft::device_matrix_view<const uint8_t, int64_t, raft::row_major> dataset)
4
-> cuvs::neighbors::ivf_pq::index<int64_t>;

Usage example:

Parameters

Returns

cuvs::neighbors::ivf_pq::index<int64_t>

Additional overload: neighbors::ivf_pq::build

Build the index from the dataset for efficient search.

1
void build(raft::resources const& handle,
2
const cuvs::neighbors::ivf_pq::index_params& index_params,
3
raft::device_matrix_view<const uint8_t, int64_t, raft::row_major> dataset,
4
cuvs::neighbors::ivf_pq::index<int64_t>* idx);

NB: Currently, the following distance metrics are supported:

• L2Expanded
• L2Unexpanded
• InnerProduct
• CosineExpanded

Usage example:

Parameters

Returns

void

Additional overload: neighbors::ivf_pq::build

Build the index from the dataset for efficient search.

1
auto build(raft::resources const& handle,
2
const cuvs::neighbors::ivf_pq::index_params& index_params,
3
raft::host_matrix_view<const float, int64_t, raft::row_major> dataset)
4
-> cuvs::neighbors::ivf_pq::index<int64_t>;

Note, if index_params.add_data_on_build is set to true, the user can set a stream pool in the input raft::resource with at least one stream to enable kernel and copy overlapping.

Usage example:

Parameters

Returns

cuvs::neighbors::ivf_pq::index<int64_t>

Additional overload: neighbors::ivf_pq::build

Build the index from the dataset for efficient search.

1
void build(raft::resources const& handle,
2
const cuvs::neighbors::ivf_pq::index_params& index_params,
3
raft::host_matrix_view<const float, int64_t, raft::row_major> dataset,
4
cuvs::neighbors::ivf_pq::index<int64_t>* idx);

NB: Currently, the following distance metrics are supported:

• L2Expanded
• L2Unexpanded
• InnerProduct
• CosineExpanded

Note, if index_params.add_data_on_build is set to true, the user can set a stream pool in the input raft::resource with at least one stream to enable kernel and copy overlapping.

Usage example:

Parameters

Returns

void

Additional overload: neighbors::ivf_pq::build

Build the index from the dataset for efficient search.

1
auto build(raft::resources const& handle,
2
const cuvs::neighbors::ivf_pq::index_params& index_params,
3
raft::host_matrix_view<const half, int64_t, raft::row_major> dataset)
4
-> cuvs::neighbors::ivf_pq::index<int64_t>;

Note, if index_params.add_data_on_build is set to true, the user can set a stream pool in the input raft::resource with at least one stream to enable kernel and copy overlapping.

Usage example:

Parameters

Returns

cuvs::neighbors::ivf_pq::index<int64_t>

Additional overload: neighbors::ivf_pq::build

Build the index from the dataset for efficient search.

1
void build(raft::resources const& handle,
2
const cuvs::neighbors::ivf_pq::index_params& index_params,
3
raft::host_matrix_view<const half, int64_t, raft::row_major> dataset,
4
cuvs::neighbors::ivf_pq::index<int64_t>* idx);

NB: Currently, the following distance metrics are supported:

• L2Expanded
• L2Unexpanded
• InnerProduct
• CosineExpanded

Usage example:

Parameters

Returns

void

Additional overload: neighbors::ivf_pq::build

Build the index from the dataset for efficient search.

1
auto build(raft::resources const& handle,
2
const cuvs::neighbors::ivf_pq::index_params& index_params,
3
raft::host_matrix_view<const int8_t, int64_t, raft::row_major> dataset)
4
-> cuvs::neighbors::ivf_pq::index<int64_t>;

Usage example:

Parameters

Returns

cuvs::neighbors::ivf_pq::index<int64_t>

Additional overload: neighbors::ivf_pq::build

Build the index from the dataset for efficient search.

1
void build(raft::resources const& handle,
2
const cuvs::neighbors::ivf_pq::index_params& index_params,
3
raft::host_matrix_view<const int8_t, int64_t, raft::row_major> dataset,
4
cuvs::neighbors::ivf_pq::index<int64_t>* idx);

NB: Currently, the following distance metrics are supported:

• L2Expanded
• L2Unexpanded
• InnerProduct
• CosineExpanded

Note, if index_params.add_data_on_build is set to true, the user can set a stream pool in the input raft::resource with at least one stream to enable kernel and copy overlapping.

Usage example:

Parameters

Returns

void

Additional overload: neighbors::ivf_pq::build

Build the index from the dataset for efficient search.

1
auto build(raft::resources const& handle,
2
const cuvs::neighbors::ivf_pq::index_params& index_params,
3
raft::host_matrix_view<const uint8_t, int64_t, raft::row_major> dataset)
4
-> cuvs::neighbors::ivf_pq::index<int64_t>;

Note, if index_params.add_data_on_build is set to true, the user can set a stream pool in the input raft::resource with at least one stream to enable kernel and copy overlapping.

Usage example:

Parameters

Returns

cuvs::neighbors::ivf_pq::index<int64_t>

Additional overload: neighbors::ivf_pq::build

Build the index from the dataset for efficient search.

1
void build(raft::resources const& handle,
2
const cuvs::neighbors::ivf_pq::index_params& index_params,
3
raft::host_matrix_view<const uint8_t, int64_t, raft::row_major> dataset,
4
cuvs::neighbors::ivf_pq::index<int64_t>* idx);

NB: Currently, the following distance metrics are supported:

• L2Expanded
• L2Unexpanded
• InnerProduct
• CosineExpanded

Note, if index_params.add_data_on_build is set to true, the user can set a stream pool in the input raft::resource with at least one stream to enable kernel and copy overlapping.

Usage example:

Parameters

Returns

void

Additional overload: neighbors::ivf_pq::build

Build a view-type IVF-PQ index from device memory centroids and codebook.

1
auto build(raft::resources const& handle,
2
const cuvs::neighbors::ivf_pq::index_params& index_params,
3
const uint32_t dim,
4
raft::device_mdspan<const float, raft::extent_3d<uint32_t>, raft::row_major> pq_centers,
5
raft::device_matrix_view<const float, uint32_t, raft::row_major> centers,
6
raft::device_matrix_view<const float, uint32_t, raft::row_major> centers_rot,
7
raft::device_matrix_view<const float, uint32_t, raft::row_major> rotation_matrix)
8
-> cuvs::neighbors::ivf_pq::index<int64_t>;

This function creates a non-owning index that stores a reference to the provided device data. All parameters must be provided with correct extents. The caller is responsible for ensuring the lifetime of the input data exceeds the lifetime of the returned index.

The index_params must be consistent with the provided matrices. Specifically:

• index_params.codebook_kind determines the expected shape of pq_centers
• index_params.metric will be stored in the index
• index_params.conservative_memory_allocation will be stored in the index The function will verify consistency between index_params, dim, and the matrix extents.

dim]

Parameters

Returns

cuvs::neighbors::ivf_pq::index<int64_t>

Additional overload: neighbors::ivf_pq::build

Build an IVF-PQ index from device memory centroids and codebook.

1
void build(raft::resources const& handle,
2
const cuvs::neighbors::ivf_pq::index_params& index_params,
3
const uint32_t dim,
4
raft::device_mdspan<const float, raft::extent_3d<uint32_t>, raft::row_major> pq_centers,
5
raft::device_matrix_view<const float, uint32_t, raft::row_major> centers,
6
raft::device_matrix_view<const float, uint32_t, raft::row_major> centers_rot,
7
raft::device_matrix_view<const float, uint32_t, raft::row_major> rotation_matrix,
8
cuvs::neighbors::ivf_pq::index<int64_t>* idx);

This function creates a non-owning index that references the provided device data directly. All parameters must be provided with correct extents. The caller is responsible for ensuring the lifetime of the input data exceeds the lifetime of the returned index.

The index_params must be consistent with the provided matrices. Specifically:

• index_params.codebook_kind determines the expected shape of pq_centers
• index_params.metric will be stored in the index
• index_params.conservative_memory_allocation will be stored in the index The function will verify consistency between index_params, dim, and the matrix extents.

dim]

Parameters

Returns

void

Additional overload: neighbors::ivf_pq::build

Build an IVF-PQ index from host memory centroids and codebook (in-place).

1
auto build(
2
raft::resources const& handle,
3
const cuvs::neighbors::ivf_pq::index_params& index_params,
4
const uint32_t dim,
5
raft::host_mdspan<const float, raft::extent_3d<uint32_t>, raft::row_major> pq_centers,
6
raft::host_matrix_view<const float, uint32_t, raft::row_major> centers,
7
std::optional<raft::host_matrix_view<const float, uint32_t, raft::row_major>> centers_rot,
8
std::optional<raft::host_matrix_view<const float, uint32_t, raft::row_major>> rotation_matrix)
9
-> cuvs::neighbors::ivf_pq::index<int64_t>;

Parameters

Returns

cuvs::neighbors::ivf_pq::index<int64_t>

Additional overload: neighbors::ivf_pq::build

Build an IVF-PQ index from host memory centroids and codebook (in-place).

1
void build(
2
raft::resources const& handle,
3
const cuvs::neighbors::ivf_pq::index_params& index_params,
4
const uint32_t dim,
5
raft::host_mdspan<const float, raft::extent_3d<uint32_t>, raft::row_major> pq_centers,
6
raft::host_matrix_view<const float, uint32_t, raft::row_major> centers,
7
std::optional<raft::host_matrix_view<const float, uint32_t, raft::row_major>> centers_rot,
8
std::optional<raft::host_matrix_view<const float, uint32_t, raft::row_major>> rotation_matrix,
9
cuvs::neighbors::ivf_pq::index<int64_t>* idx);

Parameters

Returns

void

IVF-PQ index extend

neighbors::ivf_pq::extend

Extend the index with the new data.

1
auto extend(raft::resources const& handle,
2
raft::device_matrix_view<const float, int64_t, raft::row_major> new_vectors,
3
std::optional<raft::device_vector_view<const int64_t, int64_t>> new_indices,
4
const cuvs::neighbors::ivf_pq::index<int64_t>& idx)
5
-> cuvs::neighbors::ivf_pq::index<int64_t>;

Usage example:

Parameters

Returns

cuvs::neighbors::ivf_pq::index<int64_t>

Additional overload: neighbors::ivf_pq::extend

Extend the index with the new data.

1
void extend(raft::resources const& handle,
2
raft::device_matrix_view<const float, int64_t, raft::row_major> new_vectors,
3
std::optional<raft::device_vector_view<const int64_t, int64_t>> new_indices,
4
cuvs::neighbors::ivf_pq::index<int64_t>* idx);

Usage example:

Parameters

Returns

void

Additional overload: neighbors::ivf_pq::extend

Extend the index with the new data.

1
auto extend(raft::resources const& handle,
2
raft::device_matrix_view<const half, int64_t, raft::row_major> new_vectors,
3
std::optional<raft::device_vector_view<const int64_t, int64_t>> new_indices,
4
const cuvs::neighbors::ivf_pq::index<int64_t>& idx)
5
-> cuvs::neighbors::ivf_pq::index<int64_t>;

Usage example:

Parameters

Returns

cuvs::neighbors::ivf_pq::index<int64_t>

Additional overload: neighbors::ivf_pq::extend

Extend the index with the new data.

1
void extend(raft::resources const& handle,
2
raft::device_matrix_view<const half, int64_t, raft::row_major> new_vectors,
3
std::optional<raft::device_vector_view<const int64_t, int64_t>> new_indices,
4
cuvs::neighbors::ivf_pq::index<int64_t>* idx);

Usage example:

Parameters

Returns

void

Additional overload: neighbors::ivf_pq::extend

Extend the index with the new data.

1
auto extend(raft::resources const& handle,
2
raft::device_matrix_view<const int8_t, int64_t, raft::row_major> new_vectors,
3
std::optional<raft::device_vector_view<const int64_t, int64_t>> new_indices,
4
const cuvs::neighbors::ivf_pq::index<int64_t>& idx)
5
-> cuvs::neighbors::ivf_pq::index<int64_t>;

Usage example:

Parameters

Returns

cuvs::neighbors::ivf_pq::index<int64_t>

Additional overload: neighbors::ivf_pq::extend

Extend the index with the new data.

1
void extend(raft::resources const& handle,
2
raft::device_matrix_view<const int8_t, int64_t, raft::row_major> new_vectors,
3
std::optional<raft::device_vector_view<const int64_t, int64_t>> new_indices,
4
cuvs::neighbors::ivf_pq::index<int64_t>* idx);

Usage example:

Parameters

Returns

void

Additional overload: neighbors::ivf_pq::extend

Extend the index with the new data.

1
auto extend(raft::resources const& handle,
2
raft::device_matrix_view<const uint8_t, int64_t, raft::row_major> new_vectors,
3
std::optional<raft::device_vector_view<const int64_t, int64_t>> new_indices,
4
const cuvs::neighbors::ivf_pq::index<int64_t>& idx)
5
-> cuvs::neighbors::ivf_pq::index<int64_t>;

Usage example:

Parameters

Returns

cuvs::neighbors::ivf_pq::index<int64_t>

Additional overload: neighbors::ivf_pq::extend

Extend the index with the new data.

1
void extend(raft::resources const& handle,
2
raft::device_matrix_view<const uint8_t, int64_t, raft::row_major> new_vectors,
3
std::optional<raft::device_vector_view<const int64_t, int64_t>> new_indices,
4
cuvs::neighbors::ivf_pq::index<int64_t>* idx);

Usage example:

Parameters

Returns

void

Additional overload: neighbors::ivf_pq::extend

Extend the index with the new data.

1
auto extend(raft::resources const& handle,
2
raft::host_matrix_view<const float, int64_t, raft::row_major> new_vectors,
3
std::optional<raft::host_vector_view<const int64_t, int64_t>> new_indices,
4
const cuvs::neighbors::ivf_pq::index<int64_t>& idx)
5
-> cuvs::neighbors::ivf_pq::index<int64_t>;

Note, the user can set a stream pool in the input raft::resource with at least one stream to enable kernel and copy overlapping.

Usage example:

Parameters

Returns

cuvs::neighbors::ivf_pq::index<int64_t>

Additional overload: neighbors::ivf_pq::extend

Extend the index with the new data.

1
void extend(raft::resources const& handle,
2
raft::host_matrix_view<const float, int64_t, raft::row_major> new_vectors,
3
std::optional<raft::host_vector_view<const int64_t, int64_t>> new_indices,
4
cuvs::neighbors::ivf_pq::index<int64_t>* idx);

Note, the user can set a stream pool in the input raft::resource with at least one stream to enable kernel and copy overlapping.

Usage example:

Parameters

Returns

void

Additional overload: neighbors::ivf_pq::extend

Extend the index with the new data.

1
auto extend(raft::resources const& handle,
2
raft::host_matrix_view<const half, int64_t, raft::row_major> new_vectors,
3
std::optional<raft::host_vector_view<const int64_t, int64_t>> new_indices,
4
const cuvs::neighbors::ivf_pq::index<int64_t>& idx)
5
-> cuvs::neighbors::ivf_pq::index<int64_t>;

Note, the user can set a stream pool in the input raft::resource with at least one stream to enable kernel and copy overlapping.

Usage example:

Parameters

Returns

cuvs::neighbors::ivf_pq::index<int64_t>

Additional overload: neighbors::ivf_pq::extend

Extend the index with the new data.

1
void extend(raft::resources const& handle,
2
raft::host_matrix_view<const half, int64_t, raft::row_major> new_vectors,
3
std::optional<raft::host_vector_view<const int64_t, int64_t>> new_indices,
4
cuvs::neighbors::ivf_pq::index<int64_t>* idx);

Note, the user can set a stream pool in the input raft::resource with at least one stream to enable kernel and copy overlapping.

Usage example:

Parameters

Returns

void

Additional overload: neighbors::ivf_pq::extend

Extend the index with the new data.

1
auto extend(raft::resources const& handle,
2
raft::host_matrix_view<const int8_t, int64_t, raft::row_major> new_vectors,
3
std::optional<raft::host_vector_view<const int64_t, int64_t>> new_indices,
4
const cuvs::neighbors::ivf_pq::index<int64_t>& idx)
5
-> cuvs::neighbors::ivf_pq::index<int64_t>;

Note, the user can set a stream pool in the input raft::resource with at least one stream to enable kernel and copy overlapping.

Usage example:

Parameters

Returns

cuvs::neighbors::ivf_pq::index<int64_t>

Additional overload: neighbors::ivf_pq::extend

Extend the index with the new data.

1
void extend(raft::resources const& handle,
2
raft::host_matrix_view<const int8_t, int64_t, raft::row_major> new_vectors,
3
std::optional<raft::host_vector_view<const int64_t, int64_t>> new_indices,
4
cuvs::neighbors::ivf_pq::index<int64_t>* idx);

Note, the user can set a stream pool in the input raft::resource with at least one stream to enable kernel and copy overlapping.

Usage example:

Parameters

Returns

void

Additional overload: neighbors::ivf_pq::extend

Extend the index with the new data.

1
auto extend(raft::resources const& handle,
2
raft::host_matrix_view<const uint8_t, int64_t, raft::row_major> new_vectors,
3
std::optional<raft::host_vector_view<const int64_t, int64_t>> new_indices,
4
const cuvs::neighbors::ivf_pq::index<int64_t>& idx)
5
-> cuvs::neighbors::ivf_pq::index<int64_t>;

Note, the user can set a stream pool in the input raft::resource with at least one stream to enable kernel and copy overlapping.

Usage example:

Parameters

Returns

cuvs::neighbors::ivf_pq::index<int64_t>

Additional overload: neighbors::ivf_pq::extend

Extend the index with the new data.

1
void extend(raft::resources const& handle,
2
raft::host_matrix_view<const uint8_t, int64_t, raft::row_major> new_vectors,
3
std::optional<raft::host_vector_view<const int64_t, int64_t>> new_indices,
4
cuvs::neighbors::ivf_pq::index<int64_t>* idx);

Note, the user can set a stream pool in the input raft::resource with at least one stream to enable kernel and copy overlapping.

Usage example:

Parameters

Returns

void

IVF-PQ index transform

neighbors::ivf_pq::transform

Transform a dataset by applying pq-encoding to each vector

1
void transform(raft::resources const& handle,
2
const cuvs::neighbors::ivf_pq::index<int64_t>& index,
3
raft::device_matrix_view<const float, int64_t, raft::row_major> dataset,
4
raft::device_vector_view<uint32_t, int64_t> output_labels,
5
raft::device_matrix_view<uint8_t, int64_t> output_dataset);

cluster ids (labels) for each vector in the input dataset index.pq_bits(), 8)]] that will get populated with the pq-encoded dataset

Parameters

Returns

void

Additional overload: neighbors::ivf_pq::transform

1
void transform(raft::resources const& handle,
2
const cuvs::neighbors::ivf_pq::index<int64_t>& index,
3
raft::device_matrix_view<const half, int64_t, raft::row_major> dataset,
4
raft::device_vector_view<uint32_t, int64_t> output_labels,
5
raft::device_matrix_view<uint8_t, int64_t> output_dataset);

Parameters

Returns

void

Additional overload: neighbors::ivf_pq::transform

1
void transform(raft::resources const& handle,
2
const cuvs::neighbors::ivf_pq::index<int64_t>& index,
3
raft::device_matrix_view<const int8_t, int64_t, raft::row_major> dataset,
4
raft::device_vector_view<uint32_t, int64_t> output_labels,
5
raft::device_matrix_view<uint8_t, int64_t> output_dataset);

Parameters

Returns

void

Additional overload: neighbors::ivf_pq::transform

1
void transform(raft::resources const& handle,
2
const cuvs::neighbors::ivf_pq::index<int64_t>& index,
3
raft::device_matrix_view<const uint8_t, int64_t, raft::row_major> dataset,
4
raft::device_vector_view<uint32_t, int64_t> output_labels,
5
raft::device_matrix_view<uint8_t, int64_t> output_dataset);

Parameters

Returns

void

IVF-PQ index serialize

neighbors::ivf_pq::serialize

Write the index to an output stream

1
void serialize(raft::resources const& handle,
2
std::ostream& os,
3
const cuvs::neighbors::ivf_pq::index<int64_t>& index);

Parameters

Returns

void

Additional overload: neighbors::ivf_pq::serialize

Save the index to file.

1
void serialize(raft::resources const& handle,
2
const std::string& filename,
3
const cuvs::neighbors::ivf_pq::index<int64_t>& index);

Parameters

Returns

void

neighbors::ivf_pq::deserialize

Load index from input stream

1
void deserialize(raft::resources const& handle,
2
std::istream& str,
3
cuvs::neighbors::ivf_pq::index<int64_t>* index);

Parameters

Returns

void

Additional overload: neighbors::ivf_pq::deserialize

Load index from file.

1
void deserialize(raft::resources const& handle,
2
const std::string& filename,
3
cuvs::neighbors::ivf_pq::index<int64_t>* index);

Parameters

Returns

void

IVF-PQ helper methods

neighbors::ivf_pq::helpers::codepacker::unpack

Unpack n_take consecutive records of a single list (cluster) in the compressed index

1
void unpack(raft::resources const& res,
2
raft::device_mdspan<const uint8_t,
3
list_spec_interleaved<uint32_t, uint32_t>::list_extents,
4
raft::row_major> list_data,
5
uint32_t pq_bits,
6
uint32_t offset,
7
raft::device_matrix_view<uint8_t, uint32_t, raft::row_major> codes);

starting at given offset.

Bit compression is removed, which means output will have pq_dim dimensional vectors (one code per byte, instead of ceildiv(pq_dim * pq_bits, 8) bytes of pq codes).

Usage example:

Parameters

Returns

void

neighbors::ivf_pq::helpers::codepacker::unpack_contiguous

Unpack n_rows consecutive records of a single list (cluster) in the compressed index

1
void unpack_contiguous(raft::resources const& res,
2
raft::device_mdspan<const uint8_t,
3
list_spec_interleaved<uint32_t, uint32_t>::list_extents,
4
raft::row_major> list_data,
5
uint32_t pq_bits,
6
uint32_t offset,
7
uint32_t n_rows,
8
uint32_t pq_dim,
9
uint8_t* codes);

starting at given offset. The output codes of a single vector are contiguous, not expanded to one code per byte, which means the output has ceildiv(pq_dim * pq_bits, 8) bytes per PQ encoded vector.

Usage example:

Parameters

Returns

void

neighbors::ivf_pq::helpers::codepacker::pack

Write flat PQ codes into an existing list by the given offset.

1
void pack(raft::resources const& res,
2
raft::device_matrix_view<const uint8_t, uint32_t, raft::row_major> codes,
3
uint32_t pq_bits,
4
uint32_t offset,
5
raft::device_mdspan<uint8_t,
6
list_spec_interleaved<uint32_t, uint32_t>::list_extents,
7
raft::row_major> list_data);

NB: no memory allocation happens here; the list must fit the data (offset + n_vec).

Usage example:

Parameters

Returns

void

neighbors::ivf_pq::helpers::codepacker::pack_contiguous

Write flat PQ codes into an existing list by the given offset. The input codes of a single vector

1
void pack_contiguous(raft::resources const& res,
2
const uint8_t* codes,
3
uint32_t n_rows,
4
uint32_t pq_dim,
5
uint32_t pq_bits,
6
uint32_t offset,
7
raft::device_mdspan<uint8_t,
8
list_spec_interleaved<uint32_t, uint32_t>::list_extents,
9
raft::row_major> list_data);

are contiguous (not expanded to one code per byte).

NB: no memory allocation happens here; the list must fit the data (offset + n_rows records).

Usage example:

Parameters

Returns

void

neighbors::ivf_pq::helpers::codepacker::pack_list_data

Write flat PQ codes into an existing list by the given offset.

1
void pack_list_data(raft::resources const& res,
2
index<int64_t>* index,
3
raft::device_matrix_view<const uint8_t, uint32_t, raft::row_major> codes,
4
uint32_t label,
5
uint32_t offset);

The list is identified by its label.

NB: no memory allocation happens here; the list must fit the data (offset + n_vec).

Usage example:

Parameters

Returns

void

neighbors::ivf_pq::helpers::codepacker::pack_contiguous_list_data

Write flat PQ codes into an existing list by the given offset. Use this when the input

1
void pack_contiguous_list_data(raft::resources const& res,
2
index<int64_t>* index,
3
uint8_t* codes,
4
uint32_t n_rows,
5
uint32_t label,
6
uint32_t offset);

vectors are PQ encoded and not expanded to one code per byte.

The list is identified by its label.

NB: no memory allocation happens here; the list into which the vectors are packed must fit offset + n_rows rows.

Usage example:

Parameters

Returns

void

neighbors::ivf_pq::helpers::codepacker::unpack_list_data

Unpack n_take consecutive records of a single list (cluster) in the compressed index

1
void unpack_list_data(raft::resources const& res,
2
const index<int64_t>& index,
3
raft::device_matrix_view<uint8_t, uint32_t, raft::row_major> out_codes,
4
uint32_t label,
5
uint32_t offset);

starting at given offset, one code per byte (independently of pq_bits).

Usage example:

Parameters

Returns

void

Additional overload: neighbors::ivf_pq::helpers::codepacker::unpack_list_data

Unpack a series of records of a single list (cluster) in the compressed index

1
void unpack_list_data(raft::resources const& res,
2
const index<int64_t>& index,
3
raft::device_vector_view<const uint32_t> in_cluster_indices,
4
raft::device_matrix_view<uint8_t, uint32_t, raft::row_major> out_codes,
5
uint32_t label);

by their in-list offsets, one code per byte (independently of pq_bits).

Usage example:

Parameters

Returns

void

neighbors::ivf_pq::helpers::codepacker::unpack_contiguous_list_data

Unpack n_rows consecutive PQ encoded vectors of a single list (cluster) in the

1
void unpack_contiguous_list_data(raft::resources const& res,
2
const index<int64_t>& index,
3
uint8_t* out_codes,
4
uint32_t n_rows,
5
uint32_t label,
6
uint32_t offset);

compressed index starting at given offset, not expanded to one code per byte. Each code in the output buffer occupies ceildiv(index.pq_dim() * index.pq_bits(), 8) bytes.

Usage example:

Parameters

Returns

void

neighbors::ivf_pq::helpers::codepacker::reconstruct_list_data

Decode n_take consecutive records of a single list (cluster) in the compressed index

1
void reconstruct_list_data(raft::resources const& res,
2
const index<int64_t>& index,
3
raft::device_matrix_view<float, uint32_t, raft::row_major> out_vectors,
4
uint32_t label,
5
uint32_t offset);

starting at given offset.

Usage example:

Parameters

Returns

void

Additional overload: neighbors::ivf_pq::helpers::codepacker::reconstruct_list_data

Decode a series of records of a single list (cluster) in the compressed index

1
void reconstruct_list_data(raft::resources const& res,
2
const index<int64_t>& index,
3
raft::device_vector_view<const uint32_t> in_cluster_indices,
4
raft::device_matrix_view<float, uint32_t, raft::row_major> out_vectors,
5
uint32_t label);

by their in-list offsets.

Usage example:

Parameters

Returns

void

neighbors::ivf_pq::helpers::codepacker::extend_list_with_codes

Extend one list of the index in-place, by the list label, skipping the classification and

1
void extend_list_with_codes(
2
raft::resources const& res,
3
index<int64_t>* index,
4
raft::device_matrix_view<const uint8_t, uint32_t, raft::row_major> new_codes,
5
raft::device_vector_view<const int64_t, uint32_t, raft::row_major> new_indices,
6
uint32_t label);

encoding steps.

Usage example:

Parameters

Returns

void

neighbors::ivf_pq::helpers::codepacker::extend_list_with_contiguous_codes

Extend one list of the index in-place, by the list label, skipping the classification and

1
void extend_list_with_contiguous_codes(
2
raft::resources const& res,
3
index<int64_t>* index,
4
raft::device_matrix_view<const uint8_t, uint32_t, raft::row_major> new_codes,
5
raft::device_vector_view<const int64_t, uint32_t, raft::row_major> new_indices,
6
uint32_t label);

encoding steps. Uses contiguous/packed codes format.

This is similar to extend_list_with_codes but takes codes in contiguous packed format [n_rows, ceildiv(pq_dim * pq_bits, 8)] instead of unpacked format [n_rows, pq_dim]. This works correctly with any pq_bits value.

Usage example:

Parameters

Returns

void

neighbors::ivf_pq::helpers::codepacker::extend_list

Extend one list of the index in-place, by the list label, skipping the classification

1
void extend_list(raft::resources const& res,
2
index<int64_t>* index,
3
raft::device_matrix_view<const float, uint32_t, raft::row_major> new_vectors,
4
raft::device_vector_view<const int64_t, uint32_t, raft::row_major> new_indices,
5
uint32_t label);

step.

Usage example:

Parameters

Returns

void

neighbors::ivf_pq::helpers::erase_list

Remove all data from a single list (cluster) in the index.

1
void erase_list(raft::resources const& res, index<int64_t>* index, uint32_t label);

Usage example:

Parameters

Returns

void

neighbors::ivf_pq::helpers::reset_index

Public helper API to reset the data and indices ptrs, and the list sizes. Useful for

1
void reset_index(const raft::resources& res, index<int64_t>* index);

externally modifying the index without going through the build stage. The data and indices of the IVF lists will be lost.

Usage example:

Parameters

Returns

void

neighbors::ivf_pq::helpers::pad_centers_with_norms

Pad cluster centers with their L2 norms for efficient GEMM operations.

1
void pad_centers_with_norms(
2
raft::resources const& res,
3
raft::device_matrix_view<const float, uint32_t, raft::row_major> centers,
4
raft::device_matrix_view<float, uint32_t, raft::row_major> padded_centers);

This function takes cluster centers and pads them with their L2 norms to create extended centers suitable for coarse search operations. The output has dimensions [n_centers, dim_ext] where dim_ext = round_up(dim + 1, 8).

Parameters

Returns

void

Additional overload: neighbors::ivf_pq::helpers::pad_centers_with_norms

Pad cluster centers with their L2 norms for efficient GEMM operations.

1
void pad_centers_with_norms(
2
raft::resources const& res,
3
raft::host_matrix_view<const float, uint32_t, raft::row_major> centers,
4
raft::device_matrix_view<float, uint32_t, raft::row_major> padded_centers);

This function takes cluster centers and pads them with their L2 norms to create extended centers suitable for coarse search operations. The output has dimensions [n_centers, dim_ext] where dim_ext = round_up(dim + 1, 8).

Parameters