IVF-PQ#

#include <raft/neighbors/ivf_pq.cuh>

namespace raft::neighbors::ivf_pq

enum class codebook_gen#

A type for specifying how PQ codebooks are created.

Values:

enumerator PER_SUBSPACE#
enumerator PER_CLUSTER#
template<typename IdxT, typename SizeT = uint32_t>
using list_data = ivf::list<list_spec, SizeT, IdxT>#
static constexpr uint32_t kIndexGroupSize = 32

Size of the interleaved group.

static constexpr uint32_t kIndexGroupVecLen = 16#

Stride of the interleaved group for vectorized loads.

template<typename IdxT>
static constexpr IdxT kOutOfBoundsRecord = std::numeric_limits<IdxT>::max()#

Default value returned by search when the n_probes is too small and top-k is too large. One may encounter it if the combined size of probed clusters is smaller than the requested number of results per query.

template<typename T, typename IdxT = uint32_t>
index<IdxT> build(raft::resources const &handle, const index_params &params, raft::device_matrix_view<const T, IdxT, row_major> dataset)#

Build the index from the dataset for efficient search.

NB: Currently, the following distance metrics are supported:

  • L2Expanded

  • L2Unexpanded

  • InnerProduct

Template Parameters:
  • T – data element type

  • IdxT – type of the indices in the source dataset

Parameters:
  • handle[in]

  • params[in] configure the index building

  • dataset[in] a device matrix view to a row-major matrix [n_rows, dim]

Returns:

the constructed ivf-pq index

template<typename T, typename IdxT>
index<IdxT> extend(raft::resources const &handle, raft::device_matrix_view<const T, IdxT, row_major> new_vectors, std::optional<raft::device_vector_view<const IdxT, IdxT>> new_indices, const index<IdxT> &idx)#

Extend the index with the new data. *.

Template Parameters:
  • T – data element type

  • IdxT – type of the indices in the source dataset

Parameters:
  • handle[in]

  • new_vectors[in] a device matrix view to a row-major matrix [n_rows, idx.dim()]

  • new_indices[in] a device vector view to a vector of indices [n_rows]. If the original index is empty (idx.size() == 0), you can pass std::nullopt here to imply a continuous range [0...n_rows).

  • idx[inout]

template<typename T, typename IdxT>
void extend(raft::resources const &handle, raft::device_matrix_view<const T, IdxT, row_major> new_vectors, std::optional<raft::device_vector_view<const IdxT, IdxT>> new_indices, index<IdxT> *idx)#

Extend the index with the new data. *.

Template Parameters:
  • T – data element type

  • IdxT – type of the indices in the source dataset

Parameters:
  • handle[in]

  • new_vectors[in] a device matrix view to a row-major matrix [n_rows, idx.dim()]

  • new_indices[in] a device vector view to a vector of indices [n_rows]. If the original index is empty (idx.size() == 0), you can pass std::nullopt here to imply a continuous range [0...n_rows).

  • idx[inout]

template<typename T, typename IdxT, typename IvfSampleFilterT>
void search_with_filtering(raft::resources const &handle, const search_params &params, const index<IdxT> &idx, raft::device_matrix_view<const T, uint32_t, row_major> queries, raft::device_matrix_view<IdxT, uint32_t, row_major> neighbors, raft::device_matrix_view<float, uint32_t, row_major> distances, IvfSampleFilterT sample_filter = IvfSampleFilterT{})#

Search ANN using the constructed index using the given filter.

See the ivf_pq::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. The exact size of the temporary buffer depends on multiple factors and is an implementation detail. However, you can safely specify a small initial size for the memory pool, so that only a few allocations happen to grow it during the first invocations of the search.

Template Parameters:
  • T – data element type

  • IdxT – type of the indices

Parameters:
  • handle[in]

  • params[in] configure the search

  • idx[in] ivf-pq constructed index

  • queries[in] a device matrix view to a row-major matrix [n_queries, index->dim()]

  • neighbors[out] a device matrix view to the indices of the neighbors in the source dataset [n_queries, k]

  • distances[out] a device matrix view to the distances to the selected neighbors [n_queries, k]

  • sample_filter[in] a filter the greenlights samples for a given query.

template<typename T, typename IdxT>
void search(raft::resources const &handle, const search_params &params, const index<IdxT> &idx, raft::device_matrix_view<const T, uint32_t, row_major> queries, raft::device_matrix_view<IdxT, uint32_t, row_major> neighbors, raft::device_matrix_view<float, uint32_t, row_major> distances)#

Search ANN using the constructed index.

See the ivf_pq::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. The exact size of the temporary buffer depends on multiple factors and is an implementation detail. However, you can safely specify a small initial size for the memory pool, so that only a few allocations happen to grow it during the first invocations of the search.

Template Parameters:
  • T – data element type

  • IdxT – type of the indices

Parameters:
  • handle[in]

  • params[in] configure the search

  • idx[in] ivf-pq constructed index

  • queries[in] a device matrix view to a row-major matrix [n_queries, index->dim()]

  • neighbors[out] a device matrix view to the indices of the neighbors in the source dataset [n_queries, k]

  • distances[out] a device matrix view to the distances to the selected neighbors [n_queries, k]

Serializer Methods#

#include <raft/neighbors/ivf_pq_serialize.cuh>

namespace raft::neighbors::ivf_pq

template<typename IdxT>
void serialize(raft::resources const &handle, std::ostream &os, const index<IdxT> &index)#

Write the index to an output stream

Experimental, both the API and the serialization format are subject to change.

#include <raft/core/resources.hpp>

raft::resources handle;

// create an output stream
std::ostream os(std::cout.rdbuf());
// create an index with `auto index = ivf_pq::build(...);`
raft::serialize(handle, os, index);
Template Parameters:

IdxT – type of the index

Parameters:
  • handle[in] the raft handle

  • os[in] output stream

  • index[in] IVF-PQ index

template<typename IdxT>
void serialize(raft::resources const &handle, const std::string &filename, const index<IdxT> &index)#

Save the index to file.

Experimental, both the API and the serialization format are subject to change.

#include <raft/core/resources.hpp>

raft::resources handle;

// create a string with a filepath
std::string filename("/path/to/index");
// create an index with `auto index = ivf_pq::build(...);`
raft::serialize(handle, filename, index);
Template Parameters:

IdxT – type of the index

Parameters:
  • handle[in] the raft handle

  • filename[in] the file name for saving the index

  • index[in] IVF-PQ index

template<typename IdxT>
index<IdxT> deserialize(raft::resources const &handle, std::istream &is)#

Load index from input stream

Experimental, both the API and the serialization format are subject to change.

#include <raft/core/resources.hpp>

raft::resources handle;

// create an input stream
std::istream is(std::cin.rdbuf());
using IdxT = int; // type of the index
auto index = raft::deserialize<IdxT>(handle, is);
Template Parameters:

IdxT – type of the index

Parameters:
  • handle[in] the raft handle

  • is[in] input stream

Returns:

raft::neighbors::ivf_pq::index<IdxT>

template<typename IdxT>
index<IdxT> deserialize(raft::resources const &handle, const std::string &filename)#

Load index from file.

Experimental, both the API and the serialization format are subject to change.

#include <raft/core/resources.hpp>

raft::resources handle;

// create a string with a filepath
std::string filename("/path/to/index");
using IdxT = int; // type of the index
auto index = raft::deserialize<IdxT>(handle, filename);
Template Parameters:

IdxT – type of the index

Parameters:
  • handle[in] the raft handle

  • filename[in] the name of the file that stores the index

Returns:

raft::neighbors::ivf_pq::index<IdxT>