Merge branch 'fix/sonatype' of github.com:rapidsai/cuvs into fix/sonatype

Author: jameslamb

cpp/CMakeLists.txt

@@ -1177,6 +1177,7 @@ if(NOT BUILD_CPU_ONLY)
     ${iface_pq_inst_files}
     src/neighbors/detail/cagra/topk_for_cagra/topk.cu
     ${cuvs_cagra_search_cuda_inst_files}
+    src/neighbors/detail/cagra/cagra_helpers.cpp
     src/neighbors/dynamic_batching.cu
     src/neighbors/composite/index.cu
     $<$<BOOL:${BUILD_CAGRA_HNSWLIB}>:src/neighbors/cagra.cpp>

cpp/bench/ann/src/cuvs/cuvs_cagra_hnswlib.cu

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -40,38 +40,8 @@ auto parse_build_param(const nlohmann::json& conf) ->
 
   // Reuse the CAGRA wrapper params parser
   ::parse_build_param<T, IdxT>(conf, cagra_params);
-  // If the users provides parameter M, we can use the CAGRA-HNSW heuristics to find optimal
-  // parameters for the dataset and HNSW reference.
-  if (conf.contains("M")) {
-    // Postpone the parsing of the CAGRA build params until the dataset extents are known.
-    // We the default parameters depend on the dataset extents; and we still would like to be able
-    // to override them.
-    cagra_params.cagra_params = [conf, hnsw_params](raft::matrix_extent<int64_t> extents,
-                                                    cuvs::distance::DistanceType dist_type) {
-      auto ps = cuvs::neighbors::cagra::index_params::from_hnsw_params(
-        extents,
-        conf.at("M"),
-        hnsw_params.ef_construction,
-        cuvs::neighbors::cagra::hnsw_heuristic_type::SAME_GRAPH_FOOTPRINT,
-        dist_type);
-      ps.metric = dist_type;
-      // Parse ACE parameters if provided
-      if (conf.contains("npartitions") || conf.contains("build_dir") ||
-          conf.contains("ef_construction") || conf.contains("use_disk")) {
-        auto ace_params = cuvs::neighbors::cagra::graph_build_params::ace_params();
-        if (conf.contains("npartitions")) { ace_params.npartitions = conf.at("npartitions"); }
-        if (conf.contains("build_dir")) { ace_params.build_dir = conf.at("build_dir"); }
-        if (conf.contains("ef_construction")) {
-          ace_params.ef_construction = conf.at("ef_construction");
-        }
-        if (conf.contains("use_disk")) { ace_params.use_disk = conf.at("use_disk"); }
-        ps.graph_build_params = ace_params;
-      }
-      // NB: above, we only provide the defaults. Below we parse the explicit parameters as usual.
-      ::parse_build_param<T, uint32_t>(conf, ps);
-      return ps;
-    };
-  }
+
+  if (conf.contains("M")) { hnsw_params.M = conf.at("M"); }
   return param;
 }
 

cpp/bench/ann/src/cuvs/cuvs_cagra_hnswlib_wrapper.h

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once
@@ -85,38 +85,9 @@ void cuvs_cagra_hnswlib<T, IdxT>::build(const T* dataset, size_t nrow)
   // when the data set is on host, we can pass it directly to HNSW
   bool dataset_is_on_host = raft::get_device_for_address(dataset) == -1;
 
-  // re-use the CAGRA wrapper to parse build params
-  auto bps = build_param_.cagra_build_params;
-  // Not very conveniently, the CAGRA wrapper resolves parameters after the dataset shape is known,
-  // so it takes a lambda to do it. Even though we know the shape, we want to use the wrapper as-is,
-  // so we just modify that lambda.
-  bps.cagra_params = [dataset_is_on_host, orig_cagra_params = bps.cagra_params](
-                       auto dataset_extents, auto metric) {
-    auto params                    = orig_cagra_params(dataset_extents, metric);
-    params.attach_dataset_on_build = !dataset_is_on_host;
-    return params;
-  };
-  cuvs_cagra<T, IdxT> cagra_wrapper{this->metric_, this->dim_, bps};
-
-  // build the CAGRA index
-  cagra_wrapper.build(dataset, nrow);
-  auto& cagra_index = *cagra_wrapper.get_index();
-
-  // pass the dataset directly to HNSW if it's on the host
-  std::optional<raft::host_matrix_view<const T, int64_t>> opt_dataset_view = std::nullopt;
-  if (dataset_is_on_host) {
-    opt_dataset_view.emplace(
-      raft::make_host_matrix_view<const T, int64_t>(dataset, nrow, this->dim_));
-  }
-
+  auto dataset_view = raft::make_host_matrix_view<const T, int64_t>(dataset, nrow, this->dim_);
   // convert the index to HNSW format
-  hnsw_index_ = cuvs::neighbors::hnsw::from_cagra(
-    handle_, build_param_.hnsw_index_params, cagra_index, opt_dataset_view);
-
-  // special treatment in save/serialize step
-  if (cagra_index.dataset_fd().has_value() && cagra_index.graph_fd().has_value()) {
-    cagra_ace_build_ = true;
-  }
+  hnsw_index_ = cuvs::neighbors::hnsw::build(handle_, build_param_.hnsw_index_params, dataset_view);
 }
 
 template <typename T, typename IdxT>

cpp/include/cuvs/neighbors/cagra.hpp

