Serialize Vamana index with SSD sector alignment per MSFT DiskANN format, generate quantized dataset for integration with DiskANN

cpp/include/cuvs/neighbors/common.hpp

@@ -206,10 +206,14 @@ inline constexpr bool is_strided_dataset_v = is_strided_dataset<DatasetT>::value
  * @param[in] res raft resources handle
  * @param[in] src the source mdarray or mdspan
  * @param[in] required_stride the leading dimension (in elements)
+ * @param[in] force_ownership force an owning_dataset to be returned (default: false)
  * @return maybe owning current-device-accessible strided matrix
  */
 template <typename SrcT>
-auto make_strided_dataset(const raft::resources& res, const SrcT& src, uint32_t required_stride)
+auto make_strided_dataset(const raft::resources& res,
+                          const SrcT& src,
+                          uint32_t required_stride,
+                          bool force_ownership = false)
   -> std::unique_ptr<strided_dataset<typename SrcT::value_type, typename SrcT::index_type>>
 {
   using extents_type = typename SrcT::extents_type;
@@ -231,13 +235,14 @@ auto make_strided_dataset(const raft::resources& res, const SrcT& src, uint32_t
   const bool row_major         = src.stride(1) <= 1;
   const bool stride_matches    = required_stride == src_stride;
 
-  if (device_accessible && row_major && stride_matches) {
+  if (device_accessible && row_major && stride_matches && !force_ownership) {
     // Everything matches: make a non-owning dataset
     return std::make_unique<non_owning_dataset<value_type, index_type>>(
       raft::make_device_strided_matrix_view<const value_type, index_type>(
         device_ptr, src.extent(0), src.extent(1), required_stride));
   }
-  // Something is wrong: have to make a copy and produce an owning dataset
+  // Something is wrong (or force_ownership = true): have to make a copy and produce an owning
+  // dataset
   auto out_layout =
     raft::make_strided_layout(src.extents(), std::array<index_type, 2>{required_stride, 1});
   auto out_array =
@@ -347,18 +352,22 @@ auto make_strided_dataset(
  * @param[in] res raft resources handle
  * @param[in] src the source mdarray or mdspan
  * @param[in] align_bytes the required byte alignment for the dataset rows.
+ * @param[in] force_ownership force an owning_dataset to be returned (default: false)
  * @return maybe owning current-device-accessible strided matrix
  */
 template <typename SrcT>
-auto make_aligned_dataset(const raft::resources& res, SrcT src, uint32_t align_bytes = 16)
+auto make_aligned_dataset(const raft::resources& res,
+                          SrcT src,
+                          uint32_t align_bytes = 16,
+                          bool force_ownership = false)
   -> std::unique_ptr<strided_dataset<typename SrcT::value_type, typename SrcT::index_type>>
 {
   using source_type      = std::remove_cv_t<std::remove_reference_t<SrcT>>;
   using value_type       = typename source_type::value_type;
   constexpr size_t kSize = sizeof(value_type);
   uint32_t required_stride =
     raft::round_up_safe<size_t>(src.extent(1) * kSize, std::lcm(align_bytes, kSize)) / kSize;
-  return make_strided_dataset(res, std::forward<SrcT>(src), required_stride);
+  return make_strided_dataset(res, std::forward<SrcT>(src), required_stride, force_ownership);
 }
 /**
  * @brief VPQ compressed dataset.

cpp/include/cuvs/neighbors/vamana.hpp

@@ -72,6 +72,10 @@ struct index_params : cuvs::neighbors::index_params {
   uint32_t queue_size = 127;
   /** Max batchsize of reverse edge processing (reduces memory footprint) */
   uint32_t reverse_batchsize = 1000000;
+  /** Path prefix to pq pivots and rotation matrix files. Expects pq pivots file at
+   * "${codebook_prefix}_pq_pivots.bin" and rotation matrix file at
+   * "${codebook_prefix}_pq_pivots.bin_rotation_matrix.bin". */
+  std::string codebook_prefix = "";
 };
 
 /**
@@ -127,6 +131,13 @@ struct index : cuvs::neighbors::index {
     return *dataset_;
   }
 
+  /** Quantized dataset [size, codes_rowlen] */
+  [[nodiscard]] inline auto quantized_data() const -> const cuvs::neighbors::dataset<int64_t>&
+  {
+    RAFT_EXPECTS(quantized_dataset_, "Invalid quantized dataset");
+    return *quantized_dataset_;
+  }
+
   /** vamana graph [size, graph-degree] */
   [[nodiscard]] inline auto graph() const noexcept
     -> raft::device_matrix_view<const IdxT, int64_t, raft::row_major>
@@ -212,11 +223,28 @@ struct index : cuvs::neighbors::index {
     graph_view_ = graph_.view();
   }
 
+  /**
+   * Replace the quantized dataset with a new dataset.
+   *
+   * If `force_ownership` is set, we create a copy of the quantized dataset on the device,
+   * and the index manages the lifetime of this copy.
+   * Otherwise, we store a reference to the device data in `new_quantized_dataset`, and it
+   * is the caller's responsibility to ensure that the data stays alive as long as the index.
+   */
+  void update_quantized_dataset(
+    raft::resources const& res,
+    raft::device_matrix_view<uint8_t, int64_t, raft::row_major> new_quantized_dataset,
+    bool force_ownership)
+  {
+    quantized_dataset_ = make_aligned_dataset(res, new_quantized_dataset, 16, force_ownership);
+  }
+
  private:
   cuvs::distance::DistanceType metric_;
   raft::device_matrix<IdxT, int64_t, raft::row_major> graph_;
   raft::device_matrix_view<const IdxT, int64_t, raft::row_major> graph_view_;
   std::unique_ptr<neighbors::dataset<int64_t>> dataset_;
+  std::unique_ptr<neighbors::dataset<int64_t>> quantized_dataset_;
   IdxT medoid_id_;
 };
 /**
@@ -457,13 +485,15 @@ auto build(raft::resources const& res,
  * @param[in] file_prefix prefix of path and name of index files
  * @param[in] index Vamana index
  * @param[in] include_dataset whether or not to serialize the dataset
+ * @param[in] sector_aligned whether output file should be aligned to disk sectors of 4096 bytes
  *
  */
 
 void serialize(raft::resources const& handle,
                const std::string& file_prefix,
                const cuvs::neighbors::vamana::index<float, uint32_t>& index,
-               bool include_dataset = true);
+               bool include_dataset = true,
+               bool sector_aligned  = false);
 
 /**
  * Save the index to file.
@@ -486,12 +516,14 @@ void serialize(raft::resources const& handle,
  * @param[in] file_prefix prefix of path and name of index files
  * @param[in] index Vamana index
  * @param[in] include_dataset whether or not to serialize the dataset
+ * @param[in] sector_aligned whether output file should be aligned to disk sectors of 4096 bytes
  *
  */
 void serialize(raft::resources const& handle,
                const std::string& file_prefix,
                const cuvs::neighbors::vamana::index<int8_t, uint32_t>& index,
-               bool include_dataset = true);
+               bool include_dataset = true,
+               bool sector_aligned  = false);
 
 /**
  * Save the index to file.
@@ -514,12 +546,14 @@ void serialize(raft::resources const& handle,
  * @param[in] file_prefix prefix of path and name of index files
  * @param[in] index Vamana index
  * @param[in] include_dataset whether or not to serialize the dataset
+ * @param[in] sector_aligned whether output file should be aligned to disk sectors of 4096 bytes
  *
  */
 void serialize(raft::resources const& handle,
                const std::string& file_prefix,
                const cuvs::neighbors::vamana::index<uint8_t, uint32_t>& index,
-               bool include_dataset = true);
+               bool include_dataset = true,
+               bool sector_aligned  = false);
 
 /**
  * @}