Distributed ANN

The SNMG (single-node multi-GPUs) ANN API provides a set of functions to deploy ANN indexes on multiple GPUs.

#include <cuvs/neighbors/mg.hpp>

namespace cuvs::neighbors::mg

Index build parameters

enum distribution_mode

Distribution mode

Values:

enumerator REPLICATED

Index is replicated on each device, favors throughput

enumerator SHARDED

Index is split on several devices, favors scaling

template<typename Upstream>
struct index_params : public Upstream

#include <mg.hpp>

Build parameters

Public Members

cuvs::neighbors::mg::distribution_mode mode = SHARDED

Distribution mode

Search parameters

enum replicated_search_mode

Search mode when using a replicated index

Values:

enumerator LOAD_BALANCER

Search queries are splited to maintain equal load on GPUs

enumerator ROUND_ROBIN

Each search query is processed by a single GPU in a round-robin fashion

enum sharded_merge_mode

Merge mode when using a sharded index

Values:

enumerator MERGE_ON_ROOT_RANK

Search batches are merged on the root rank

enumerator TREE_MERGE

Search batches are merged in a tree reduction fashion

template<typename Upstream>
struct search_params : public Upstream

#include <mg.hpp>

Search parameters

Public Members

cuvs::neighbors::mg::replicated_search_mode search_mode = LOAD_BALANCER

Replicated search mode

cuvs::neighbors::mg::sharded_merge_mode merge_mode = TREE_MERGE

Sharded merge mode

Index build

auto build(const raft::device_resources &handle, const mg::index_params<ivf_flat::index_params> &index_params, raft::host_matrix_view<const float, int64_t, row_major> index_dataset) -> index<ivf_flat::index<float, int64_t>, float, int64_t>

Builds a multi-GPU index.

Usage example:

raft::handle_t handle;
cuvs::neighbors::mg::index_params<ivf_flat::index_params> index_params;
auto index = cuvs::neighbors::mg::build(handle, index_params, index_dataset);

Parameters:

• handle – [in]

• index_params – configure the index building

• index_dataset – [in] a row-major matrix on host [n_rows, dim]

Returns:

the constructed IVF-Flat MG index

auto build(const raft::device_resources &handle, const mg::index_params<ivf_flat::index_params> &index_params, raft::host_matrix_view<const int8_t, int64_t, row_major> index_dataset) -> index<ivf_flat::index<int8_t, int64_t>, int8_t, int64_t>

Builds a multi-GPU index.

Usage example:

raft::handle_t handle;
cuvs::neighbors::mg::index_params<ivf_flat::index_params> index_params;
auto index = cuvs::neighbors::mg::build(handle, index_params, index_dataset);

Parameters:

• handle – [in]

• index_params – configure the index building

• index_dataset – [in] a row-major matrix on host [n_rows, dim]

Returns:

the constructed IVF-Flat MG index

auto build(const raft::device_resources &handle, const mg::index_params<ivf_flat::index_params> &index_params, raft::host_matrix_view<const uint8_t, int64_t, row_major> index_dataset) -> index<ivf_flat::index<uint8_t, int64_t>, uint8_t, int64_t>

Builds a multi-GPU index.

Usage example:

raft::handle_t handle;
cuvs::neighbors::mg::index_params<ivf_flat::index_params> index_params;
auto index = cuvs::neighbors::mg::build(handle, index_params, index_dataset);

Parameters:

• handle – [in]

• index_params – configure the index building

• index_dataset – [in] a row-major matrix on host [n_rows, dim]

Returns:

the constructed IVF-Flat MG index

auto build(const raft::device_resources &handle, const mg::index_params<ivf_pq::index_params> &index_params, raft::host_matrix_view<const float, int64_t, row_major> index_dataset) -> index<ivf_pq::index<int64_t>, float, int64_t>

Builds a multi-GPU index.

Usage example:

raft::handle_t handle;
cuvs::neighbors::mg::index_params<ivf_pq::index_params> index_params;
auto index = cuvs::neighbors::mg::build(handle, index_params, index_dataset);

Parameters:

• handle – [in]

• index_params – configure the index building

• index_dataset – [in] a row-major matrix on host [n_rows, dim]

Returns:

the constructed IVF-PQ MG index

auto build(const raft::device_resources &handle, const mg::index_params<ivf_pq::index_params> &index_params, raft::host_matrix_view<const half, int64_t, row_major> index_dataset) -> index<ivf_pq::index<int64_t>, half, int64_t>

Builds a multi-GPU index.

Usage example:

raft::handle_t handle;
cuvs::neighbors::mg::index_params<ivf_pq::index_params> index_params;
auto index = cuvs::neighbors::mg::build(handle, index_params, index_dataset);

Parameters:

• handle – [in]

• index_params – configure the index building

• index_dataset – [in] a row-major matrix on host [n_rows, dim]

Returns:

the constructed IVF-PQ MG index

auto build(const raft::device_resources &handle, const mg::index_params<ivf_pq::index_params> &index_params, raft::host_matrix_view<const int8_t, int64_t, row_major> index_dataset) -> index<ivf_pq::index<int64_t>, int8_t, int64_t>

Builds a multi-GPU index.

Usage example:

raft::handle_t handle;
cuvs::neighbors::mg::index_params<ivf_pq::index_params> index_params;
auto index = cuvs::neighbors::mg::build(handle, index_params, index_dataset);

Parameters:

• handle – [in]

• index_params – configure the index building

