[WIP] Vpq dataset serialisation

c/src/neighbors/cagra.cpp

@@ -1011,7 +1011,7 @@ static auto read_serialized_header(cuvsResources_t res, const char *filename)
       "serialization version mismatch, expected %d, got %d",
       cuvs::neighbors::cagra::cagra_serialization_version, version);
   using kind = cuvs::neighbors::cagra::serialized_dataset_kind;
-  RAFT_EXPECTS(dataset_kind_raw <= static_cast<std::uint32_t>(kind::host_standard),
+  RAFT_EXPECTS(dataset_kind_raw <= static_cast<std::uint32_t>(kind::device_vpq_f16),
                "Invalid serialized dataset kind %u in file %s",
                dataset_kind_raw, filename);
   return {output_dtype, static_cast<kind>(dataset_kind_raw)};
@@ -1058,6 +1058,12 @@ void dispatch_serialized_dataset_kind(
       fn.template operator()<
           cuvs::neighbors::device_padded_dataset_view<T, int64_t>>();
       break;
+    case serialized_kind::device_vpq_f16:
+      // A recognised file the C API has no index layout for, as opposed to an unreadable one.
+      // cuvsDatasetLayout_t covers standard and padded only, and every C entry point dispatches
+      // on that layout, so there is nothing here to hand a VPQ index to yet.
+      RAFT_FAIL("File holds a VPQ-compressed (CAGRA-Q) dataset, which the C API has no dataset "
+                "layout for; load it through the C++ API");
   }
 }
 

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/neighbors/cagra.hpp

@@ -2252,7 +2252,7 @@ void search(
  * @{
  */
 
-/** Dense dataset storage kind recorded in a serialized CAGRA index. */
+/** Dataset storage kind recorded in a serialized CAGRA index. */
 enum class serialized_dataset_kind : std::uint32_t {
   /** The serialized index does not contain a dataset payload. */
   none = 0,
@@ -2264,16 +2264,19 @@ enum class serialized_dataset_kind : std::uint32_t {
   host_padded = 3,
   /** Host-resident dataset using its standard row layout. */
   host_standard = 4,
+  /** Device-resident VPQ-compressed dataset with f16 codebooks (CAGRA-Q). */
+  device_vpq_f16 = 5,
 };
 
 /** Current experimental CAGRA serialization format version. */
 inline constexpr int cagra_serialization_version = 6;
 
-// Serialize and deserialize are overloaded for device/host and padded/standard dense indexes.
-// They use the same strided dataset payload; the serialized dataset kind selects the matching
-// owning dataset type during deserialization. To support a new dataset kind (e.g. vpq_f16_index),
-// add matching overloads here and a corresponding deserialize_<kind> in
-// detail/dataset_serialize.hpp (dense views use serialize_cagra_dense_dataset).
+// Serialize and deserialize are overloaded for device/host and padded/standard dense indexes,
+// which share the same strided dataset payload, and for vpq_f16_index, which writes a VPQ payload
+// instead. The serialized dataset kind selects the matching owning dataset type during
+// deserialization. To support a further kind, add matching overloads here and a corresponding
+// serialize_/deserialize_<kind> in detail/dataset_serialize.hpp (dense views use
+// serialize_cagra_dense_dataset, VPQ ones serialize_vpq_dataset).
 
 /**
  * Save the index to file.
@@ -2824,6 +2827,106 @@ void deserialize(raft::resources const& handle,
                  std::unique_ptr<cuvs::neighbors::device_standard_dataset<uint8_t, int64_t>>*
                    out_dataset = nullptr);
 
+/* vpq_f16_index overloads (CAGRA-Q).
+ *
+ * The compressed rows travel with the index, so that a deserialized index can be searched without
+ * the dense dataset it was compressed from and without retraining the codebooks. As everywhere
+ * else, the index holds a view: `deserialize` returns the owning dataset through `out_dataset`,
+ * which the caller has to keep alive for as long as the index is used.
+ *
+ * Unlike the dense overloads, `out_dataset` is required. Nothing can be searched in a VPQ index
+ * whose rows were dropped, so there is no use for a graph-only load, and asking for one is an
+ * error rather than a silently unusable index. For the same reason `include_dataset = false`
+ * produces an index that only `update_dataset` can make searchable again.
+ */
+void serialize(raft::resources const& handle,
+               const std::string& filename,
+               const cuvs::neighbors::cagra::vpq_f16_index<float>& index,
+               bool include_dataset = true);
+
+void deserialize(
+  raft::resources const& handle,
+  const std::string& filename,
+  cuvs::neighbors::cagra::vpq_f16_index<float>* index,
+  std::unique_ptr<cuvs::neighbors::device_vpq_dataset<half, int64_t>>* out_dataset);
+
+void serialize(raft::resources const& handle,
+               std::ostream& os,
+               const cuvs::neighbors::cagra::vpq_f16_index<float>& index,
+               bool include_dataset = true);
+
+void deserialize(
+  raft::resources const& handle,
+  std::istream& is,
+  cuvs::neighbors::cagra::vpq_f16_index<float>* index,
+  std::unique_ptr<cuvs::neighbors::device_vpq_dataset<half, int64_t>>* out_dataset);
+
+void serialize(raft::resources const& handle,
+               const std::string& filename,
+               const cuvs::neighbors::cagra::vpq_f16_index<half>& index,
+               bool include_dataset = true);
+
+void deserialize(
+  raft::resources const& handle,
+  const std::string& filename,
+  cuvs::neighbors::cagra::vpq_f16_index<half>* index,
+  std::unique_ptr<cuvs::neighbors::device_vpq_dataset<half, int64_t>>* out_dataset);
+
+void serialize(raft::resources const& handle,
+               std::ostream& os,
+               const cuvs::neighbors::cagra::vpq_f16_index<half>& index,
+               bool include_dataset = true);
+
+void deserialize(
+  raft::resources const& handle,
+  std::istream& is,
+  cuvs::neighbors::cagra::vpq_f16_index<half>* index,
+  std::unique_ptr<cuvs::neighbors::device_vpq_dataset<half, int64_t>>* out_dataset);
+
+void serialize(raft::resources const& handle,
+               const std::string& filename,
+               const cuvs::neighbors::cagra::vpq_f16_index<int8_t>& index,
+               bool include_dataset = true);
+
+void deserialize(
+  raft::resources const& handle,
+  const std::string& filename,
+  cuvs::neighbors::cagra::vpq_f16_index<int8_t>* index,
+  std::unique_ptr<cuvs::neighbors::device_vpq_dataset<half, int64_t>>* out_dataset);
+
+void serialize(raft::resources const& handle,
+               std::ostream& os,
+               const cuvs::neighbors::cagra::vpq_f16_index<int8_t>& index,
+               bool include_dataset = true);
+
+void deserialize(
+  raft::resources const& handle,
+  std::istream& is,
+  cuvs::neighbors::cagra::vpq_f16_index<int8_t>* index,
+  std::unique_ptr<cuvs::neighbors::device_vpq_dataset<half, int64_t>>* out_dataset);
+
+void serialize(raft::resources const& handle,
+               const std::string& filename,
+               const cuvs::neighbors::cagra::vpq_f16_index<uint8_t>& index,
+               bool include_dataset = true);
+
+void deserialize(
+  raft::resources const& handle,
+  const std::string& filename,
+  cuvs::neighbors::cagra::vpq_f16_index<uint8_t>* index,
+  std::unique_ptr<cuvs::neighbors::device_vpq_dataset<half, int64_t>>* out_dataset);
+
+void serialize(raft::resources const& handle,
+               std::ostream& os,
+               const cuvs::neighbors::cagra::vpq_f16_index<uint8_t>& index,
+               bool include_dataset = true);
+
+void deserialize(
+  raft::resources const& handle,
+  std::istream& is,
+  cuvs::neighbors::cagra::vpq_f16_index<uint8_t>* index,
+  std::unique_ptr<cuvs::neighbors::device_vpq_dataset<half, int64_t>>* out_dataset);
+
 /** @copydoc serialize */
 void serialize(raft::resources const& handle,
                const std::string& filename,

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>
 
@@ -331,6 +334,82 @@ template <typename SrcT>
   }
 }
 
+/** Current VPQ dataset serialization format version. */
+inline constexpr int vpq_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: the encoded rows are what CAGRA-Q builds and
+ * searches over, so a stored VPQ dataset removes the need to keep the dense vectors around or
+ * re-quantize them on every run.
+ *
+ * The file opens with the same preamble as `cagra::serialize` — a 4-byte NumPy dtype prefix then
+ * `vpq_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.
+ *
+ * @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, "base.vpq", vpq);
+ *
+ * // Later, per run: load the compressed rows and build a CAGRA-Q 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] os output stream, opened in binary mode
+ * @param[in] dataset the VPQ dataset to write
+ */
+void serialize(raft::resources const& res,
+               std::ostream& os,
+               const cuvs::neighbors::device_vpq_dataset<half, int64_t>& dataset);
+
+/**
+ * @copydoc serialize
+ *
+ * @param[in] res raft resource
+ * @param[in] filename path to write, truncated if it exists
+ * @param[in] dataset the VPQ dataset to write
+ */
+void serialize(raft::resources const& res,
+               const std::string& filename,
+               const cuvs::neighbors::device_vpq_dataset<half, int64_t>& dataset);
+
+/**
+ * @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_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,7 @@ 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>;
+using inst_vpq_f16_view_t         = cuvs::neighbors::device_vpq_dataset_view<half, int64_t>;
 
 }  // namespace
 
@@ -21,6 +22,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);