Skip to content

Commit

Permalink
Migrate feature diff for NN Descent from RAFT to cuVS (#421)
Browse files Browse the repository at this point in the history 
This PR is an amalgamation of the diff of 3 PRs in RAFT:

1. rapidsai/raft#2345
2. rapidsai/raft#2380
3. rapidsai/raft#2403

This PR also addresses part 1 and 2 of #419, closes #391 and makes CAGRA use the compiled headers of NN Descent, which seemed to have been a pending TODO https://github.com/rapidsai/cuvs/blob/009bb8de03ce9708d4d797166187250f77a59a36/cpp/src/neighbors/detail/cagra/cagra_build.cuh#L36-L37

Also, batch tests are disabled in this PR due to issue rapidsai/raft#2450. PR #424 will attempt to re-enable them.

Authors:
  - Divye Gala (https://github.com/divyegala)
  - Corey J. Nolet (https://github.com/cjnolet)

Approvers:
  - Corey J. Nolet (https://github.com/cjnolet)

URL: #421
  • Loading branch information
@divyegala
divyegala authored Nov 8, 2024
1 parent e559d58 commit 29c0f5b
Show file tree
Hide file tree
Showing 13 changed files with 1,361 additions and 224 deletions.
92 changes: 70 additions & 22 deletions cpp/include/cuvs/neighbors/nn_descent.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,8 @@ struct index_params : cuvs::neighbors::index_params {
size_t intermediate_graph_degree = 128; // Degree of input graph for pruning.
size_t max_iterations = 20; // Number of nn-descent iterations.
float termination_threshold = 0.0001; // Termination threshold of nn-descent.
bool return_distances = true; // return distances if true
size_t n_clusters = 1; // defaults to not using any batching

/** @brief Construct NN descent parameters for a specific kNN graph degree
*
Expand Down Expand Up @@ -100,14 +102,20 @@ struct index : cuvs::neighbors::index {
* @param res raft::resources is an object mangaging resources
* @param n_rows number of rows in knn-graph
* @param n_cols number of cols in knn-graph
* @param return_distances whether to return distances
*/
index(raft::resources const& res, int64_t n_rows, int64_t n_cols)
index(raft::resources const& res, int64_t n_rows, int64_t n_cols, bool return_distances = false)
: cuvs::neighbors::index(),
res_{res},
metric_{cuvs::distance::DistanceType::L2Expanded},
graph_{raft::make_host_matrix<IdxT, int64_t, raft::row_major>(n_rows, n_cols)},
graph_view_{graph_.view()}
graph_view_{graph_.view()},
return_distances_{return_distances}
{
if (return_distances) {
distances_ = raft::make_device_matrix<float, int64_t>(res_, n_rows, n_cols);
distances_view_ = distances_.value().view();
}
}

/**
Expand All @@ -119,14 +127,20 @@ struct index : cuvs::neighbors::index {
*
* @param res raft::resources is an object mangaging resources
* @param graph_view raft::host_matrix_view<IdxT, int64_t, raft::row_major> for storing knn-graph
* @param distances_view optional raft::device_matrix_view<float, int64_t, row_major> for storing
* distances
*/
index(raft::resources const& res,
raft::host_matrix_view<IdxT, int64_t, raft::row_major> graph_view)
raft::host_matrix_view<IdxT, int64_t, raft::row_major> graph_view,
std::optional<raft::device_matrix_view<float, int64_t, row_major>> distances_view =
std::nullopt)
: cuvs::neighbors::index(),
res_{res},
metric_{cuvs::distance::DistanceType::L2Expanded},
graph_{raft::make_host_matrix<IdxT, int64_t, raft::row_major>(0, 0)},
graph_view_{graph_view}
graph_view_{graph_view},
distances_view_{distances_view},
return_distances_{distances_view.has_value()}
{
}

Expand Down Expand Up @@ -155,6 +169,13 @@ struct index : cuvs::neighbors::index {
return graph_view_;
}

/** neighborhood graph distances [size, graph-degree] */
[[nodiscard]] inline auto distances() noexcept
-> std::optional<device_matrix_view<float, int64_t, row_major>>
{
return distances_view_;
}

// Don't allow copying the index for performance reasons (try avoiding copying data)
index(const index&) = delete;
index(index&&) = default;
Expand All @@ -166,8 +187,11 @@ struct index : cuvs::neighbors::index {
raft::resources const& res_;
cuvs::distance::DistanceType metric_;
raft::host_matrix<IdxT, int64_t, raft::row_major> graph_; // graph to return for non-int IdxT
std::optional<raft::device_matrix<float, int64_t, row_major>> distances_;
raft::host_matrix_view<IdxT, int64_t, raft::row_major>
graph_view_; // view of graph for user provided matrix
std::optional<raft::device_matrix_view<float, int64_t, row_major>> distances_view_;
bool return_distances_;
};

/** @} */
Expand Down Expand Up @@ -200,12 +224,15 @@ struct index : cuvs::neighbors::index {
* to run the nn-descent algorithm
* @param[in] dataset raft::device_matrix_view input dataset expected to be located
* in device memory
* @param[in] graph optional raft::host_matrix_view<uint32_t, int64_t, raft::row_major> for owning
* the output graph
* @return index<IdxT> index containing all-neighbors knn graph in host memory
*/
auto build(raft::resources const& res,
index_params const& params,
raft::device_matrix_view<const float, int64_t, raft::row_major> dataset)
-> cuvs::neighbors::nn_descent::index<uint32_t>;
raft::device_matrix_view<const float, int64_t, raft::row_major> dataset,
std::optional<raft::host_matrix_view<uint32_t, int64_t, raft::row_major>> graph =
std::nullopt) -> cuvs::neighbors::nn_descent::index<uint32_t>;

/**
* @brief Build nn-descent Index with dataset in host memory
Expand All @@ -232,12 +259,15 @@ auto build(raft::resources const& res,
* to run the nn-descent algorithm
* @param[in] dataset raft::host_matrix_view input dataset expected to be located
* in host memory
* @param[in] graph optional raft::host_matrix_view<uint32_t, int64_t, raft::row_major> for owning
* the output graph
* @return index<IdxT> index containing all-neighbors knn graph in host memory
*/
auto build(raft::resources const& res,
index_params const& params,
raft::host_matrix_view<const float, int64_t, raft::row_major> dataset)
-> cuvs::neighbors::nn_descent::index<uint32_t>;
raft::host_matrix_view<const float, int64_t, raft::row_major> dataset,
std::optional<raft::host_matrix_view<uint32_t, int64_t, raft::row_major>> graph =
std::nullopt) -> cuvs::neighbors::nn_descent::index<uint32_t>;

/**
* @brief Build nn-descent Index with dataset in device memory
Expand All @@ -262,12 +292,15 @@ auto build(raft::resources const& res,
* to run the nn-descent algorithm
* @param[in] dataset raft::device_matrix_view input dataset expected to be located
* in device memory
* @param[in] graph optional raft::host_matrix_view<uint32_t, int64_t, raft::row_major> for owning
* the output graph
* @return index<IdxT> index containing all-neighbors knn graph in host memory
*/
auto build(raft::resources const& res,
index_params const& params,
raft::device_matrix_view<const half, int64_t, raft::row_major> dataset)
-> cuvs::neighbors::nn_descent::index<uint32_t>;
raft::device_matrix_view<const half, int64_t, raft::row_major> dataset,
std::optional<raft::host_matrix_view<uint32_t, int64_t, raft::row_major>> graph =
std::nullopt) -> cuvs::neighbors::nn_descent::index<uint32_t>;

/**
* @brief Build nn-descent Index with dataset in host memory
Expand All @@ -294,12 +327,15 @@ auto build(raft::resources const& res,
* to run the nn-descent algorithm
* @param[in] dataset raft::host_matrix_view input dataset expected to be located
* in host memory
* @param[in] graph optional raft::host_matrix_view<uint32_t, int64_t, raft::row_major> for owning
* the output graph
* @return index<IdxT> index containing all-neighbors knn graph in host memory
*/
auto build(raft::resources const& res,
index_params const& params,
raft::host_matrix_view<const half, int64_t, raft::row_major> dataset)
-> cuvs::neighbors::nn_descent::index<uint32_t>;
raft::host_matrix_view<const half, int64_t, raft::row_major> dataset,
std::optional<raft::host_matrix_view<uint32_t, int64_t, raft::row_major>> graph =
std::nullopt) -> cuvs::neighbors::nn_descent::index<uint32_t>;

/**
* @brief Build nn-descent Index with dataset in device memory
Expand All @@ -324,12 +360,15 @@ auto build(raft::resources const& res,
* to run the nn-descent algorithm
* @param[in] dataset raft::device_matrix_view input dataset expected to be located
* in device memory
* @param[in] graph optional raft::host_matrix_view<uint32_t, int64_t, raft::row_major> for owning
* the output graph
* @return index<IdxT> index containing all-neighbors knn graph in host memory
*/
auto build(raft::resources const& res,
index_params const& params,
raft::device_matrix_view<const int8_t, int64_t, raft::row_major> dataset)
-> cuvs::neighbors::nn_descent::index<uint32_t>;
raft::device_matrix_view<const int8_t, int64_t, raft::row_major> dataset,
std::optional<raft::host_matrix_view<uint32_t, int64_t, raft::row_major>> graph =
std::nullopt) -> cuvs::neighbors::nn_descent::index<uint32_t>;