• index_dataset – [in] a row-major matrix on host [n_rows, dim]

Returns:

the constructed IVF-PQ MG index

auto build(const raft::device_resources &handle, const mg::index_params<ivf_pq::index_params> &index_params, raft::host_matrix_view<const uint8_t, int64_t, row_major> index_dataset) -> index<ivf_pq::index<int64_t>, uint8_t, int64_t>

Builds a multi-GPU index.

Usage example:

raft::handle_t handle;
cuvs::neighbors::mg::index_params<ivf_pq::index_params> index_params;
auto index = cuvs::neighbors::mg::build(handle, index_params, index_dataset);

Parameters:

• handle – [in]

• index_params – configure the index building

• index_dataset – [in] a row-major matrix on host [n_rows, dim]

Returns:

the constructed IVF-PQ MG index

auto build(const raft::device_resources &handle, const mg::index_params<cagra::index_params> &index_params, raft::host_matrix_view<const float, int64_t, row_major> index_dataset) -> index<cagra::index<float, uint32_t>, float, uint32_t>

Builds a multi-GPU index.

Usage example:

raft::handle_t handle;
cuvs::neighbors::mg::index_params<cagra::index_params> index_params;
auto index = cuvs::neighbors::mg::build(handle, index_params, index_dataset);

Parameters:

• handle – [in]

• index_params – configure the index building

• index_dataset – [in] a row-major matrix on host [n_rows, dim]

Returns:

the constructed CAGRA MG index

auto build(const raft::device_resources &handle, const mg::index_params<cagra::index_params> &index_params, raft::host_matrix_view<const half, int64_t, row_major> index_dataset) -> index<cagra::index<half, uint32_t>, half, uint32_t>

Builds a multi-GPU index.

Usage example:

raft::handle_t handle;
cuvs::neighbors::mg::index_params<cagra::index_params> index_params;
auto index = cuvs::neighbors::mg::build(handle, index_params, index_dataset);

Parameters:

• handle – [in]

• index_params – configure the index building

• index_dataset – [in] a row-major matrix on host [n_rows, dim]

Returns:

the constructed CAGRA MG index

auto build(const raft::device_resources &handle, const mg::index_params<cagra::index_params> &index_params, raft::host_matrix_view<const int8_t, int64_t, row_major> index_dataset) -> index<cagra::index<int8_t, uint32_t>, int8_t, uint32_t>

Builds a multi-GPU index.

Usage example:

raft::handle_t handle;
cuvs::neighbors::mg::index_params<cagra::index_params> index_params;
auto index = cuvs::neighbors::mg::build(handle, index_params, index_dataset);

Parameters:

• handle – [in]

• index_params – configure the index building

• index_dataset – [in] a row-major matrix on host [n_rows, dim]

Returns:

the constructed CAGRA MG index

auto build(const raft::device_resources &handle, const mg::index_params<cagra::index_params> &index_params, raft::host_matrix_view<const uint8_t, int64_t, row_major> index_dataset) -> index<cagra::index<uint8_t, uint32_t>, uint8_t, uint32_t>

Builds a multi-GPU index.

Usage example:

raft::handle_t handle;
cuvs::neighbors::mg::index_params<cagra::index_params> index_params;
auto index = cuvs::neighbors::mg::build(handle, index_params, index_dataset);

Parameters:

• handle – [in]

• index_params – configure the index building

• index_dataset – [in] a row-major matrix on host [n_rows, dim]

Returns:

the constructed CAGRA MG index

Index extend

void extend(const raft::device_resources &handle, index<ivf_flat::index<float, int64_t>, float, int64_t> &index, raft::host_matrix_view<const float, int64_t, row_major> new_vectors, std::optional<raft::host_vector_view<const int64_t, int64_t>> new_indices)

Extends a multi-GPU index.

Usage example:

raft::handle_t handle;
cuvs::neighbors::mg::index_params<ivf_flat::index_params> index_params;
auto index = cuvs::neighbors::mg::build(handle, index_params, index_dataset);
cuvs::neighbors::mg::extend(handle, index, new_vectors, std::nullopt);

Parameters:

• handle – [in]

• index – [in] the pre-built index

• new_vectors – [in] a row-major matrix on host [n_rows, dim]

• new_indices – [in] optional vector on host [n_rows], std::nullopt means default continuous range [0...n_rows)

void extend(const raft::device_resources &handle, index<ivf_flat::index<int8_t, int64_t>, int8_t, int64_t> &index, raft::host_matrix_view<const int8_t, int64_t, row_major> new_vectors, std::optional<raft::host_vector_view<const int64_t, int64_t>> new_indices)

Extends a multi-GPU index.

Usage example:

raft::handle_t handle;
cuvs::neighbors::mg::index_params<ivf_flat::index_params> index_params;
auto index = cuvs::neighbors::mg::build(handle, index_params, index_dataset);
cuvs::neighbors::mg::extend(handle, index, new_vectors, std::nullopt);

Parameters:

• handle – [in]

• index – [in] the pre-built index

• new_vectors – [in] a row-major matrix on host [n_rows, dim]

• new_indices – [in] optional vector on host [n_rows], std::nullopt means default continuous range [0...n_rows)

void extend(const raft::device_resources &handle, index<ivf_flat::index<uint8_t, int64_t>, uint8_t, int64_t> &index, raft::host_matrix_view<const uint8_t, int64_t, row_major> new_vectors, std::optional<raft::host_vector_view<const int64_t, int64_t>> new_indices)