@@ -3245,6 +3245,38 @@ namespace neighbors {
 namespace cagra {
 namespace helpers {
 
+/** Calculates the workspace for graph optimization
+ *
+ * @param[in] n_rows number of rows in the dataset (or number of points in the graph)
+ * @param[in] graph_degree degree of the output graph
+ * @param[in] intermediate_graph_degree degree of the input graph for the optimization process
+ * @param[in] index_size
+ * @param[in] mst_optimize whether to use MST optimization
+ * @return tuple of [host_size, device_size, host_fixed_size, device_fixed_size] memory sizes in
+ * bytes
+ */
+std::tuple<size_t, size_t, size_t, size_t> optimize_workspace_size(size_t n_rows,
+                                                                   size_t graph_degree,
+                                                                   size_t intermediate_degree,
+                                                                   size_t index_size,
+                                                                   bool mst_optimize = false);
+
+/**
+ * Calculate memory usage of CAGRA build.
+ *
+ * @param[in] res raft resource
+ * @param[in] dataset shape of the dataset
+ * @param[in] dtype element type of the dataset
+ *            (e.g. `CUDA_R_32F`, `CUDA_R_16F`, `CUDA_R_8I`, `CUDA_R_8U`)
+ * @param[in] cparams CAGRA index building parameters
+ *
+ * @return pair of [host_size, device_size] memory sizes in bytes
+ */
+std::pair<size_t, size_t> cagra_build_mem_usage(raft::resources const& res,
+                                                raft::matrix_extent<int64_t> dataset,
+                                                cudaDataType_t dtype,
+                                                cuvs::neighbors::cagra::index_params cparams);
+
 /**
  * @brief Optimize a KNN graph into a CAGRA graph.
  *

cpp/include/cuvs/neighbors/ivf_pq.hpp

@@ -3294,6 +3294,18 @@ void make_rotation_matrix(
   raft::device_matrix_view<float, uint32_t, raft::row_major> rotation_matrix,
   bool force_random_rotation);
 
+/** Calculate the size of the compressed dataset.
+ *
+ * @param[in] res raft resource
+ * @param[in] dataset shape of the dataset
+ * @param[in] param ivf-pq compression params
+ *
+ * @return compressed dataset size in bytes
+ */
+size_t compressed_dataset_size(raft::resources const& res,
+                               raft::matrix_extent<int64_t> dataset,
+                               cuvs::neighbors::ivf_pq::index_params params);
+
 /**
  * @brief Resize an IVF-PQ list with flat layout.
  *
@@ -3355,6 +3367,7 @@ void resize_list(raft::resources const& res,
                  const list_spec_interleaved<uint32_t, int64_t>& spec,
                  uint32_t new_used_size,
                  uint32_t old_used_size);
+
 /**
  * @}
  */

cpp/include/cuvs/neighbors/nn_descent.hpp

@@ -19,6 +19,8 @@
 #include <cuvs/core/cuda_fp16.hpp>
 #include <cuvs/core/export.hpp>
 
+#include <utility>
+
 namespace CUVS_EXPORT cuvs {
 namespace neighbors {
 namespace nn_descent {
@@ -535,6 +537,24 @@ bool has_enough_device_memory(raft::resources const& res,
                               raft::matrix_extent<int64_t> dataset,
                               size_t idx_size = 4);
 
+/**
+ * @brief Estimate the host and device memory (in bytes) required to build an
+ * all-neighbors kNN graph with NN-descent for a dataset of the given shape.
+ *
+ * The estimate mirrors the persistent allocations performed by the GNND solver.
+ * It does not include the input dataset nor the output graph.
+ *
+ * @param res
+ * @param dataset shape of the dataset
+ * @param graph_degree the degree of the all-neighbors graph
+ * @param idx_size the size of the graph index type in bytes
+ * @return std::pair<host_bytes, device_bytes>
+ */
+std::pair<size_t, size_t> build_mem_usage(raft::resources const& res,
+                                          raft::matrix_extent<int64_t> dataset,
+                                          size_t graph_degree,
+                                          size_t idx_size = 4);
+
 }  // namespace nn_descent
 }  // namespace neighbors
 }  // namespace CUVS_EXPORT cuvs

cpp/src/cluster/detail/kmeans_balanced.cuh

@@ -743,7 +743,6 @@ void build_clusters(const raft::resources& handle,
                     const MathT* dataset_norm = nullptr)
 {
   auto stream = raft::resource::get_cuda_stream(handle);
-
   // "randomly" initialize labels
   auto labels_view = raft::make_device_vector_view<LabelT, IdxT>(cluster_labels, n_rows);
   raft::linalg::map_offset(
@@ -886,7 +885,10 @@ auto build_fine_clusters(const raft::resources& handle,
 {
   auto stream = raft::resource::get_cuda_stream(handle);
   rmm::device_uvector<IdxT> mc_trainset_ids_buf(mesocluster_size_max, stream, managed_memory);
-  rmm::device_uvector<MathT> mc_trainset_buf(mesocluster_size_max * dim, stream, device_memory);
+  // for small cluster counts the maximum mesocluster size is proportional to the number of rows, so
+  // we use large workspace
+  auto large_ws = raft::resource::get_large_workspace_resource_ref(handle);
+  rmm::device_uvector<MathT> mc_trainset_buf(mesocluster_size_max * dim, stream, large_ws);
   rmm::device_uvector<MathT> mc_trainset_norm_buf(mesocluster_size_max, stream, device_memory);
   auto mc_trainset_ids  = mc_trainset_ids_buf.data();
   auto mc_trainset      = mc_trainset_buf.data();