diff --git a/c/src/neighbors/cagra.cpp b/c/src/neighbors/cagra.cpp index 99e456e23c..42200258b6 100644 --- a/c/src/neighbors/cagra.cpp +++ b/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(kind::host_standard), + RAFT_EXPECTS(dataset_kind_raw <= static_cast(kind::device_vpq_f16), "Invalid serialized dataset kind %u in file %s", dataset_kind_raw, filename); return {output_dtype, static_cast(dataset_kind_raw)}; @@ -1058,6 +1058,12 @@ void dispatch_serialized_dataset_kind( fn.template operator()< cuvs::neighbors::device_padded_dataset_view>(); 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"); } } 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..56394b5398 100644 --- a/cpp/include/cuvs/neighbors/cagra.hpp +++ b/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_ 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_ 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>* 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& index, + bool include_dataset = true); + +void deserialize( + raft::resources const& handle, + const std::string& filename, + cuvs::neighbors::cagra::vpq_f16_index* index, + std::unique_ptr>* out_dataset); + +void serialize(raft::resources const& handle, + std::ostream& os, + const cuvs::neighbors::cagra::vpq_f16_index& index, + bool include_dataset = true); + +void deserialize( + raft::resources const& handle, + std::istream& is, + cuvs::neighbors::cagra::vpq_f16_index* index, + std::unique_ptr>* out_dataset); + +void serialize(raft::resources const& handle, + const std::string& filename, + const cuvs::neighbors::cagra::vpq_f16_index& index, + bool include_dataset = true); + +void deserialize( + raft::resources const& handle, + const std::string& filename, + cuvs::neighbors::cagra::vpq_f16_index* index, + std::unique_ptr>* out_dataset); + +void serialize(raft::resources const& handle, + std::ostream& os, + const cuvs::neighbors::cagra::vpq_f16_index& index, + bool include_dataset = true); + +void deserialize( + raft::resources const& handle, + std::istream& is, + cuvs::neighbors::cagra::vpq_f16_index* index, + std::unique_ptr>* out_dataset); + +void serialize(raft::resources const& handle, + const std::string& filename, + const cuvs::neighbors::cagra::vpq_f16_index& index, + bool include_dataset = true); + +void deserialize( + raft::resources const& handle, + const std::string& filename, + cuvs::neighbors::cagra::vpq_f16_index* index, + std::unique_ptr>* out_dataset); + +void serialize(raft::resources const& handle, + std::ostream& os, + const cuvs::neighbors::cagra::vpq_f16_index& index, + bool include_dataset = true); + +void deserialize( + raft::resources const& handle, + std::istream& is, + cuvs::neighbors::cagra::vpq_f16_index* index, + std::unique_ptr>* out_dataset); + +void serialize(raft::resources const& handle, + const std::string& filename, + const cuvs::neighbors::cagra::vpq_f16_index& index, + bool include_dataset = true); + +void deserialize( + raft::resources const& handle, + const std::string& filename, + cuvs::neighbors::cagra::vpq_f16_index* index, + std::unique_ptr>* out_dataset); + +void serialize(raft::resources const& handle, + std::ostream& os, + const cuvs::neighbors::cagra::vpq_f16_index& index, + bool include_dataset = true); + +void deserialize( + raft::resources const& handle, + std::istream& is, + cuvs::neighbors::cagra::vpq_f16_index* index, + std::unique_ptr>* out_dataset); + /** @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..c26a24fea6 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 @@ -331,6 +334,82 @@ template } } +/** 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 + * #include + * + * // 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> 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& 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& 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>* 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_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..58e555d17e 100644 --- a/cpp/src/neighbors/cagra_serialize_inst.cu.in +++ b/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; using inst_device_standard_view_t = cuvs::neighbors::device_standard_dataset_view; +using inst_vpq_f16_view_t = cuvs::neighbors::device_vpq_dataset_view; } // namespace @@ -21,6 +22,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..add4bb2532 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh @@ -60,6 +60,8 @@ 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_f16_dataset_view_v) { + return kind::device_vpq_f16; } else { static_assert(sizeof(DatasetViewT) == 0, "serialized_dataset_kind_for_view: unsupported dataset view type"); @@ -69,7 +71,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_vpq_f16); } /** @@ -123,9 +125,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_vpq_dataset(res, os, index_.dataset().dset()); } 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 +408,16 @@ void deserialize( std::unique_ptr dataset_owner{}; if (has_dataset) { if (out_dataset == nullptr) { - cuvs::neighbors::detail::skip_dense_dataset(res, is); + // Dropping the rows leaves a searchable index for a dense view, whose dataset can be + // reattached from the caller's own copy, but not for a VPQ one: the compressed rows exist + // nowhere else. Refuse rather than hand back an index that cannot answer a query, and skip + // the dense payload only when it is in fact dense. + if constexpr (cuvs::neighbors::is_vpq_dataset_view_v) { + RAFT_FAIL( + "cagra::deserialize: a VPQ index cannot be loaded without its dataset; pass out_dataset"); + } else { + cuvs::neighbors::detail::skip_dense_dataset(res, is); + } } else { auto const expected_kind = serialized_dataset_kind_for_view(); RAFT_EXPECTS( @@ -419,6 +435,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_vpq_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..6e73f36d10 100644 --- a/cpp/src/neighbors/detail/dataset_serialize.hpp +++ b/cpp/src/neighbors/detail/dataset_serialize.hpp @@ -279,6 +279,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 vpq_wire_dtype() -> cudaDataType_t +{ + static_assert(std::is_same_v || std::is_same_v, + "serialize_vpq: 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_vpq(raft::resources const& res, + std::ostream& os, + device_vpq_dataset const& dataset) +{ + 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 +339,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_vpq`, 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_vpq_dataset(raft::resources const& res, + std::ostream& os, + device_vpq_dataset const& dataset) +{ + raft::serialize_scalar(res, os, kSerializeVPQDataset); + raft::serialize_scalar(res, os, vpq_wire_dtype()); + serialize_vpq(res, os, dataset); +} + +/** Read a blob written by `serialize_vpq_dataset`, validating the tag and codebook dtype. */ +template +auto deserialize_vpq_dataset(raft::resources const& res, std::istream& is) + -> std::unique_ptr> +{ + const auto tag = raft::deserialize_scalar(res, is); + RAFT_EXPECTS(tag == kSerializeVPQDataset, + "deserialize_vpq_dataset: expected VPQ tag (%u), got %u", + static_cast(kSerializeVPQDataset), + static_cast(tag)); + const auto dtype = raft::deserialize_scalar(res, is); + RAFT_EXPECTS(dtype == vpq_wire_dtype(), + "deserialize_vpq_dataset: codebook dtype (%d) does not match expected (%d)", + static_cast(dtype), + static_cast(vpq_wire_dtype())); + return deserialize_vpq(res, is); +} + template auto deserialize_dense_dataset(raft::resources const& res, std::istream& is) -> std::unique_ptr diff --git a/cpp/src/preprocessing/quantize/pq.cu b/cpp/src/preprocessing/quantize/pq.cu index 20b8f21d36..673e49759b 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, + std::ostream& os, + const cuvs::neighbors::device_vpq_dataset& dataset) +{ + // 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, vpq_serialization_version); + ::cuvs::neighbors::detail::serialize_vpq_dataset(res, os, dataset); +} + +void serialize(raft::resources const& res, + const std::string& filename, + const cuvs::neighbors::device_vpq_dataset& dataset) +{ + 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, os, dataset); +} + +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 == vpq_serialization_version, + "pq::deserialize: serialization version mismatch, expected %d, got %d", + vpq_serialization_version, + version); + *out_dataset = ::cuvs::neighbors::detail::deserialize_vpq_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..26300b1f8d 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_vpq_serialize.cu GPUS 1 PERCENT 100 ) @@ -415,6 +415,7 @@ ConfigureTest( preprocessing/binary_quantization.cu preprocessing/spectral_embedding.cu preprocessing/product_quantization.cu + preprocessing/vpq_serialization.cu preprocessing/pca.cu GPUS 1 PERCENT 100 diff --git a/cpp/tests/neighbors/ann_cagra/test_vpq_serialize.cu b/cpp/tests/neighbors/ann_cagra/test_vpq_serialize.cu new file mode 100644 index 0000000000..41ff37381c --- /dev/null +++ b/cpp/tests/neighbors/ann_cagra/test_vpq_serialize.cu @@ -0,0 +1,262 @@ +/* + * 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 (CAGRA-Q). + * + * 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, and that the + * cases where they cannot travel fail loudly. Fidelity of the dataset payload itself is covered by + * preprocessing/vpq_serialization.cu. + */ + +#include + +#include +#include + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +namespace cuvs::neighbors::cagra { + +using vpq_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) -> vpq_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 CagraVpqSerializeTest : 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)); + } + + 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(CagraVpqSerializeTest, RoundTripsThroughAFileWithItsDataset) +{ + auto compressed = compress(res_, dataset(), pq_dim); + auto graph_source = build_graph_source(res_, dataset()); + vpq_f16_index idx{res_, + cuvs::distance::DistanceType::L2Expanded, + compressed.as_dataset_view(), + graph_source.graph()}; + + 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(CagraVpqSerializeTest, RefusesToLoadWithoutItsDataset) +{ + auto compressed = compress(res_, dataset(), pq_dim); + auto graph_source = build_graph_source(res_, dataset()); + vpq_f16_index idx{res_, + cuvs::distance::DistanceType::L2Expanded, + compressed.as_dataset_view(), + graph_source.graph()}; + + std::stringstream stored; + cagra::serialize(res_, stored, idx); + + // Dropping the rows on load is fine for a dense index, whose caller can attach its own copy, but + // it would leave a VPQ index unsearchable with no way back: the rows exist nowhere else. + vpq_f16_index restored{res_}; + EXPECT_THROW(cagra::deserialize(res_, stored, &restored, nullptr), raft::exception); +} + +TEST_F(CagraVpqSerializeTest, SerializesTheGraphAloneWhenAsked) +{ + auto compressed = compress(res_, dataset(), pq_dim); + auto graph_source = build_graph_source(res_, dataset()); + vpq_f16_index idx{res_, + cuvs::distance::DistanceType::L2Expanded, + compressed.as_dataset_view(), + graph_source.graph()}; + + 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_dataset() can make searchable again. + EXPECT_EQ(owner, nullptr); + EXPECT_EQ(restored.size(), idx.size()); + EXPECT_EQ(restored.graph_degree(), idx.graph_degree()); +} + +TEST_F(CagraVpqSerializeTest, RejectsLoadingACompressedIndexAsDense) +{ + auto compressed = compress(res_, dataset(), pq_dim); + auto graph_source = build_graph_source(res_, dataset()); + vpq_f16_index idx{res_, + cuvs::distance::DistanceType::L2Expanded, + compressed.as_dataset_view(), + graph_source.graph()}; + + 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/vpq_serialization.cu b/cpp/tests/preprocessing/vpq_serialization.cu new file mode 100644 index 0000000000..ac36cb6f23 --- /dev/null +++ b/cpp/tests/preprocessing/vpq_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 vpq_dataset_t = cuvs::neighbors::device_vpq_dataset; + +struct VpqSerializationInputs { + 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 VpqSerializationInputs& 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 VpqSerializationTest : public ::testing::TestWithParam { + public: + VpqSerializationTest() + : 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() -> vpq_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 vpq_dataset_t& expected, const vpq_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 vpq_dataset_t& expected, const vpq_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_; + VpqSerializationInputs params_; + raft::device_matrix dataset_; +}; + +TEST_P(VpqSerializationTest, RoundTrip) +{ + auto original = compress(); + + { + SCOPED_TRACE("through a stream"); + std::stringstream stream; + serialize(res_, stream, original); + 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_vpq_serialization_test.bin"; + serialize(res_, path, original); + 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 vpq_serialization_inputs = { + // pq_len = dim / pq_dim of 2, 4 and 8: the three values CAGRA-Q 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(VpqSerializationTests, + VpqSerializationTest, + ::testing::ValuesIn(vpq_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(VpqSerialization, RejectsEmptyStream) +{ + raft::resources res; + std::stringstream stream; + std::unique_ptr restored; + EXPECT_THROW(deserialize(res, stream, &restored), raft::exception); +} + +TEST(VpqSerialization, 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, vpq_serialization_version); + + std::unique_ptr restored; + EXPECT_THROW(deserialize(res, stream, &restored), raft::exception); +} + +TEST(VpqSerialization, RejectsFutureVersion) +{ + raft::resources res; + std::stringstream stream; + write_preamble(res, stream, vpq_serialization_version + 1); + + std::unique_ptr restored; + EXPECT_THROW(deserialize(res, stream, &restored), raft::exception); +} + +TEST(VpqSerialization, RejectsTruncatedPayload) +{ + raft::resources res; + std::stringstream stream; + write_preamble(res, stream, vpq_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(VpqSerialization, RejectsNullOutParameter) +{ + raft::resources res; + std::stringstream stream; + write_preamble(res, stream, vpq_serialization_version); + EXPECT_THROW(deserialize(res, stream, nullptr), raft::exception); +} + +} // namespace cuvs::preprocessing::quantize::pq