Extends a multi-GPU index.

Usage example:

raft::handle_t handle;
cuvs::neighbors::mg::index_params<ivf_flat::index_params> index_params;
auto index = cuvs::neighbors::mg::build(handle, index_params, index_dataset);
cuvs::neighbors::mg::extend(handle, index, new_vectors, std::nullopt);

Parameters:

• handle – [in]

• index – [in] the pre-built index

• new_vectors – [in] a row-major matrix on host [n_rows, dim]

• new_indices – [in] optional vector on host [n_rows], std::nullopt means default continuous range [0...n_rows)

void extend(const raft::device_resources &handle, index<ivf_pq::index<int64_t>, float, int64_t> &index, raft::host_matrix_view<const float, int64_t, row_major> new_vectors, std::optional<raft::host_vector_view<const int64_t, int64_t>> new_indices)

Extends a multi-GPU index.

Usage example:

raft::handle_t handle;
cuvs::neighbors::mg::index_params<ivf_pq::index_params> index_params;
auto index = cuvs::neighbors::mg::build(handle, index_params, index_dataset);
cuvs::neighbors::mg::extend(handle, index, new_vectors, std::nullopt);

Parameters:

• handle – [in]

• index – [in] the pre-built index

• new_vectors – [in] a row-major matrix on host [n_rows, dim]

• new_indices – [in] optional vector on host [n_rows], std::nullopt means default continuous range [0...n_rows)

void extend(const raft::device_resources &handle, index<ivf_pq::index<int64_t>, half, int64_t> &index, raft::host_matrix_view<const half, int64_t, row_major> new_vectors, std::optional<raft::host_vector_view<const int64_t, int64_t>> new_indices)

Extends a multi-GPU index.

Usage example:

raft::handle_t handle;
cuvs::neighbors::mg::index_params<ivf_pq::index_params> index_params;
auto index = cuvs::neighbors::mg::build(handle, index_params, index_dataset);
cuvs::neighbors::mg::extend(handle, index, new_vectors, std::nullopt);

Parameters:

• handle – [in]

• index – [in] the pre-built index

• new_vectors – [in] a row-major matrix on host [n_rows, dim]

• new_indices – [in] optional vector on host [n_rows], std::nullopt means default continuous range [0...n_rows)

void extend(const raft::device_resources &handle, index<ivf_pq::index<int64_t>, int8_t, int64_t> &index, raft::host_matrix_view<const int8_t, int64_t, row_major> new_vectors, std::optional<raft::host_vector_view<const int64_t, int64_t>> new_indices)

Extends a multi-GPU index.

Usage example:

raft::handle_t handle;
cuvs::neighbors::mg::index_params<ivf_pq::index_params> index_params;
auto index = cuvs::neighbors::mg::build(handle, index_params, index_dataset);
cuvs::neighbors::mg::extend(handle, index, new_vectors, std::nullopt);

Parameters:

• handle – [in]

• index – [in] the pre-built index

• new_vectors – [in] a row-major matrix on host [n_rows, dim]

• new_indices – [in] optional vector on host [n_rows], std::nullopt means default continuous range [0...n_rows)

void extend(const raft::device_resources &handle, index<ivf_pq::index<int64_t>, uint8_t, int64_t> &index, raft::host_matrix_view<const uint8_t, int64_t, row_major> new_vectors, std::optional<raft::host_vector_view<const int64_t, int64_t>> new_indices)

Extends a multi-GPU index.

Usage example:

raft::handle_t handle;
cuvs::neighbors::mg::index_params<ivf_pq::index_params> index_params;
auto index = cuvs::neighbors::mg::build(handle, index_params, index_dataset);
cuvs::neighbors::mg::extend(handle, index, new_vectors, std::nullopt);

Parameters:

• handle – [in]

• index – [in] the pre-built index

• new_vectors – [in] a row-major matrix on host [n_rows, dim]

• new_indices – [in] optional vector on host [n_rows], std::nullopt means default continuous range [0...n_rows)

void extend(const raft::device_resources &handle, index<cagra::index<float, uint32_t>, float, uint32_t> &index, raft::host_matrix_view<const float, int64_t, row_major> new_vectors, std::optional<raft::host_vector_view<const uint32_t, int64_t>> new_indices)

Extends a multi-GPU index.

Usage example:

raft::handle_t handle;
cuvs::neighbors::mg::index_params<cagra::index_params> index_params;
auto index = cuvs::neighbors::mg::build(handle, index_params, index_dataset);
cuvs::neighbors::mg::extend(handle, index, new_vectors, std::nullopt);

Parameters:

• handle – [in]

• index – [in] the pre-built index

• new_vectors – [in] a row-major matrix on host [n_rows, dim]

• new_indices – [in] optional vector on host [n_rows], std::nullopt means default continuous range [0...n_rows)

void extend(const raft::device_resources &handle, index<cagra::index<half, uint32_t>, half, uint32_t> &index, raft::host_matrix_view<const half, int64_t, row_major> new_vectors, std::optional<raft::host_vector_view<const uint32_t, int64_t>> new_indices)

Extends a multi-GPU index.

Usage example:

raft::handle_t handle;
cuvs::neighbors::mg::index_params<cagra::index_params> index_params;
auto index = cuvs::neighbors::mg::build(handle, index_params, index_dataset);
cuvs::neighbors::mg::extend(handle, index, new_vectors, std::nullopt);

Parameters:

• handle – [in]

• index – [in] the pre-built index

• new_vectors – [in] a row-major matrix on host [n_rows, dim]

• new_indices – [in] optional vector on host [n_rows], std::nullopt means default continuous range [0...n_rows)

void extend(const raft::device_resources &handle, index<cagra::index<int8_t, uint32_t>, int8_t, uint32_t> &index, raft::host_matrix_view<const int8_t, int64_t, row_major> new_vectors, std::optional<raft::host_vector_view<const uint32_t, int64_t>> new_indices)

Extends a multi-GPU index.

Usage example:

raft::handle_t handle;
cuvs::neighbors::mg::index_params<cagra::index_params> index_params;
auto index = cuvs::neighbors::mg::build(handle, index_params, index_dataset);
cuvs::neighbors::mg::extend(handle, index, new_vectors, std::nullopt);

Parameters:

• handle – [in]

• index – [in] the pre-built index

• new_vectors – [in] a row-major matrix on host [n_rows, dim]

• new_indices – [in] optional vector on host [n_rows], std::nullopt means default continuous range [0...n_rows)

void extend(const raft::device_resources &handle, index<cagra::index<uint8_t, uint32_t>, uint8_t, uint32_t> &index, raft::host_matrix_view<const uint8_t, int64_t, row_major> new_vectors, std::optional<raft::host_vector_view<const uint32_t, int64_t>> new_indices)

Extends a multi-GPU index.

Usage example:

raft::handle_t handle;
cuvs::neighbors::mg::index_params<cagra::index_params> index_params;
auto index = cuvs::neighbors::mg::build(handle, index_params, index_dataset);
cuvs::neighbors::mg::extend(handle, index, new_vectors, std::nullopt);

Parameters:

• handle – [in]

• index – [in] the pre-built index

• new_vectors – [in] a row-major matrix on host [n_rows, dim]

• new_indices – [in] optional vector on host [n_rows], std::nullopt means default continuous range [0...n_rows)

Index search

void search(const raft::device_resources &handle, const index<ivf_flat::index<float, int64_t>, float, int64_t> &index, const mg::search_params<ivf_flat::search_params> &search_params, raft::host_matrix_view<const float, int64_t, row_major> queries, raft::host_matrix_view<int64_t, int64_t, row_major> neighbors, raft::host_matrix_view<float, int64_t, row_major> distances, int64_t n_rows_per_batch = DEFAULT_SEARCH_BATCH_SIZE)

Searches a multi-GPU index.

Usage example:

raft::handle_t handle;
cuvs::neighbors::mg::index_params<ivf_flat::index_params> index_params;
auto index = cuvs::neighbors::mg::build(handle, index_params, index_dataset);
cuvs::neighbors::mg::search_params<ivf_flat::search_params> search_params;
cuvs::neighbors::mg::search(handle, index, search_params, queries, neighbors,
distances);

Parameters:

• handle – [in]

• index – [in] the pre-built index

• search_params – configure the index search

• queries – [in] a row-major matrix on host [n_rows, dim]

• neighbors – [out] a row-major matrix on host [n_rows, n_neighbors]

• distances – [out] a row-major matrix on host [n_rows, n_neighbors]

• n_rows_per_batch – [in] (optional) search batch size

void search(const raft::device_resources &handle, const index<ivf_flat::index<int8_t, int64_t>, int8_t, int64_t> &index, const mg::search_params<ivf_flat::search_params> &search_params, raft::host_matrix_view<const int8_t, int64_t, row_major> queries, raft::host_matrix_view<int64_t, int64_t, row_major> neighbors, raft::host_matrix_view<float, int64_t, row_major> distances, int64_t n_rows_per_batch = DEFAULT_SEARCH_BATCH_SIZE)

Searches a multi-GPU index.

Usage example:

raft::handle_t handle;
cuvs::neighbors::mg::index_params<ivf_flat::index_params> index_params;
auto index = cuvs::neighbors::mg::build(handle, index_params, index_dataset);
cuvs::neighbors::mg::search_params<ivf_flat::search_params> search_params;
cuvs::neighbors::mg::search(handle, index, search_params, queries, neighbors,
distances);

Parameters:

• handle – [in]

• index – [in] the pre-built index

• search_params – configure the index search

• queries – [in] a row-major matrix on host [n_rows, dim]

• neighbors – [out] a row-major matrix on host [n_rows, n_neighbors]

• distances – [out] a row-major matrix on host [n_rows, n_neighbors]

• n_rows_per_batch – [in] (optional) search batch size

void search(const raft::device_resources &handle, const index<ivf_flat::index<uint8_t, int64_t>, uint8_t, int64_t> &index, const mg::search_params<ivf_flat::search_params> &search_params, raft::host_matrix_view<const uint8_t, int64_t, row_major> queries, raft::host_matrix_view<int64_t, int64_t, row_major> neighbors, raft::host_matrix_view<float, int64_t, row_major> distances, int64_t n_rows_per_batch = DEFAULT_SEARCH_BATCH_SIZE)

Searches a multi-GPU index.

Usage example:

raft::handle_t handle;
cuvs::neighbors::mg::index_params<ivf_flat::index_params> index_params;
auto index = cuvs::neighbors::mg::build(handle, index_params, index_dataset);
cuvs::neighbors::mg::search_params<ivf_flat::search_params> search_params;
cuvs::neighbors::mg::search(handle, index, search_params, queries, neighbors,
distances);

Parameters:

