IVF-Flat#

The IVF-Flat method is an ANN algorithm. It uses an inverted file index (IVF) with unmodified (that is, flat) vectors. This algorithm provides simple knobs to reduce the overall search space and to trade-off accuracy for speed.

#include <cuvs/neighbors/ivf_flat.hpp>

namespace cuvs::neighbors::ivf_flat

Index build parameters#

struct index_params : public cuvs::neighbors::ann::index_params#

#include <ivf_flat.hpp>

Public Functions

inline operator raft::neighbors::ivf_flat::index_params() const#: Build a raft IVF_FLAT index params from an existing cuvs IVF_FLAT index params.

Public Members

uint32_t n_lists = 1024#: The number of inverted lists (clusters)

uint32_t kmeans_n_iters = 20#: The number of iterations searching for kmeans centers (index building).

double kmeans_trainset_fraction = 0.5#: The fraction of data to use during iterative kmeans building.

bool adaptive_centers = false#

By default (adaptive_centers = false), the cluster centers are trained in ivf_flat::build, and never modified in ivf_flat::extend. As a result, you may need to retrain the index from scratch after invoking (ivf_flat::extend) a few times with new data, the distribution of which is no longer representative of the original training set.

The alternative behavior (adaptive_centers = true) is to update the cluster centers for new data when it is added. In this case, index.centers() are always exactly the centroids of the data in the corresponding clusters. The drawback of this behavior is that the centroids depend on the order of adding new data (through the classification of the added data); that is, index.centers() “drift” together with the changing distribution of the newly added data.

bool conservative_memory_allocation = false#

By default, the algorithm allocates more space than necessary for individual clusters (list_data). This allows to amortize the cost of memory allocation and reduce the number of data copies during repeated calls to extend (extending the database).

The alternative is the conservative allocation behavior; when enabled, the algorithm always allocates the minimum amount of memory required to store the given number of records. Set this flag to true if you prefer to use as little GPU memory for the database as possible.

Index search parameters#

struct search_params : public cuvs::neighbors::ann::search_params#

#include <ivf_flat.hpp>

Public Functions

inline operator raft::neighbors::ivf_flat::search_params() const#: Build a raft IVF_FLAT search params from an existing cuvs IVF_FLAT search params.

Public Members

uint32_t n_probes = 20#: The number of clusters to search.

Index#

template<typename T, typename IdxT> struct index : public cuvs::neighbors::ann::index#

#include <ivf_flat.hpp>

IVF-flat index.

Template Parameters:

T – data element type
IdxT – type of the indices in the source dataset

Public Functions

index(raft::resources const &res, const index_params &params, uint32_t dim)#: Construct an empty index. It needs to be trained and then populated.

index(raft::resources const &res, cuvs::distance::DistanceType metric, uint32_t n_lists, bool adaptive_centers, bool conservative_memory_allocation, uint32_t dim)#: Construct an empty index. It needs to be trained and then populated.

uint32_t veclen() const noexcept#: Vectorized load/store size in elements, determines the size of interleaved data chunks.

cuvs::distance::DistanceType metric() const noexcept#: Distance metric used for clustering.

bool adaptive_centers() const noexcept#: Whether centers() change upon extending the index (ivf_flat::extend).

raft::device_vector_view<uint32_t, uint32_t> list_sizes() noexcept#

Inverted list data [size, dim].

The data consists of the dataset rows, grouped by their labels (into clusters/lists). Within each list (cluster), the data is grouped into blocks of kIndexGroupSize interleaved vectors. Note, the total index length is slightly larger than the source dataset length, because each cluster is padded by kIndexGroupSize elements.

Interleaving pattern: within groups of kIndexGroupSize rows, the data is interleaved with the block size equal to veclen * sizeof(T). That is, a chunk of veclen consecutive components of one row is followed by a chunk of the same size of the next row, and so on.

