Skip to content

Commit

Permalink
Iteratively build graph index (#612)
Browse files Browse the repository at this point in the history 
This PR is about how CAGRA's search() and optimize() can be used to iteratively create and improve graph index.

Currently, IVFPQ and NND are used to create the initial kNN graph, which is then optimized to create the CAGRA search graph. So, for example, if you want to support a new data type in CAGRA, you need to create an initial kNN graph with that data type, and IVFPQ or NND must also support that new data type. This is a bit of hassle.

This PR is one solution to that problem. With functionality of this PR, once the CAGRA search supports the new data type, it can be used to create a graph index with it.

Authors:
  - Akira Naruse (https://github.com/anaruse)
  - Corey J. Nolet (https://github.com/cjnolet)

Approvers:
  - Artem M. Chirkin (https://github.com/achirkin)
  - Tamas Bela Feher (https://github.com/tfeher)

URL: #612
  • Loading branch information
@anaruse
anaruse authored Feb 2, 2025
1 parent 995b244 commit 2b7dce6
Show file tree
Hide file tree
Showing 3 changed files with 234 additions and 37 deletions.
17 changes: 13 additions & 4 deletions cpp/include/cuvs/neighbors/cagra.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -72,6 +72,10 @@ struct ivf_pq_params {
};

using nn_descent_params = cuvs::neighbors::nn_descent::index_params;

// **** Experimental ****
using iterative_search_params = cuvs::neighbors::search_params;

} // namespace graph_build_params

struct index_params : cuvs::neighbors::index_params {
Expand All @@ -87,9 +91,9 @@ struct index_params : cuvs::neighbors::index_params {

/** Parameters for graph building.
*
* Set ivf_pq_params or nn_descent_params to select the graph build algorithm and control their
* parameters. The default (std::monostate) is to use a heuristic to decide the algorithm and its
* parameters.
* Set ivf_pq_params, nn_descent_params, or iterative_search_params to select the graph build
* algorithm and control their parameters. The default (std::monostate) is to use a heuristic
* to decide the algorithm and its parameters.
*
* @code{.cpp}
* cagra::index_params params;
Expand All @@ -100,11 +104,16 @@ struct index_params : cuvs::neighbors::index_params {
* // 2. Choose NN Descent algorithm for kNN graph construction
* params.graph_build_params =
* cagra::graph_build_params::nn_descent_params(params.intermediate_graph_degree);
*
* // 3. Choose iterative graph building using CAGRA's search() and optimize() [Experimental]
* params.graph_build_params =
* cagra::graph_build_params::iterative_search_params();
* @endcode
*/
std::variant<std::monostate,
graph_build_params::ivf_pq_params,
graph_build_params::nn_descent_params>
graph_build_params::nn_descent_params,
graph_build_params::iterative_search_params>
graph_build_params;

/**
Expand Down
224 changes: 196 additions & 28 deletions cpp/src/neighbors/detail/cagra/cagra_build.cuh
View file
Original file line number Diff line number Diff line change
Expand Up @@ -401,6 +401,167 @@ void optimize(
res, knn_graph_internal, new_graph_internal, guarantee_connectivity);
}

template <typename T,
typename IdxT = uint32_t,
typename Accessor = raft::host_device_accessor<std::experimental::default_accessor<T>,
raft::memory_type::host>>
auto iterative_build_graph(
raft::resources const& res,
const index_params& params,
raft::mdspan<const T, raft::matrix_extent<int64_t>, raft::row_major, Accessor> dataset)
{
size_t intermediate_degree = params.intermediate_graph_degree;
size_t graph_degree = params.graph_degree;

auto cagra_graph = raft::make_host_matrix<IdxT, int64_t>(0, 0);

// Iteratively improve the accuracy of the graph by repeatedly running
// CAGRA's search() and optimize(). As for the size of the graph, instead
// of targeting all nodes from the beginning, the number of nodes is
// initially small, and the number of nodes is doubled with each iteration.
RAFT_LOG_INFO("Iteratively creating/improving graph index using CAGRA's search() and optimize()");

// If dataset is a host matrix, change it to a device matrix. Also, if the
// dimensionality of the dataset does not meet the alighnemt restriction,
// add extra dimensions and change it to a strided matrix.
std::unique_ptr<strided_dataset<T, int64_t>> dev_aligned_dataset;
try {
dev_aligned_dataset = make_aligned_dataset(res, dataset);
} catch (raft::logic_error& e) {
RAFT_LOG_ERROR("Iterative CAGRA graph build requires the dataset to fit GPU memory");
throw e;
}
auto dev_aligned_dataset_view = dev_aligned_dataset.get()->view();

// If the matrix stride and extent do no match, the extra dimensions are
// also as extent since it cannot be used as query matrix.
auto dev_dataset =
raft::make_device_matrix_view<const T, int64_t>(dev_aligned_dataset_view.data_handle(),
dev_aligned_dataset_view.extent(0),
dev_aligned_dataset_view.stride(0));

// Determine initial graph size.
uint64_t final_graph_size = (uint64_t)dataset.extent(0);
uint64_t initial_graph_size = (final_graph_size + 1) / 2;
while (initial_graph_size > graph_degree * 64) {
initial_graph_size = (initial_graph_size + 1) / 2;
}
RAFT_LOG_DEBUG("# initial graph size = %lu", (uint64_t)initial_graph_size);

// Allocate memory for search results.
constexpr uint64_t max_chunk_size = 8192;
auto topk = intermediate_degree;
auto dev_neighbors = raft::make_device_matrix<IdxT, int64_t>(res, max_chunk_size, topk);
auto dev_distances = raft::make_device_matrix<float, int64_t>(res, max_chunk_size, topk);

// Determine graph degree and number of search results while increasing
// graph size.
auto small_graph_degree = std::max(graph_degree / 2, std::min(graph_degree, (uint64_t)32));
auto small_topk = topk * small_graph_degree / graph_degree;
RAFT_LOG_DEBUG("# graph_degree = %lu", (uint64_t)graph_degree);
RAFT_LOG_DEBUG("# small_graph_degree = %lu", (uint64_t)small_graph_degree);
RAFT_LOG_DEBUG("# topk = %lu", (uint64_t)topk);
RAFT_LOG_DEBUG("# small_topk = %lu", (uint64_t)small_topk);

// Create an initial graph. The initial graph created here is not suitable for
// searching, but connectivity is guaranteed.
auto offset = raft::make_host_vector<IdxT, int64_t>(small_graph_degree);
const double base = sqrt((double)2.0);
for (uint64_t j = 0; j < small_graph_degree; j++) {
if (j == 0) {
offset(j) = 1;
} else {
offset(j) = offset(j - 1) + 1;
}
IdxT ofst = initial_graph_size * pow(base, (double)j - small_graph_degree - 1);
if (offset(j) < ofst) { offset(j) = ofst; }
RAFT_LOG_DEBUG("# offset(%lu) = %lu\n", (uint64_t)j, (uint64_t)offset(j));
}
cagra_graph = raft::make_host_matrix<IdxT, int64_t>(initial_graph_size, small_graph_degree);
for (uint64_t i = 0; i < initial_graph_size; i++) {
for (uint64_t j = 0; j < small_graph_degree; j++) {
cagra_graph(i, j) = (i + offset(j)) % initial_graph_size;
}
}

auto curr_graph_size = initial_graph_size;
while (true) {
RAFT_LOG_DEBUG("# graph_size = %lu (%.3lf)",
(uint64_t)curr_graph_size,
(double)curr_graph_size / final_graph_size);

auto curr_query_size = std::min(2 * curr_graph_size, final_graph_size);
auto curr_topk = small_topk;
auto curr_itopk_size = small_topk * 3 / 2;
auto curr_graph_degree = small_graph_degree;
if (curr_query_size == final_graph_size) {
curr_topk = topk;
curr_itopk_size = topk * 2;
curr_graph_degree = graph_degree;
}

cuvs::neighbors::cagra::search_params search_params;
search_params.algo = cuvs::neighbors::cagra::search_algo::AUTO;
search_params.max_queries = max_chunk_size;
search_params.itopk_size = curr_itopk_size;

// Create an index (idx), a query view (dev_query_view), and a mdarray for
// search results (neighbors).
auto dev_dataset_view = raft::make_device_matrix_view<const T, int64_t>(
dev_dataset.data_handle(), (int64_t)curr_graph_size, dev_dataset.extent(1));

auto idx = index<T, IdxT>(
res, params.metric, dev_dataset_view, raft::make_const_mdspan(cagra_graph.view()));

auto dev_query_view = raft::make_device_matrix_view<const T, int64_t>(
dev_dataset.data_handle(), (int64_t)curr_query_size, dev_dataset.extent(1));
auto neighbors = raft::make_host_matrix<IdxT, int64_t>(curr_query_size, curr_topk);

// Search.
// Since there are many queries, divide them into batches and search them.
cuvs::spatial::knn::detail::utils::batch_load_iterator<T> query_batch(
dev_query_view.data_handle(),
curr_query_size,
dev_query_view.extent(1),
max_chunk_size,
raft::resource::get_cuda_stream(res),
raft::resource::get_workspace_resource(res));
for (const auto& batch : query_batch) {
auto batch_dev_query_view = raft::make_device_matrix_view<const T, int64_t>(
batch.data(), batch.size(), dev_query_view.extent(1));
auto batch_dev_neighbors_view = raft::make_device_matrix_view<IdxT, int64_t>(
dev_neighbors.data_handle(), batch.size(), curr_topk);
auto batch_dev_distances_view = raft::make_device_matrix_view<float, int64_t>(
dev_distances.data_handle(), batch.size(), curr_topk);

cuvs::neighbors::cagra::search(res,
search_params,
idx,
batch_dev_query_view,
batch_dev_neighbors_view,
batch_dev_distances_view);

auto batch_neighbors_view = raft::make_host_matrix_view<IdxT, int64_t>(
neighbors.data_handle() + batch.offset() * curr_topk, batch.size(), curr_topk);
raft::copy(batch_neighbors_view.data_handle(),
batch_dev_neighbors_view.data_handle(),
batch_neighbors_view.size(),
raft::resource::get_cuda_stream(res));
}

// Optimize graph
bool flag_last = (curr_graph_size == final_graph_size);
curr_graph_size = curr_query_size;
cagra_graph = raft::make_host_matrix<IdxT, int64_t>(0, 0); // delete existing grahp
cagra_graph = raft::make_host_matrix<IdxT, int64_t>(curr_graph_size, curr_graph_degree);
optimize<IdxT>(
res, neighbors.view(), cagra_graph.view(), flag_last ? params.guarantee_connectivity : 0);
if (flag_last) { break; }
}

return cagra_graph;
}

template <typename T,
typename IdxT = uint32_t,
typename Accessor = raft::host_device_accessor<std::experimental::default_accessor<T>,
Expand All @@ -427,9 +588,6 @@ index<T, IdxT> build(
graph_degree = intermediate_degree;
}

std::optional<raft::host_matrix<IdxT, int64_t>> knn_graph(
raft::make_host_matrix<IdxT, int64_t>(dataset.extent(0), intermediate_degree));

// Set default value in case knn_build_params is not defined.
auto knn_build_params = params.graph_build_params;
if (std::holds_alternative<std::monostate>(params.graph_build_params)) {
Expand All @@ -445,38 +603,48 @@ index<T, IdxT> build(
}
}

auto cagra_graph = raft::make_host_matrix<IdxT, int64_t>(0, 0);

// Dispatch based on graph_build_params
if (std::holds_alternative<cagra::graph_build_params::ivf_pq_params>(knn_build_params)) {
auto ivf_pq_params =
std::get<cuvs::neighbors::cagra::graph_build_params::ivf_pq_params>(knn_build_params);
build_knn_graph(res, dataset, knn_graph->view(), ivf_pq_params);
if (std::holds_alternative<cagra::graph_build_params::iterative_search_params>(
knn_build_params)) {
cagra_graph = iterative_build_graph<T, IdxT, Accessor>(res, params, dataset);
} else {
auto nn_descent_params =
std::get<cagra::graph_build_params::nn_descent_params>(knn_build_params);
std::optional<raft::host_matrix<IdxT, int64_t>> knn_graph(
raft::make_host_matrix<IdxT, int64_t>(dataset.extent(0), intermediate_degree));

if (nn_descent_params.graph_degree != intermediate_degree) {
RAFT_LOG_WARN(
"Graph degree (%lu) for nn-descent needs to match cagra intermediate graph degree (%lu), "
"aligning "
"nn-descent graph_degree.",
nn_descent_params.graph_degree,
intermediate_degree);
nn_descent_params =
cagra::graph_build_params::nn_descent_params(intermediate_degree, params.metric);
}
if (std::holds_alternative<cagra::graph_build_params::ivf_pq_params>(knn_build_params)) {
auto ivf_pq_params =
std::get<cuvs::neighbors::cagra::graph_build_params::ivf_pq_params>(knn_build_params);
build_knn_graph(res, dataset, knn_graph->view(), ivf_pq_params);
} else {
auto nn_descent_params =
std::get<cagra::graph_build_params::nn_descent_params>(knn_build_params);

if (nn_descent_params.graph_degree != intermediate_degree) {
RAFT_LOG_WARN(
"Graph degree (%lu) for nn-descent needs to match cagra intermediate graph degree (%lu), "
"aligning "
"nn-descent graph_degree.",
nn_descent_params.graph_degree,
intermediate_degree);
nn_descent_params =
cagra::graph_build_params::nn_descent_params(intermediate_degree, params.metric);
}

// Use nn-descent to build CAGRA knn graph
nn_descent_params.return_distances = false;
build_knn_graph<T, IdxT>(res, dataset, knn_graph->view(), nn_descent_params);
}
// Use nn-descent to build CAGRA knn graph
nn_descent_params.return_distances = false;
build_knn_graph<T, IdxT>(res, dataset, knn_graph->view(), nn_descent_params);
}

auto cagra_graph = raft::make_host_matrix<IdxT, int64_t>(dataset.extent(0), graph_degree);
cagra_graph = raft::make_host_matrix<IdxT, int64_t>(dataset.extent(0), graph_degree);

RAFT_LOG_INFO("optimizing graph");
optimize<IdxT>(res, knn_graph->view(), cagra_graph.view(), params.guarantee_connectivity);
RAFT_LOG_INFO("optimizing graph");
optimize<IdxT>(res, knn_graph->view(), cagra_graph.view(), params.guarantee_connectivity);

// free intermediate graph before trying to create the index
knn_graph.reset();
// free intermediate graph before trying to create the index
knn_graph.reset();
}

RAFT_LOG_INFO("Graph optimized, creating index");

Expand Down
30 changes: 25 additions & 5 deletions cpp/tests/neighbors/ann_cagra.cuh
View file
Original file line number Diff line number Diff line change
Expand Up @@ -254,6 +254,8 @@ enum class graph_build_algo {
IVF_PQ,
/* Experimental, use NN-Descent to build all-neighbors knn graph */
NN_DESCENT,
/* Experimental, iteratively execute CAGRA's search() and optimize() */
ITERATIVE_CAGRA_SEARCH,
/* Choose default automatically */
AUTO
};
Expand Down Expand Up @@ -294,7 +296,7 @@ inline ::std::ostream& operator<<(::std::ostream& os, const AnnCagraInputs& p)
};

std::vector<std::string> algo = {"single-cta", "multi_cta", "multi_kernel", "auto"};
std::vector<std::string> build_algo = {"IVF_PQ", "NN_DESCENT", "AUTO"};
std::vector<std::string> build_algo = {"IVF_PQ", "NN_DESCENT", "ITERATIVE_CAGRA_SEARCH", "AUTO"};
os << "{n_queries=" << p.n_queries << ", dataset shape=" << p.n_rows << "x" << p.dim
<< ", k=" << p.k << ", " << algo.at((int)p.algo) << ", max_queries=" << p.max_queries
<< ", itopk_size=" << p.itopk_size << ", search_width=" << p.search_width
Expand Down Expand Up @@ -374,6 +376,10 @@ class AnnCagraTest : public ::testing::TestWithParam<AnnCagraInputs> {
index_params.intermediate_graph_degree, index_params.metric);
break;
}
case graph_build_algo::ITERATIVE_CAGRA_SEARCH: {
index_params.graph_build_params = graph_build_params::iterative_search_params();
break;
}
case graph_build_algo::AUTO:
// do nothing
break;
Expand Down Expand Up @@ -551,6 +557,10 @@ class AnnCagraAddNodesTest : public ::testing::TestWithParam<AnnCagraInputs> {
graph_build_params::nn_descent_params(index_params.intermediate_graph_degree);
break;
}
case graph_build_algo::ITERATIVE_CAGRA_SEARCH: {
index_params.graph_build_params = graph_build_params::iterative_search_params();
break;
}
case graph_build_algo::AUTO:
// do nothing
break;
Expand Down Expand Up @@ -752,6 +762,10 @@ class AnnCagraFilterTest : public ::testing::TestWithParam<AnnCagraInputs> {
graph_build_params::nn_descent_params(index_params.intermediate_graph_degree);
break;
}
case graph_build_algo::ITERATIVE_CAGRA_SEARCH: {
index_params.graph_build_params = graph_build_params::iterative_search_params();
break;
}
case graph_build_algo::AUTO:
// do nothing
break;
Expand Down Expand Up @@ -882,7 +896,7 @@ inline std::vector<AnnCagraInputs> generate_inputs()
{1000},
{1, 8, 17},
{16}, // k
{graph_build_algo::NN_DESCENT},
{graph_build_algo::NN_DESCENT, graph_build_algo::ITERATIVE_CAGRA_SEARCH},
{search_algo::SINGLE_CTA, search_algo::MULTI_CTA, search_algo::MULTI_KERNEL},
{0, 10}, // query size
{0},
Expand Down Expand Up @@ -917,7 +931,9 @@ inline std::vector<AnnCagraInputs> generate_inputs()
{1000},
{1, 3, 5, 7, 8, 17, 64, 128, 137, 192, 256, 512, 1024}, // dim
{16}, // k
{graph_build_algo::IVF_PQ, graph_build_algo::NN_DESCENT},
{graph_build_algo::IVF_PQ,
graph_build_algo::NN_DESCENT,
graph_build_algo::ITERATIVE_CAGRA_SEARCH},
{search_algo::AUTO},
{10},
{0},
Expand All @@ -935,7 +951,9 @@ inline std::vector<AnnCagraInputs> generate_inputs()
{1000},
{64},
{16},
{graph_build_algo::IVF_PQ, graph_build_algo::NN_DESCENT},
{graph_build_algo::IVF_PQ,
graph_build_algo::NN_DESCENT,
graph_build_algo::ITERATIVE_CAGRA_SEARCH},
{search_algo::AUTO},
{10},
{0, 8, 16, 32}, // team_size
Expand All @@ -953,7 +971,9 @@ inline std::vector<AnnCagraInputs> generate_inputs()
{1000},
{64},
{16},
{graph_build_algo::IVF_PQ, graph_build_algo::NN_DESCENT},
{graph_build_algo::IVF_PQ,
graph_build_algo::NN_DESCENT,
graph_build_algo::ITERATIVE_CAGRA_SEARCH},
{search_algo::AUTO},
{10},
{0}, // team_size
Expand Down

0 comments on commit 2b7dce6

Please sign in to comment.