• handle – [in]

• index – [in] the pre-built index

• search_params – configure the index search

• queries – [in] a row-major matrix on host [n_rows, dim]

• neighbors – [out] a row-major matrix on host [n_rows, n_neighbors]

• distances – [out] a row-major matrix on host [n_rows, n_neighbors]

• n_rows_per_batch – [in] (optional) search batch size

void search(const raft::device_resources &handle, const index<ivf_pq::index<int64_t>, float, int64_t> &index, const mg::search_params<ivf_pq::search_params> &search_params, raft::host_matrix_view<const float, int64_t, row_major> queries, raft::host_matrix_view<int64_t, int64_t, row_major> neighbors, raft::host_matrix_view<float, int64_t, row_major> distances, int64_t n_rows_per_batch = DEFAULT_SEARCH_BATCH_SIZE)

Searches a multi-GPU index.

Usage example:

raft::handle_t handle;
cuvs::neighbors::mg::index_params<ivf_pq::index_params> index_params;
auto index = cuvs::neighbors::mg::build(handle, index_params, index_dataset);
cuvs::neighbors::mg::search_params<ivf_pq::search_params> search_params;
cuvs::neighbors::mg::search(handle, index, search_params, queries, neighbors,
distances);

Parameters:

• handle – [in]

• index – [in] the pre-built index

• search_params – configure the index search

• queries – [in] a row-major matrix on host [n_rows, dim]

• neighbors – [out] a row-major matrix on host [n_rows, n_neighbors]

• distances – [out] a row-major matrix on host [n_rows, n_neighbors]

• n_rows_per_batch – [in] (optional) search batch size

void search(const raft::device_resources &handle, const index<ivf_pq::index<int64_t>, half, int64_t> &index, const mg::search_params<ivf_pq::search_params> &search_params, raft::host_matrix_view<const half, int64_t, row_major> queries, raft::host_matrix_view<int64_t, int64_t, row_major> neighbors, raft::host_matrix_view<float, int64_t, row_major> distances, int64_t n_rows_per_batch = DEFAULT_SEARCH_BATCH_SIZE)

Searches a multi-GPU index.

Usage example:

raft::handle_t handle;
cuvs::neighbors::mg::index_params<ivf_pq::index_params> index_params;
auto index = cuvs::neighbors::mg::build(handle, index_params, index_dataset);
cuvs::neighbors::mg::search_params<ivf_pq::search_params> search_params;
cuvs::neighbors::mg::search(handle, index, search_params, queries, neighbors,
distances);

Parameters:

• handle – [in]

• index – [in] the pre-built index

• search_params – configure the index search

• queries – [in] a row-major matrix on host [n_rows, dim]

• neighbors – [out] a row-major matrix on host [n_rows, n_neighbors]

• distances – [out] a row-major matrix on host [n_rows, n_neighbors]

• n_rows_per_batch – [in] (optional) search batch size

void search(const raft::device_resources &handle, const index<ivf_pq::index<int64_t>, int8_t, int64_t> &index, const mg::search_params<ivf_pq::search_params> &search_params, raft::host_matrix_view<const int8_t, int64_t, row_major> queries, raft::host_matrix_view<int64_t, int64_t, row_major> neighbors, raft::host_matrix_view<float, int64_t, row_major> distances, int64_t n_rows_per_batch = DEFAULT_SEARCH_BATCH_SIZE)

Searches a multi-GPU index.

Usage example:

raft::handle_t handle;
cuvs::neighbors::mg::index_params<ivf_pq::index_params> index_params;
auto index = cuvs::neighbors::mg::build(handle, index_params, index_dataset);
cuvs::neighbors::mg::search_params<ivf_pq::search_params> search_params;
cuvs::neighbors::mg::search(handle, index, search_params, queries, neighbors,
distances);

Parameters:

• handle – [in]

• index – [in] the pre-built index

• search_params – configure the index search

• queries – [in] a row-major matrix on host [n_rows, dim]

• neighbors – [out] a row-major matrix on host [n_rows, n_neighbors]

• distances – [out] a row-major matrix on host [n_rows, n_neighbors]

• n_rows_per_batch – [in] (optional) search batch size

void search(const raft::device_resources &handle, const index<ivf_pq::index<int64_t>, uint8_t, int64_t> &index, const mg::search_params<ivf_pq::search_params> &search_params, raft::host_matrix_view<const uint8_t, int64_t, row_major> queries, raft::host_matrix_view<int64_t, int64_t, row_major> neighbors, raft::host_matrix_view<float, int64_t, row_major> distances, int64_t n_rows_per_batch = DEFAULT_SEARCH_BATCH_SIZE)

Searches a multi-GPU index.

Usage example:

raft::handle_t handle;
cuvs::neighbors::mg::index_params<ivf_pq::index_params> index_params;
auto index = cuvs::neighbors::mg::build(handle, index_params, index_dataset);
cuvs::neighbors::mg::search_params<ivf_pq::search_params> search_params;
cuvs::neighbors::mg::search(handle, index, search_params, queries, neighbors,
distances);

Parameters:

• handle – [in]

• index – [in] the pre-built index

• search_params – configure the index search

• queries – [in] a row-major matrix on host [n_rows, dim]

• neighbors – [out] a row-major matrix on host [n_rows, n_neighbors]

• distances – [out] a row-major matrix on host [n_rows, n_neighbors]

• n_rows_per_batch – [in] (optional) search batch size

void search(const raft::device_resources &handle, const index<cagra::index<float, uint32_t>, float, uint32_t> &index, const mg::search_params<cagra::search_params> &search_params, raft::host_matrix_view<const float, int64_t, row_major> queries, raft::host_matrix_view<uint32_t, int64_t, row_major> neighbors, raft::host_matrix_view<float, int64_t, row_major> distances, int64_t n_rows_per_batch = DEFAULT_SEARCH_BATCH_SIZE)

Searches a multi-GPU index.

Usage example:

raft::handle_t handle;
cuvs::neighbors::mg::index_params<cagra::index_params> index_params;
auto index = cuvs::neighbors::mg::build(handle, index_params, index_dataset);
cuvs::neighbors::mg::search_params<cagra::search_params> search_params;
cuvs::neighbors::mg::search(handle, index, search_params, queries, neighbors,
distances);

Parameters:

• handle – [in]

• index – [in] the pre-built index

• search_params – configure the index search

• queries – [in] a row-major matrix on host [n_rows, dim]

• neighbors – [out] a row-major matrix on host [n_rows, n_neighbors]

• distances – [out] a row-major matrix on host [n_rows, n_neighbors]

• n_rows_per_batch – [in] (optional) search batch size

void search(const raft::device_resources &handle, const index<cagra::index<half, uint32_t>, half, uint32_t> &index, const mg::search_params<cagra::search_params> &search_params, raft::host_matrix_view<const half, int64_t, row_major> queries, raft::host_matrix_view<uint32_t, int64_t, row_major> neighbors, raft::host_matrix_view<float, int64_t, row_major> distances, int64_t n_rows_per_batch = DEFAULT_SEARCH_BATCH_SIZE)

Searches a multi-GPU index.

Usage example:

raft::handle_t handle;
cuvs::neighbors::mg::index_params<cagra::index_params> index_params;
auto index = cuvs::neighbors::mg::build(handle, index_params, index_dataset);
cuvs::neighbors::mg::search_params<cagra::search_params> search_params;
cuvs::neighbors::mg::search(handle, index, search_params, queries, neighbors,
distances);

Parameters:

• handle – [in]

• index – [in] the pre-built index

• search_params – configure the index search

• queries – [in] a row-major matrix on host [n_rows, dim]

• neighbors – [out] a row-major matrix on host [n_rows, n_neighbors]

• distances – [out] a row-major matrix on host [n_rows, n_neighbors]

• n_rows_per_batch – [in] (optional) search batch size

void search(const raft::device_resources &handle, const index<cagra::index<int8_t, uint32_t>, int8_t, uint32_t> &index, const mg::search_params<cagra::search_params> &search_params, raft::host_matrix_view<const int8_t, int64_t, row_major> queries, raft::host_matrix_view<uint32_t, int64_t, row_major> neighbors, raft::host_matrix_view<float, int64_t, row_major> distances, int64_t n_rows_per_batch = DEFAULT_SEARCH_BATCH_SIZE)

Searches a multi-GPU index.

Usage example:

raft::handle_t handle;
cuvs::neighbors::mg::index_params<cagra::index_params> index_params;
auto index = cuvs::neighbors::mg::build(handle, index_params, index_dataset);
cuvs::neighbors::mg::search_params<cagra::search_params> search_params;
cuvs::neighbors::mg::search(handle, index, search_params, queries, neighbors,
distances);

Parameters:

• handle – [in]

• index – [in] the pre-built index

• search_params – configure the index search

• queries – [in] a row-major matrix on host [n_rows, dim]

• neighbors – [out] a row-major matrix on host [n_rows, n_neighbors]

• distances – [out] a row-major matrix on host [n_rows, n_neighbors]

• n_rows_per_batch – [in] (optional) search batch size

void search(const raft::device_resources &handle, const index<cagra::index<uint8_t, uint32_t>, uint8_t, uint32_t> &index, const mg::search_params<cagra::search_params> &search_params, raft::host_matrix_view<const uint8_t, int64_t, row_major> queries, raft::host_matrix_view<uint32_t, int64_t, row_major> neighbors, raft::host_matrix_view<float, int64_t, row_major> distances, int64_t n_rows_per_batch = DEFAULT_SEARCH_BATCH_SIZE)

Searches a multi-GPU index.

Usage example:

raft::handle_t handle;
cuvs::neighbors::mg::index_params<cagra::index_params> index_params;
auto index = cuvs::neighbors::mg::build(handle, index_params, index_dataset);
cuvs::neighbors::mg::search_params<cagra::search_params> search_params;
cuvs::neighbors::mg::search(handle, index, search_params, queries, neighbors,
distances);

Parameters:

• handle – [in]

• index – [in] the pre-built index

• search_params – configure the index search

• queries – [in] a row-major matrix on host [n_rows, dim]

• neighbors – [out] a row-major matrix on host [n_rows, n_neighbors]

• distances – [out] a row-major matrix on host [n_rows, n_neighbors]

• n_rows_per_batch – [in] (optional) search batch size

Index serialize

void serialize(const raft::device_resources &handle, const index<ivf_flat::index<float, int64_t>, float, int64_t> &index, const std::string &filename)

Serializes a multi-GPU index.

Usage example:

raft::handle_t handle;
cuvs::neighbors::mg::index_params<ivf_flat::index_params> index_params;
auto index = cuvs::neighbors::mg::build(handle, index_params, index_dataset);
const std::string filename = "mg_index.cuvs";
cuvs::neighbors::mg::serialize(handle, index, filename);

