[WIP] Vpq dataset serialisation

c/src/neighbors/cagra.cpp

@@ -1058,6 +1058,10 @@ void dispatch_serialized_dataset_kind(
       fn.template operator()<
           cuvs::neighbors::device_padded_dataset_view<T, int64_t>>();
       break;
+    // Unreachable: read_serialized_header rejects this kind before the dispatch, since
+    // cuvsDatasetLayout_t has no PQ-compressed layout to hand back. Listed only because the switch
+    // is exhaustive and -Wswitch is an error. Delete it when the C API gains the layout.
+    case serialized_kind::device_pq: break;
   }
 }
 

c/tests/CMakeLists.txt

@@ -89,7 +89,9 @@ ConfigureTest(NAME IVF_FLAT_C_TEST PATH neighbors/run_ivf_flat_c.c neighbors/ann
 ConfigureTest(NAME IVF_PQ_C_TEST PATH neighbors/run_ivf_pq_c.c neighbors/ann_ivf_pq_c.cu)
 ConfigureTest(NAME IVF_SQ_C_TEST PATH neighbors/run_ivf_sq_c.c neighbors/ann_ivf_sq_c.cu)
 ConfigureTest(NAME CAGRA_C_TEST PATH neighbors/ann_cagra_c.cu)
-ConfigureTest(NAME MG_C_TEST PATH neighbors/run_mg_c.c neighbors/ann_mg_c.cu)
+if(BUILD_MG_ALGOS)
+  ConfigureTest(NAME MG_C_TEST PATH neighbors/run_mg_c.c neighbors/ann_mg_c.cu)
+endif()
 ConfigureTest(
   NAME ALL_NEIGHBORS_C_TEST PATH neighbors/run_all_neighbors_c.c neighbors/all_neighbors_c.cu
 )

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

@@ -14,6 +14,9 @@
 
 #include <cuda_runtime.h>
 #include <cuvs/core/export.hpp>
+#include <iosfwd>
+#include <memory>
+#include <string>
 #include <type_traits>
 #include <variant>
 
@@ -279,6 +282,15 @@ namespace detail {
  * dense dataset is never staged on the device in full; they must be tightly packed. Empty sources
  * are rejected. The element type must be `float`, `half`, `int8_t` or `uint8_t`.
  *
+ * Only the input streams. The result is a single device allocation of `n_rows` encoded rows, so the
+ * compressed dataset has to fit in whatever the current device memory resource can serve, and there
+ * is no host-resident output to fall back on: nothing produces, searches or serializes the
+ * `host_vpq_dataset` type today. A row is `sizeof(uint32_t) + pq_dim * pq_bits / 8` bytes rounded
+ * up to a multiple of 4, so at `pq_bits = 8` and `pq_dim = 384` a hundred million rows come to
+ * about 39 GB, and a billion rows exceed any single device. Past that point the options are an
+ * oversubscribed (managed) memory resource, which is enough to encode and serialize but not to
+ * search, or sharding the rows and merging the search results.
+ *
  * Typical **CAGRA** usage: build the graph on dense vectors, then attach VPQ for search (metric
  * must remain `L2Expanded` for this path). Train VPQ from the same CAGRA-padded device layout you
  * used for graph build, keep the `device_vpq_dataset` alive, and call
@@ -331,6 +343,88 @@ template <typename SrcT>
   }
 }
 
