diff --git a/c/src/neighbors/cagra.cpp b/c/src/neighbors/cagra.cpp index 99e456e23c..f65f571713 100644 --- a/c/src/neighbors/cagra.cpp +++ b/c/src/neighbors/cagra.cpp @@ -1058,6 +1058,10 @@ void dispatch_serialized_dataset_kind( fn.template operator()< cuvs::neighbors::device_padded_dataset_view>(); 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; } } diff --git a/c/tests/CMakeLists.txt b/c/tests/CMakeLists.txt index 7d6c588bd9..ff8f807a6a 100644 --- a/c/tests/CMakeLists.txt +++ b/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 ) diff --git a/cpp/include/cuvs/neighbors/cagra.hpp b/cpp/include/cuvs/neighbors/cagra.hpp index 43ae7a6235..b1a046a2c2 100644 --- a/cpp/include/cuvs/neighbors/cagra.hpp +++ b/cpp/include/cuvs/neighbors/cagra.hpp @@ -25,6 +25,7 @@ #include #include #include +#include #include #include #include @@ -469,6 +470,7 @@ struct CUVS_EXPORT index : cuvs::neighbors::index { using value_type = T; using dataset_index_type = int64_t; using graph_index_type = uint32_t; + using dataset_view_type = DatasetViewT; static_assert(!raft::is_narrowing_v, "IdxT must be able to represent all values of uint32_t"); @@ -918,6 +920,18 @@ using cagra_index_t = index>; +/** + * @brief The dataset type `deserialize` produces for `IndexT`: what that index views, made owning. + * + * Saves a caller from restating it, which is a mouthful once the index type has already said it: + * `std::unique_ptr> rows;` + * + * @tparam IndexT a `cagra::index` instantiation, or one of its aliases + */ +template +using owning_dataset_for_index_t = + cuvs::neighbors::owning_dataset_for_view_t; + /** * @} */ @@ -1493,7 +1507,7 @@ void search(raft::resources const& res, const cuvs::neighbors::filtering::base_filter& sample_filter = cuvs::neighbors::filtering::none_sample_filter{}); -// vpq_f16_index overloads (uint32_t neighbor indices) +// Indexes over a PQ-compressed dataset with f16 codebooks (uint32_t neighbor indices) /** * @brief Search ANN using the constructed index. * @@ -1501,8 +1515,8 @@ void search(raft::resources const& res, * * @param[in] res raft resources * @param[in] params configure the search - * @param[in] index pre-built vpq_f16_index (CAGRA-Q, VPQ f16-compressed dataset) with uint32_t - * neighbor indices + * @param[in] index a pre-built CAGRA index over a PQ-compressed dataset with f16 codebooks, + * returning uint32_t neighbor indices * @param[in] queries a device matrix view to a row-major matrix [n_queries, index.dim()] * @param[out] neighbors a device matrix view to the indices of the neighbors in the source dataset * [n_queries, k] @@ -1527,8 +1541,8 @@ void search(raft::resources const& res, * * @param[in] res raft resources * @param[in] params configure the search - * @param[in] index pre-built vpq_f16_index (CAGRA-Q, VPQ f16-compressed dataset) with uint32_t - * neighbor indices + * @param[in] index a pre-built CAGRA index over a PQ-compressed dataset with f16 codebooks, + * returning uint32_t neighbor indices * @param[in] queries a device matrix view to a row-major matrix [n_queries, index.dim()] * @param[out] neighbors a device matrix view to the indices of the neighbors in the source dataset * [n_queries, k] @@ -1553,8 +1567,8 @@ void search(raft::resources const& res, * * @param[in] res raft resources * @param[in] params configure the search - * @param[in] index pre-built vpq_f16_index (CAGRA-Q, VPQ f16-compressed dataset) with uint32_t - * neighbor indices + * @param[in] index a pre-built CAGRA index over a PQ-compressed dataset with f16 codebooks, + * returning uint32_t neighbor indices * @param[in] queries a device matrix view to a row-major matrix [n_queries, index.dim()] * @param[out] neighbors a device matrix view to the indices of the neighbors in the source dataset * [n_queries, k] @@ -1579,8 +1593,8 @@ void search(raft::resources const& res, * * @param[in] res raft resources * @param[in] params configure the search - * @param[in] index pre-built vpq_f16_index (CAGRA-Q, VPQ f16-compressed dataset) with uint32_t - * neighbor indices + * @param[in] index a pre-built CAGRA index over a PQ-compressed dataset with f16 codebooks, + * returning uint32_t neighbor indices * @param[in] queries a device matrix view to a row-major matrix [n_queries, index.dim()] * @param[out] neighbors a device matrix view to the indices of the neighbors in the source dataset * [n_queries, k] @@ -1598,7 +1612,7 @@ void search(raft::resources const& res, const cuvs::neighbors::filtering::base_filter& sample_filter = cuvs::neighbors::filtering::none_sample_filter{}); -// vpq_f16_index overloads (int64_t neighbor indices) +// Indexes over a PQ-compressed dataset with f16 codebooks (int64_t neighbor indices) /** * @brief Search ANN using the constructed index. * @@ -1606,8 +1620,8 @@ void search(raft::resources const& res, * * @param[in] res raft resources * @param[in] params configure the search - * @param[in] index pre-built vpq_f16_index (CAGRA-Q, VPQ f16-compressed dataset) with int64_t - * neighbor indices + * @param[in] index a pre-built CAGRA index over a PQ-compressed dataset with f16 codebooks, + * returning int64_t neighbor indices * @param[in] queries a device matrix view to a row-major matrix [n_queries, index.dim()] * @param[out] neighbors a device matrix view to the indices of the neighbors in the source dataset * [n_queries, k] @@ -1632,8 +1646,8 @@ void search(raft::resources const& res, * * @param[in] res raft resources * @param[in] params configure the search - * @param[in] index pre-built vpq_f16_index (CAGRA-Q, VPQ f16-compressed dataset) with int64_t - * neighbor indices + * @param[in] index a pre-built CAGRA index over a PQ-compressed dataset with f16 codebooks, + * returning int64_t neighbor indices * @param[in] queries a device matrix view to a row-major matrix [n_queries, index.dim()] * @param[out] neighbors a device matrix view to the indices of the neighbors in the source dataset * [n_queries, k] @@ -1658,8 +1672,8 @@ void search(raft::resources const& res, * * @param[in] res raft resources * @param[in] params configure the search - * @param[in] index pre-built vpq_f16_index (CAGRA-Q, VPQ f16-compressed dataset) with int64_t - * neighbor indices + * @param[in] index a pre-built CAGRA index over a PQ-compressed dataset with f16 codebooks, + * returning int64_t neighbor indices * @param[in] queries a device matrix view to a row-major matrix [n_queries, index.dim()] * @param[out] neighbors a device matrix view to the indices of the neighbors in the source dataset * [n_queries, k] @@ -1684,8 +1698,8 @@ void search(raft::resources const& res, * * @param[in] res raft resources * @param[in] params configure the search - * @param[in] index pre-built vpq_f16_index (CAGRA-Q, VPQ f16-compressed dataset) with int64_t - * neighbor indices + * @param[in] index a pre-built CAGRA index over a PQ-compressed dataset with f16 codebooks, + * returning int64_t neighbor indices * @param[in] queries a device matrix view to a row-major matrix [n_queries, index.dim()] * @param[out] neighbors a device matrix view to the indices of the neighbors in the source dataset * [n_queries, k] @@ -1703,7 +1717,7 @@ void search(raft::resources const& res, const cuvs::neighbors::filtering::base_filter& sample_filter = cuvs::neighbors::filtering::none_sample_filter{}); -// vpq_f32_index overloads (uint32_t neighbor indices) +// Indexes over a PQ-compressed dataset with f32 codebooks (uint32_t neighbor indices) /** * @brief Search ANN using the constructed index. * @@ -1711,8 +1725,8 @@ void search(raft::resources const& res, * * @param[in] res raft resources * @param[in] params configure the search - * @param[in] index pre-built vpq_f32_index (CAGRA-Q, VPQ f32-compressed dataset) with uint32_t - * neighbor indices + * @param[in] index a pre-built CAGRA index over a PQ-compressed dataset with f32 codebooks, + * returning uint32_t neighbor indices * @param[in] queries a device matrix view to a row-major matrix [n_queries, index.dim()] * @param[out] neighbors a device matrix view to the indices of the neighbors in the source dataset * [n_queries, k] @@ -1739,8 +1753,8 @@ void search(raft::resources const& res, * * @param[in] res raft resources * @param[in] params configure the search - * @param[in] index pre-built vpq_f32_index (CAGRA-Q, VPQ f32-compressed dataset) with uint32_t - * neighbor indices + * @param[in] index a pre-built CAGRA index over a PQ-compressed dataset with f32 codebooks, + * returning uint32_t neighbor indices * @param[in] queries a device matrix view to a row-major matrix [n_queries, index.dim()] * @param[out] neighbors a device matrix view to the indices of the neighbors in the source dataset * [n_queries, k] @@ -1767,8 +1781,8 @@ void search(raft::resources const& res, * * @param[in] res raft resources * @param[in] params configure the search - * @param[in] index pre-built vpq_f32_index (CAGRA-Q, VPQ f32-compressed dataset) with uint32_t - * neighbor indices + * @param[in] index a pre-built CAGRA index over a PQ-compressed dataset with f32 codebooks, + * returning uint32_t neighbor indices * @param[in] queries a device matrix view to a row-major matrix [n_queries, index.dim()] * @param[out] neighbors a device matrix view to the indices of the neighbors in the source dataset * [n_queries, k] @@ -1795,8 +1809,8 @@ void search(raft::resources const& res, * * @param[in] res raft resources * @param[in] params configure the search - * @param[in] index pre-built vpq_f32_index (CAGRA-Q, VPQ f32-compressed dataset) with uint32_t - * neighbor indices + * @param[in] index a pre-built CAGRA index over a PQ-compressed dataset with f32 codebooks, + * returning uint32_t neighbor indices * @param[in] queries a device matrix view to a row-major matrix [n_queries, index.dim()] * @param[out] neighbors a device matrix view to the indices of the neighbors in the source dataset * [n_queries, k] @@ -1816,7 +1830,7 @@ void search(raft::resources const& res, const cuvs::neighbors::filtering::base_filter& sample_filter = cuvs::neighbors::filtering::none_sample_filter{}); -// vpq_f32_index overloads (int64_t neighbor indices) +// Indexes over a PQ-compressed dataset with f32 codebooks (int64_t neighbor indices) /** * @brief Search ANN using the constructed index. * @@ -1824,8 +1838,8 @@ void search(raft::resources const& res, * * @param[in] res raft resources * @param[in] params configure the search - * @param[in] index pre-built vpq_f32_index (CAGRA-Q, VPQ f32-compressed dataset) with int64_t - * neighbor indices + * @param[in] index a pre-built CAGRA index over a PQ-compressed dataset with f32 codebooks, + * returning int64_t neighbor indices * @param[in] queries a device matrix view to a row-major matrix [n_queries, index.dim()] * @param[out] neighbors a device matrix view to the indices of the neighbors in the source dataset * [n_queries, k] @@ -1852,8 +1866,8 @@ void search(raft::resources const& res, * * @param[in] res raft resources * @param[in] params configure the search - * @param[in] index pre-built vpq_f32_index (CAGRA-Q, VPQ f32-compressed dataset) with int64_t - * neighbor indices + * @param[in] index a pre-built CAGRA index over a PQ-compressed dataset with f32 codebooks, + * returning int64_t neighbor indices * @param[in] queries a device matrix view to a row-major matrix [n_queries, index.dim()] * @param[out] neighbors a device matrix view to the indices of the neighbors in the source dataset * [n_queries, k] @@ -1880,8 +1894,8 @@ void search(raft::resources const& res, * * @param[in] res raft resources * @param[in] params configure the search - * @param[in] index pre-built vpq_f32_index (CAGRA-Q, VPQ f32-compressed dataset) with int64_t - * neighbor indices + * @param[in] index a pre-built CAGRA index over a PQ-compressed dataset with f32 codebooks, + * returning int64_t neighbor indices * @param[in] queries a device matrix view to a row-major matrix [n_queries, index.dim()] * @param[out] neighbors a device matrix view to the indices of the neighbors in the source dataset * [n_queries, k] @@ -1908,8 +1922,8 @@ void search(raft::resources const& res, * * @param[in] res raft resources * @param[in] params configure the search - * @param[in] index pre-built vpq_f32_index (CAGRA-Q, VPQ f32-compressed dataset) with int64_t - * neighbor indices + * @param[in] index a pre-built CAGRA index over a PQ-compressed dataset with f32 codebooks, + * returning int64_t neighbor indices * @param[in] queries a device matrix view to a row-major matrix [n_queries, index.dim()] * @param[out] neighbors a device matrix view to the indices of the neighbors in the source dataset * [n_queries, k] @@ -2252,7 +2266,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 +2278,76 @@ enum class serialized_dataset_kind : std::uint32_t { host_padded = 3, /** Host-resident dataset using its standard row layout. */ host_standard = 4, + /** + * Device-resident PQ-compressed dataset, which CAGRA searches with quantized vectors. + * + * One kind for any codebook element type: the payload records its own, so f16 and f32 codebooks + * are told apart by the blob rather than by a second enumerator. + */ + device_pq = 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_ in -// detail/dataset_serialize.hpp (dense views use serialize_cagra_dense_dataset). +/** + * @brief What a serialized CAGRA index says about itself. + * + * @see read_serialized_header + */ +struct serialized_index_header { + /** Element type of the index that wrote the file, i.e. the `T` of its `index`. */ + cudaDataType_t dtype; + /** Which dataset, if any, travels with the graph. */ + serialized_dataset_kind dataset_kind; +}; + +/** + * Read what a serialized index holds, without loading it. + * + * An index carries its dataset kind in its type, so a caller loading a file someone else wrote has + * to know what is in it before it can name the type to load it into — `deserialize` rejects a file + * whose dataset kind does not match the index it was handed. This answers that question first. + * + * @code{.cpp} + * auto header = cagra::read_serialized_header(res, "index.bin"); + * if (header.dataset_kind == cagra::serialized_dataset_kind::device_pq) { + * cagra::vpq_f16_index index{res}; + * std::unique_ptr> rows; + * cagra::deserialize(res, "index.bin", &index, &rows); + * ... + * } + * @endcode + * + * A caller who only wants the graph does not need this: `deserialize` without an `out_dataset` + * skips whatever dataset the file holds, for any index type. + * + * @param[in] res raft resources + * @param[in] filename the file to inspect + * @return what the file records about itself + */ +auto read_serialized_header(raft::resources const& res, const std::string& filename) + -> serialized_index_header; + +/** + * @copydoc read_serialized_header + * + * The stream is left where it was found, so it can be passed straight to `deserialize`. It has to + * be seekable for that reason. + * + * @param[in] res raft resources + * @param[in] is input stream, opened in binary mode and positioned at the start of the index + * @return what the stream records about itself + */ +auto read_serialized_header(raft::resources const& res, std::istream& is) + -> serialized_index_header; + +// 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_ in detail/dataset_serialize.hpp (dense views use +// serialize_cagra_dense_dataset, PQ-compressed ones serialize_pq_dataset). /** * Save the index to file. @@ -2824,6 +2898,138 @@ void deserialize(raft::resources const& handle, std::unique_ptr>* out_dataset = nullptr); +/** + * @brief Save a CAGRA index over a PQ-compressed dataset, with or without those rows. + * + * The compressed rows can travel with the index, so that a deserialized index searches without + * the dense dataset it was compressed from and without retraining the codebooks. Passing + * `include_dataset = false` leaves them out, for a caller who keeps their own PQ-quantized dataset + * and wants nothing from the file but the graph. + * + * @param[in] handle the raft handle + * @param[in] filename the file to write, truncated if it exists + * @param[in] index the index to save + * @param[in] include_dataset whether the compressed rows are written with the graph + */ +void serialize(raft::resources const& handle, + const std::string& filename, + const cuvs::neighbors::cagra::vpq_f16_index& index, + bool include_dataset = true); + +/** + * @brief Load a CAGRA index over a PQ-compressed dataset, with those rows if the file has + * them. + * + * As everywhere else the index holds a view, so `deserialize` hands back the owning dataset through + * `out_dataset`, which has to stay alive for as long as the index is used. + * + * `out_dataset` is optional here as it is everywhere else. Leaving it out loads the graph and skips + * whatever rows the file holds; the index then has nothing to search until + * `update_device_dataset_same_layout` gives it some, from the file or from the caller's own + * compressed dataset. + * + * @param[in] handle the raft handle + * @param[in] filename the file that stores the index + * @param[out] index the index to load into + * @param[out] out_dataset receives the compressed rows the file held, if any and if not null + */ +void deserialize( + raft::resources const& handle, + const std::string& filename, + cuvs::neighbors::cagra::vpq_f16_index* index, + std::unique_ptr>* out_dataset = nullptr); + +/** @copydoc serialize */ +void serialize(raft::resources const& handle, + std::ostream& os, + const cuvs::neighbors::cagra::vpq_f16_index& index, + bool include_dataset = true); + +/** @copydoc deserialize */ +void deserialize( + raft::resources const& handle, + std::istream& is, + cuvs::neighbors::cagra::vpq_f16_index* index, + std::unique_ptr>* out_dataset = nullptr); + +/** @copydoc serialize */ +void serialize(raft::resources const& handle, + const std::string& filename, + const cuvs::neighbors::cagra::vpq_f16_index& index, + bool include_dataset = true); + +/** @copydoc deserialize */ +void deserialize( + raft::resources const& handle, + const std::string& filename, + cuvs::neighbors::cagra::vpq_f16_index* index, + std::unique_ptr>* out_dataset = nullptr); + +/** @copydoc serialize */ +void serialize(raft::resources const& handle, + std::ostream& os, + const cuvs::neighbors::cagra::vpq_f16_index& index, + bool include_dataset = true); + +/** @copydoc deserialize */ +void deserialize( + raft::resources const& handle, + std::istream& is, + cuvs::neighbors::cagra::vpq_f16_index* index, + std::unique_ptr>* out_dataset = nullptr); + +/** @copydoc serialize */ +void serialize(raft::resources const& handle, + const std::string& filename, + const cuvs::neighbors::cagra::vpq_f16_index& index, + bool include_dataset = true); + +/** @copydoc deserialize */ +void deserialize( + raft::resources const& handle, + const std::string& filename, + cuvs::neighbors::cagra::vpq_f16_index* index, + std::unique_ptr>* out_dataset = nullptr); + +/** @copydoc serialize */ +void serialize(raft::resources const& handle, + std::ostream& os, + const cuvs::neighbors::cagra::vpq_f16_index& index, + bool include_dataset = true); + +/** @copydoc deserialize */ +void deserialize( + raft::resources const& handle, + std::istream& is, + cuvs::neighbors::cagra::vpq_f16_index* index, + std::unique_ptr>* out_dataset = nullptr); + +/** @copydoc serialize */ +void serialize(raft::resources const& handle, + const std::string& filename, + const cuvs::neighbors::cagra::vpq_f16_index& index, + bool include_dataset = true); + +/** @copydoc deserialize */ +void deserialize( + raft::resources const& handle, + const std::string& filename, + cuvs::neighbors::cagra::vpq_f16_index* index, + std::unique_ptr>* out_dataset = nullptr); + +/** @copydoc serialize */ +void serialize(raft::resources const& handle, + std::ostream& os, + const cuvs::neighbors::cagra::vpq_f16_index& index, + bool include_dataset = true); + +/** @copydoc deserialize */ +void deserialize( + raft::resources const& handle, + std::istream& is, + cuvs::neighbors::cagra::vpq_f16_index* index, + std::unique_ptr>* out_dataset = nullptr); + /** @copydoc serialize */ void serialize(raft::resources const& handle, const std::string& filename, diff --git a/cpp/include/cuvs/preprocessing/quantize/pq.hpp b/cpp/include/cuvs/preprocessing/quantize/pq.hpp index 112341f2ad..232ecc61aa 100644 --- a/cpp/include/cuvs/preprocessing/quantize/pq.hpp +++ b/cpp/include/cuvs/preprocessing/quantize/pq.hpp @@ -14,6 +14,9 @@ #include #include +#include +#include +#include #include #include @@ -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 } } +/** 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 + * #include + * + * // 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> 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& 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& 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>* 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>* out_dataset); + /** @} */ // end of group product } // namespace pq diff --git a/cpp/src/neighbors/cagra.cpp b/cpp/src/neighbors/cagra.cpp index 1de6d5a669..3c03abc0e5 100644 --- a/cpp/src/neighbors/cagra.cpp +++ b/cpp/src/neighbors/cagra.cpp @@ -8,6 +8,15 @@ #include +#include +#include +#include + +#include +#include +#include +#include + namespace cuvs::neighbors::cagra { graph_build_params_t index_params::graph_build_heuristic(raft::matrix_extent dataset, @@ -67,4 +76,75 @@ cagra::index_params index_params::from_hnsw_params(raft::matrix_extent 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()`, 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(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(res, is); + RAFT_EXPECTS(kind_raw <= static_cast(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(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 diff --git a/cpp/src/neighbors/cagra_serialize.cuh b/cpp/src/neighbors/cagra_serialize.cuh index 83d047560b..9d7614e498 100644 --- a/cpp/src/neighbors/cagra_serialize.cuh +++ b/cpp/src/neighbors/cagra_serialize.cuh @@ -155,6 +155,43 @@ namespace cuvs::neighbors::cagra { cuvs::neighbors::cagra::detail::deserialize(handle, is, index, out_dataset); \ } \ \ + void serialize(raft::resources const& handle, \ + const std::string& filename, \ + const cuvs::neighbors::cagra::vpq_f16_index& index, \ + bool include_dataset) \ + { \ + cuvs::neighbors::cagra::detail::serialize( \ + handle, filename, index, include_dataset); \ + } \ + \ + void deserialize( \ + raft::resources const& handle, \ + const std::string& filename, \ + cuvs::neighbors::cagra::vpq_f16_index* index, \ + std::unique_ptr>* out_dataset) \ + { \ + cuvs::neighbors::cagra::detail::deserialize( \ + handle, filename, index, out_dataset); \ + } \ + \ + void serialize(raft::resources const& handle, \ + std::ostream& os, \ + const cuvs::neighbors::cagra::vpq_f16_index& index, \ + bool include_dataset) \ + { \ + cuvs::neighbors::cagra::detail::serialize( \ + handle, os, index, include_dataset); \ + } \ + \ + void deserialize( \ + raft::resources const& handle, \ + std::istream& is, \ + cuvs::neighbors::cagra::vpq_f16_index* index, \ + std::unique_ptr>* out_dataset) \ + { \ + cuvs::neighbors::cagra::detail::deserialize(handle, is, index, out_dataset); \ + } \ + \ void serialize_to_hnswlib( \ raft::resources const& handle, \ std::ostream& os, \ diff --git a/cpp/src/neighbors/cagra_serialize_inst.cu.in b/cpp/src/neighbors/cagra_serialize_inst.cu.in index 3d34adb36f..21993e22a5 100644 --- a/cpp/src/neighbors/cagra_serialize_inst.cu.in +++ b/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; using inst_device_standard_view_t = cuvs::neighbors::device_standard_dataset_view; +// `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; } // namespace @@ -21,6 +24,8 @@ extern template void index::compute raft::resources const&); extern template void index::compute_dataset_norms_( raft::resources const&); +extern template void index::compute_dataset_norms_( + raft::resources const&); CUVS_INST_CAGRA_SERIALIZE(data_t); diff --git a/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh b/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh index f2e0c4f07b..e985a6bdd8 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh +++ b/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) { return kind::host_standard; + } else if constexpr (cuvs::neighbors::is_device_vpq_dataset_view_v) { + // 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(kind::host_standard); + return raw <= static_cast(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) { neighbors::detail::serialize_cagra_dense_dataset(res, os, index_.dataset()); + } else if constexpr (cuvs::neighbors::is_device_vpq_f16_dataset_view_v) { + // 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(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 dataset_owner{}; if (has_dataset) { if (out_dataset == nullptr) { - cuvs::neighbors::detail::skip_dense_dataset(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(res, is); } else { auto const expected_kind = serialized_dataset_kind_for_view(); RAFT_EXPECTS( @@ -419,6 +432,8 @@ void deserialize( } else if constexpr (cuvs::neighbors::is_host_standard_dataset_view_v) { dataset_owner = cuvs::neighbors::detail::deserialize_host_standard_dataset(res, is); + } else if constexpr (cuvs::neighbors::is_device_vpq_f16_dataset_view_v) { + dataset_owner = cuvs::neighbors::detail::deserialize_pq_dataset(res, is); } else { static_assert(sizeof(DatasetViewT) == 0, "deserialize: dataset deserialization is not implemented for this view"); diff --git a/cpp/src/neighbors/detail/dataset_serialize.hpp b/cpp/src/neighbors/detail/dataset_serialize.hpp index 05c71e0213..36cc982efb 100644 --- a/cpp/src/neighbors/detail/dataset_serialize.hpp +++ b/cpp/src/neighbors/detail/dataset_serialize.hpp @@ -188,45 +188,70 @@ auto deserialize_dense_payload_metadata(raft::resources const& res, std::istream return {n_rows, dim, stride, elements}; } -template -void skip_dense_payload(raft::resources const& res, std::istream& is) +/** Advance past `count` bytes, by seeking where the stream allows it and by reading them off where + * it does not. */ +inline void skip_bytes(std::istream& is, std::size_t count, char const* context) { - auto const metadata = deserialize_dense_payload_metadata(res, is); - RAFT_EXPECTS(metadata.elements <= std::numeric_limits::max() / sizeof(DataT), - "skip_dense_payload: byte count overflow"); - auto remaining = metadata.elements * sizeof(DataT); - using pos_type = std::istream::pos_type; using off_type = std::istream::off_type; auto* buffer = is.rdbuf(); auto const invalid_position = pos_type{off_type{-1}}; auto const current = buffer->pubseekoff(0, std::ios_base::cur, std::ios_base::in); if (current != invalid_position && - remaining <= static_cast(std::numeric_limits::max())) { + count <= static_cast(std::numeric_limits::max())) { auto const end = buffer->pubseekoff(0, std::ios_base::end, std::ios_base::in); if (end != invalid_position) { auto const available = end - current; - RAFT_EXPECTS(available >= 0 && static_cast(available) >= remaining, - "skip_dense_payload: truncated payload"); + RAFT_EXPECTS(available >= 0 && static_cast(available) >= count, + "%s: truncated payload", + context); auto const next = - buffer->pubseekpos(current + static_cast(remaining), std::ios_base::in); - RAFT_EXPECTS(next != invalid_position, "skip_dense_payload: failed to seek past payload"); + buffer->pubseekpos(current + static_cast(count), std::ios_base::in); + RAFT_EXPECTS(next != invalid_position, "%s: failed to seek past payload", context); return; } RAFT_EXPECTS(buffer->pubseekpos(current, std::ios_base::in) != invalid_position, - "skip_dense_payload: failed to restore stream position"); + "%s: failed to restore stream position", + context); } std::array discard_buffer{}; - while (remaining > 0) { - auto const chunk = std::min(remaining, discard_buffer.size()); + while (count > 0) { + auto const chunk = std::min(count, discard_buffer.size()); is.read(discard_buffer.data(), static_cast(chunk)); - RAFT_EXPECTS(static_cast(is.gcount()) == chunk, - "skip_dense_payload: truncated payload"); - remaining -= chunk; + RAFT_EXPECTS(static_cast(is.gcount()) == chunk, "%s: truncated payload", context); + count -= chunk; } } +/** Advance past one `raft::serialize_mdspan` payload, whose NumPy header states its own size. */ +inline void skip_serialized_mdspan(std::istream& is, char const* context) +{ + auto const header = raft::numpy_serializer::read_header(is); + std::size_t elements = 1; + for (auto const extent : header.shape) { + auto const dimension = static_cast(extent); + RAFT_EXPECTS(dimension == 0 || elements <= std::numeric_limits::max() / dimension, + "%s: element count overflow", + context); + elements *= dimension; + } + auto const itemsize = static_cast(header.dtype.itemsize); + RAFT_EXPECTS(itemsize == 0 || elements <= std::numeric_limits::max() / itemsize, + "%s: byte count overflow", + context); + skip_bytes(is, elements * itemsize, context); +} + +template +void skip_dense_payload(raft::resources const& res, std::istream& is) +{ + auto const metadata = deserialize_dense_payload_metadata(res, is); + RAFT_EXPECTS(metadata.elements <= std::numeric_limits::max() / sizeof(DataT), + "skip_dense_payload: byte count overflow"); + skip_bytes(is, metadata.elements * sizeof(DataT), "skip_dense_payload"); +} + template auto deserialize_device_dense(raft::resources const& res, std::istream& is) -> std::unique_ptr @@ -279,6 +304,40 @@ auto deserialize_host_dense(raft::resources const& res, std::istream& is) return std::make_unique(std::move(storage), metadata.dim); } +/** VPQ codebooks are floating point; the encoded rows are always uint8 and carry no dtype. */ +template +constexpr auto pq_wire_dtype() -> cudaDataType_t +{ + static_assert(std::is_same_v || std::is_same_v, + "serialize_pq: codebook element type must be float or half"); + return std::is_same_v ? CUDA_R_16F : CUDA_R_32F; +} + +/** + * Write the payload of a VPQ dataset: six scalars followed by the two codebooks and the encoded + * rows. + * + * Stays on `raft::serialize_mdspan` rather than the `write_dense_bytes` scheme used by the dense + * path above, because `deserialize_vpq` reads with `raft::deserialize_mdspan`, which expects the + * NumPy header that helper embeds per matrix. The scalar types must also match the reader exactly: + * `n_rows` is `IdxT` and the remaining five are `uint32_t`. + */ +template +void serialize_pq(raft::resources const& res, + device_vpq_dataset const& dataset, + std::ostream& os) +{ + raft::serialize_scalar(res, os, dataset.n_rows()); + raft::serialize_scalar(res, os, dataset.dim()); + raft::serialize_scalar(res, os, dataset.vq_n_centers()); + raft::serialize_scalar(res, os, dataset.pq_n_centers()); + raft::serialize_scalar(res, os, dataset.pq_len()); + raft::serialize_scalar(res, os, dataset.encoded_row_length()); + raft::serialize_mdspan(res, os, raft::make_const_mdspan(dataset.vq_code_book.view())); + raft::serialize_mdspan(res, os, raft::make_const_mdspan(dataset.pq_code_book.view())); + raft::serialize_mdspan(res, os, raft::make_const_mdspan(dataset.data.view())); +} + template auto deserialize_vpq(raft::resources const& res, std::istream& is) -> std::unique_ptr> @@ -305,6 +364,41 @@ auto deserialize_vpq(raft::resources const& res, std::istream& is) std::move(vq_code_book), std::move(pq_code_book), std::move(data)); } +/** + * Write a self-describing VPQ dataset blob: tag + codebook dtype + payload. + * + * The tag and dtype are deliberately written here rather than inside `serialize_pq`, mirroring how + * `serialize_cagra_dense_dataset` wraps the dense payload, so that a reader can identify the blob + * before committing to a `DataT`. + */ +template +void serialize_pq_dataset(raft::resources const& res, + device_vpq_dataset const& dataset, + std::ostream& os) +{ + raft::serialize_scalar(res, os, kSerializeVPQDataset); + raft::serialize_scalar(res, os, pq_wire_dtype()); + serialize_pq(res, dataset, os); +} + +/** Read a blob written by `serialize_pq_dataset`, validating the tag and codebook dtype. */ +template +auto deserialize_pq_dataset(raft::resources const& res, std::istream& is) + -> std::unique_ptr> +{ + const auto tag = raft::deserialize_scalar(res, is); + RAFT_EXPECTS(tag == kSerializeVPQDataset, + "deserialize_pq_dataset: expected VPQ tag (%u), got %u", + static_cast(kSerializeVPQDataset), + static_cast(tag)); + const auto dtype = raft::deserialize_scalar(res, is); + RAFT_EXPECTS(dtype == pq_wire_dtype(), + "deserialize_pq_dataset: codebook dtype (%d) does not match expected (%d)", + static_cast(dtype), + static_cast(pq_wire_dtype())); + return deserialize_vpq(res, is); +} + template auto deserialize_dense_dataset(raft::resources const& res, std::istream& is) -> std::unique_ptr @@ -336,23 +430,64 @@ auto deserialize_dense_dataset(raft::resources const& res, std::istream& is) } } -template -void skip_dense_dataset(raft::resources const& res, std::istream& is) +/** Advance past the element dtype and strided payload that follow a dense dataset tag. */ +template +void skip_dense_dtype_and_payload(raft::resources const& res, std::istream& is) +{ + const auto dtype = raft::deserialize_scalar(res, is); + switch (dtype) { + case CUDA_R_32F: return skip_dense_payload(res, is); + case CUDA_R_16F: return skip_dense_payload(res, is); + case CUDA_R_8I: return skip_dense_payload(res, is); + case CUDA_R_8U: return skip_dense_payload(res, is); + default: + RAFT_FAIL("skip_dataset: unsupported dense element dtype (%d)", static_cast(dtype)); + } +} + +/** + * Advance past the codebook dtype and payload that follow a VPQ dataset tag. + * + * The payload is the six scalars, then the two codebooks and the encoded rows. Their element types + * are not needed: each of the three matrices carries a NumPy header stating its own size. + */ +template +void skip_pq_dtype_and_payload(raft::resources const& res, std::istream& is) +{ + const auto dtype = raft::deserialize_scalar(res, is); + RAFT_EXPECTS(dtype == CUDA_R_16F || dtype == CUDA_R_32F, + "skip_dataset: unsupported VPQ codebook dtype (%d)", + static_cast(dtype)); + static_cast(raft::deserialize_scalar(res, is)); // n_rows + static_cast(raft::deserialize_scalar(res, is)); // dim + static_cast(raft::deserialize_scalar(res, is)); // vq_n_centers + static_cast(raft::deserialize_scalar(res, is)); // pq_n_centers + static_cast(raft::deserialize_scalar(res, is)); // pq_len + static_cast(raft::deserialize_scalar(res, is)); // encoded_row_length + skip_serialized_mdspan(is, "skip_dataset: VPQ vq_code_book"); + skip_serialized_mdspan(is, "skip_dataset: VPQ pq_code_book"); + skip_serialized_mdspan(is, "skip_dataset: VPQ encoded rows"); +} + +/** + * Advance past a dataset blob of any kind, leaving the stream on whatever follows it. + * + * How many bytes to skip is worked out from the tag and dtype the blob begins with, not from a type + * the caller names, so that dropping a dataset does not require knowing what it was. `IdxT` is the + * index type the blob was written with, which for a CAGRA index file is always int64_t. + */ +template +void skip_dataset(raft::resources const& res, std::istream& is) { const auto tag = raft::deserialize_scalar(res, is); - RAFT_EXPECTS(tag == kSerializeStridedDataset, - "skip_dense_dataset: expected strided tag, got %u", - static_cast(tag)); - const auto dtype = raft::deserialize_scalar(res, is); - constexpr cudaDataType_t expected_dtype = std::is_same_v ? CUDA_R_32F - : std::is_same_v ? CUDA_R_16F - : std::is_same_v ? CUDA_R_8I - : CUDA_R_8U; - RAFT_EXPECTS(dtype == expected_dtype, - "skip_dense_dataset: serialized dtype (%d) does not match expected (%d)", - static_cast(dtype), - static_cast(expected_dtype)); - skip_dense_payload(res, is); + switch (tag) { + case kSerializeEmptyDataset: + static_cast(raft::deserialize_scalar(res, is)); // suggested_dim + return; + case kSerializeStridedDataset: return skip_dense_dtype_and_payload(res, is); + case kSerializeVPQDataset: return skip_pq_dtype_and_payload(res, is); + default: RAFT_FAIL("skip_dataset: unknown dataset tag %u", static_cast(tag)); + } } // Reads tag + dtype prefix, validates they match DataT, and returns the requested concrete diff --git a/cpp/src/preprocessing/quantize/pq.cu b/cpp/src/preprocessing/quantize/pq.cu index 20b8f21d36..fc2c16d317 100644 --- a/cpp/src/preprocessing/quantize/pq.cu +++ b/cpp/src/preprocessing/quantize/pq.cu @@ -3,13 +3,20 @@ * SPDX-License-Identifier: Apache-2.0 */ +#include "../../neighbors/detail/dataset_serialize.hpp" +#include "../../util/serialize_validation.hpp" #include "./detail/pq.cuh" #include +#include #include #include +#include +#include +#include + namespace cuvs::preprocessing::quantize::pq { #define CUVS_INST_QUANTIZATION(T, QuantI) \ @@ -76,6 +83,55 @@ CUVS_INST_VPQ_BUILD(uint8_t); #undef CUVS_INST_VPQ_BUILD +void serialize(raft::resources const& res, + const cuvs::neighbors::device_vpq_dataset& dataset, + std::ostream& os) +{ + // Same file preamble as cagra::serialize. The nested blob carries only a kind tag and dtype, + // matching serialize_cagra_dense_dataset, because a nested blob relies on its enclosing file for + // the version; a standalone .vpq has no enclosing file, so the version is written here. + std::string dtype_string = raft::numpy_serializer::get_numpy_dtype().to_string(); + dtype_string.resize(4); + os << dtype_string; + raft::serialize_scalar(res, os, pq_serialization_version); + ::cuvs::neighbors::detail::serialize_pq_dataset(res, dataset, os); +} + +void serialize(raft::resources const& res, + const cuvs::neighbors::device_vpq_dataset& dataset, + const std::string& filename) +{ + std::ofstream os(filename, std::ios::out | std::ios::binary | std::ios::trunc); + RAFT_EXPECTS(os.good(), "pq::serialize: cannot open %s for writing", filename.c_str()); + serialize(res, dataset, os); +} + +void deserialize(raft::resources const& res, + std::istream& is, + std::unique_ptr>* out_dataset) +{ + RAFT_EXPECTS(out_dataset != nullptr, "pq::deserialize: out_dataset must not be null"); + char dtype_string[4]; + RAFT_EXPECTS(is.read(dtype_string, 4), "pq::deserialize: failed to read the dtype prefix"); + RAFT_EXPECTS(cuvs::util::validate_serialized_dtype(dtype_string, sizeof(dtype_string)), + "pq::deserialize: dtype prefix does not match a VPQ dataset with half codebooks"); + auto const version = raft::deserialize_scalar(res, is); + RAFT_EXPECTS(version == pq_serialization_version, + "pq::deserialize: serialization version mismatch, expected %d, got %d", + pq_serialization_version, + version); + *out_dataset = ::cuvs::neighbors::detail::deserialize_pq_dataset(res, is); +} + +void deserialize(raft::resources const& res, + const std::string& filename, + std::unique_ptr>* out_dataset) +{ + std::ifstream is(filename, std::ios::in | std::ios::binary); + RAFT_EXPECTS(is.good(), "pq::deserialize: cannot open %s for reading", filename.c_str()); + deserialize(res, is, out_dataset); +} + namespace detail { template diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index b4a657c90c..b5df68f9b1 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -199,7 +199,7 @@ ConfigureTest( ConfigureTest( NAME NEIGHBORS_ANN_CAGRA_FLOAT_UINT32_TEST - PATH neighbors/ann_cagra/test_float_uint32_t.cu + PATH neighbors/ann_cagra/test_float_uint32_t.cu neighbors/ann_cagra/test_pq_serialize.cu GPUS 1 PERCENT 100 ) @@ -415,6 +415,7 @@ ConfigureTest( preprocessing/binary_quantization.cu preprocessing/spectral_embedding.cu preprocessing/product_quantization.cu + preprocessing/pq_serialization.cu preprocessing/pca.cu GPUS 1 PERCENT 100 diff --git a/cpp/tests/neighbors/ann_cagra/test_pq_serialize.cu b/cpp/tests/neighbors/ann_cagra/test_pq_serialize.cu new file mode 100644 index 0000000000..7152c0a0f7 --- /dev/null +++ b/cpp/tests/neighbors/ann_cagra/test_pq_serialize.cu @@ -0,0 +1,333 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +/* + * Serializing a CAGRA index whose dataset is PQ-compressed. + * + * Such an index cannot be saved by the dtype-templated suites in ann_cagra.cuh: its rows are VPQ + * codes rather than values of `DataT`, it only searches with `L2Expanded`, `pq_bits == 8` and + * `pq_len` in {2, 4, 8}, and it is assembled rather than built, since `cagra::build` produces dense + * indices only. The assembly here is the usual one: a graph from a dense build, a dataset + * compressed separately, and an index that views both. + * + * What is checked is that the compressed rows travel with the index, so a loaded index searches on + * its own without the dense dataset it came from and without retraining codebooks; that a caller + * after the graph alone can leave them in the file; that a caller can find out what a file holds + * before naming the index type to load it into; and that a dense reader refuses the file rather + * than misreading it. Fidelity of the dataset payload itself is covered by + * preprocessing/pq_serialization.cu. + */ + +#include + +#include +#include + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace cuvs::neighbors::cagra { + +using pq_dataset_t = cuvs::neighbors::device_vpq_dataset; + +namespace { + +constexpr int64_t kSearchK = 10; +constexpr uint32_t kGraphDegree = 32; + +auto compress(const raft::resources& res, + raft::device_matrix_view dataset, + uint32_t pq_dim) -> pq_dataset_t +{ + cuvs::neighbors::vpq_params params; + params.pq_dim = pq_dim; + params.pq_bits = 8; + params.vq_n_centers = 32; + params.kmeans_n_iters = 5; // Codebooks need to be well defined here, not optimal. + return cuvs::preprocessing::quantize::pq::make_vpq_dataset(res, params, dataset); +} + +/** A dense index built over the same rows, kept alive only to lend its graph. */ +auto build_graph_source(const raft::resources& res, + raft::device_matrix_view dataset) + -> device_standard_index +{ + index_params params; + params.metric = cuvs::distance::DistanceType::L2Expanded; + params.graph_degree = kGraphDegree; + params.intermediate_graph_degree = kGraphDegree * 2; + return cagra::build(res, params, cuvs::neighbors::make_device_standard_dataset_view(dataset)); +} + +/** Neighbour ids for `queries`, row-major [n_queries, kSearchK]. */ +template +auto neighbor_ids(const raft::resources& res, + const IndexT& idx, + raft::device_matrix_view queries) -> std::vector +{ + const auto n_queries = queries.extent(0); + auto neighbors = raft::make_device_matrix(res, n_queries, kSearchK); + auto distances = raft::make_device_matrix(res, n_queries, kSearchK); + + search_params params; + params.itopk_size = 64; + search(res, params, idx, queries, neighbors.view(), distances.view()); + + std::vector ids(static_cast(n_queries * kSearchK)); + raft::copy(ids.data(), neighbors.data_handle(), ids.size(), raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + return ids; +} + +/** + * Fraction of queries that retrieve their own row, where the queries are dataset rows. + * + * A sanity signal rather than a quality metric: it is here so that comparing neighbour ids before + * and after a round trip compares useful answers rather than two copies of the same nonsense. + */ +auto self_recall_at_1(const std::vector& ids) -> double +{ + const size_t n_queries = ids.size() / kSearchK; + size_t hits = 0; + for (size_t q = 0; q < n_queries; q++) { + hits += static_cast(ids[q * kSearchK] == static_cast(q)); + } + return static_cast(hits) / static_cast(n_queries); +} + +} // namespace + +/** + * An index over compressed rows is serialized with those rows. The ownership split is the usual + * one: the file yields an owning dataset, the index only views it. + */ +class CagraPqSerializeTest : public ::testing::Test { + protected: + void SetUp() override + { + dataset_.emplace(raft::make_device_matrix(res_, n_rows, dim)); + auto labels = raft::make_device_vector(res_, n_rows); + raft::random::make_blobs(res_, + dataset_->view(), + labels.view(), + 5, // clusters + std::nullopt, // random centers + std::nullopt, // scalar std + 1.0F, // cluster std + true, // shuffle + -10.0F, // center box min + 10.0F, // center box max + 1234ULL); + raft::resource::sync_stream(res_); + } + + void TearDown() override + { + dataset_.reset(); + raft::resource::sync_stream(res_); + } + + auto dataset() -> raft::device_matrix_view + { + return raft::make_const_mdspan(dataset_->view()); + } + + /** The first rows of the dataset, reused as queries. */ + auto queries(int64_t n_queries) -> raft::device_matrix_view + { + return raft::make_device_matrix_view( + dataset_->data_handle(), std::min(n_queries, dataset_->extent(0)), dataset_->extent(1)); + } + + /** An index viewing `compressed` and the graph of `graph_source`; both must outlive it. */ + auto assemble(const pq_dataset_t& compressed, + const device_standard_index& graph_source) -> vpq_f16_index + { + return vpq_f16_index{res_, + cuvs::distance::DistanceType::L2Expanded, + compressed.as_dataset_view(), + graph_source.graph()}; + } + + static constexpr int64_t n_rows = 2000; + static constexpr int64_t dim = 128; + static constexpr uint32_t pq_dim = 32; // pq_len 4 + + raft::resources res_; + std::optional> dataset_ = std::nullopt; +}; + +TEST_F(CagraPqSerializeTest, RoundTripsThroughAFileWithItsDataset) +{ + auto compressed = compress(res_, dataset(), pq_dim); + auto graph_source = build_graph_source(res_, dataset()); + auto idx = assemble(compressed, graph_source); + + auto before = neighbor_ids(res_, idx, queries(500)); + ASSERT_GT(self_recall_at_1(before), 0.5); + + std::stringstream stored; + cagra::serialize(res_, stored, idx); + + vpq_f16_index restored{res_}; + std::unique_ptr owner; + cagra::deserialize(res_, stored, &restored, &owner); + + ASSERT_NE(owner, nullptr); + EXPECT_EQ(owner->n_rows(), compressed.n_rows()); + EXPECT_EQ(owner->dim(), compressed.dim()); + EXPECT_EQ(owner->pq_len(), compressed.pq_len()); + EXPECT_EQ(owner->pq_bits(), compressed.pq_bits()); + EXPECT_EQ(owner->vq_n_centers(), compressed.vq_n_centers()); + EXPECT_EQ(owner->encoded_row_length(), compressed.encoded_row_length()); + + ASSERT_EQ(restored.size(), idx.size()); + ASSERT_EQ(restored.dim(), idx.dim()); + ASSERT_EQ(restored.graph_degree(), idx.graph_degree()); + EXPECT_EQ(restored.metric(), idx.metric()); + + // Same graph over the same rows, so the results are identical rather than merely comparable. + auto after = neighbor_ids(res_, restored, queries(500)); + ASSERT_EQ(after.size(), before.size()); + size_t mismatches = 0; + for (size_t i = 0; i < before.size(); i++) { + mismatches += static_cast(after[i] != before[i]); + } + EXPECT_EQ(mismatches, 0u) << mismatches << " of " << before.size() << " neighbour ids changed"; +} + +TEST_F(CagraPqSerializeTest, LoadsTheGraphAloneWhenNoOwnerIsAskedFor) +{ + auto compressed = compress(res_, dataset(), pq_dim); + auto graph_source = build_graph_source(res_, dataset()); + auto idx = assemble(compressed, graph_source); + auto before = neighbor_ids(res_, idx, queries(500)); + + std::stringstream stored; + cagra::serialize(res_, stored, idx); + + // A caller who wants the graph alone says nothing about the rows: not their type, not even + // whether the file has any. They are skipped here, and the graph arrives without them. + vpq_f16_index restored{res_}; + cagra::deserialize(res_, stored, &restored); + EXPECT_EQ(restored.size(), idx.size()); + EXPECT_EQ(restored.graph_degree(), idx.graph_degree()); + + // Skipping has to consume the payload to the byte, or anything the format writes after the rows + // would be read as garbage. Nothing follows them here, so the stream has to be spent. + EXPECT_EQ(stored.peek(), std::char_traits::eof()); + + // And the graph is intact: give it rows again and it answers exactly as it did before. + restored.update_device_dataset_same_layout(res_, compressed.as_dataset_view()); + auto after = neighbor_ids(res_, restored, queries(500)); + ASSERT_EQ(after.size(), before.size()); + size_t mismatches = 0; + for (size_t i = 0; i < before.size(); i++) { + mismatches += static_cast(after[i] != before[i]); + } + EXPECT_EQ(mismatches, 0u) << mismatches << " of " << before.size() << " neighbour ids changed"; +} + +TEST_F(CagraPqSerializeTest, SkipsWhicheverDatasetTheFileHappensToHold) +{ + // The graph-only load asks nothing about the rows, so it also does not care that these are dense + // floats rather than the compressed rows this index type would view. + auto graph_source = build_graph_source(res_, dataset()); + std::stringstream stored; + cagra::serialize(res_, stored, graph_source); + + vpq_f16_index restored{res_}; + cagra::deserialize(res_, stored, &restored); + EXPECT_EQ(restored.size(), graph_source.size()); + EXPECT_EQ(restored.graph_degree(), graph_source.graph_degree()); + EXPECT_EQ(stored.peek(), std::char_traits::eof()); +} + +TEST_F(CagraPqSerializeTest, SerializesTheGraphAloneWhenAsked) +{ + auto compressed = compress(res_, dataset(), pq_dim); + auto graph_source = build_graph_source(res_, dataset()); + auto idx = assemble(compressed, graph_source); + + std::stringstream stored; + cagra::serialize(res_, stored, idx, /* include_dataset */ false); + + vpq_f16_index restored{res_}; + std::unique_ptr owner; + cagra::deserialize(res_, stored, &restored, &owner); + + // Nothing to own, and a graph that only update_device_dataset_same_layout() makes searchable + // again. + EXPECT_EQ(owner, nullptr); + EXPECT_EQ(restored.size(), idx.size()); + EXPECT_EQ(restored.graph_degree(), idx.graph_degree()); +} + +TEST_F(CagraPqSerializeTest, SaysWhatItHoldsBeforeItIsLoaded) +{ + auto compressed = compress(res_, dataset(), pq_dim); + auto graph_source = build_graph_source(res_, dataset()); + auto idx = assemble(compressed, graph_source); + + std::stringstream stored; + cagra::serialize(res_, stored, idx); + + // An index carries its dataset kind in its type, so a caller loading someone else's file has to + // be able to ask what is in it before naming the type to load it into. + auto header = cagra::read_serialized_header(res_, stored); + EXPECT_EQ(header.dtype, CUDA_R_32F); + EXPECT_EQ(header.dataset_kind, serialized_dataset_kind::device_pq); + + // Asking rewinds, so the same stream still loads. The index type names the dataset type, which is + // what owning_dataset_for_index_t is for. + vpq_f16_index restored{res_}; + std::unique_ptr> rows; + cagra::deserialize(res_, stored, &restored, &rows); + ASSERT_NE(rows, nullptr); + EXPECT_EQ(restored.size(), idx.size()); + + // A dense index records its own layout, and a graph-only file records no dataset at all. + std::stringstream dense; + cagra::serialize(res_, dense, graph_source); + EXPECT_EQ(cagra::read_serialized_header(res_, dense).dataset_kind, + serialized_dataset_kind::device_standard); + + std::stringstream graph_only; + cagra::serialize(res_, graph_only, idx, /* include_dataset */ false); + auto graph_only_header = cagra::read_serialized_header(res_, graph_only); + EXPECT_EQ(graph_only_header.dataset_kind, serialized_dataset_kind::none); + // Still float: the dtype is the index's element type, not its rows' storage. + EXPECT_EQ(graph_only_header.dtype, CUDA_R_32F); +} + +TEST_F(CagraPqSerializeTest, RejectsLoadingACompressedIndexAsDense) +{ + auto compressed = compress(res_, dataset(), pq_dim); + auto graph_source = build_graph_source(res_, dataset()); + auto idx = assemble(compressed, graph_source); + + std::stringstream stored; + cagra::serialize(res_, stored, idx); + + // The dtype prefix says float either way, so it is the recorded dataset kind that has to stop the + // dense reader from interpreting VPQ codes as rows of floats. + device_padded_index dense{res_}; + std::unique_ptr> dense_owner; + EXPECT_THROW(cagra::deserialize(res_, stored, &dense, &dense_owner), raft::exception); +} + +} // namespace cuvs::neighbors::cagra diff --git a/cpp/tests/preprocessing/pq_serialization.cu b/cpp/tests/preprocessing/pq_serialization.cu new file mode 100644 index 0000000000..90dd237271 --- /dev/null +++ b/cpp/tests/preprocessing/pq_serialization.cu @@ -0,0 +1,261 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "../neighbors/vpq_utils.cuh" +#include "../test_utils.cuh" + +#include + +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace cuvs::preprocessing::quantize::pq { + +using pq_dataset_t = cuvs::neighbors::device_vpq_dataset; + +struct PqSerializationInputs { + int64_t n_rows; + int64_t dim; + uint32_t pq_bits; + uint32_t pq_dim; + uint32_t vq_n_centers; // 0 lets the heuristic choose + uint64_t seed; +}; + +std::ostream& operator<<(std::ostream& os, const PqSerializationInputs& in) +{ + return os << "n_rows:" << in.n_rows << " dim:" << in.dim << " pq_bits:" << in.pq_bits + << " pq_dim:" << in.pq_dim << " vq_n_centers:" << in.vq_n_centers + << " seed:" << in.seed; +} + +template +auto to_host(const raft::resources& res, raft::device_matrix_view m) + -> std::vector +{ + std::vector host(static_cast(m.extent(0)) * static_cast(m.extent(1))); + raft::copy(host.data(), m.data_handle(), host.size(), raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + return host; +} + +/** Bitwise, not approximate: serialization is expected not to perturb a single bit. */ +template +void expect_same_bits(const raft::resources& res, + raft::device_matrix_view expected, + raft::device_matrix_view actual, + const char* what) +{ + ASSERT_EQ(expected.extent(0), actual.extent(0)) << what; + ASSERT_EQ(expected.extent(1), actual.extent(1)) << what; + const auto lhs = to_host(res, expected); + const auto rhs = to_host(res, actual); + EXPECT_EQ(0, std::memcmp(lhs.data(), rhs.data(), lhs.size() * sizeof(T))) << what; +} + +class PqSerializationTest : public ::testing::TestWithParam { + public: + PqSerializationTest() + : params_(::testing::TestWithParam::GetParam()), + dataset_(raft::make_device_matrix(res_, params_.n_rows, params_.dim)) + { + } + + protected: + void SetUp() override + { + auto labels = raft::make_device_vector(res_, params_.n_rows); + raft::random::make_blobs(res_, + dataset_.view(), + labels.view(), + 5, // clusters + std::nullopt, // random centers + std::nullopt, // scalar std + 1.0F, // cluster std + true, // shuffle + -10.0F, // center box min + 10.0F, // center box max + params_.seed); + raft::resource::sync_stream(res_); + } + + auto compress() -> pq_dataset_t + { + cuvs::neighbors::vpq_params vpq; + vpq.pq_bits = params_.pq_bits; + vpq.pq_dim = params_.pq_dim; + vpq.vq_n_centers = params_.vq_n_centers; + // The codebooks only have to be well defined here, not good, so keep training short. + vpq.kmeans_n_iters = 5; + return make_vpq_dataset(res_, vpq, raft::make_const_mdspan(dataset_.view())); + } + + void expect_equivalent(const pq_dataset_t& expected, const pq_dataset_t& actual) + { + ASSERT_EQ(expected.n_rows(), actual.n_rows()); + ASSERT_EQ(expected.dim(), actual.dim()); + ASSERT_EQ(expected.vq_n_centers(), actual.vq_n_centers()); + ASSERT_EQ(expected.pq_n_centers(), actual.pq_n_centers()); + ASSERT_EQ(expected.pq_len(), actual.pq_len()); + ASSERT_EQ(expected.encoded_row_length(), actual.encoded_row_length()); + ASSERT_EQ(expected.pq_bits(), actual.pq_bits()); + ASSERT_EQ(expected.pq_dim(), actual.pq_dim()); + + expect_same_bits(res_, + raft::make_const_mdspan(expected.vq_code_book.view()), + raft::make_const_mdspan(actual.vq_code_book.view()), + "vq_code_book"); + expect_same_bits(res_, + raft::make_const_mdspan(expected.pq_code_book.view()), + raft::make_const_mdspan(actual.pq_code_book.view()), + "pq_code_book"); + expect_same_bits(res_, + raft::make_const_mdspan(expected.data.view()), + raft::make_const_mdspan(actual.data.view()), + "encoded rows"); + } + + /** + * Decodes both datasets and compares the reconstructions, which checks that a kernel can consume + * the deserialized extents and strides rather than only that the numbers match. + */ + void expect_same_decoded(const pq_dataset_t& expected, const pq_dataset_t& actual) + { + if (expected.pq_bits() != 8) { return; } // decode_vpq_dataset implements pq_bits == 8 only + auto stream = raft::resource::get_cuda_stream(res_); + auto lhs = raft::make_device_matrix(res_, expected.n_rows(), expected.dim()); + auto rhs = raft::make_device_matrix(res_, actual.n_rows(), actual.dim()); + cuvs::neighbors::decode_vpq_dataset(lhs.view(), expected, stream); + cuvs::neighbors::decode_vpq_dataset(rhs.view(), actual, stream); + raft::resource::sync_stream(res_); + expect_same_bits(res_, + raft::make_const_mdspan(lhs.view()), + raft::make_const_mdspan(rhs.view()), + "decoded rows"); + } + + raft::resources res_; + PqSerializationInputs params_; + raft::device_matrix dataset_; +}; + +TEST_P(PqSerializationTest, RoundTrip) +{ + auto original = compress(); + + { + SCOPED_TRACE("through a stream"); + std::stringstream stream; + serialize(res_, original, stream); + std::unique_ptr restored; + deserialize(res_, stream, &restored); + ASSERT_NE(restored, nullptr); + expect_equivalent(original, *restored); + expect_same_decoded(original, *restored); + } + + { + SCOPED_TRACE("through a file"); + const std::string path = "cuvs_pq_serialization_test.bin"; + serialize(res_, original, path); + std::unique_ptr restored; + deserialize(res_, path, &restored); + std::remove(path.c_str()); + ASSERT_NE(restored, nullptr); + expect_equivalent(original, *restored); + } +} + +// Named for this suite rather than `inputs`: product_quantization.cu declares a variable of that +// name in this same namespace, which would collide under a unity build. +const std::vector pq_serialization_inputs = { + // pq_len = dim / pq_dim of 2, 4 and 8: the three values a CAGRA search accepts. + {1000, 64, 8, 32, 0, 42ULL}, + {1000, 128, 8, 32, 0, 42ULL}, + {1000, 256, 8, 32, 0, 42ULL}, + // An explicit VQ codebook size rather than the heuristic. + {2000, 128, 8, 64, 64, 42ULL}, + // pq_bits below 8 packs several codes per byte, so encoded_row_length stops being pq_dim. + {500, 96, 6, 24, 0, 42ULL}, + {500, 32, 4, 16, 0, 42ULL}, +}; + +INSTANTIATE_TEST_CASE_P(PqSerializationTests, + PqSerializationTest, + ::testing::ValuesIn(pq_serialization_inputs)); + +/** Writes the preamble that `serialize` emits, so only the field under test differs. */ +static void write_preamble(const raft::resources& res, std::ostream& os, int version) +{ + std::string dtype_string = raft::numpy_serializer::get_numpy_dtype().to_string(); + dtype_string.resize(4); + os << dtype_string; + raft::serialize_scalar(res, os, version); +} + +TEST(PqSerialization, RejectsEmptyStream) +{ + raft::resources res; + std::stringstream stream; + std::unique_ptr restored; + EXPECT_THROW(deserialize(res, stream, &restored), raft::exception); +} + +TEST(PqSerialization, RejectsForeignDtypePrefix) +{ + raft::resources res; + std::stringstream stream; + std::string dtype_string = raft::numpy_serializer::get_numpy_dtype().to_string(); + dtype_string.resize(4); + stream << dtype_string; + raft::serialize_scalar(res, stream, pq_serialization_version); + + std::unique_ptr restored; + EXPECT_THROW(deserialize(res, stream, &restored), raft::exception); +} + +TEST(PqSerialization, RejectsFutureVersion) +{ + raft::resources res; + std::stringstream stream; + write_preamble(res, stream, pq_serialization_version + 1); + + std::unique_ptr restored; + EXPECT_THROW(deserialize(res, stream, &restored), raft::exception); +} + +TEST(PqSerialization, RejectsTruncatedPayload) +{ + raft::resources res; + std::stringstream stream; + write_preamble(res, stream, pq_serialization_version); + // A correct preamble followed by nothing: the payload reader must fail rather than return a + // dataset built from whatever the scalars happened to deserialize to. + std::unique_ptr restored; + EXPECT_THROW(deserialize(res, stream, &restored), raft::exception); +} + +TEST(PqSerialization, RejectsNullOutParameter) +{ + raft::resources res; + std::stringstream stream; + write_preamble(res, stream, pq_serialization_version); + EXPECT_THROW(deserialize(res, stream, nullptr), raft::exception); +} + +} // namespace cuvs::preprocessing::quantize::pq