Parameters:

• handle – [in]

• index – [in] the pre-built index

• filename – [in] path to the file to be serialized

void serialize(const raft::device_resources &handle, const index<ivf_flat::index<int8_t, int64_t>, int8_t, int64_t> &index, const std::string &filename)

Serializes a multi-GPU index.

Usage example:

raft::handle_t handle;
cuvs::neighbors::mg::index_params<ivf_flat::index_params> index_params;
auto index = cuvs::neighbors::mg::build(handle, index_params, index_dataset);
const std::string filename = "mg_index.cuvs";
cuvs::neighbors::mg::serialize(handle, index, filename);

Parameters:

• handle – [in]

• index – [in] the pre-built index

• filename – [in] path to the file to be serialized

void serialize(const raft::device_resources &handle, const index<ivf_flat::index<uint8_t, int64_t>, uint8_t, int64_t> &index, const std::string &filename)

Serializes a multi-GPU index.

Usage example:

raft::handle_t handle;
cuvs::neighbors::mg::index_params<ivf_flat::index_params> index_params;
auto index = cuvs::neighbors::mg::build(handle, index_params, index_dataset);
const std::string filename = "mg_index.cuvs";
cuvs::neighbors::mg::serialize(handle, index, filename);

Parameters:

• handle – [in]

• index – [in] the pre-built index

• filename – [in] path to the file to be serialized

void serialize(const raft::device_resources &handle, const index<ivf_pq::index<int64_t>, float, int64_t> &index, const std::string &filename)

Serializes a multi-GPU index.

Usage example:

raft::handle_t handle;
cuvs::neighbors::mg::index_params<ivf_pq::index_params> index_params;
auto index = cuvs::neighbors::mg::build(handle, index_params, index_dataset);
const std::string filename = "mg_index.cuvs";
cuvs::neighbors::mg::serialize(handle, index, filename);

Parameters:

• handle – [in]

• index – [in] the pre-built index

• filename – [in] path to the file to be serialized

void serialize(const raft::device_resources &handle, const index<ivf_pq::index<int64_t>, half, int64_t> &index, const std::string &filename)

Serializes a multi-GPU index.

Usage example:

raft::handle_t handle;
cuvs::neighbors::mg::index_params<ivf_pq::index_params> index_params;
auto index = cuvs::neighbors::mg::build(handle, index_params, index_dataset);
const std::string filename = "mg_index.cuvs";
cuvs::neighbors::mg::serialize(handle, index, filename);

Parameters:

• handle – [in]

• index – [in] the pre-built index

• filename – [in] path to the file to be serialized

void serialize(const raft::device_resources &handle, const index<ivf_pq::index<int64_t>, int8_t, int64_t> &index, const std::string &filename)

Serializes a multi-GPU index.

Usage example:

raft::handle_t handle;
cuvs::neighbors::mg::index_params<ivf_pq::index_params> index_params;
auto index = cuvs::neighbors::mg::build(handle, index_params, index_dataset);
const std::string filename = "mg_index.cuvs";
cuvs::neighbors::mg::serialize(handle, index, filename);

Parameters:

• handle – [in]

• index – [in] the pre-built index

• filename – [in] path to the file to be serialized

void serialize(const raft::device_resources &handle, const index<ivf_pq::index<int64_t>, uint8_t, int64_t> &index, const std::string &filename)

Serializes a multi-GPU index.

Usage example:

raft::handle_t handle;
cuvs::neighbors::mg::index_params<ivf_pq::index_params> index_params;
auto index = cuvs::neighbors::mg::build(handle, index_params, index_dataset);
const std::string filename = "mg_index.cuvs";
cuvs::neighbors::mg::serialize(handle, index, filename);

Parameters:

• handle – [in]

• index – [in] the pre-built index

• filename – [in] path to the file to be serialized

void serialize(const raft::device_resources &handle, const index<cagra::index<float, uint32_t>, float, uint32_t> &index, const std::string &filename)

Serializes a multi-GPU index.

Usage example:

raft::handle_t handle;
cuvs::neighbors::mg::index_params<cagra::index_params> index_params;
auto index = cuvs::neighbors::mg::build(handle, index_params, index_dataset);
const std::string filename = "mg_index.cuvs";
cuvs::neighbors::mg::serialize(handle, index, filename);

Parameters:

• handle – [in]

• index – [in] the pre-built index

• filename – [in] path to the file to be serialized

void serialize(const raft::device_resources &handle, const index<cagra::index<half, uint32_t>, half, uint32_t> &index, const std::string &filename)

Serializes a multi-GPU index.

Usage example:

raft::handle_t handle;
cuvs::neighbors::mg::index_params<cagra::index_params> index_params;
auto index = cuvs::neighbors::mg::build(handle, index_params, index_dataset);
const std::string filename = "mg_index.cuvs";
cuvs::neighbors::mg::serialize(handle, index, filename);

Parameters:

• handle – [in]

• index – [in] the pre-built index

• filename – [in] path to the file to be serialized

void serialize(const raft::device_resources &handle, const index<cagra::index<int8_t, uint32_t>, int8_t, uint32_t> &index, const std::string &filename)

Serializes a multi-GPU index.

Usage example:

raft::handle_t handle;
cuvs::neighbors::mg::index_params<cagra::index_params> index_params;
auto index = cuvs::neighbors::mg::build(handle, index_params, index_dataset);
const std::string filename = "mg_index.cuvs";
cuvs::neighbors::mg::serialize(handle, index, filename);