+/** Current VPQ dataset serialization format version. */
+inline constexpr int pq_serialization_version = 1;
+
+/**
+ * @brief Write a VPQ dataset (both codebooks plus the encoded rows) to a stream.
+ *
+ * Lets compression be done once, offline, and reused: a CAGRA graph over a compressed dataset
+ * builds and searches on the encoded rows, so storing them removes the need to keep the dense
+ * vectors around or to re-quantize them on every run.
+ *
+ * The file opens with the same preamble as `cagra::serialize` — a 4-byte NumPy dtype prefix then
+ * `pq_serialization_version` — followed by a dataset kind tag and the codebook element type. A
+ * file of the wrong kind, or one written by an older format, is rejected rather than misread. Bump
+ * the version whenever the encoded row layout changes, since that layout is a library convention
+ * and is not otherwise described by the file.
+ *
+ * Writing copies the encoded rows to the host in one piece, as `raft::serialize_mdspan` does for
+ * any device matrix: it allocates a host buffer the size of those rows alongside the device copy
+ * it reads from, and frees it afterwards. The two codebooks go the same way and are small. Reading
+ * is the mirror image, host buffer first and then a copy to the device. So a file costs the encoded
+ * rows twice while it is being written or read, once on each side, and neither direction streams.
+ *
+ * @code{.cpp}
+ * #include <cuvs/neighbors/cagra.hpp>
+ * #include <cuvs/preprocessing/quantize/pq.hpp>
+ *
+ * // Offline, once.
+ * auto vpq = cuvs::preprocessing::quantize::pq::make_vpq_dataset(res, vpq_params, rows);
+ * cuvs::preprocessing::quantize::pq::serialize(res, vpq, "base.vpq");
+ *
+ * // Later, per run: load the compressed rows and build a CAGRA graph over them.
+ * std::unique_ptr<cuvs::neighbors::device_vpq_dataset<half, int64_t>> loaded;
+ * cuvs::preprocessing::quantize::pq::deserialize(res, "base.vpq", &loaded);
+ * auto index = cuvs::neighbors::cagra::build(res, index_params, loaded->as_dataset_view());
+ * // `loaded` must outlive `index`, which only holds a view of it.
+ * @endcode
+ *
+ * @param[in] res raft resource
+ * @param[in] dataset the VPQ dataset to write
+ * @param[out] os output stream, opened in binary mode
+ */
+void serialize(raft::resources const& res,
+               const cuvs::neighbors::device_vpq_dataset<half, int64_t>& dataset,
+               std::ostream& os);
+
+/**
+ * @copydoc serialize
+ *
+ * @param[in] res raft resource
+ * @param[in] dataset the VPQ dataset to write
+ * @param[out] filename path to write, truncated if it exists
+ */
+void serialize(raft::resources const& res,
+               const cuvs::neighbors::device_vpq_dataset<half, int64_t>& dataset,
+               const std::string& filename);
+
+/**
+ * @brief Read a VPQ dataset written by `serialize`.
+ *
+ * Returned through an out-parameter because the dataset owns device allocations and has no default
+ * constructor, matching how `cagra::deserialize` hands back its dataset. Throws if the blob was not
+ * written by `serialize` or holds codebooks of a different element type.
+ *
+ * @param[in] res raft resource
+ * @param[in] is input stream, opened in binary mode
+ * @param[out] out_dataset receives the loaded dataset; must not be null
+ */
+void deserialize(raft::resources const& res,
+                 std::istream& is,
+                 std::unique_ptr<cuvs::neighbors::device_vpq_dataset<half, int64_t>>* out_dataset);
+
+/**
+ * @copydoc deserialize
+ *
+ * @param[in] res raft resource
+ * @param[in] filename path to read
+ * @param[out] out_dataset receives the loaded dataset; must not be null
+ */
+void deserialize(raft::resources const& res,
+                 const std::string& filename,
+                 std::unique_ptr<cuvs::neighbors::device_vpq_dataset<half, int64_t>>* out_dataset);
+
 /** @} */  // end of group product
 
 }  // namespace pq

cpp/src/neighbors/cagra.cpp

@@ -8,6 +8,15 @@
 
 #include <cuvs/neighbors/common.hpp>
 
+#include <raft/core/error.hpp>
+#include <raft/core/numpy_serializer.hpp>
+#include <raft/core/serialize.hpp>
+
+#include <cstdint>
+#include <fstream>
+#include <istream>
+#include <string>
+
 namespace cuvs::neighbors::cagra {
 
 graph_build_params_t index_params::graph_build_heuristic(raft::matrix_extent<int64_t> dataset,
@@ -67,4 +76,75 @@ cagra::index_params index_params::from_hnsw_params(raft::matrix_extent<int64_t>
   return params;
 }
 
+namespace {
+
+/**
+ * Map the file's 4-byte NumPy dtype descriptor back to the element type that wrote it.
+ *
+ * Parses the descriptor rather than comparing against `get_numpy_dtype<T>()`, which has no answer
+ * for `half` outside a CUDA translation unit. 'e' is how raft spells a half, as the C API's reader
+ * also has to know.
+ */
+auto element_dtype_of(const char (&prefix)[4], const char* source) -> cudaDataType_t
+{
+  auto const dtype = raft::numpy_serializer::parse_descr(std::string(prefix, sizeof(prefix)));
+  if (dtype.kind == 'f' && dtype.itemsize == 4) { return CUDA_R_32F; }
+  if (dtype.kind == 'e' && dtype.itemsize == 2) { return CUDA_R_16F; }
+  if (dtype.kind == 'i' && dtype.itemsize == 1) { return CUDA_R_8I; }
+  if (dtype.kind == 'u' && dtype.itemsize == 1) { return CUDA_R_8U; }
+  RAFT_FAIL("cagra::read_serialized_header: %s holds an index whose element type (%s) is not one "
+            "CAGRA writes",
+            source,
+            dtype.to_string().c_str());
+}
+
+auto read_header(raft::resources const& res, std::istream& is, const char* source)
+  -> serialized_index_header
+{
+  using pos_type   = std::istream::pos_type;
+  using off_type   = std::istream::off_type;
+  auto const start = is.tellg();
+  RAFT_EXPECTS(start != pos_type{off_type{-1}},
+               "cagra::read_serialized_header: %s is not seekable",
+               source);
+
+  char dtype_prefix[4];
+  RAFT_EXPECTS(is.read(dtype_prefix, sizeof(dtype_prefix)),
+               "cagra::read_serialized_header: failed to read the dtype prefix of %s",
+               source);
+  auto const dtype = element_dtype_of(dtype_prefix, source);
+
+  auto const version = raft::deserialize_scalar<int>(res, is);
+  RAFT_EXPECTS(version == cagra_serialization_version,
+               "cagra::read_serialized_header: serialization version mismatch, expected %d, got %d",
+               cagra_serialization_version,
+               version);
+
+  // Read after the version check: an older or newer format need not put the kind here at all.
+  auto const kind_raw = raft::deserialize_scalar<std::uint32_t>(res, is);
+  RAFT_EXPECTS(kind_raw <= static_cast<std::uint32_t>(serialized_dataset_kind::device_pq),
+               "cagra::read_serialized_header: invalid serialized dataset kind %u in %s",
+               kind_raw,
+               source);
+
+  // Rewind, so that the caller can hand the same stream to deserialize.
+  is.seekg(start);
+  return {dtype, static_cast<serialized_dataset_kind>(kind_raw)};
+}
+
+}  // namespace
+
+auto read_serialized_header(raft::resources const& res, std::istream& is) -> serialized_index_header
+{
+  return read_header(res, is, "the stream");
+}
+
+auto read_serialized_header(raft::resources const& res, const std::string& filename)
+  -> serialized_index_header
+{
+  std::ifstream is(filename, std::ios::in | std::ios::binary);
+  RAFT_EXPECTS(is, "cagra::read_serialized_header: cannot open %s", filename.c_str());
+  return read_header(res, is, filename.c_str());
+}
+
 }  // namespace cuvs::neighbors::cagra

cpp/src/neighbors/cagra_serialize.cuh

@@ -155,6 +155,43 @@ namespace cuvs::neighbors::cagra {
     cuvs::neighbors::cagra::detail::deserialize<DTYPE, uint32_t>(handle, is, index, out_dataset); \
   }                                                                                               \
                                                                                                   \
+  void serialize(raft::resources const& handle,                                                   \
+                 const std::string& filename,                                                     \
+                 const cuvs::neighbors::cagra::vpq_f16_index<DTYPE, uint32_t>& index,             \
+                 bool include_dataset)                                                            \
+  {                                                                                               \
+    cuvs::neighbors::cagra::detail::serialize<DTYPE, uint32_t>(                                   \
+      handle, filename, index, include_dataset);                                                  \
+  }                                                                                               \
+                                                                                                  \
+  void deserialize(                                                                               \
+    raft::resources const& handle,                                                                \
+    const std::string& filename,                                                                  \
+    cuvs::neighbors::cagra::vpq_f16_index<DTYPE, uint32_t>* index,                                \
+    std::unique_ptr<cuvs::neighbors::device_vpq_dataset<half, int64_t>>* out_dataset)             \
+  {                                                                                               \
+    cuvs::neighbors::cagra::detail::deserialize<DTYPE, uint32_t>(                                 \
+      handle, filename, index, out_dataset);                                                      \
+  }                                                                                               \
+                                                                                                  \
+  void serialize(raft::resources const& handle,                                                   \
+                 std::ostream& os,                                                                \
+                 const cuvs::neighbors::cagra::vpq_f16_index<DTYPE, uint32_t>& index,             \
+                 bool include_dataset)                                                            \
+  {                                                                                               \
+    cuvs::neighbors::cagra::detail::serialize<DTYPE, uint32_t>(                                   \
+      handle, os, index, include_dataset);                                                        \
+  }                                                                                               \
+                                                                                                  \
+  void deserialize(                                                                               \
+    raft::resources const& handle,                                                                \
+    std::istream& is,                                                                             \
+    cuvs::neighbors::cagra::vpq_f16_index<DTYPE, uint32_t>* index,                                \
+    std::unique_ptr<cuvs::neighbors::device_vpq_dataset<half, int64_t>>* out_dataset)             \
+  {                                                                                               \
+    cuvs::neighbors::cagra::detail::deserialize<DTYPE, uint32_t>(handle, is, index, out_dataset); \
+  }                                                                                               \
+                                                                                                  \
   void serialize_to_hnswlib(                                                                      \
     raft::resources const& handle,                                                                \
     std::ostream& os,                                                                             \

cpp/src/neighbors/cagra_serialize_inst.cu.in

@@ -12,6 +12,9 @@ namespace {
 using data_t                      = @data_type@;
 using inst_device_padded_view_t   = cuvs::neighbors::device_padded_dataset_view<data_t, int64_t>;
 using inst_device_standard_view_t = cuvs::neighbors::device_standard_dataset_view<data_t, int64_t>;
+// `vpq` rather than `pq`, to keep the same name as the identical alias in cagra_search_inst.cu.in
+// and as the dataset type it stands for. It changes when those do.
+using inst_vpq_f16_view_t = cuvs::neighbors::device_vpq_dataset_view<half, int64_t>;
 
 }  // namespace
 
@@ -21,6 +24,8 @@ extern template void index<data_t, uint32_t, inst_device_padded_view_t>::compute
   raft::resources const&);
 extern template void index<data_t, uint32_t, inst_device_standard_view_t>::compute_dataset_norms_(
   raft::resources const&);
+extern template void index<data_t, uint32_t, inst_vpq_f16_view_t>::compute_dataset_norms_(
+  raft::resources const&);
 
 CUVS_INST_CAGRA_SERIALIZE(data_t);
 

cpp/src/neighbors/detail/cagra/cagra_serialize.cuh

@@ -60,6 +60,10 @@ constexpr auto serialized_dataset_kind_for_view() -> cuvs::neighbors::cagra::ser
     return kind::host_padded;
   } else if constexpr (cuvs::neighbors::is_host_standard_dataset_view_v<DatasetViewT>) {
     return kind::host_standard;
+  } else if constexpr (cuvs::neighbors::is_device_vpq_dataset_view_v<DatasetViewT>) {
+    // Any codebook element type maps to the one kind, since the payload records which it is. Only
+    // f16 codebooks are written today, and the branches below say so.
+    return kind::device_pq;
   } else {
     static_assert(sizeof(DatasetViewT) == 0,
                   "serialized_dataset_kind_for_view: unsupported dataset view type");
@@ -69,7 +73,7 @@ constexpr auto serialized_dataset_kind_for_view() -> cuvs::neighbors::cagra::ser
 constexpr bool is_valid_serialized_dataset_kind(std::uint32_t raw)
 {
   using kind = cuvs::neighbors::cagra::serialized_dataset_kind;
-  return raw <= static_cast<std::uint32_t>(kind::host_standard);
+  return raw <= static_cast<std::uint32_t>(kind::device_pq);
 }
 
 /**
@@ -123,9 +127,14 @@ void serialize(raft::resources const& res,
     RAFT_LOG_DEBUG("Saving CAGRA index with dataset");
     if constexpr (cuvs::neighbors::is_dense_row_major_dataset_view_v<DatasetViewT>) {
       neighbors::detail::serialize_cagra_dense_dataset<T, int64_t>(res, os, index_.dataset());
+    } else if constexpr (cuvs::neighbors::is_device_vpq_f16_dataset_view_v<DatasetViewT>) {
+      // The payload describes its own codebook type, which is `half` here regardless of T: the
+      // dtype prefix written above is the type of the queries this index answers, not of its rows.
+      // `dset()` is safe to call because a view over no rows left include_dataset false above.
+      neighbors::detail::serialize_pq_dataset<half, int64_t>(res, index_.dataset().dset(), os);
     } else {
-      // Future dataset types (e.g. VPQ) require a new branch here and a corresponding
-      // deserialize overload. Use static_assert to catch unsupported types at compile time.
+      // A further dataset type requires a new branch here and a corresponding deserialize branch.
+      // Use static_assert to catch unsupported types at compile time.
       static_assert(
         sizeof(DatasetViewT) == 0,
         "serialize: dataset serialization is not yet implemented for this DatasetViewT");
@@ -401,7 +410,11 @@ void deserialize(
   std::unique_ptr<owner_t> dataset_owner{};
   if (has_dataset) {
     if (out_dataset == nullptr) {
-      cuvs::neighbors::detail::skip_dense_dataset<T, int64_t>(res, is);
+      // No out_dataset means the caller wants the graph alone. The dataset bytes still have to be
+      // stepped over to reach the source indices that follow them, and the payload starts with a
+      // tag naming its kind, so skipping it needs nothing from the caller. The index comes back
+      // with no rows, and cannot be searched until update_device_dataset_same_layout gives it some.
+      cuvs::neighbors::detail::skip_dataset<int64_t>(res, is);
     } else {
       auto const expected_kind = serialized_dataset_kind_for_view<DatasetViewT>();
       RAFT_EXPECTS(
@@ -419,6 +432,8 @@ void deserialize(
       } else if constexpr (cuvs::neighbors::is_host_standard_dataset_view_v<DatasetViewT>) {
         dataset_owner =
           cuvs::neighbors::detail::deserialize_host_standard_dataset<T, int64_t>(res, is);
+      } else if constexpr (cuvs::neighbors::is_device_vpq_f16_dataset_view_v<DatasetViewT>) {
+        dataset_owner = cuvs::neighbors::detail::deserialize_pq_dataset<half, int64_t>(res, is);
       } else {
         static_assert(sizeof(DatasetViewT) == 0,
                       "deserialize: dataset deserialization is not implemented for this view");