CAGRA#

CAGRA is a graph-based nearest neighbors implementation with state-of-the art query performance for both small- and large-batch sized search.

Please note that the CAGRA implementation is currently experimental and the API is subject to change from release to release. We are currently working on promoting CAGRA to a top-level stable API within RAFT.

#include <raft/neighbors/cagra.cuh>

namespace raft::neighbors::cagra

template<typename DataT, typename IdxT, typename accessor>
void build_knn_graph(raft::resources const &res, mdspan<const DataT, matrix_extent<int64_t>, row_major, accessor> dataset, raft::host_matrix_view<IdxT, int64_t, row_major> knn_graph, std::optional<float> refine_rate = std::nullopt, std::optional<ivf_pq::index_params> build_params = std::nullopt, std::optional<ivf_pq::search_params> search_params = std::nullopt)#

Build a kNN graph.

The kNN graph is the first building block for CAGRA index. This function uses the IVF-PQ method to build a kNN graph.

The output is a dense matrix that stores the neighbor indices for each pont in the dataset. Each point has the same number of neighbors.

See cagra::build for an alternative method.

The following distance metrics are supported:

  • L2Expanded

Usage example: 

  using namespace raft::neighbors;
  // use default index parameters
  cagra::index_params build_params;
  cagra::search_params search_params
  auto knn_graph      = raft::make_host_matrix<IdxT, IdxT>(dataset.extent(0), 128);
  // create knn graph
  cagra::build_knn_graph(res, dataset, knn_graph.view(), 2, build_params, search_params);
  auto optimized_gaph      = raft::make_host_matrix<IdxT, IdxT>(dataset.extent(0), 64);
  cagra::optimize(res, dataset, knn_graph.view(), optimized_graph.view());
  // Construct an index from dataset and optimized knn_graph
  auto index = cagra::index<T, IdxT>(res, build_params.metric(), dataset,
optimized_graph.view());

Template Parameters:

  • DataT – data element type

  • IdxT – type of the dataset vector indices

Parameters:

  • res – [in] raft resources

  • dataset – [in] a matrix view (host or device) to a row-major matrix [n_rows, dim]

  • knn_graph – [out] a host matrix view to store the output knn graph [n_rows, graph_degree]

  • refine_rate – [in] refinement rate for ivf-pq search

  • build_params – [in] (optional) ivf_pq index building parameters for knn graph

  • search_params – (optional) ivf_pq search parameters

template<typename DataT, typename IdxT = uint32_t, typename d_accessor = host_device_accessor<std::experimental::default_accessor<DataT>, memory_type::device>, typename g_accessor = host_device_accessor<std::experimental::default_accessor<IdxT>, memory_type::host>>
void sort_knn_graph(raft::resources const &res, mdspan<const DataT, matrix_extent<int64_t>, row_major, d_accessor> dataset, mdspan<IdxT, matrix_extent<int64_t>, row_major, g_accessor> knn_graph)#

Sort a KNN graph index. Preprocessing step for cagra::optimize: If a KNN graph is not built using cagra::build_knn_graph, then it is necessary to call this function before calling cagra::optimize. If the graph is built by cagra::build_knn_graph, it is already sorted and you do not need to call this function.

Usage example: 

  using namespace raft::neighbors;
  cagra::index_params build_params;
  auto knn_graph      = raft::make_host_matrix<IdxT, IdxT>(dataset.extent(0), 128);
  // build KNN graph not using `cagra::build_knn_graph`
  // build(knn_graph, dataset, ...);
  // sort graph index
  sort_knn_graph(res, dataset.view(), knn_graph.view());
  // optimize graph
  cagra::optimize(res, dataset, knn_graph.view(), optimized_graph.view());
  // Construct an index from dataset and optimized knn_graph
  auto index = cagra::index<T, IdxT>(res, build_params.metric(), dataset,
optimized_graph.view());

Template Parameters:

  • DataT – type of the data in the source dataset

  • IdxT – type of the dataset vector indices

Parameters:

  • res – [in] raft resources

  • dataset – [in] a matrix view (host or device) to a row-major matrix [n_rows, dim]

  • knn_graph – [inout] a matrix view (host or device) of the input knn graph [n_rows, knn_graph_degree]

template<typename IdxT = uint32_t, typename g_accessor = host_device_accessor<std::experimental::default_accessor<IdxT>, memory_type::host>>
void optimize(raft::resources const &res, mdspan<IdxT, matrix_extent<int64_t>, row_major, g_accessor> knn_graph, raft::host_matrix_view<IdxT, int64_t, row_major> new_graph)#

Prune a KNN graph.

Decrease the number of neighbors for each node.

See cagra::build_knn_graph for usage example

Template Parameters:

IdxT – type of the indices in the source dataset

Parameters:

  • res – [in] raft resources

  • knn_graph – [in] a matrix view (host or device) of the input knn graph [n_rows, knn_graph_degree]

  • new_graph – [out] a host matrix view of the optimized knn graph [n_rows, graph_degree]

template<typename T, typename IdxT = uint32_t, typename Accessor = host_device_accessor<std::experimental::default_accessor<T>, memory_type::host>>
index<T, IdxT> build(raft::resources const &res, const index_params &params, mdspan<const T, matrix_extent<int64_t>, row_major, Accessor> dataset)#

Build the index from the dataset for efficient search.

The build consist of two steps: build an intermediate knn-graph, and optimize it to create the final graph. The index_params struct controls the node degree of these graphs.

It is required that dataset and the optimized graph fit the GPU memory.

To customize the parameters for knn-graph building and pruning, and to reuse the intermediate results, you could build the index in two steps using cagra::build_knn_graph and cagra::optimize.

The following distance metrics are supported:

  • L2

Usage example: 

using namespace raft::neighbors;
// use default index parameters
cagra::index_params index_params;
// create and fill the index from a [N, D] dataset
auto index = cagra::build(res, index_params, dataset);
// use default search parameters
cagra::search_params search_params;
// search K nearest neighbours
auto neighbors = raft::make_device_matrix<uint32_t>(res, n_queries, k);
auto distances = raft::make_device_matrix<float>(res, n_queries, k);
cagra::search(res, search_params, index, queries, neighbors, distances);

Template Parameters:

  • T – data element type

  • IdxT – type of the indices in the source dataset

Parameters:

  • res – [in]

  • params – [in] parameters for building the index

  • dataset – [in] a matrix view (host or device) to a row-major matrix [n_rows, dim]

Returns:

the constructed cagra index

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

Search ANN using the constructed index.

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

Template Parameters:

  • T – data element type

  • IdxT – type of the indices

Parameters:

  • res – [in] raft resources

  • params – [in] configure the search

  • idx – [in] cagra 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]

struct index_params : public raft::neighbors::ann::index_params#
#include <cagra_types.hpp>