Example: veclen = 2, dim = 6, kIndexGroupSize = 32, list_size = 31

x[ 0, 0], x[ 0, 1], x[ 1, 0], x[ 1, 1], ... x[14, 0], x[14, 1], x[15, 0], x[15, 1],
x[16, 0], x[16, 1], x[17, 0], x[17, 1], ... x[30, 0], x[30, 1],    -    ,    -    ,
x[ 0, 2], x[ 0, 3], x[ 1, 2], x[ 1, 3], ... x[14, 2], x[14, 3], x[15, 2], x[15, 3],
x[16, 2], x[16, 3], x[17, 2], x[17, 3], ... x[30, 2], x[30, 3],    -    ,    -    ,
x[ 0, 4], x[ 0, 5], x[ 1, 4], x[ 1, 5], ... x[14, 4], x[14, 5], x[15, 4], x[15, 5],
x[16, 4], x[16, 5], x[17, 4], x[17, 5], ... x[30, 4], x[30, 5],    -    ,    -    ,

Sizes of the lists (clusters) [n_lists] NB: This may differ from the actual list size if the shared lists have been extended by another index

raft::device_matrix_view<float, uint32_t, raft::row_major> centers() noexcept#: k-means cluster centers corresponding to the lists [n_lists, dim]

std::optional<raft::device_vector_view<float, uint32_t>> center_norms() noexcept#

(Optional) Precomputed norms of the centers w.r.t. the chosen distance metric [n_lists].

NB: this may be empty if the index is empty or if the metric does not require the center norms calculation.

IdxT size() const noexcept#: Total length of the index.

uint32_t dim() const noexcept#: Dimensionality of the data.

uint32_t n_lists() const noexcept#: Number of clusters/inverted lists.

raft::device_vector_view<IdxT*, uint32_t> inds_ptrs() noexcept#: Pointers to the inverted lists (clusters) indices [n_lists].

bool conservative_memory_allocation() const noexcept#: Whether to use convervative memory allocation when extending the list (cluster) data (see index_params.conservative_memory_allocation).

std::vector<std::shared_ptr<raft::neighbors::ivf_flat::list_data<T, IdxT>>> &lists() noexcept#: Lists’ data and indices.

Index build#

auto build(raft::resources const &handle, const cuvs::neighbors::ivf_flat::index_params &index_params, raft::device_matrix_view<const float, int64_t, raft::row_major> dataset) -> cuvs::neighbors::ivf_flat::index<float, int64_t>#

Build the index from the dataset for efficient search.

Usage example:

using namespace cuvs::neighbors;
// use default index parameters
ivf_flat::index_params index_params;
// create and fill the index from a [N, D] dataset
auto index = ivf_flat::build(handle, dataset, index_params);

Parameters:

handle – [in]
index_params – configure the index building
dataset – [in] a device pointer to a row-major matrix [n_rows, dim]

Returns:

the constructed ivf-flat index

void build(raft::resources const &handle, const cuvs::neighbors::ivf_flat::index_params &index_params, raft::device_matrix_view<const float, int64_t, raft::row_major> dataset, cuvs::neighbors::ivf_flat::index<float, int64_t> &idx)#

Build the index from the dataset for efficient search.

Usage example:

using namespace cuvs::neighbors;
// use default index parameters
ivf_flat::index_params index_params;
// create and fill the index from a [N, D] dataset
ivf_flat::index<decltype(dataset::value_type), decltype(dataset::index_type)> index;
ivf_flat::build(handle, dataset, index_params, index);

Parameters:

handle – [in]
index_params – configure the index building
dataset – [in] raft::device_matrix_view to a row-major matrix [n_rows, dim]
idx – [out] reference to ivf_flat::index

auto build(raft::resources const &handle, const cuvs::neighbors::ivf_flat::index_params &index_params, raft::device_matrix_view<const int8_t, int64_t, raft::row_major> dataset) -> cuvs::neighbors::ivf_flat::index<int8_t, int64_t>#

