All-neighbors API docs (rapidsai#944)

docs for all-neighbors API

Authors:
  - Jinsol Park (https://github.com/jinsolp)

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

URL: rapidsai#944

Author: jinsolp

cpp/include/cuvs/neighbors/all_neighbors.hpp

@@ -22,6 +22,10 @@
 #include <variant>
 
 namespace cuvs::neighbors::all_neighbors {
+/**
+ * @defgroup all_neighbors_cpp_params The all-neighbors algorithm parameters.
+ * @{
+ */
 
 /**
  * @brief Parameters used to build an all-neighbors knn graph (find nearest neighbors for all the
@@ -43,20 +47,21 @@ using GraphBuildParams =
   std::variant<graph_build_params::ivf_pq_params, graph_build_params::nn_descent_params>;
 
 /**
- * @brief Parameters used to build an all-neighbors graph  (find nearest neighbors for all the
- * training vectors)
- *
- * graph_build_params: graph building parameters for the given graph building algorithm. defaults
- * to ivfpq.
- * n_nearest_clusters: number of nearest clusters each data point will be assigned to in
- * the batching algorithm
- * n_clusters: number of total clusters (aka batches) to split the data into. If set to 1, algorithm
- * creates an all-neighbors graph without batching
- * metric: metric type
+ * @brief Parameters used to build an all-neighbors graph (find nearest neighbors for all the
+ * training vectors).
+ * For scalability, the all-neighbors graph construction algorithm partitions a set of training
+ * vectors into overlapping clusters, computes a local knn graph on each cluster, and merges the
+ * local graphs into a single global graph.
+ * Device memory usage and accuracy can be configured by changing the `overlap_factor` and
+ * `n_clusters`.
+ * The algorithm used to build each local graph is also configurable.
  *
  */
 struct all_neighbors_params {
   /** Parameters for knn graph building algorithm
+   * Approximate nearest neighbors methods are used to build the knn graph. Currently supported
+   * options are 'IVF-PQ' and 'NN Descent'. IVF-PQ is more accurate, but slower compared to NN
+   * Descent.
    *
    * Set ivf_pq_params, or nn_descent_params to select the graph build
    * algorithm and control their parameters.
@@ -74,35 +79,55 @@ struct all_neighbors_params {
   GraphBuildParams graph_build_params;
 
   /**
-   * Usage of n_nearest_clusters and n_clusters
+   * Number of nearest clusters each data point will be assigned to in the batching algorithm.
+   * Start with `overlap_factor = 2` and gradually increase (2->3->4 ...) for better accuracy at the
+   * cost of device memory usage.
+   */
+  size_t overlap_factor = 2;
+
+  /**
+   * Number of total clusters (aka batches) to split the data into. If set to 1, algorithm creates
+   * an all-neighbors graph without batching.
+   * Start with `n_clusters = 4` and increase (4 → 8 → 16...) for less device memory usage at the
+   * cost of accuracy. This is independent from `overlap_factor` as long as `overlap_factor` <
+   * `n_clusters`.
    *
-   * The ratio of n_nearest_clusters / n_clusters determines device memory usage.
-   * Approximately (n_nearest_clusters / n_clusters) * num_rows_in_entire_data number of rows will
+   * The ratio of `overlap_factor / n_clusters` determines device memory usage.
+   * Approximately `(overlap_factor / n_clusters) * num_rows_in_entire_data` number of rows will
    * be put on device memory at once.
-   * E.g. between (n_nearest_clusters / n_clusters) = 2/10 and 2/20, the latter will use less device
+   * E.g. between `(overlap_factor / n_clusters)` = 2/10 and 2/20, the latter will use less device
    * memory.
    *
-   * Larger n_nearest_clusters results in better accuracy of the final all-neighbors knn
-   * graph. E.g. With the similar device memory usages, (n_nearest_clusters / n_clusters) = 4/20
+   * Larger `overlap_factor` results in better accuracy of the final all-neighbors knn
+   * graph. E.g. While using similar device memory, `(overlap_factor / n_clusters)` = 4/20
    * will have better accuracy than 2/10 at the cost of performance.
+   *
    */
-  size_t n_nearest_clusters           = 2;
-  size_t n_clusters                   = 1;  // defaults to not batching
+  size_t n_clusters = 1;  // defaults to not batching
+
+  /** Metric used. */
   cuvs::distance::DistanceType metric = cuvs::distance::DistanceType::L2Expanded;
 };
 
+/** @} */
+
+/**
+ * @defgroup all_neighbors_cpp_build The all-neighbors knn graph build
+ * @{
+ */
+
 /**
  * @brief Builds an approximate all-neighbors knn graph  (find nearest neighbors for all the
  * training vectors)
  *
  * Usage example:
  * @code{.cpp}
- *   using namespace cuvs::neighbors;
- *   // use default index parameters
- *   all_neighbors::all_neighbors_params params;
+ *  using namespace cuvs::neighbors;
+ *  // use default index parameters
+ *  all_neighbors::all_neighbors_params params;
  *  auto indices = raft::make_device_matrix<int64_t, int64_t>(handle, n_row, k);
  *  auto distances = raft::make_device_matrix<float, int64_t>(handle, n_row, k);
- *   all_neighbors::build(res, params, dataset, indices.view(), distances.view());
+ *  all_neighbors::build(res, params, dataset, indices.view(), distances.view());
  * @endcode
  *
  * @param[in] handle raft::resources is an object mangaging resources
@@ -127,12 +152,12 @@ void build(
  *
  * Usage example:
  * @code{.cpp}
- *   using namespace cuvs::neighbors;
- *   // use default index parameters
- *   all_neighbors::all_neighbors_params params;
+ *  using namespace cuvs::neighbors;
+ *  // use default index parameters
+ *  all_neighbors::all_neighbors_params params;
  *  auto indices = raft::make_device_matrix<int64_t, int64_t>(handle, n_row, k);
  *  auto distances = raft::make_device_matrix<float, int64_t>(handle, n_row, k);
- *   all_neighbors::build(res, params, dataset, indices.view(), distances.view());
+ *  all_neighbors::build(res, params, dataset, indices.view(), distances.view());
  * @endcode
  *
  * @param[in] handle raft::resources is an object mangaging resources
@@ -149,4 +174,6 @@ void build(
   raft::device_matrix_view<const float, int64_t, row_major> dataset,
   raft::device_matrix_view<int64_t, int64_t, row_major> indices,
   std::optional<raft::device_matrix_view<float, int64_t, row_major>> distances = std::nullopt);
+
+/** @} */
 }  // namespace cuvs::neighbors::all_neighbors

cpp/src/neighbors/all_neighbors/all_neighbors_batched.cuh

@@ -71,7 +71,7 @@ void get_centroids_on_data_subsample(raft::resources const& res,
 template <typename T, typename IdxT>
 void single_gpu_assign_clusters(
   raft::resources const& res,
-  size_t n_nearest_clusters,
+  size_t overlap_factor,
   size_t n_clusters,
   size_t n_rows_per_batch,
   size_t base_row_offset,
@@ -89,10 +89,10 @@ void single_gpu_assign_clusters(
   auto dataset_batch_d =
     raft::make_device_matrix<T, IdxT, raft::row_major>(res, n_rows_per_batch, num_cols);
 
-  auto nearest_clusters_idx_d = raft::make_device_matrix<IdxT, int64_t, raft::row_major>(
-    res, n_rows_per_batch, n_nearest_clusters);
-  auto nearest_clusters_dist_d = raft::make_device_matrix<T, int64_t, raft::row_major>(
-    res, n_rows_per_batch, n_nearest_clusters);
+  auto nearest_clusters_idx_d =
+    raft::make_device_matrix<IdxT, int64_t, raft::row_major>(res, n_rows_per_batch, overlap_factor);
+  auto nearest_clusters_dist_d =
+    raft::make_device_matrix<T, int64_t, raft::row_major>(res, n_rows_per_batch, overlap_factor);
 
   std::optional<raft::device_vector_view<const T, int64_t>> norms_view;
   cuvs::neighbors::brute_force::index<T> brute_force_index(res, centroids, norms_view, metric);
@@ -111,21 +111,21 @@ void single_gpu_assign_clusters(
                                          raft::make_const_mdspan(dataset_batch_d.view()),
                                          nearest_clusters_idx_d.view(),
                                          nearest_clusters_dist_d.view());
-    raft::copy(global_nearest_cluster.data_handle() + row_offset * n_nearest_clusters,
+    raft::copy(global_nearest_cluster.data_handle() + row_offset * overlap_factor,
                nearest_clusters_idx_d.data_handle(),
-               n_rows_of_current_batch * n_nearest_clusters,
+               n_rows_of_current_batch * overlap_factor,
                resource::get_cuda_stream(res));
   }
 }
 
 /**
- * Assign each data point to top n_nearest_clusters number of clusters. Loads the data in batches
+ * Assign each data point to top overlap_factor number of clusters. Loads the data in batches
  * onto device for efficiency. Arguments:
  * - [in] res: raft resource
  * - [in] params: params for graph building
  * - [in] dataset [num_rows x num_cols]: entire dataset located on host memory
  * - [in] centroids [n_clusters x num_cols] : centroid vectors
- * - [out] global_nearest_cluster [num_rows X n_nearest_clusters] : top n_nearest_clusters closest
+ * - [out] global_nearest_cluster [num_rows X overlap_factor] : top overlap_factor closest
  * clusters for each data point
  */
 template <typename T, typename IdxT>
@@ -169,7 +169,7 @@ void assign_clusters(raft::resources const& res,
       size_t base_row_offset_for_this_rank = n_rows_per_cluster * base_cluster_idx;
 
       single_gpu_assign_clusters(dev_res,
-                                 params.n_nearest_clusters,
+                                 params.overlap_factor,
                                  n_clusters_for_this_rank,
                                  n_rows_per_cluster,
                                  base_row_offset_for_this_rank,
@@ -180,7 +180,7 @@ void assign_clusters(raft::resources const& res,
     }
   } else {
     single_gpu_assign_clusters(res,
-                               params.n_nearest_clusters,
+                               params.overlap_factor,
                                params.n_clusters,
                                n_rows_per_cluster,
                                0,
@@ -195,9 +195,9 @@ void assign_clusters(raft::resources const& res,
  * Getting data indices that belong to cluster
  * Arguments:
  * - [in] res: raft resource
- * - [in] global_nearest_cluster [num_rows X n_nearest_clusters] : top n_nearest_clusters closest
+ * - [in] global_nearest_cluster [num_rows X overlap_factor] : top overlap_factor closest
  * clusters for each data point
- * - [out] inverted_indices [num_rows x n_nearest_clusters sized vector] : vector for data indices
+ * - [out] inverted_indices [num_rows x overlap_factor sized vector] : vector for data indices
  * for each cluster
  * - [out] cluster_sizes [n_cluster] : cluster size for each cluster
  * - [out] cluster_offsets [n_cluster] : offset in inverted_indices for each cluster
@@ -210,17 +210,17 @@ void get_inverted_indices(raft::resources const& res,
                           raft::host_vector_view<IdxT, IdxT> cluster_offsets)
 {
   // build sparse inverted indices and get number of data points for each cluster
-  size_t num_rows           = global_nearest_cluster.extent(0);
-  size_t n_nearest_clusters = global_nearest_cluster.extent(1);
-  size_t n_clusters         = cluster_sizes.extent(0);
+  size_t num_rows       = global_nearest_cluster.extent(0);
+  size_t overlap_factor = global_nearest_cluster.extent(1);
+  size_t n_clusters     = cluster_sizes.extent(0);
 
   auto local_offsets = raft::make_host_vector<IdxT>(n_clusters);
 
   std::fill(cluster_sizes.data_handle(), cluster_sizes.data_handle() + n_clusters, 0);
   std::fill(local_offsets.data_handle(), local_offsets.data_handle() + n_clusters, 0);
 
   for (size_t i = 0; i < num_rows; i++) {
-    for (size_t j = 0; j < n_nearest_clusters; j++) {
+    for (size_t j = 0; j < overlap_factor; j++) {
       IdxT cluster_id = global_nearest_cluster(i, j);
       cluster_sizes(cluster_id) += 1;
     }
@@ -231,7 +231,7 @@ void get_inverted_indices(raft::resources const& res,
     cluster_offsets(i) = cluster_offsets(i - 1) + cluster_sizes(i - 1);
   }
   for (size_t i = 0; i < num_rows; i++) {
-    for (size_t j = 0; j < n_nearest_clusters; j++) {
+    for (size_t j = 0; j < overlap_factor; j++) {
       IdxT cluster_id = global_nearest_cluster(i, j);
       inverted_indices(cluster_offsets(cluster_id) + local_offsets(cluster_id)) = i;
       local_offsets(cluster_id) += 1;
@@ -389,20 +389,19 @@ void batch_build(
   size_t num_cols = static_cast<size_t>(dataset.extent(1));
   size_t k        = indices.extent(1);
 
-  RAFT_EXPECTS(params.n_clusters > params.n_nearest_clusters,
-               "n_nearest_clusters should be smaller than n_clusters. We recommend starting from "
-               "n_nearest_clusters=2 and gradually increase it for better knn graph recall.");
+  RAFT_EXPECTS(params.n_clusters > params.overlap_factor,
+               "overlap_factor should be smaller than n_clusters. We recommend starting from "
+               "overlap_factor=2 and gradually increase it for better knn graph recall.");
 
   auto centroids = raft::make_device_matrix<T, IdxT>(handle, params.n_clusters, num_cols);
   get_centroids_on_data_subsample<T, IdxT>(handle, params.metric, dataset, centroids.view());
 
-  auto global_nearest_cluster =
-    raft::make_host_matrix<IdxT, IdxT>(num_rows, params.n_nearest_clusters);
+  auto global_nearest_cluster = raft::make_host_matrix<IdxT, IdxT>(num_rows, params.overlap_factor);
   assign_clusters<T, IdxT>(
     handle, params, dataset, centroids.view(), global_nearest_cluster.view());
 
   auto inverted_indices =
-    raft::make_host_vector<IdxT, IdxT, raft::row_major>(num_rows * params.n_nearest_clusters);
+    raft::make_host_vector<IdxT, IdxT, raft::row_major>(num_rows * params.overlap_factor);
   auto cluster_sizes   = raft::make_host_vector<IdxT, IdxT, raft::row_major>(params.n_clusters);
   auto cluster_offsets = raft::make_host_vector<IdxT, IdxT, raft::row_major>(params.n_clusters);
   get_inverted_indices(handle,

cpp/tests/neighbors/all_neighbors.cuh

@@ -58,7 +58,7 @@ inline ::std::ostream& operator<<(::std::ostream& os, const AllNeighborsInputs&
   os << "dataset shape=" << p.n_rows << "x" << p.dim << ", k=" << p.k
      << ", metric=" << static_cast<int>(std::get<1>(p.build_algo_metric_recall))
      << ", clusters=" << std::get<0>(p.cluster_nearestcluster)
-     << ", num nearest clusters=" << std::get<1>(p.cluster_nearestcluster) << std::endl;
+     << ", overlap_factor=" << std::get<1>(p.cluster_nearestcluster) << std::endl;
   return os;
 }
 
@@ -73,9 +73,9 @@ void get_graphs(raft::resources& handle,
                 size_t queries_size)
 {
   all_neighbors_params params;
-  params.n_clusters         = std::get<0>(ps.cluster_nearestcluster);
-  params.n_nearest_clusters = std::get<1>(ps.cluster_nearestcluster);
-  params.metric             = std::get<1>(ps.build_algo_metric_recall);
+  params.n_clusters     = std::get<0>(ps.cluster_nearestcluster);
+  params.overlap_factor = std::get<1>(ps.cluster_nearestcluster);
+  params.metric         = std::get<1>(ps.build_algo_metric_recall);
 
   auto build_algo = std::get<0>(ps.build_algo_metric_recall);
 
@@ -91,8 +91,7 @@ void get_graphs(raft::resources& handle,
     ivfq_build_params.build_params.metric = params.metric;
     // heuristically good ivfpq n_lists
     ivfq_build_params.build_params.n_lists = std::max(
-      5u,
-      static_cast<uint32_t>(ps.n_rows * params.n_nearest_clusters / (5000 * params.n_clusters)));
+      5u, static_cast<uint32_t>(ps.n_rows * params.overlap_factor / (5000 * params.n_clusters)));
     params.graph_build_params = ivfq_build_params;
   }
 
@@ -209,7 +208,7 @@ const std::vector<AllNeighborsInputs> inputsSingle =
      std::make_tuple(NN_DESCENT, cuvs::distance::DistanceType::L2SqrtExpanded, 0.9),
      std::make_tuple(NN_DESCENT, cuvs::distance::DistanceType::CosineExpanded, 0.9),
      std::make_tuple(NN_DESCENT, cuvs::distance::DistanceType::InnerProduct, 0.8)},
-    {std::make_tuple(1lu, 2lu)},  // min_recall, n_clusters, num_nearest_cluster
+    {std::make_tuple(1lu, 2lu)},  // min_recall, n_clusters, overlap_factor
     {5000, 7151},                 // n_rows
     {64, 137},                    // dim
     {16, 23},                     // graph_degree
@@ -227,7 +226,7 @@ const std::vector<AllNeighborsInputs> inputsBatch =
       std::make_tuple(4lu, 2lu),
       std::make_tuple(7lu, 2lu),
       std::make_tuple(10lu, 2lu),
-    },             // min_recall, n_clusters, num_nearest_cluster
+    },             // min_recall, n_clusters, overlap_factor
     {5000, 7151},  // n_rows
     {64, 137},     // dim
     {16, 23},      // graph_degree

docs/source/cpp_api/neighbors.rst

@@ -9,6 +9,7 @@ Nearest Neighbors
    :maxdepth: 2
    :caption: Contents:
 
+   neighbors_all_neighbors.rst
    neighbors_bruteforce.rst
    neighbors_cagra.rst
    neighbors_dynamic_batching.rst

docs/source/cpp_api/neighbors_all_neighbors.rst

@@ -0,0 +1,29 @@
+All-neighbors
+==========
+
+All-neighbors allows building an approximate all-neighbors knn graph. Given a full dataset, it finds nearest neighbors for all the training vectors in the dataset.
+
+.. role:: py(code)
+   :language: c++
+   :class: highlight
+
+``#include <cuvs/neighbors/all_neighbors.hpp>``
+
+namespace *cuvs::neighbors::all_neighbors*
+
+All neighbors knn graph build parameters
+----------------------
+
+.. doxygengroup:: all_neighbors_cpp_params
+    :project: cuvs
+    :members:
+    :content-only:
+
+
+Build
+-----------
+
+.. doxygengroup:: all_neighbors_cpp_build
+    :project: cuvs
+    :members:
+    :content-only: