Merge branch 'main' into cutile-python-to-cpp

Author: divyegala

cpp/CMakeLists.txt

@@ -1367,6 +1367,8 @@ if(NOT BUILD_CPU_ONLY)
     cuvs_objs OBJECT
     src/cluster/detail/minClusterDistanceCompute.cu
     src/cluster/agglomerative.cu
+    src/cluster/gmm_double.cu
+    src/cluster/gmm_float.cu
     src/cluster/kmeans_cluster_cost.cu
     src/cluster/kmeans_fit_double.cu
     src/cluster/kmeans_fit_float.cu

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

@@ -416,6 +416,14 @@ void parse_build_param(const nlohmann::json& conf,
       throw std::runtime_error("invalid value for merge_type");
     }
   }
+
+  nlohmann::json comp_search_conf = collect_conf_with_prefix(conf, "compression_");
+  if (!comp_search_conf.empty()) {
+    auto vpq_pams = param.compression.value_or(cuvs::neighbors::vpq_params{});
+    parse_build_param(comp_search_conf, vpq_pams);
+    param.compression.emplace(vpq_pams);
+  }
+
   param.cagra_params = [conf](raft::matrix_extent<int64_t> extents,
                               cuvs::distance::DistanceType dist_type) {
     // Delayed parsing/initialization of cagra_params - it's called once the dataset shape is known

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

@@ -8,12 +8,18 @@
 #include <cuvs/neighbors/hnsw.hpp>
 #include <raft/core/logger.hpp>
 
+#include <algorithm>
 #include <chrono>
+#include <filesystem>
+#include <fstream>
+#include <iterator>
 #include <memory>
 #include <optional>
 #include <variant>
 
 #include "../common/ann_types.hpp"
+#include "../common/blob.hpp"
+#include "../common/conf.hpp"
 #include "../diskann/diskann_wrapper.h"
 #include "cuvs_ann_bench_utils.h"
 #include <cuvs/neighbors/vamana.hpp>
@@ -165,51 +171,35 @@ void cuvs_cagra_diskann<T, IdxT>::save(const std::string& file) const
   index_of.close();
   if (!index_of) { RAFT_FAIL("Error writing output %s", file.c_str()); }
 
-  // try allocating a buffer for the dataset on host
-  try {
-    auto const* idx_ptr                                    = cagra_build_.get_index();
-    std::optional<raft::host_matrix<T, int64_t>> h_dataset = std::nullopt;
-    auto const& data_view                                  = idx_ptr->dataset();
-    if constexpr (cuvs::neighbors::is_padded_dataset_view_v<std::decay_t<decltype(data_view)>>) {
-      auto const& v = data_view;
-      auto n_rows   = v.n_rows();
-      auto dim      = v.dim();
-      auto stride   = v.stride();
-      h_dataset.emplace(raft::make_host_matrix<T, int64_t>(n_rows, dim));
-      raft::copy_matrix(h_dataset->data_handle(),
-                        dim,
-                        v.view().data_handle(),
-                        stride,
-                        dim,
-                        n_rows,
-                        raft::resource::get_cuda_stream(handle_));
-    } else {
-      RAFT_LOG_DEBUG("dataset serialization: index dataset is not device_padded_dataset_view");
-    }
-
-    if (h_dataset.has_value()) {
-      raft::resource::sync_stream(handle_);
-      std::string dataset_base_file = file + ".data";
-      std::ofstream dataset_of(dataset_base_file, std::ios::out | std::ios::binary);
-      if (!dataset_of) { RAFT_FAIL("Cannot open file %s", dataset_base_file.c_str()); }
-      size_t dataset_file_offset = 0;
-      int size                   = static_cast<int>(cagra_build_.get_index()->size());
-      int dim                    = static_cast<int>(cagra_build_.get_index()->dim());
-      dataset_of.seekp(dataset_file_offset, dataset_of.beg);
-      dataset_of.write((char*)&size, sizeof(int));
-      dataset_of.write((char*)&dim, sizeof(int));
-      for (int i = 0; i < size; i++) {
-        dataset_of.write((char*)(h_dataset->data_handle() + i * h_dataset->extent(1)),
-                         dim * sizeof(T));
-      }
-      dataset_of.close();
-      if (!dataset_of) { RAFT_FAIL("Error writing output %s", dataset_base_file.c_str()); }
-    }
-  } catch (std::bad_alloc& e) {
-    RAFT_LOG_INFO("Failed to serialize dataset");
-  } catch (raft::logic_error& e) {
-    RAFT_LOG_INFO("Failed to serialize dataset");
-  }
+  // Write the rows next to the graph; diskann::Index::load() reads them from `<file>.data`.
+  // The benchmark base file is already in the same bin format, so copy it rather than pull the
+  // rows out of memory - this way `save()` does not care where the dataset was allocated.
+  const auto& ds_conf = configuration::singleton().get_dataset_conf();
+  blob_file<T> base{ds_conf.base_file, ds_conf.subset_first_row, ds_conf.subset_size};
+  int size = static_cast<int>(base.rows_limit());
+  int dim  = static_cast<int>(base.n_cols());
+  RAFT_EXPECTS(dim == this->dim_, "base_file dimensionality does not match the index");
+
+  size_t header_bytes = 2 * sizeof(uint32_t);
+  size_t skip_bytes   = sizeof(T) * static_cast<size_t>(base.rows_offset()) * dim;
+  size_t copy_bytes   = sizeof(T) * static_cast<size_t>(size) * dim;
+  RAFT_EXPECTS(std::filesystem::file_size(base.path()) >= header_bytes + skip_bytes + copy_bytes,
+               "base_file is shorter than its header claims");
+
+  std::ifstream base_in(base.path(), std::ios::in | std::ios::binary);
+  if (!base_in) { RAFT_FAIL("Cannot open file %s", base.path().c_str()); }
+  base_in.seekg(header_bytes + skip_bytes);
+
+  std::string dataset_base_file = file + ".data";
+  std::ofstream dataset_of(dataset_base_file, std::ios::out | std::ios::binary);
+  if (!dataset_of) { RAFT_FAIL("Cannot open file %s", dataset_base_file.c_str()); }
+  dataset_of.write((char*)&size, sizeof(int));
+  dataset_of.write((char*)&dim, sizeof(int));
+  std::copy_n(std::istreambuf_iterator<char>(base_in),
+              copy_bytes,
+              std::ostreambuf_iterator<char>(dataset_of));
+  dataset_of.close();
+  if (!base_in || !dataset_of) { RAFT_FAIL("Error writing output %s", dataset_base_file.c_str()); }
 }
 
 template <typename T, typename IdxT>