/**
* @brief Build nn-descent Index with dataset in host memory
Expand All @@ -356,12 +395,15 @@ auto build(raft::resources const& res,
* to run the nn-descent algorithm
* @param[in] dataset raft::host_matrix_view input dataset expected to be located
* in host memory
* @param[in] graph optional raft::host_matrix_view<uint32_t, int64_t, raft::row_major> for owning
* the output graph
* @return index<IdxT> index containing all-neighbors knn graph in host memory
*/
auto build(raft::resources const& res,
index_params const& params,
raft::host_matrix_view<const int8_t, int64_t, raft::row_major> dataset)
-> cuvs::neighbors::nn_descent::index<uint32_t>;
raft::host_matrix_view<const int8_t, int64_t, raft::row_major> dataset,
std::optional<raft::host_matrix_view<uint32_t, int64_t, raft::row_major>> graph =
std::nullopt) -> cuvs::neighbors::nn_descent::index<uint32_t>;

/**
* @brief Build nn-descent Index with dataset in device memory
Expand All @@ -386,14 +428,15 @@ auto build(raft::resources const& res,
* to run the nn-descent algorithm
* @param[in] dataset raft::device_matrix_view input dataset expected to be located
* in device memory
* @param[in] graph optional raft::host_matrix_view<uint32_t, int64_t, raft::row_major> for owning
* the output graph
* @return index<IdxT> index containing all-neighbors knn graph in host memory
*/
auto build(raft::resources const& res,
index_params const& params,
raft::device_matrix_view<const uint8_t, int64_t, raft::row_major> dataset)
-> cuvs::neighbors::nn_descent::index<uint32_t>;