Build the index from the dataset for efficient search.

Usage example:

using namespace cuvs::neighbors;
// use default index parameters
ivf_flat::index_params index_params;
// create and fill the index from a [N, D] dataset
auto index = ivf_flat::build(handle, dataset, index_params);

Parameters:

handle – [in]
index_params – configure the index building
dataset – [in] a device pointer to a row-major matrix [n_rows, dim]

Returns:

the constructed ivf-flat index

void build(raft::resources const &handle, const cuvs::neighbors::ivf_flat::index_params &index_params, raft::device_matrix_view<const int8_t, int64_t, raft::row_major> dataset, cuvs::neighbors::ivf_flat::index<int8_t, int64_t> &idx)#

Build the index from the dataset for efficient search.

Usage example:

using namespace cuvs::neighbors;
// use default index parameters
ivf_flat::index_params index_params;
// create and fill the index from a [N, D] dataset
ivf_flat::index<decltype(dataset::value_type), decltype(dataset::index_type)> index;
ivf_flat::build(handle, dataset, index_params, index);

Parameters:

handle – [in]
index_params – configure the index building
dataset – [in] raft::device_matrix_view to a row-major matrix [n_rows, dim]
idx – [out] reference to ivf_flat::index

auto build(raft::resources const &handle, const cuvs::neighbors::ivf_flat::index_params &index_params, raft::device_matrix_view<const uint8_t, int64_t, raft::row_major> dataset) -> cuvs::neighbors::ivf_flat::index<uint8_t, int64_t>#

Build the index from the dataset for efficient search.

Usage example:

using namespace cuvs::neighbors;
// use default index parameters
ivf_flat::index_params index_params;
// create and fill the index from a [N, D] dataset
auto index = ivf_flat::build(handle, dataset, index_params);

Parameters:

handle – [in]
index_params – configure the index building
dataset – [in] a device pointer to a row-major matrix [n_rows, dim]

Returns:

the constructed ivf-flat index

void build(raft::resources const &handle, const cuvs::neighbors::ivf_flat::index_params &index_params, raft::device_matrix_view<const uint8_t, int64_t, raft::row_major> dataset, cuvs::neighbors::ivf_flat::index<uint8_t, int64_t> &idx)#

Build the index from the dataset for efficient search.

Usage example:

using namespace cuvs::neighbors;
// use default index parameters
ivf_flat::index_params index_params;
// create and fill the index from a [N, D] dataset
ivf_flat::index<decltype(dataset::value_type), decltype(dataset::index_type)> index;
ivf_flat::build(handle, dataset, index_params, index);

Parameters:

handle – [in]
index_params – configure the index building
dataset – [in] raft::device_matrix_view to a row-major matrix [n_rows, dim]
idx – [out] reference to ivf_flat::index

Index extend#

auto extend(raft::resources const &handle, raft::device_matrix_view<const float, int64_t, raft::row_major> new_vectors, std::optional<raft::device_vector_view<const int64_t, int64_t>> new_indices, const cuvs::neighbors::ivf_flat::index<float, int64_t> &idx) -> cuvs::neighbors::ivf_flat::index<float, int64_t>#

Build a new index containing the data of the original plus new extra vectors.

Implementation note: The new data is clustered according to existing kmeans clusters, then the cluster centers are adjusted to match the newly labeled data.

Usage example:

using namespace cuvs::neighbors;
ivf_flat::index_params index_params;
index_params.add_data_on_build = false;      // don't populate index on build
index_params.kmeans_trainset_fraction = 1.0; // use whole dataset for kmeans training
// train the index from a [N, D] dataset
auto index_empty = ivf_flat::build(handle, index_params, dataset);
// fill the index with the data
std::optional<raft::device_vector_view<const IdxT, IdxT>> no_op = std::nullopt;
auto index = ivf_flat::extend(handle, new_vectors, no_op, index_empty);

