Improved CAGRA build parameter heuristics

c/include/cuvs/neighbors/cagra.h

@@ -156,6 +156,48 @@ cuvsError_t cuvsCagraCompressionParamsCreate(cuvsCagraCompressionParams_t* param
  */
 cuvsError_t cuvsCagraCompressionParamsDestroy(cuvsCagraCompressionParams_t params);
 
+/**
+ * @brief Create CAGRA index parameters similar to an HNSW index with hard M constraint
+ *
+ * This factory function creates CAGRA parameters that yield a graph of the same size
+ * as an HNSW graph with the given parameters.
+ *
+ * @param[out] params The CAGRA index params to populate
+ * @param[in] n_rows Number of rows in the dataset
+ * @param[in] dim Number of dimensions in the dataset
+ * @param[in] M HNSW index parameter M
+ * @param[in] ef_construction HNSW index parameter ef_construction
+ * @param[in] metric Distance metric to use
+ * @return cuvsError_t
+ */
+cuvsError_t cuvsCagraIndexParamsFromHnswHardM(cuvsCagraIndexParams_t params,
+                                               int64_t n_rows,
+                                               int64_t dim,
+                                               int M,
+                                               int ef_construction,
+                                               cuvsDistanceType metric);
+
+/**
+ * @brief Create CAGRA index parameters similar to an HNSW index with soft M constraint
+ *
+ * This factory function creates CAGRA parameters that yield a graph with similar search
+ * performance as an HNSW graph with the given parameters.
+ *
+ * @param[out] params The CAGRA index params to populate
+ * @param[in] n_rows Number of rows in the dataset
+ * @param[in] dim Number of dimensions in the dataset
+ * @param[in] M HNSW index parameter M
+ * @param[in] ef_construction HNSW index parameter ef_construction
+ * @param[in] metric Distance metric to use
+ * @return cuvsError_t
+ */
+cuvsError_t cuvsCagraIndexParamsFromHnswSoftM(cuvsCagraIndexParams_t params,
+                                               int64_t n_rows,
+                                               int64_t dim,
+                                               int M,
+                                               int ef_construction,
+                                               cuvsDistanceType metric);
+
 /**
  * @}
  */

c/src/neighbors/cagra.cpp

@@ -676,6 +676,94 @@ extern "C" cuvsError_t cuvsCagraCompressionParamsDestroy(cuvsCagraCompressionPar
   return cuvs::core::translate_exceptions([=] { delete params; });
 }
 
+// Helper function to populate C IVF-PQ params from C++ params
+static void populate_c_ivf_pq_params(cuvsIvfPqParams* c_ivf_pq,
+                                       const cuvs::neighbors::cagra::graph_build_params::ivf_pq_params& cpp_ivf_pq)
+{
+  // Populate the IVF-PQ build params
+  auto& bp = cpp_ivf_pq.build_params;
+  c_ivf_pq->ivf_pq_build_params->metric = static_cast<cuvsDistanceType>(bp.metric);
+  c_ivf_pq->ivf_pq_build_params->metric_arg = bp.metric_arg;
+  c_ivf_pq->ivf_pq_build_params->add_data_on_build = bp.add_data_on_build;
+  c_ivf_pq->ivf_pq_build_params->n_lists = bp.n_lists;
+  c_ivf_pq->ivf_pq_build_params->kmeans_n_iters = bp.kmeans_n_iters;
+  c_ivf_pq->ivf_pq_build_params->kmeans_trainset_fraction = bp.kmeans_trainset_fraction;
+  c_ivf_pq->ivf_pq_build_params->pq_bits = bp.pq_bits;
+  c_ivf_pq->ivf_pq_build_params->pq_dim = bp.pq_dim;
+  c_ivf_pq->ivf_pq_build_params->codebook_kind = static_cast<int>(bp.codebook_kind);
+  c_ivf_pq->ivf_pq_build_params->force_random_rotation = bp.force_random_rotation;
+  c_ivf_pq->ivf_pq_build_params->conservative_memory_allocation = bp.conservative_memory_allocation;
+  c_ivf_pq->ivf_pq_build_params->max_train_points_per_pq_code = bp.max_train_points_per_pq_code;
+
+  // Populate the IVF-PQ search params
+  auto& sp = cpp_ivf_pq.search_params;
+  c_ivf_pq->ivf_pq_search_params->n_probes = sp.n_probes;
+  c_ivf_pq->ivf_pq_search_params->lut_dtype = static_cast<int>(sp.lut_dtype);
+  c_ivf_pq->ivf_pq_search_params->internal_distance_dtype = static_cast<int>(sp.internal_distance_dtype);
+  c_ivf_pq->ivf_pq_search_params->preferred_shmem_carveout = sp.preferred_shmem_carveout;
+
+  c_ivf_pq->refinement_rate = cpp_ivf_pq.refinement_rate;
+}
+
+// Helper function to populate C struct from C++ index_params
+static void populate_cagra_index_params_from_cpp(cuvsCagraIndexParams_t c_params,
+                                                   const cuvs::neighbors::cagra::index_params& cpp_params)
+{
+  c_params->metric = static_cast<cuvsDistanceType>(cpp_params.metric);
+  c_params->intermediate_graph_degree = cpp_params.intermediate_graph_degree;
+  c_params->graph_degree = cpp_params.graph_degree;
+
+  // Set build algo and parameters based on the variant
+  if (std::holds_alternative<cuvs::neighbors::cagra::graph_build_params::nn_descent_params>(
+        cpp_params.graph_build_params)) {
+    c_params->build_algo = NN_DESCENT;
+    auto nn_params =
+      std::get<cuvs::neighbors::cagra::graph_build_params::nn_descent_params>(
+        cpp_params.graph_build_params);
+    c_params->nn_descent_niter = nn_params.max_iterations;
+  } else if (std::holds_alternative<cuvs::neighbors::cagra::graph_build_params::ivf_pq_params>(
+               cpp_params.graph_build_params)) {
+    c_params->build_algo = IVF_PQ;
+    auto ivf_pq_params =
+      std::get<cuvs::neighbors::cagra::graph_build_params::ivf_pq_params>(
+        cpp_params.graph_build_params);
+
+    populate_c_ivf_pq_params(c_params->graph_build_params, ivf_pq_params);
+  }
+}
+
+extern "C" cuvsError_t cuvsCagraIndexParamsFromHnswHardM(cuvsCagraIndexParams_t params,
+                                                          int64_t n_rows,
+                                                          int64_t dim,
+                                                          int M,
+                                                          int ef_construction,
+                                                          cuvsDistanceType metric)
+{
+  return cuvs::core::translate_exceptions([=] {
+    auto cpp_metric = static_cast<cuvs::distance::DistanceType>((int)metric);
+    auto cpp_params = cuvs::neighbors::cagra::index_params::from_hnsw_hard_m(
+      raft::matrix_extent<int64_t>(n_rows, dim), M, ef_construction, cpp_metric);
+
+    populate_cagra_index_params_from_cpp(params, cpp_params);
+  });
+}
+
+extern "C" cuvsError_t cuvsCagraIndexParamsFromHnswSoftM(cuvsCagraIndexParams_t params,
+                                                          int64_t n_rows,
+                                                          int64_t dim,
+                                                          int M,
+                                                          int ef_construction,
+                                                          cuvsDistanceType metric)
+{
+  return cuvs::core::translate_exceptions([=] {
+    auto cpp_metric = static_cast<cuvs::distance::DistanceType>((int)metric);
+    auto cpp_params = cuvs::neighbors::cagra::index_params::from_hnsw_soft_m(
+      raft::matrix_extent<int64_t>(n_rows, dim), M, ef_construction, cpp_metric);
+
+    populate_cagra_index_params_from_cpp(params, cpp_params);
+  });
+}
+
 extern "C" cuvsError_t cuvsCagraExtendParamsCreate(cuvsCagraExtendParams_t* params)
 {
   return cuvs::core::translate_exceptions(

cpp/include/cuvs/neighbors/cagra.hpp

@@ -116,6 +116,74 @@ struct index_params : cuvs::neighbors::index_params {
    * @endcode
    */
   bool attach_dataset_on_build = true;
+
+  /**
+   * @brief Create a CAGRA index parameters compatible with HNSW index
+   *
+   * @param dataset The shape of the input dataset.
+   * @param M HNSW index parameter M (graph_degree = 2M).
+   * @param ef_construction HNSW index parameter ef_construction.
+   * @param metric The distance metric to search.
+   *
+   *
+   * * IMPORTANT NOTE *
+   *
+   * The reference HNSW index and the corresponding from-CAGRA generated HNSW index will NOT produce
+   * the same recalls and QPS for the same parameter `ef`. The graphs are different internally. For
+   * the same `ef`, the from-CAGRA index likely has a slightly higher recall and slightly lower QPS.
+   * However, the Recall-QPS curves should be similar (i.e. the points are just shifted along the
+   * curve).
+   *
+   * Usage example:
+   * @code{.cpp}
+   *   using namespace cuvs::neighbors;
+   *   raft::resources res;
+   *   auto dataset = raft::make_device_matrix<float, int64_t>(res, N, D);
+   *   auto cagra_params = cagra::index_params::from_hnsw_hard_m(dataset.extents(), M, efc);
+   *   auto cagra_index = cagra::build(res, cagra_params, dataset);
+   *   auto hnsw_index = hnsw::from_cagra(res, hnsw_params, cagra_index);
+   * @endcode
+   */
+  static cagra::index_params from_hnsw_hard_m(
+    raft::matrix_extent<int64_t> dataset,
+    int M,
+    int ef_construction,
+    cuvs::distance::DistanceType metric = cuvs::distance::DistanceType::L2Expanded);
+
+  /**
+   * @brief Create a CAGRA index parameters similar to an HNSW index
+   *
+   * @param dataset The shape of the input dataset.
+   * @param M HNSW index parameter M.
+   * @param ef_construction HNSW index parameter ef_construction.
+   * @param metric The distance metric to search.
+   *
+   *
+   * This variant of parameter heuristic yields a graph that is very similar to an HNSW graph in
+   * terms of the number of nodes and search performance. Since HNSW produces a variable-degree
+   * graph (2M being the max graph degree) and CAGRA produces a fixed-degree graph, there's always a
+   * difference in the performance of the two.
+   *
+   * This function attempts to produce such a graph that the QPS and recall of the two graphs being
+   * searched by HNSW are close for any search parameter combination. The CAGRA-produced graph tends
+   * to have a "longer tail" on the low recall side (that is being slightly faster and less
+   * precise).
+   *
+   * Usage example:
+   * @code{.cpp}
+   *   using namespace cuvs::neighbors;
+   *   raft::resources res;
+   *   auto dataset = raft::make_device_matrix<float, int64_t>(res, N, D);
+   *   auto cagra_params = cagra::index_params::from_hnsw_soft_m(dataset.extents(), M, efc);
+   *   auto cagra_index = cagra::build(res, cagra_params, dataset);
+   *   auto hnsw_index = hnsw::from_cagra(res, hnsw_params, cagra_index);
+   * @endcode
+   */
+  static cagra::index_params from_hnsw_soft_m(
+    raft::matrix_extent<int64_t> dataset,
+    int M,
+    int ef_construction,
+    cuvs::distance::DistanceType metric = cuvs::distance::DistanceType::L2Expanded);
 };
 
 /**

cpp/include/cuvs/neighbors/hnsw.hpp

@@ -73,13 +73,15 @@ struct index_params : cuvs::neighbors::index_params {
  * @param metric The distance metric to search.
  *
  *
- * * IMPORTANT NOTE *
+ * This parameter heuristic yields a graph that is very similar to an HNSW graph in
+ * terms of the number of nodes and search performance. Since HNSW produces a variable-degree
+ * graph (2M being the max graph degree) and CAGRA produces a fixed-degree graph, there's always a
+ * difference in the performance of the two.
  *
- * The reference HNSW index and the corresponding from-CAGRA generated HNSW index will NOT produce
- * the same recalls and QPS for the same parameter `ef`. The graphs are different internally. For
- * the same `ef`, the from-CAGRA index likely has a slightly higher recall and slightly lower QPS.
- * However, the Recall-QPS curves should be similar (i.e. the points are just shifted along the
- * curve).
+ * This function attempts to produce such a graph that the QPS and recall of the two graphs being
+ * searched by HNSW are close for any search parameter combination. The CAGRA-produced graph tends
+ * to have a "longer tail" on the low recall side (that is being slightly faster and less
+ * precise).
  *
  * Usage example:
  * @code{.cpp}
@@ -91,6 +93,7 @@ struct index_params : cuvs::neighbors::index_params {
  *   auto hnsw_index = hnsw::from_cagra(res, hnsw_params, cagra_index);
  * @endcode
  */
+[[deprecated("Use cagra::index_params::from_hnsw_soft_m instead")]]
 cuvs::neighbors::cagra::index_params to_cagra_params(
   raft::matrix_extent<int64_t> dataset,
   int M,

cpp/src/neighbors/cagra.cuh

@@ -416,4 +416,51 @@ index<T, IdxT> merge(raft::resources const& handle,
 
 /** @} */  // end group cagra
 
+inline auto graph_params_heuristic(raft::matrix_extent<int64_t> dataset,
+                                   int intermediate_graph_degree,
+                                   int ef_construction,
+                                   cuvs::distance::DistanceType metric)
+  -> decltype(index_params::graph_build_params)
+{
+  if (dataset.extent(0) < int64_t(1e6)) {
+    // Use NN descent for smaller datasets
+    auto nn_descent_params =
+      graph_build_params::nn_descent_params(intermediate_graph_degree, metric);
+    nn_descent_params.max_iterations = 5 + ef_construction / 16;
+    return nn_descent_params;
+  } else {
+    // Otherwise, use IVF-PQ
+    auto ivf_pq_params = cuvs::neighbors::graph_build_params::ivf_pq_params(dataset, metric);
+    ivf_pq_params.search_params.n_probes =
+      std::round(2 + std::sqrt(ivf_pq_params.build_params.n_lists) / 20 + ef_construction / 16);
+    return ivf_pq_params;
+  }
+}
+
+inline cagra::index_params index_params::from_hnsw_hard_m(raft::matrix_extent<int64_t> dataset,
+                                                          int M,
+                                                          int ef_construction,
+                                                          cuvs::distance::DistanceType metric)
+{
+  cagra::index_params params;
+  params.graph_degree              = M * 2;
+  params.intermediate_graph_degree = M * 3;
+  params.graph_build_params =
+    graph_params_heuristic(dataset, params.intermediate_graph_degree, ef_construction, metric);
+  return params;
+}
+
+inline cagra::index_params index_params::from_hnsw_soft_m(raft::matrix_extent<int64_t> dataset,
+                                                          int M,
+                                                          int ef_construction,
+                                                          cuvs::distance::DistanceType metric)
+{
+  cagra::index_params params;
+  params.graph_degree              = 2 + M * 2 / 3;
+  params.intermediate_graph_degree = M + M * ef_construction / 256;
+  params.graph_build_params =
+    graph_params_heuristic(dataset, params.intermediate_graph_degree, ef_construction, metric);
+  return params;
+}
+
 }  // namespace cuvs::neighbors::cagra

cpp/src/neighbors/hnsw.cpp

@@ -15,8 +15,11 @@
  */
 
 #include "detail/hnsw.hpp"
-#include <cstdint>
+
+#include <cuvs/neighbors/cagra.hpp>
 #include <cuvs/neighbors/hnsw.hpp>
+
+#include <cstdint>
 #include <sys/types.h>
 
 namespace cuvs::neighbors::hnsw {
@@ -26,16 +29,8 @@ auto to_cagra_params(raft::matrix_extent<int64_t> dataset,
                      int ef_construction,
                      cuvs::distance::DistanceType metric) -> cuvs::neighbors::cagra::index_params
 {
-  auto ivf_pq_params = cuvs::neighbors::graph_build_params::ivf_pq_params(dataset, metric);
-  ivf_pq_params.search_params.n_probes =
-    std::round(std::sqrt(ivf_pq_params.build_params.n_lists) / 20 + ef_construction / 16);
-
-  cagra::index_params params;
-  params.graph_build_params        = ivf_pq_params;
-  params.graph_degree              = M * 2;
-  params.intermediate_graph_degree = M * 3;
-
-  return params;
+  return cuvs::neighbors::cagra::index_params::from_hnsw_soft_m(
+    dataset, M, ef_construction, metric);
 }
 
 #define CUVS_INST_HNSW_FROM_CAGRA(T)                                                  \