Parameters:

• handle – [in]

• index – [in] the pre-built index

• filename – [in] path to the file to be serialized

void serialize(const raft::device_resources &handle, const index<cagra::index<uint8_t, uint32_t>, uint8_t, uint32_t> &index, const std::string &filename)

Serializes a multi-GPU index.

Usage example:

raft::handle_t handle;
cuvs::neighbors::mg::index_params<cagra::index_params> index_params;
auto index = cuvs::neighbors::mg::build(handle, index_params, index_dataset);
const std::string filename = "mg_index.cuvs";
cuvs::neighbors::mg::serialize(handle, index, filename);

Parameters:

• handle – [in]

• index – [in] the pre-built index

• filename – [in] path to the file to be serialized

Index deserialize

template<typename T, typename IdxT>
auto deserialize_flat(const raft::device_resources &handle, const std::string &filename) -> index<ivf_flat::index<T, IdxT>, T, IdxT>

Deserializes an IVF-Flat multi-GPU index.

Usage example:

raft::handle_t handle;
cuvs::neighbors::mg::index_params<ivf_flat::index_params> index_params;
auto index = cuvs::neighbors::mg::build(handle, index_params, index_dataset);
const std::string filename = "mg_index.cuvs";
cuvs::neighbors::mg::serialize(handle, index, filename);
auto new_index = cuvs::neighbors::mg::deserialize_flat<float, int64_t>(handle, filename);

Parameters:

• handle – [in]

• filename – [in] path to the file to be deserialized

template<typename T, typename IdxT>
auto deserialize_pq(const raft::device_resources &handle, const std::string &filename) -> index<ivf_pq::index<IdxT>, T, IdxT>

Deserializes an IVF-PQ multi-GPU index.

Usage example:

raft::handle_t handle;
cuvs::neighbors::mg::index_params<ivf_pq::index_params> index_params;
auto index = cuvs::neighbors::mg::build(handle, index_params, index_dataset);
const std::string filename = "mg_index.cuvs";
cuvs::neighbors::mg::serialize(handle, index, filename);
auto new_index = cuvs::neighbors::mg::deserialize_pq<float, int64_t>(handle, filename);

Parameters:

• handle – [in]

• filename – [in] path to the file to be deserialized

template<typename T, typename IdxT>
auto deserialize_cagra(const raft::device_resources &handle, const std::string &filename) -> index<cagra::index<T, IdxT>, T, IdxT>

Deserializes a CAGRA multi-GPU index.

Usage example:

raft::handle_t handle;
cuvs::neighbors::mg::index_params<cagra::index_params> index_params;
auto index = cuvs::neighbors::mg::build(handle, index_params, index_dataset);
const std::string filename = "mg_index.cuvs";
cuvs::neighbors::mg::serialize(handle, index, filename);
auto new_index = cuvs::neighbors::mg::deserialize_cagra<float, uint32_t>(handle, filename);

Parameters:

• handle – [in]

• filename – [in] path to the file to be deserialized

Distribute pre-built local index

template<typename T, typename IdxT>
auto distribute_flat(const raft::device_resources &handle, const std::string &filename) -> index<ivf_flat::index<T, IdxT>, T, IdxT>

Replicates a locally built and serialized IVF-Flat index to all GPUs to form a distributed multi-GPU index.

Usage example:

raft::handle_t handle;
cuvs::neighbors::ivf_flat::index_params index_params;
auto index = cuvs::neighbors::ivf_flat::build(handle, index_params, index_dataset);
const std::string filename = "local_index.cuvs";
cuvs::neighbors::ivf_flat::serialize(handle, filename, index);
auto new_index = cuvs::neighbors::mg::distribute_flat<float, int64_t>(handle, filename);

Parameters:

• handle – [in]

• filename – [in] path to the file to be deserialized : a local index

template<typename T, typename IdxT>
auto distribute_pq(const raft::device_resources &handle, const std::string &filename) -> index<ivf_pq::index<IdxT>, T, IdxT>

Replicates a locally built and serialized IVF-PQ index to all GPUs to form a distributed multi-GPU index.

Usage example:

raft::handle_t handle;
cuvs::neighbors::ivf_pq::index_params index_params;
auto index = cuvs::neighbors::ivf_pq::build(handle, index_params, index_dataset);
const std::string filename = "local_index.cuvs";
cuvs::neighbors::ivf_pq::serialize(handle, filename, index);
auto new_index = cuvs::neighbors::mg::distribute_pq<float, int64_t>(handle, filename);

Parameters:

• handle – [in]

• filename – [in] path to the file to be deserialized : a local index

template<typename T, typename IdxT>
auto distribute_cagra(const raft::device_resources &handle, const std::string &filename) -> index<cagra::index<T, IdxT>, T, IdxT>

Replicates a locally built and serialized CAGRA index to all GPUs to form a distributed multi-GPU index.

Usage example:

raft::handle_t handle;
cuvs::neighbors::cagra::index_params index_params;
auto index = cuvs::neighbors::cagra::build(handle, index_params, index_dataset);
const std::string filename = "local_index.cuvs";
cuvs::neighbors::cagra::serialize(handle, filename, index);
auto new_index = cuvs::neighbors::mg::distribute_cagra<float, uint32_t>(handle, filename);

Parameters:

• handle – [in]

• filename – [in] path to the file to be deserialized : a local index