Parameters:

handle – [in]
new_vectors – [in] raft::device_matrix_view to a row-major matrix [n_rows, index.dim()]
new_indices – [in] optional raft::device_vector_view to a vector of indices [n_rows]. If the original index is empty (orig_index.size() == 0), you can pass std::nullopt here to imply a continuous range [0...n_rows).
idx – [in] original index

Returns:

the constructed extended ivf-flat index

void extend(raft::resources const &handle, raft::device_matrix_view<const float, int64_t, raft::row_major> new_vectors, std::optional<raft::device_vector_view<const int64_t, int64_t>> new_indices, cuvs::neighbors::ivf_flat::index<float, int64_t> *idx)#

Extend the index in-place with the new data.

Usage example:

using namespace cuvs::neighbors;
ivf_flat::index_params index_params;
index_params.add_data_on_build = false;      // don't populate index on build
index_params.kmeans_trainset_fraction = 1.0; // use whole dataset for kmeans training
// train the index from a [N, D] dataset
auto index_empty = ivf_flat::build(handle, index_params, dataset);
// fill the index with the data
std::optional<raft::device_vector_view<const IdxT, IdxT>> no_op = std::nullopt;
ivf_flat::extend(handle, dataset, no_opt, &index_empty);

Parameters:

handle – [in]
new_vectors – [in] raft::device_matrix_view to a row-major matrix [n_rows, index.dim()]
new_indices – [in] optional raft::device_vector_view to a vector of indices [n_rows]. If the original index is empty (orig_index.size() == 0), you can pass std::nullopt here to imply a continuous range [0...n_rows).
idx – [inout] pointer to index, to be overwritten in-place

auto extend(raft::resources const &handle, raft::device_matrix_view<const int8_t, int64_t, raft::row_major> new_vectors, std::optional<raft::device_vector_view<const int64_t, int64_t>> new_indices, const cuvs::neighbors::ivf_flat::index<int8_t, int64_t> &idx) -> cuvs::neighbors::ivf_flat::index<int8_t, int64_t>#

Build a new index containing the data of the original plus new extra vectors.

Implementation note: The new data is clustered according to existing kmeans clusters, then the cluster centers are adjusted to match the newly labeled data.

Usage example:

using namespace cuvs::neighbors;
ivf_flat::index_params index_params;
index_params.add_data_on_build = false;      // don't populate index on build
index_params.kmeans_trainset_fraction = 1.0; // use whole dataset for kmeans training
// train the index from a [N, D] dataset
auto index_empty = ivf_flat::build(handle, dataset, index_params, dataset);
// fill the index with the data
std::optional<raft::device_vector_view<const IdxT, IdxT>> no_op = std::nullopt;
auto index = ivf_flat::extend(handle, new_vectors, no_op, index_empty);

Parameters:

handle – [in]
new_vectors – [in] raft::device_matrix_view to a row-major matrix [n_rows, index.dim()]
new_indices – [in] optional raft::device_vector_view to a vector of indices [n_rows]. If the original index is empty (orig_index.size() == 0), you can pass std::nullopt here to imply a continuous range [0...n_rows).
idx – [in] original index

Returns:

the constructed extended ivf-flat index

void extend(raft::resources const &handle, raft::device_matrix_view<const int8_t, int64_t, raft::row_major> new_vectors, std::optional<raft::device_vector_view<const int64_t, int64_t>> new_indices, cuvs::neighbors::ivf_flat::index<int8_t, int64_t> *idx)#

Extend the index in-place with the new data.

Usage example:

using namespace cuvs::neighbors;
ivf_flat::index_params index_params;
index_params.add_data_on_build = false;      // don't populate index on build
index_params.kmeans_trainset_fraction = 1.0; // use whole dataset for kmeans training
// train the index from a [N, D] dataset
auto index_empty = ivf_flat::build(handle, index_params, dataset);
// fill the index with the data
std::optional<raft::device_vector_view<const IdxT, IdxT>> no_op = std::nullopt;
ivf_flat::extend(handle, dataset, no_opt, &index_empty);

