Merge branch 'main' into cagra_optimize

Author: mfoerste4

c/src/preprocessing/quantize/pq.cpp

@@ -62,18 +62,10 @@ void* _build(cuvsResources_t res,
   auto dataset = dataset_tensor->dl_tensor;
 
   auto res_ptr = reinterpret_cast<raft::resources*>(res);
-
-  auto quantizer_params = cuvs::preprocessing::quantize::pq::params{
-    .pq_bits = params->pq_bits,
-    .pq_dim = params->pq_dim,
-    .use_subspaces = params->use_subspaces,
-    .use_vq = params->use_vq,
-    .vq_n_centers = params->vq_n_centers,
-    .kmeans_n_iters = params->kmeans_n_iters,
-    .pq_kmeans_type = static_cast<cuvs::cluster::kmeans::kmeans_type>(params->pq_kmeans_type),
-    .max_train_points_per_pq_code = params->max_train_points_per_pq_code,
-    .max_train_points_per_vq_cluster = params->max_train_points_per_vq_cluster
-  };
+  cuvs::preprocessing::quantize::pq::params quantizer_params(
+    params->pq_bits, params->pq_dim, params->use_subspaces, params->use_vq, params->vq_n_centers,
+    params->kmeans_n_iters,  static_cast<cuvs::cluster::kmeans::kmeans_type>(params->pq_kmeans_type), params->max_train_points_per_pq_code,
+     params->max_train_points_per_vq_cluster);
   cuvs::preprocessing::quantize::pq::quantizer<T>* ret = nullptr;
 
   if (cuvs::core::is_dlpack_device_compatible(dataset)) {

cpp/include/cuvs/preprocessing/quantize/pq.hpp

@@ -5,22 +5,78 @@
 
 #pragma once
 
+#include <cuvs/cluster/kmeans.hpp>
 #include <cuvs/neighbors/common.hpp>
 #include <raft/core/device_mdspan.hpp>
 #include <raft/core/handle.hpp>
 #include <raft/core/host_mdspan.hpp>
 
+#include <variant>
+
 namespace cuvs::preprocessing::quantize::pq {
 
 /**
  * @defgroup pq Product Quantizer utilities
  * @{
  */
 
+/** Alias for the variant holding either balanced or regular k-means parameters. */
+using kmeans_params_variant =
+  std::variant<cuvs::cluster::kmeans::balanced_params, cuvs::cluster::kmeans::params>;
+
 /**
  * @brief Product Quantizer parameters.
  */
 struct params {
+  /**
+   * Simplified constructor that will build an appropriate kmeans params object.
+   */
+  params(uint32_t pq_bits,
+         uint32_t pq_dim,
+         bool use_subspaces,
+         bool use_vq,
+         uint32_t vq_n_centers,
+         uint32_t kmeans_n_iters,
+         cuvs::cluster::kmeans::kmeans_type pq_kmeans_type =
+           cuvs::cluster::kmeans::kmeans_type::KMeansBalanced,
+         uint32_t max_train_points_per_pq_code    = 256,
+         uint32_t max_train_points_per_vq_cluster = 1024)
+    : pq_bits(pq_bits),
+      pq_dim(pq_dim),
+      use_subspaces(use_subspaces),
+      use_vq(use_vq),
+      vq_n_centers(vq_n_centers),
+      kmeans_params(
+        pq_kmeans_type == cuvs::cluster::kmeans::kmeans_type::KMeansBalanced
+          ? kmeans_params_variant{cuvs::cluster::kmeans::balanced_params{.n_iters = kmeans_n_iters}}
+          : kmeans_params_variant{cuvs::cluster::kmeans::params{
+              .n_clusters = 1 << pq_bits, .max_iter = static_cast<int>(kmeans_n_iters)}}),
+      max_train_points_per_pq_code(max_train_points_per_pq_code),
+      max_train_points_per_vq_cluster(max_train_points_per_vq_cluster)
+  {
+  }
+
+  params(uint32_t pq_bits,
+         uint32_t pq_dim,
+         bool use_subspaces,
+         bool use_vq,
+         uint32_t vq_n_centers,
+         kmeans_params_variant kmeans_params,
+         uint32_t max_train_points_per_pq_code    = 256,
+         uint32_t max_train_points_per_vq_cluster = 1024)
+    : pq_bits(pq_bits),
+      pq_dim(pq_dim),
+      use_subspaces(use_subspaces),
+      use_vq(use_vq),
+      vq_n_centers(vq_n_centers),
+      kmeans_params(kmeans_params),
+      max_train_points_per_pq_code(max_train_points_per_pq_code),
+      max_train_points_per_vq_cluster(max_train_points_per_vq_cluster)
+  {
+  }
+
+  params() = default;
+
   /**
    * The bit length of the vector element after compression by PQ.
    *
@@ -32,7 +88,7 @@ struct params {
   uint32_t pq_bits = 8;
   /**
    * The dimensionality of the vector after compression by PQ.
-   * When zero, an optimal value is selected using a heuristic.
+   * When zero, dim / 4 is used as default.
    *
    * TODO: at the moment `dim` must be a multiple `pq_dim`.
    */
@@ -50,19 +106,19 @@ struct params {
   bool use_vq = false;
   /**
    * Vector Quantization (VQ) codebook size - number of "coarse cluster centers".
-   * When zero, an optimal value is selected using a heuristic.
+   * When zero, an optimal value is selected using a heuristic. (sqrt(n_rows))
    */
   uint32_t vq_n_centers = 0;
-  /** The number of iterations searching for kmeans centers (both VQ & PQ phases). */
-  uint32_t kmeans_n_iters = 25;
   /**
-   * Type of k-means algorithm for PQ training.
-   * Balanced k-means tends to be faster than regular k-means for PQ training, for
-   * problem sets where the number of points per cluster are approximately equal.
-   * Regular k-means may be better for skewed cluster distributions.
+   * K-means parameters for PQ codebook training.
+   *
+   * Set to cuvs::cluster::kmeans::balanced_params for balanced k-means (default),
+   * or cuvs::cluster::kmeans::params for regular k-means.
+   * The active variant type selects the algorithm; balanced k-means tends to be faster
+   * for PQ training where cluster sizes are approximately equal.
+   * Only L2Expanded metric is supported. The number of clusters is always set to 1 << pq_bits.
    */
-  cuvs::cluster::kmeans::kmeans_type pq_kmeans_type =
-    cuvs::cluster::kmeans::kmeans_type::KMeansBalanced;
+  kmeans_params_variant kmeans_params = cuvs::cluster::kmeans::balanced_params{};
   /**
    * The max number of data points to use per PQ code during PQ codebook training. Using more data
    * points per PQ code may increase the quality of PQ codebook but may also increase the build

cpp/src/neighbors/detail/vpq_dataset.cuh

@@ -5,6 +5,7 @@
 #pragma once
 
 #include <cuvs/neighbors/common.hpp>
+#include <cuvs/preprocessing/quantize/pq.hpp>
 
 #include "../../cluster/kmeans_balanced.cuh"
 #include "../../preprocessing/quantize/detail/pq_codepacking.cuh"  // pq_bits-bitfield
@@ -74,50 +75,49 @@ namespace cuvs::neighbors::detail {
 template <typename MathT, typename IdxT>
 void train_pq_centers(
   const raft::resources& res,
-  const cuvs::neighbors::vpq_params& params,
+  const cuvs::preprocessing::quantize::pq::kmeans_params_variant& kmeans_params,
   const raft::device_matrix_view<const MathT, IdxT, raft::row_major> pq_trainset_view,
   const raft::device_matrix_view<MathT, uint32_t, raft::row_major> pq_centers_view,
   raft::device_vector_view<uint32_t, IdxT> sub_labels_view,
   raft::device_vector_view<uint32_t, IdxT> pq_cluster_sizes_view)
 {
-  if (params.pq_kmeans_type == cuvs::cluster::kmeans::kmeans_type::KMeansBalanced) {
-    cuvs::cluster::kmeans::balanced_params kmeans_params;
-    kmeans_params.n_iters = params.kmeans_n_iters;
-    kmeans_params.metric  = cuvs::distance::DistanceType::L2Expanded;
-
-    cuvs::cluster::kmeans_balanced::helpers::build_clusters<
-      MathT,
-      MathT,
-      IdxT,
-      uint32_t,
-      uint32_t,
-      cuvs::spatial::knn::detail::utils::mapping<MathT>>(
-      res,
-      kmeans_params,
-      pq_trainset_view,
-      pq_centers_view,
-      sub_labels_view,
-      pq_cluster_sizes_view,
-      cuvs::spatial::knn::detail::utils::mapping<MathT>{});
-  } else {
-    const auto pq_n_centers = pq_centers_view.extent(0);
-    cuvs::cluster::kmeans::params kmeans_params;
-    kmeans_params.n_clusters = pq_n_centers;
-    kmeans_params.max_iter   = params.kmeans_n_iters;
-    kmeans_params.metric     = cuvs::distance::DistanceType::L2Expanded;
-    kmeans_params.init       = cuvs::cluster::kmeans::params::InitMethod::Random;
-
-    std::optional<raft::device_vector_view<const MathT, IdxT>> sample_weight = std::nullopt;
-    MathT inertia;
-    IdxT n_iter;
-    cuvs::cluster::kmeans::fit(res,
-                               kmeans_params,
-                               pq_trainset_view,
-                               sample_weight,
-                               pq_centers_view,
-                               raft::make_host_scalar_view<MathT>(&inertia),
-                               raft::make_host_scalar_view<IdxT>(&n_iter));
-  }
+  std::visit(
+    [&](auto const& base_kmeans_params) {
+      using KP = std::decay_t<decltype(base_kmeans_params)>;
+      if constexpr (std::is_same_v<KP, cuvs::cluster::kmeans::balanced_params>) {
+        auto bal_params   = base_kmeans_params;
+        bal_params.metric = cuvs::distance::DistanceType::L2Expanded;
+        cuvs::cluster::kmeans_balanced::helpers::build_clusters<
+          MathT,
+          MathT,
+          IdxT,
+          uint32_t,
+          uint32_t,
+          cuvs::spatial::knn::detail::utils::mapping<MathT>>(
+          res,
+          bal_params,
+          pq_trainset_view,
+          pq_centers_view,
+          sub_labels_view,
+          pq_cluster_sizes_view,
+          cuvs::spatial::knn::detail::utils::mapping<MathT>{});
+      } else {
+        auto classic_params       = base_kmeans_params;
+        classic_params.n_clusters = pq_centers_view.extent(0);
+        classic_params.metric     = cuvs::distance::DistanceType::L2Expanded;
+        std::optional<raft::device_vector_view<const MathT, IdxT>> sample_weight = std::nullopt;
+        MathT inertia;
+        IdxT n_iter;
+        cuvs::cluster::kmeans::fit(res,
+                                   classic_params,
+                                   pq_trainset_view,
+                                   sample_weight,
+                                   pq_centers_view,
+                                   raft::make_host_scalar_view<MathT>(&inertia),
+                                   raft::make_host_scalar_view<IdxT>(&n_iter));
+      }
+    },
+    kmeans_params);
 }
 
 template <typename DatasetT>
@@ -219,7 +219,7 @@ auto predict_vq(const raft::resources& res,
 
 template <typename MathT, typename DatasetT>
 auto train_pq(const raft::resources& res,
-              const vpq_params& params,
+              const cuvs::preprocessing::quantize::pq::params& params,
               const DatasetT& dataset,
               const raft::device_matrix_view<const MathT, uint32_t, raft::row_major> vq_centers)
   -> raft::device_matrix<MathT, uint32_t, raft::row_major>
@@ -230,8 +230,8 @@ auto train_pq(const raft::resources& res,
   const ix_t pq_bits      = params.pq_bits;
   const ix_t pq_n_centers = ix_t{1} << pq_bits;
   const ix_t pq_len       = raft::div_rounding_up_safe(dim, pq_dim);
-  const ix_t n_rows_train = std::min((ix_t)(n_rows * params.pq_kmeans_trainset_fraction),
-                                     params.max_train_points_per_pq_code * pq_n_centers);
+  const ix_t n_rows_train =
+    std::min<ix_t>(n_rows, params.max_train_points_per_pq_code * pq_n_centers);
   RAFT_EXPECTS(
     n_rows_train >= pq_n_centers,
     "The number of training samples must be greater than or equal to the number of PQ centers");
@@ -261,8 +261,12 @@ auto train_pq(const raft::resources& res,
     pq_trainset.data_handle(), n_rows_train * pq_dim, pq_len);
   auto sub_labels       = raft::make_device_vector<uint32_t, ix_t>(res, pq_trainset_view.extent(0));
   auto pq_cluster_sizes = raft::make_device_vector<uint32_t, ix_t>(res, pq_centers.extent(0));
-  train_pq_centers<MathT, ix_t>(
-    res, params, pq_trainset_view, pq_centers.view(), sub_labels.view(), pq_cluster_sizes.view());
+  train_pq_centers<MathT, ix_t>(res,
+                                params.kmeans_params,
+                                pq_trainset_view,
+                                pq_centers.view(),
+                                sub_labels.view(),
+                                pq_cluster_sizes.view());
 
   return pq_centers;
 }

cpp/src/neighbors/scann/detail/scann_build.cuh

@@ -160,14 +160,16 @@ index<T, IdxT> build(
   int dim_per_subspace = params.pq_dim;
   int num_clusters     = 1 << params.pq_bits;
 
-  cuvs::preprocessing::quantize::pq::params pq_build_params;
-  pq_build_params.pq_bits                      = params.pq_bits;
-  pq_build_params.pq_dim                       = num_subspaces;
-  pq_build_params.use_subspaces                = true;
-  pq_build_params.use_vq                       = false;  // We already computed residuals
-  pq_build_params.kmeans_n_iters               = params.pq_train_iters;
-  pq_build_params.max_train_points_per_pq_code = pq_n_rows_train / num_clusters;
-  pq_build_params.pq_kmeans_type               = cuvs::cluster::kmeans::kmeans_type::KMeansBalanced;
+  cuvs::preprocessing::quantize::pq::params pq_build_params(
+    params.pq_bits,
+    num_subspaces,
+    true,
+    false,
+    0,
+    params.pq_train_iters,
+    cuvs::cluster::kmeans::kmeans_type::KMeansBalanced,
+    pq_n_rows_train / num_clusters,
+    1024);
 
   auto pq_quantizer = cuvs::preprocessing::quantize::pq::build(
     res, pq_build_params, raft::make_const_mdspan(trainset_residuals.view()));