Merge branch 'branch-25.06' into get-version-jni-c-api

Author: mythrocks

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/include/cuvs/neighbors/ivf_flat.h

@@ -164,6 +164,25 @@ cuvsError_t cuvsIvfFlatIndexCreate(cuvsIvfFlatIndex_t* index);
  * @param[in] index cuvsIvfFlatIndex_t to de-allocate
  */
 cuvsError_t cuvsIvfFlatIndexDestroy(cuvsIvfFlatIndex_t index);
+
+/** Get the number of clusters/inverted lists */
+uint32_t cuvsIvfFlatIndexGetNLists(cuvsIvfFlatIndex_t index);
+
+/** Get the dimensionality of the data */
+uint32_t cuvsIvfFlatIndexGetDim(cuvsIvfFlatIndex_t index);
+
+/**
+ * @brief Get the cluster centers corresponding to the lists [n_lists, dim]
+ *
+ * @param[in] res cuvsResources_t opaque C handle
+ * @param[in] index cuvsIvfFlatIndex_t Built Ivf-Flat Index
+ * @param[out] centers Preallocated array on host or device memory to store output, [n_lists, dim]
+ * @return cuvsError_t
+ */
+cuvsError_t cuvsIvfFlatIndexGetCenters(cuvsResources_t res,
+                                       cuvsIvfFlatIndex_t index,
+                                       DLManagedTensor* centers);
+
 /**
  * @}
  */

cpp/include/cuvs/neighbors/ivf_pq.h

@@ -264,15 +264,16 @@ cuvsError_t cuvsIvfPqIndexDestroy(cuvsIvfPqIndex_t index);
 /** Get the number of clusters/inverted lists */
 uint32_t cuvsIvfPqIndexGetNLists(cuvsIvfPqIndex_t index);
 
-/** Get the dimensionality of the cluster centers */
-uint32_t cuvsIvfPqIndexGetDimExt(cuvsIvfPqIndex_t index);
+/** Get the dimensionality */
+uint32_t cuvsIvfPqIndexGetDim(cuvsIvfPqIndex_t index);
 
 /**
  * @brief Get the cluster centers corresponding to the lists in the original space
  *
  * @param[in] res cuvsResources_t opaque C handle
- * @param[in] index cuvsIvfPqIndex_t Built NN-Descent index
- * @param[out] centers Preallocated array on host memory to store output, [n_lists, dim_ext]
+ * @param[in] index cuvsIvfPqIndex_t Built Ivf-Pq index
+ * @param[out] centers Preallocated array on host or device memory to store output,
+ * dimensions [n_lists, dim]
  * @return cuvsError_t
  */
 cuvsError_t cuvsIvfPqIndexGetCenters(cuvsResources_t res,

cpp/include/cuvs/neighbors/ivf_pq.hpp

@@ -2861,6 +2861,10 @@ void extract_centers(raft::resources const& res,
                      const index<int64_t>& index,
                      raft::device_matrix_view<float, uint32_t, raft::row_major> cluster_centers);
 
+/** @copydoc extract_centers */
+void extract_centers(raft::resources const& res,
+                     const index<int64_t>& index,
+                     raft::host_matrix_view<float, uint32_t, raft::row_major> cluster_centers);
 /**
  * @brief Helper exposing the re-computation of list sizes and related arrays if IVF lists have been
  * modified externally.

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,