Parameters:

handle – [in]
new_vectors – [in] raft::device_matrix_view to a row-major matrix [n_rows, index.dim()]
new_indices – [in] optional raft::device_vector_view to a vector of indices [n_rows]. If the original index is empty (orig_index.size() == 0), you can pass std::nullopt here to imply a continuous range [0...n_rows).
idx – [inout] pointer to index, to be overwritten in-place

auto extend(raft::resources const &handle, raft::device_matrix_view<const uint8_t, int64_t, raft::row_major> new_vectors, std::optional<raft::device_vector_view<const int64_t, int64_t>> new_indices, const cuvs::neighbors::ivf_flat::index<uint8_t, int64_t> &idx) -> cuvs::neighbors::ivf_flat::index<uint8_t, int64_t>#

Build a new index containing the data of the original plus new extra vectors.

Implementation note: The new data is clustered according to existing kmeans clusters, then the cluster centers are adjusted to match the newly labeled data.

Usage example:

using namespace cuvs::neighbors;
ivf_flat::index_params index_params;
index_params.add_data_on_build = false;      // don't populate index on build
index_params.kmeans_trainset_fraction = 1.0; // use whole dataset for kmeans training
// train the index from a [N, D] dataset
auto index_empty = ivf_flat::build(handle, dataset, index_params, dataset);
// fill the index with the data
std::optional<raft::device_vector_view<const IdxT, IdxT>> no_op = std::nullopt;
auto index = ivf_flat::extend(handle, new_vectors, no_op, index_empty);

Parameters:

handle – [in]
new_vectors – [in] raft::device_matrix_view to a row-major matrix [n_rows, index.dim()]
new_indices – [in] optional raft::device_vector_view to a vector of indices [n_rows]. If the original index is empty (orig_index.size() == 0), you can pass std::nullopt here to imply a continuous range [0...n_rows).
idx – [in] original index

Returns:

the constructed extended ivf-flat index

void extend(raft::resources const &handle, raft::device_matrix_view<const uint8_t, int64_t, raft::row_major> new_vectors, std::optional<raft::device_vector_view<const int64_t, int64_t>> new_indices, cuvs::neighbors::ivf_flat::index<uint8_t, int64_t> *idx)#

Extend the index in-place with the new data.

Usage example:

using namespace cuvs::neighbors;
ivf_flat::index_params index_params;
index_params.add_data_on_build = false;      // don't populate index on build
index_params.kmeans_trainset_fraction = 1.0; // use whole dataset for kmeans training
// train the index from a [N, D] dataset
auto index_empty = ivf_flat::build(handle, index_params, dataset);
// fill the index with the data
std::optional<raft::device_vector_view<const IdxT, IdxT>> no_op = std::nullopt;
ivf_flat::extend(handle, dataset, no_opt, &index_empty);

Parameters:

handle – [in]
new_vectors – [in] raft::device_matrix_view to a row-major matrix [n_rows, index.dim()]
new_indices – [in] optional raft::device_vector_view to a vector of indices [n_rows]. If the original index is empty (orig_index.size() == 0), you can pass std::nullopt here to imply a continuous range [0...n_rows).
idx – [inout] pointer to index, to be overwritten in-place

Index search#

void search(raft::resources const &handle, const cuvs::neighbors::ivf_flat::search_params &params, cuvs::neighbors::ivf_flat::index<float, int64_t> &index, raft::device_matrix_view<const float, int64_t, raft::row_major> queries, raft::device_matrix_view<int64_t, int64_t, raft::row_major> neighbors, raft::device_matrix_view<float, int64_t, raft::row_major> distances)#

Search ANN using the constructed index.

See the ivf_flat::build documentation for a usage example.