/** @} */
raft::device_matrix_view<const uint8_t, int64_t, raft::row_major> dataset,
std::optional<raft::host_matrix_view<uint32_t, int64_t, raft::row_major>> graph =
std::nullopt) -> cuvs::neighbors::nn_descent::index<uint32_t>;

/**
* @brief Build nn-descent Index with dataset in host memory
Expand All @@ -420,12 +463,17 @@ auto build(raft::resources const& res,
* to run the nn-descent algorithm
* @param[in] dataset raft::host_matrix_view input dataset expected to be located
* in host memory
* @param[in] graph optional raft::host_matrix_view<uint32_t, int64_t, raft::row_major> for owning
* the output graph
* @return index<IdxT> index containing all-neighbors knn graph in host memory
*/
auto build(raft::resources const& res,
index_params const& params,
raft::host_matrix_view<const uint8_t, int64_t, raft::row_major> dataset)
-> cuvs::neighbors::nn_descent::index<uint32_t>;
raft::host_matrix_view<const uint8_t, int64_t, raft::row_major> dataset,
std::optional<raft::host_matrix_view<uint32_t, int64_t, raft::row_major>> graph =
std::nullopt) -> cuvs::neighbors::nn_descent::index<uint32_t>;

/** @} */

/**
* @brief Test if we have enough GPU memory to run NN descent algorithm.
Expand Down
8 changes: 4 additions & 4 deletions cpp/src/neighbors/detail/cagra/cagra_build.cuh
View file
Original file line number Diff line number Diff line change
Expand Up @@ -33,8 +33,7 @@
#include <cuvs/neighbors/ivf_pq.hpp>
#include <cuvs/neighbors/refine.hpp>

// TODO: Fixme- this needs to be migrated
#include "../../nn_descent.cuh"
#include <cuvs/neighbors/nn_descent.hpp>

// TODO: This shouldn't be calling spatial/knn APIs
#include "../ann_utils.cuh"
Expand Down Expand Up @@ -356,8 +355,8 @@ void build_knn_graph(
raft::host_matrix_view<IdxT, int64_t, raft::row_major> knn_graph,
cuvs::neighbors::nn_descent::index_params build_params)
{
auto nn_descent_idx = cuvs::neighbors::nn_descent::index<IdxT>(res, knn_graph);
cuvs::neighbors::nn_descent::build<DataT, IdxT>(res, build_params, dataset, nn_descent_idx);
std::optional<raft::host_matrix_view<IdxT, int64_t, row_major>> graph_view = knn_graph;
auto nn_descent_idx = cuvs::neighbors::nn_descent::build(res, build_params, dataset, graph_view);

using internal_IdxT = typename std::make_unsigned<IdxT>::type;
using g_accessor = typename decltype(nn_descent_idx.graph())::accessor_type;
Expand Down Expand Up @@ -471,6 +470,7 @@ index<T, IdxT> build(
}

// 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);
}

Expand Down
Loading

0 comments on commit 29c0f5b

Please sign in to comment.