Note, this function requires a temporary buffer to store intermediate results between cuda kernel calls, which may lead to undesirable allocations and slowdown. To alleviate the problem, you can pass a pool memory resource or a large enough pre-allocated memory resource to reduce or eliminate entirely allocations happening within search:

...
// use default search parameters
ivf_flat::search_params search_params;
// Use the same allocator across multiple searches to reduce the number of
// cuda memory allocations
ivf_flat::search(handle, search_params, index, queries1, out_inds1, out_dists1);
ivf_flat::search(handle, search_params, index, queries2, out_inds2, out_dists2);
ivf_flat::search(handle, search_params, index, queries3, out_inds3, out_dists3);
...

Parameters:

handle – [in]
params – [in] configure the search
index – [in] ivf-flat constructed index
queries – [in] raft::device_matrix_view to a row-major matrix [n_queries, index->dim()]
neighbors – [out] raft::device_matrix_view to the indices of the neighbors in the source dataset [n_queries, k]
distances – [out] raft::device_matrix_view to the distances to the selected neighbors [n_queries, k]

void search(raft::resources const &handle, const cuvs::neighbors::ivf_flat::search_params &params, cuvs::neighbors::ivf_flat::index<int8_t, int64_t> &index, raft::device_matrix_view<const int8_t, int64_t, raft::row_major> queries, raft::device_matrix_view<int64_t, int64_t, raft::row_major> neighbors, raft::device_matrix_view<float, int64_t, raft::row_major> distances)#

Search ANN using the constructed index.

See the ivf_flat::build documentation for a usage example.

...
// use default search parameters
ivf_flat::search_params search_params;
// Use the same allocator across multiple searches to reduce the number of
// cuda memory allocations
ivf_flat::search(handle, search_params, index, queries1, out_inds1, out_dists1);
ivf_flat::search(handle, search_params, index, queries2, out_inds2, out_dists2);
ivf_flat::search(handle, search_params, index, queries3, out_inds3, out_dists3);
...

Parameters:

handle – [in]
params – [in] configure the search
index – [in] ivf-flat constructed index
queries – [in] raft::device_matrix_view to a row-major matrix [n_queries, index->dim()]
neighbors – [out] raft::device_matrix_view to the indices of the neighbors in the source dataset [n_queries, k]
distances – [out] raft::device_matrix_view to the distances to the selected neighbors [n_queries, k]

void search(raft::resources const &handle, const cuvs::neighbors::ivf_flat::search_params &params, cuvs::neighbors::ivf_flat::index<uint8_t, int64_t> &index, raft::device_matrix_view<const uint8_t, int64_t, raft::row_major> queries, raft::device_matrix_view<int64_t, int64_t, raft::row_major> neighbors, raft::device_matrix_view<float, int64_t, raft::row_major> distances)#

Search ANN using the constructed index.

See the ivf_flat::build documentation for a usage example.

...
// use default search parameters
ivf_flat::search_params search_params;
// Use the same allocator across multiple searches to reduce the number of
// cuda memory allocations
ivf_flat::search(handle, search_params, index, queries1, out_inds1, out_dists1);
ivf_flat::search(handle, search_params, index, queries2, out_inds2, out_dists2);
ivf_flat::search(handle, search_params, index, queries3, out_inds3, out_dists3);
...

Parameters:

handle – [in]
params – [in] configure the search
index – [in] ivf-flat constructed index
queries – [in] raft::device_matrix_view to a row-major matrix [n_queries, index->dim()]
neighbors – [out] raft::device_matrix_view to the indices of the neighbors in the source dataset [n_queries, k]
distances – [out] raft::device_matrix_view to the distances to the selected neighbors [n_queries, k]

Index serialize#

void serialize_file(raft::resources const &handle, const std::string &filename, const cuvs::neighbors::ivf_flat::index<float, int64_t> &index)#

Save the index to file.