From 8b176088b93c011b9cd64fe4de60a5f84770e277 Mon Sep 17 00:00:00 2001 From: Irina Reshodko Date: Tue, 18 Aug 2026 01:22:40 -0700 Subject: [PATCH 1/8] Restore the VPQ writer in the dataset serializer --- .../neighbors/detail/dataset_serialize.hpp | 69 +++++++++++++++++++ 1 file changed, 69 insertions(+) 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 From f4e2e84e3e84c822613f49cd352b7c636d420857 Mon Sep 17 00:00:00 2001 From: Irina Reshodko Date: Tue, 18 Aug 2026 01:24:59 -0700 Subject: [PATCH 2/8] Add a public on-disk format for VPQ-compressed datasets --- .../cuvs/preprocessing/quantize/pq.hpp | 79 ++++++ cpp/src/preprocessing/quantize/pq.cu | 56 ++++ cpp/tests/preprocessing/vpq_serialization.cu | 261 ++++++++++++++++++ 3 files changed, 396 insertions(+) create mode 100644 cpp/tests/preprocessing/vpq_serialization.cu 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/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/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 From f24e9e36d80016703b66c68a50d6c3a249449801 Mon Sep 17 00:00:00 2001 From: Irina Reshodko Date: Tue, 18 Aug 2026 01:26:35 -0700 Subject: [PATCH 3/8] Serialize a CAGRA-Q index together with its compressed rows --- c/src/neighbors/cagra.cpp | 8 +- cpp/include/cuvs/neighbors/cagra.hpp | 115 +++++++- cpp/src/neighbors/cagra_serialize.cuh | 37 +++ cpp/src/neighbors/cagra_serialize_inst.cu.in | 3 + .../detail/cagra/cagra_serialize.cuh | 26 +- cpp/tests/CMakeLists.txt | 3 +- .../neighbors/ann_cagra/test_vpq_serialize.cu | 262 ++++++++++++++++++ 7 files changed, 442 insertions(+), 12 deletions(-) create mode 100644 cpp/tests/neighbors/ann_cagra/test_vpq_serialize.cu 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/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/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/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 From bc6190f6c36cfce5ee5e8c1aae5db6f80a2480c4 Mon Sep 17 00:00:00 2001 From: Irina Reshodko Date: Tue, 18 Aug 2026 02:22:25 -0700 Subject: [PATCH 4/8] Register MG_C_TEST only when the multi-GPU algorithms are built --- c/tests/CMakeLists.txt | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) 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 ) From 8fe4f9f7a868554ce04acb1285107d081e71af04 Mon Sep 17 00:00:00 2001 From: Irina Reshodko Date: Wed, 19 Aug 2026 06:29:52 -0700 Subject: [PATCH 5/8] Load a CAGRA index without its dataset, or ask the file what it holds first --- c/src/neighbors/cagra.cpp | 70 ++++------ cpp/include/cuvs/neighbors/cagra.hpp | 90 +++++++++--- .../cuvs/preprocessing/quantize/pq.hpp | 15 ++ cpp/src/neighbors/cagra.cpp | 80 +++++++++++ .../detail/cagra/cagra_serialize.cuh | 15 +- .../neighbors/detail/dataset_serialize.hpp | 132 +++++++++++++----- .../neighbors/ann_cagra/test_vpq_serialize.cu | 117 +++++++++++++--- 7 files changed, 397 insertions(+), 122 deletions(-) diff --git a/c/src/neighbors/cagra.cpp b/c/src/neighbors/cagra.cpp index 42200258b6..f8f115cd50 100644 --- a/c/src/neighbors/cagra.cpp +++ b/c/src/neighbors/cagra.cpp @@ -15,11 +15,8 @@ #include #include #include -#include #include -#include #include -#include #include "../core/exceptions.hpp" #include "../core/interop.hpp" @@ -974,47 +971,40 @@ struct serialized_cagra_header { cuvs::neighbors::cagra::serialized_dataset_kind dataset_kind; }; +/** + * What the file holds, in the terms the C entry points dispatch on. + * + * The reading is the C++ API's: `cagra::read_serialized_header` parses the same preamble, validates + * the format version and the dataset kind, and leaves this with nothing to do but restate the + * element type in DLPack terms. + */ static auto read_serialized_header(cuvsResources_t res, const char *filename) -> serialized_cagra_header { auto res_ptr = reinterpret_cast(res); - std::ifstream is(filename, std::ios::in | std::ios::binary); - if (!is) { - RAFT_FAIL("Cannot open file %s", filename); - } - - char dtype_string[4]{}; - if (!is.read(dtype_string, sizeof(dtype_string))) { - RAFT_FAIL("Invalid or truncated index header in file %s", filename); - } - - auto const dtype = raft::numpy_serializer::parse_descr( - std::string(dtype_string, sizeof(dtype_string))); - DLDataType output_dtype{ - .code = 0, .bits = static_cast(dtype.itemsize * 8), .lanes = 1}; - if (dtype.kind == 'f' && dtype.itemsize == 4) { - output_dtype.code = kDLFloat; - } else if (dtype.kind == 'e' && dtype.itemsize == 2) { - output_dtype.code = kDLFloat; - } else if (dtype.kind == 'i' && dtype.itemsize == 1) { - output_dtype.code = kDLInt; - } else if (dtype.kind == 'u' && dtype.itemsize == 1) { - output_dtype.code = kDLUInt; - } else { - RAFT_FAIL("Unsupported dtype in file %s", filename); + auto const header = cuvs::neighbors::cagra::read_serialized_header( + *res_ptr, std::string(filename)); + DLDataType dtype{.code = 0, .bits = 0, .lanes = 1}; + switch (header.dtype) { + case CUDA_R_32F: + dtype.code = kDLFloat; + dtype.bits = 32; + break; + case CUDA_R_16F: + dtype.code = kDLFloat; + dtype.bits = 16; + break; + case CUDA_R_8I: + dtype.code = kDLInt; + dtype.bits = 8; + break; + case CUDA_R_8U: + dtype.code = kDLUInt; + dtype.bits = 8; + break; + default: + RAFT_FAIL("Unsupported dtype in file %s", filename); } - - auto const version = raft::deserialize_scalar(*res_ptr, is); - auto const dataset_kind_raw = - raft::deserialize_scalar(*res_ptr, is); - RAFT_EXPECTS( - version == cuvs::neighbors::cagra::cagra_serialization_version, - "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::device_vpq_f16), - "Invalid serialized dataset kind %u in file %s", - dataset_kind_raw, filename); - return {output_dtype, static_cast(dataset_kind_raw)}; + return {dtype, header.dataset_kind}; } template diff --git a/cpp/include/cuvs/neighbors/cagra.hpp b/cpp/include/cuvs/neighbors/cagra.hpp index 56394b5398..d065eda04c 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,16 @@ using cagra_index_t = index>; +/** + * 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;` + */ +template +using owning_dataset_for_index_t = + cuvs::neighbors::owning_dataset_for_view_t; + /** * @} */ @@ -2271,6 +2283,50 @@ enum class serialized_dataset_kind : std::uint32_t { /** Current experimental CAGRA serialization format version. */ inline constexpr int cagra_serialization_version = 6; +/** 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_vpq_f16) { + * 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. + */ +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 @@ -2829,15 +2885,17 @@ void deserialize(raft::resources const& handle, /* 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. + * 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. As everywhere else the index holds a view: + * `deserialize` returns the owning dataset through `out_dataset`, which has to stay 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. + * `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. */ void serialize(raft::resources const& handle, const std::string& filename, @@ -2848,7 +2906,7 @@ void deserialize( raft::resources const& handle, const std::string& filename, cuvs::neighbors::cagra::vpq_f16_index* index, - std::unique_ptr>* out_dataset); + std::unique_ptr>* out_dataset = nullptr); void serialize(raft::resources const& handle, std::ostream& os, @@ -2859,7 +2917,7 @@ void deserialize( raft::resources const& handle, std::istream& is, cuvs::neighbors::cagra::vpq_f16_index* index, - std::unique_ptr>* out_dataset); + std::unique_ptr>* out_dataset = nullptr); void serialize(raft::resources const& handle, const std::string& filename, @@ -2870,7 +2928,7 @@ void deserialize( raft::resources const& handle, const std::string& filename, cuvs::neighbors::cagra::vpq_f16_index* index, - std::unique_ptr>* out_dataset); + std::unique_ptr>* out_dataset = nullptr); void serialize(raft::resources const& handle, std::ostream& os, @@ -2881,7 +2939,7 @@ void deserialize( raft::resources const& handle, std::istream& is, cuvs::neighbors::cagra::vpq_f16_index* index, - std::unique_ptr>* out_dataset); + std::unique_ptr>* out_dataset = nullptr); void serialize(raft::resources const& handle, const std::string& filename, @@ -2892,7 +2950,7 @@ void deserialize( raft::resources const& handle, const std::string& filename, cuvs::neighbors::cagra::vpq_f16_index* index, - std::unique_ptr>* out_dataset); + std::unique_ptr>* out_dataset = nullptr); void serialize(raft::resources const& handle, std::ostream& os, @@ -2903,7 +2961,7 @@ void deserialize( raft::resources const& handle, std::istream& is, cuvs::neighbors::cagra::vpq_f16_index* index, - std::unique_ptr>* out_dataset); + std::unique_ptr>* out_dataset = nullptr); void serialize(raft::resources const& handle, const std::string& filename, @@ -2914,7 +2972,7 @@ void deserialize( raft::resources const& handle, const std::string& filename, cuvs::neighbors::cagra::vpq_f16_index* index, - std::unique_ptr>* out_dataset); + std::unique_ptr>* out_dataset = nullptr); void serialize(raft::resources const& handle, std::ostream& os, @@ -2925,7 +2983,7 @@ void deserialize( raft::resources const& handle, std::istream& is, cuvs::neighbors::cagra::vpq_f16_index* index, - std::unique_ptr>* out_dataset); + std::unique_ptr>* out_dataset = nullptr); /** @copydoc serialize */ void serialize(raft::resources const& handle, diff --git a/cpp/include/cuvs/preprocessing/quantize/pq.hpp b/cpp/include/cuvs/preprocessing/quantize/pq.hpp index c26a24fea6..51bee23d12 100644 --- a/cpp/include/cuvs/preprocessing/quantize/pq.hpp +++ b/cpp/include/cuvs/preprocessing/quantize/pq.hpp @@ -282,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 @@ -350,6 +359,12 @@ inline constexpr int vpq_serialization_version = 1; * 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 diff --git a/cpp/src/neighbors/cagra.cpp b/cpp/src/neighbors/cagra.cpp index 1de6d5a669..2a0ed3b029 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_vpq_f16), + "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/detail/cagra/cagra_serialize.cuh b/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh index add4bb2532..c548e8c9b2 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh @@ -408,16 +408,11 @@ void deserialize( std::unique_ptr dataset_owner{}; if (has_dataset) { if (out_dataset == nullptr) { - // 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); - } + // 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( diff --git a/cpp/src/neighbors/detail/dataset_serialize.hpp b/cpp/src/neighbors/detail/dataset_serialize.hpp index 6e73f36d10..4a563a7d32 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 @@ -405,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_vpq_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_vpq_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/tests/neighbors/ann_cagra/test_vpq_serialize.cu b/cpp/tests/neighbors/ann_cagra/test_vpq_serialize.cu index 41ff37381c..2b2a537384 100644 --- a/cpp/tests/neighbors/ann_cagra/test_vpq_serialize.cu +++ b/cpp/tests/neighbors/ann_cagra/test_vpq_serialize.cu @@ -13,8 +13,10 @@ * 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 + * 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/vpq_serialization.cu. */ @@ -35,6 +37,7 @@ #include #include #include +#include #include namespace cuvs::neighbors::cagra { @@ -150,6 +153,16 @@ class CagraVpqSerializeTest : public ::testing::Test { dataset_->data_handle(), std::min(n_queries, dataset_->extent(0)), dataset_->extent(1)); } + /** A CAGRA-Q index viewing `compressed` and the graph of `graph_source`; both must outlive it. */ + auto assemble(const vpq_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 @@ -162,10 +175,7 @@ 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 idx = assemble(compressed, graph_source); auto before = neighbor_ids(res_, idx, queries(500)); ASSERT_GT(self_recall_at_1(before), 0.5); @@ -200,32 +210,58 @@ TEST_F(CagraVpqSerializeTest, RoundTripsThroughAFileWithItsDataset) EXPECT_EQ(mismatches, 0u) << mismatches << " of " << before.size() << " neighbour ids changed"; } -TEST_F(CagraVpqSerializeTest, RefusesToLoadWithoutItsDataset) +TEST_F(CagraVpqSerializeTest, LoadsTheGraphAloneWhenNoOwnerIsAskedFor) { 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 idx = assemble(compressed, graph_source); + auto before = neighbor_ids(res_, idx, queries(500)); 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. + // 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(CagraVpqSerializeTest, 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_}; - EXPECT_THROW(cagra::deserialize(res_, stored, &restored, nullptr), raft::exception); + 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(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()}; + auto idx = assemble(compressed, graph_source); std::stringstream stored; cagra::serialize(res_, stored, idx, /* include_dataset */ false); @@ -234,20 +270,55 @@ TEST_F(CagraVpqSerializeTest, SerializesTheGraphAloneWhenAsked) std::unique_ptr owner; cagra::deserialize(res_, stored, &restored, &owner); - // Nothing to own, and a graph that only update_dataset() can make searchable again. + // 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(CagraVpqSerializeTest, 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_vpq_f16); + + // 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(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()}; + auto idx = assemble(compressed, graph_source); std::stringstream stored; cagra::serialize(res_, stored, idx); From be9b186f3c1bf8e89573f564ec0396baad99ece4 Mon Sep 17 00:00:00 2001 From: Irina Reshodko Date: Wed, 19 Aug 2026 07:32:54 -0700 Subject: [PATCH 6/8] Changed argument order; Renamed vpq to pq, removed mentioning f16 where possible --- c/src/neighbors/cagra.cpp | 6 +- cpp/include/cuvs/neighbors/cagra.hpp | 13 ++-- .../cuvs/preprocessing/quantize/pq.hpp | 18 ++--- cpp/src/neighbors/cagra.cpp | 2 +- cpp/src/neighbors/cagra_serialize_inst.cu.in | 4 +- .../detail/cagra/cagra_serialize.cuh | 12 ++-- .../neighbors/detail/dataset_serialize.hpp | 38 +++++------ cpp/src/preprocessing/quantize/pq.cu | 20 +++--- cpp/tests/CMakeLists.txt | 4 +- ..._vpq_serialize.cu => test_pq_serialize.cu} | 28 ++++---- ...q_serialization.cu => pq_serialization.cu} | 66 +++++++++---------- 11 files changed, 110 insertions(+), 101 deletions(-) rename cpp/tests/neighbors/ann_cagra/{test_vpq_serialize.cu => test_pq_serialize.cu} (94%) rename cpp/tests/preprocessing/{vpq_serialization.cu => pq_serialization.cu} (82%) diff --git a/c/src/neighbors/cagra.cpp b/c/src/neighbors/cagra.cpp index f8f115cd50..677968250d 100644 --- a/c/src/neighbors/cagra.cpp +++ b/c/src/neighbors/cagra.cpp @@ -1048,11 +1048,11 @@ void dispatch_serialized_dataset_kind( fn.template operator()< cuvs::neighbors::device_padded_dataset_view>(); break; - case serialized_kind::device_vpq_f16: + case serialized_kind::device_pq: // 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 " + // on that layout, so there is nothing here to hand a PQ-compressed index to yet. + RAFT_FAIL("File holds a PQ-compressed (CAGRA-Q) dataset, which the C API has no dataset " "layout for; load it through the C++ API"); } } diff --git a/cpp/include/cuvs/neighbors/cagra.hpp b/cpp/include/cuvs/neighbors/cagra.hpp index d065eda04c..ff7ee9e0d1 100644 --- a/cpp/include/cuvs/neighbors/cagra.hpp +++ b/cpp/include/cuvs/neighbors/cagra.hpp @@ -2276,8 +2276,13 @@ 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, + /** + * Device-resident PQ-compressed dataset (CAGRA-Q). + * + * 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. */ @@ -2300,7 +2305,7 @@ struct serialized_index_header { * * @code{.cpp} * auto header = cagra::read_serialized_header(res, "index.bin"); - * if (header.dataset_kind == cagra::serialized_dataset_kind::device_vpq_f16) { + * 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); @@ -2332,7 +2337,7 @@ auto read_serialized_header(raft::resources const& res, std::istream& is) // 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). +// serialize_cagra_dense_dataset, PQ-compressed ones serialize_pq_dataset). /** * Save the index to file. diff --git a/cpp/include/cuvs/preprocessing/quantize/pq.hpp b/cpp/include/cuvs/preprocessing/quantize/pq.hpp index 51bee23d12..44f120fffd 100644 --- a/cpp/include/cuvs/preprocessing/quantize/pq.hpp +++ b/cpp/include/cuvs/preprocessing/quantize/pq.hpp @@ -344,7 +344,7 @@ template } /** Current VPQ dataset serialization format version. */ -inline constexpr int vpq_serialization_version = 1; +inline constexpr int pq_serialization_version = 1; /** * @brief Write a VPQ dataset (both codebooks plus the encoded rows) to a stream. @@ -354,7 +354,7 @@ inline constexpr int vpq_serialization_version = 1; * 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 + * `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. @@ -371,7 +371,7 @@ inline constexpr int vpq_serialization_version = 1; * * // Offline, once. * auto vpq = cuvs::preprocessing::quantize::pq::make_vpq_dataset(res, vpq_params, rows); - * cuvs::preprocessing::quantize::pq::serialize(res, "base.vpq", vpq); + * cuvs::preprocessing::quantize::pq::serialize(res, vpq, "base.vpq"); * * // Later, per run: load the compressed rows and build a CAGRA-Q graph over them. * std::unique_ptr> loaded; @@ -381,23 +381,23 @@ inline constexpr int vpq_serialization_version = 1; * @endcode * * @param[in] res raft resource - * @param[in] os output stream, opened in binary mode * @param[in] dataset the VPQ dataset to write + * @param[out] os output stream, opened in binary mode */ void serialize(raft::resources const& res, - std::ostream& os, - const cuvs::neighbors::device_vpq_dataset& dataset); + const cuvs::neighbors::device_vpq_dataset& dataset, + std::ostream& os); /** * @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 + * @param[out] filename path to write, truncated if it exists */ void serialize(raft::resources const& res, - const std::string& filename, - const cuvs::neighbors::device_vpq_dataset& dataset); + const cuvs::neighbors::device_vpq_dataset& dataset, + const std::string& filename); /** * @brief Read a VPQ dataset written by `serialize`. diff --git a/cpp/src/neighbors/cagra.cpp b/cpp/src/neighbors/cagra.cpp index 2a0ed3b029..3c03abc0e5 100644 --- a/cpp/src/neighbors/cagra.cpp +++ b/cpp/src/neighbors/cagra.cpp @@ -122,7 +122,7 @@ auto read_header(raft::resources const& res, std::istream& is, const char* sourc // 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_vpq_f16), + 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); diff --git a/cpp/src/neighbors/cagra_serialize_inst.cu.in b/cpp/src/neighbors/cagra_serialize_inst.cu.in index 58e555d17e..21993e22a5 100644 --- a/cpp/src/neighbors/cagra_serialize_inst.cu.in +++ b/cpp/src/neighbors/cagra_serialize_inst.cu.in @@ -12,7 +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; -using inst_vpq_f16_view_t = cuvs::neighbors::device_vpq_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 diff --git a/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh b/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh index c548e8c9b2..e985a6bdd8 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh @@ -60,8 +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_f16_dataset_view_v) { - return kind::device_vpq_f16; + } 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"); @@ -71,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::device_vpq_f16); + return raw <= static_cast(kind::device_pq); } /** @@ -129,7 +131,7 @@ void serialize(raft::resources const& res, // 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()); + neighbors::detail::serialize_pq_dataset(res, index_.dataset().dset(), os); } else { // A further dataset type requires a new branch here and a corresponding deserialize branch. // Use static_assert to catch unsupported types at compile time. @@ -431,7 +433,7 @@ void deserialize( 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); + 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 4a563a7d32..36cc982efb 100644 --- a/cpp/src/neighbors/detail/dataset_serialize.hpp +++ b/cpp/src/neighbors/detail/dataset_serialize.hpp @@ -306,10 +306,10 @@ auto deserialize_host_dense(raft::resources const& res, std::istream& is) /** VPQ codebooks are floating point; the encoded rows are always uint8 and carry no dtype. */ template -constexpr auto vpq_wire_dtype() -> cudaDataType_t +constexpr auto pq_wire_dtype() -> cudaDataType_t { static_assert(std::is_same_v || std::is_same_v, - "serialize_vpq: codebook element type must be float or half"); + "serialize_pq: codebook element type must be float or half"); return std::is_same_v ? CUDA_R_16F : CUDA_R_32F; } @@ -323,9 +323,9 @@ constexpr auto vpq_wire_dtype() -> cudaDataType_t * `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) +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()); @@ -367,35 +367,35 @@ auto deserialize_vpq(raft::resources const& res, std::istream& is) /** * 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 + * 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_vpq_dataset(raft::resources const& res, - std::ostream& os, - device_vpq_dataset const& dataset) +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, vpq_wire_dtype()); - serialize_vpq(res, os, dataset); + raft::serialize_scalar(res, os, pq_wire_dtype()); + serialize_pq(res, dataset, os); } -/** Read a blob written by `serialize_vpq_dataset`, validating the tag and codebook dtype. */ +/** Read a blob written by `serialize_pq_dataset`, validating the tag and codebook dtype. */ template -auto deserialize_vpq_dataset(raft::resources const& res, std::istream& is) +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_vpq_dataset: expected VPQ tag (%u), got %u", + "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 == vpq_wire_dtype(), - "deserialize_vpq_dataset: codebook dtype (%d) does not match expected (%d)", + RAFT_EXPECTS(dtype == pq_wire_dtype(), + "deserialize_pq_dataset: codebook dtype (%d) does not match expected (%d)", static_cast(dtype), - static_cast(vpq_wire_dtype())); + static_cast(pq_wire_dtype())); return deserialize_vpq(res, is); } @@ -452,7 +452,7 @@ void skip_dense_dtype_and_payload(raft::resources const& res, std::istream& is) * are not needed: each of the three matrices carries a NumPy header stating its own size. */ template -void skip_vpq_dtype_and_payload(raft::resources const& res, std::istream& is) +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, @@ -485,7 +485,7 @@ void skip_dataset(raft::resources const& res, std::istream& is) static_cast(raft::deserialize_scalar(res, is)); // suggested_dim return; case kSerializeStridedDataset: return skip_dense_dtype_and_payload(res, is); - case kSerializeVPQDataset: return skip_vpq_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)); } } diff --git a/cpp/src/preprocessing/quantize/pq.cu b/cpp/src/preprocessing/quantize/pq.cu index 673e49759b..fc2c16d317 100644 --- a/cpp/src/preprocessing/quantize/pq.cu +++ b/cpp/src/preprocessing/quantize/pq.cu @@ -84,8 +84,8 @@ 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) + 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 @@ -93,17 +93,17 @@ void serialize(raft::resources const& res, 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); + raft::serialize_scalar(res, os, pq_serialization_version); + ::cuvs::neighbors::detail::serialize_pq_dataset(res, dataset, os); } void serialize(raft::resources const& res, - const std::string& filename, - const cuvs::neighbors::device_vpq_dataset& dataset) + 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, os, dataset); + serialize(res, dataset, os); } void deserialize(raft::resources const& res, @@ -116,11 +116,11 @@ void deserialize(raft::resources const& res, 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, + RAFT_EXPECTS(version == pq_serialization_version, "pq::deserialize: serialization version mismatch, expected %d, got %d", - vpq_serialization_version, + pq_serialization_version, version); - *out_dataset = ::cuvs::neighbors::detail::deserialize_vpq_dataset(res, is); + *out_dataset = ::cuvs::neighbors::detail::deserialize_pq_dataset(res, is); } void deserialize(raft::resources const& res, diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index 26300b1f8d..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 neighbors/ann_cagra/test_vpq_serialize.cu + PATH neighbors/ann_cagra/test_float_uint32_t.cu neighbors/ann_cagra/test_pq_serialize.cu GPUS 1 PERCENT 100 ) @@ -415,7 +415,7 @@ ConfigureTest( preprocessing/binary_quantization.cu preprocessing/spectral_embedding.cu preprocessing/product_quantization.cu - preprocessing/vpq_serialization.cu + preprocessing/pq_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_pq_serialize.cu similarity index 94% rename from cpp/tests/neighbors/ann_cagra/test_vpq_serialize.cu rename to cpp/tests/neighbors/ann_cagra/test_pq_serialize.cu index 2b2a537384..2e23b3196b 100644 --- a/cpp/tests/neighbors/ann_cagra/test_vpq_serialize.cu +++ b/cpp/tests/neighbors/ann_cagra/test_pq_serialize.cu @@ -17,7 +17,7 @@ * 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/vpq_serialization.cu. + * preprocessing/pq_serialization.cu. */ #include @@ -42,7 +42,7 @@ namespace cuvs::neighbors::cagra { -using vpq_dataset_t = cuvs::neighbors::device_vpq_dataset; +using pq_dataset_t = cuvs::neighbors::device_vpq_dataset; namespace { @@ -51,7 +51,7 @@ constexpr uint32_t kGraphDegree = 32; auto compress(const raft::resources& res, raft::device_matrix_view dataset, - uint32_t pq_dim) -> vpq_dataset_t + uint32_t pq_dim) -> pq_dataset_t { cuvs::neighbors::vpq_params params; params.pq_dim = pq_dim; @@ -115,7 +115,7 @@ auto self_recall_at_1(const std::vector& ids) -> double * 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 { +class CagraPqSerializeTest : public ::testing::Test { protected: void SetUp() override { @@ -154,7 +154,7 @@ class CagraVpqSerializeTest : public ::testing::Test { } /** A CAGRA-Q index viewing `compressed` and the graph of `graph_source`; both must outlive it. */ - auto assemble(const vpq_dataset_t& compressed, + auto assemble(const pq_dataset_t& compressed, const device_standard_index& graph_source) -> vpq_f16_index { return vpq_f16_index{res_, @@ -171,7 +171,7 @@ class CagraVpqSerializeTest : public ::testing::Test { std::optional> dataset_ = std::nullopt; }; -TEST_F(CagraVpqSerializeTest, RoundTripsThroughAFileWithItsDataset) +TEST_F(CagraPqSerializeTest, RoundTripsThroughAFileWithItsDataset) { auto compressed = compress(res_, dataset(), pq_dim); auto graph_source = build_graph_source(res_, dataset()); @@ -184,7 +184,7 @@ TEST_F(CagraVpqSerializeTest, RoundTripsThroughAFileWithItsDataset) cagra::serialize(res_, stored, idx); vpq_f16_index restored{res_}; - std::unique_ptr owner; + std::unique_ptr owner; cagra::deserialize(res_, stored, &restored, &owner); ASSERT_NE(owner, nullptr); @@ -210,7 +210,7 @@ TEST_F(CagraVpqSerializeTest, RoundTripsThroughAFileWithItsDataset) EXPECT_EQ(mismatches, 0u) << mismatches << " of " << before.size() << " neighbour ids changed"; } -TEST_F(CagraVpqSerializeTest, LoadsTheGraphAloneWhenNoOwnerIsAskedFor) +TEST_F(CagraPqSerializeTest, LoadsTheGraphAloneWhenNoOwnerIsAskedFor) { auto compressed = compress(res_, dataset(), pq_dim); auto graph_source = build_graph_source(res_, dataset()); @@ -242,7 +242,7 @@ TEST_F(CagraVpqSerializeTest, LoadsTheGraphAloneWhenNoOwnerIsAskedFor) EXPECT_EQ(mismatches, 0u) << mismatches << " of " << before.size() << " neighbour ids changed"; } -TEST_F(CagraVpqSerializeTest, SkipsWhicheverDatasetTheFileHappensToHold) +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. @@ -257,7 +257,7 @@ TEST_F(CagraVpqSerializeTest, SkipsWhicheverDatasetTheFileHappensToHold) EXPECT_EQ(stored.peek(), std::char_traits::eof()); } -TEST_F(CagraVpqSerializeTest, SerializesTheGraphAloneWhenAsked) +TEST_F(CagraPqSerializeTest, SerializesTheGraphAloneWhenAsked) { auto compressed = compress(res_, dataset(), pq_dim); auto graph_source = build_graph_source(res_, dataset()); @@ -267,7 +267,7 @@ TEST_F(CagraVpqSerializeTest, SerializesTheGraphAloneWhenAsked) cagra::serialize(res_, stored, idx, /* include_dataset */ false); vpq_f16_index restored{res_}; - std::unique_ptr owner; + 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 @@ -277,7 +277,7 @@ TEST_F(CagraVpqSerializeTest, SerializesTheGraphAloneWhenAsked) EXPECT_EQ(restored.graph_degree(), idx.graph_degree()); } -TEST_F(CagraVpqSerializeTest, SaysWhatItHoldsBeforeItIsLoaded) +TEST_F(CagraPqSerializeTest, SaysWhatItHoldsBeforeItIsLoaded) { auto compressed = compress(res_, dataset(), pq_dim); auto graph_source = build_graph_source(res_, dataset()); @@ -290,7 +290,7 @@ TEST_F(CagraVpqSerializeTest, SaysWhatItHoldsBeforeItIsLoaded) // 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_vpq_f16); + 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. @@ -314,7 +314,7 @@ TEST_F(CagraVpqSerializeTest, SaysWhatItHoldsBeforeItIsLoaded) EXPECT_EQ(graph_only_header.dtype, CUDA_R_32F); } -TEST_F(CagraVpqSerializeTest, RejectsLoadingACompressedIndexAsDense) +TEST_F(CagraPqSerializeTest, RejectsLoadingACompressedIndexAsDense) { auto compressed = compress(res_, dataset(), pq_dim); auto graph_source = build_graph_source(res_, dataset()); diff --git a/cpp/tests/preprocessing/vpq_serialization.cu b/cpp/tests/preprocessing/pq_serialization.cu similarity index 82% rename from cpp/tests/preprocessing/vpq_serialization.cu rename to cpp/tests/preprocessing/pq_serialization.cu index ac36cb6f23..2fe38c432e 100644 --- a/cpp/tests/preprocessing/vpq_serialization.cu +++ b/cpp/tests/preprocessing/pq_serialization.cu @@ -26,9 +26,9 @@ namespace cuvs::preprocessing::quantize::pq { -using vpq_dataset_t = cuvs::neighbors::device_vpq_dataset; +using pq_dataset_t = cuvs::neighbors::device_vpq_dataset; -struct VpqSerializationInputs { +struct PqSerializationInputs { int64_t n_rows; int64_t dim; uint32_t pq_bits; @@ -37,7 +37,7 @@ struct VpqSerializationInputs { uint64_t seed; }; -std::ostream& operator<<(std::ostream& os, const VpqSerializationInputs& in) +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 @@ -68,10 +68,10 @@ void expect_same_bits(const raft::resources& res, EXPECT_EQ(0, std::memcmp(lhs.data(), rhs.data(), lhs.size() * sizeof(T))) << what; } -class VpqSerializationTest : public ::testing::TestWithParam { +class PqSerializationTest : public ::testing::TestWithParam { public: - VpqSerializationTest() - : params_(::testing::TestWithParam::GetParam()), + PqSerializationTest() + : params_(::testing::TestWithParam::GetParam()), dataset_(raft::make_device_matrix(res_, params_.n_rows, params_.dim)) { } @@ -94,7 +94,7 @@ class VpqSerializationTest : public ::testing::TestWithParam vpq_dataset_t + auto compress() -> pq_dataset_t { cuvs::neighbors::vpq_params vpq; vpq.pq_bits = params_.pq_bits; @@ -105,7 +105,7 @@ class VpqSerializationTest : public ::testing::TestWithParam dataset_; }; -TEST_P(VpqSerializationTest, RoundTrip) +TEST_P(PqSerializationTest, RoundTrip) { auto original = compress(); { SCOPED_TRACE("through a stream"); std::stringstream stream; - serialize(res_, stream, original); - std::unique_ptr restored; + serialize(res_, original, stream); + std::unique_ptr restored; deserialize(res_, stream, &restored); ASSERT_NE(restored, nullptr); expect_equivalent(original, *restored); @@ -171,9 +171,9 @@ TEST_P(VpqSerializationTest, RoundTrip) { SCOPED_TRACE("through a file"); - const std::string path = "cuvs_vpq_serialization_test.bin"; - serialize(res_, path, original); - std::unique_ptr restored; + 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); @@ -183,7 +183,7 @@ TEST_P(VpqSerializationTest, RoundTrip) // 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 = { +const std::vector pq_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}, @@ -195,9 +195,9 @@ const std::vector vpq_serialization_inputs = { {500, 32, 4, 16, 0, 42ULL}, }; -INSTANTIATE_TEST_CASE_P(VpqSerializationTests, - VpqSerializationTest, - ::testing::ValuesIn(vpq_serialization_inputs)); +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) @@ -208,53 +208,53 @@ static void write_preamble(const raft::resources& res, std::ostream& os, int ver raft::serialize_scalar(res, os, version); } -TEST(VpqSerialization, RejectsEmptyStream) +TEST(PqSerialization, RejectsEmptyStream) { raft::resources res; std::stringstream stream; - std::unique_ptr restored; + std::unique_ptr restored; EXPECT_THROW(deserialize(res, stream, &restored), raft::exception); } -TEST(VpqSerialization, RejectsForeignDtypePrefix) +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, vpq_serialization_version); + raft::serialize_scalar(res, stream, pq_serialization_version); - std::unique_ptr restored; + std::unique_ptr restored; EXPECT_THROW(deserialize(res, stream, &restored), raft::exception); } -TEST(VpqSerialization, RejectsFutureVersion) +TEST(PqSerialization, RejectsFutureVersion) { raft::resources res; std::stringstream stream; - write_preamble(res, stream, vpq_serialization_version + 1); + write_preamble(res, stream, pq_serialization_version + 1); - std::unique_ptr restored; + std::unique_ptr restored; EXPECT_THROW(deserialize(res, stream, &restored), raft::exception); } -TEST(VpqSerialization, RejectsTruncatedPayload) +TEST(PqSerialization, RejectsTruncatedPayload) { raft::resources res; std::stringstream stream; - write_preamble(res, stream, vpq_serialization_version); + 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; + std::unique_ptr restored; EXPECT_THROW(deserialize(res, stream, &restored), raft::exception); } -TEST(VpqSerialization, RejectsNullOutParameter) +TEST(PqSerialization, RejectsNullOutParameter) { raft::resources res; std::stringstream stream; - write_preamble(res, stream, vpq_serialization_version); + write_preamble(res, stream, pq_serialization_version); EXPECT_THROW(deserialize(res, stream, nullptr), raft::exception); } From 43d1acedd26e4331367481d1c96e0c460907b10f Mon Sep 17 00:00:00 2001 From: Irina Reshodko Date: Wed, 19 Aug 2026 07:54:50 -0700 Subject: [PATCH 7/8] Updated comments and doxygenised the function annotations --- c/src/neighbors/cagra.cpp | 2 +- cpp/include/cuvs/neighbors/cagra.hpp | 134 ++++++++++++------ .../cuvs/preprocessing/quantize/pq.hpp | 8 +- .../neighbors/ann_cagra/test_pq_serialize.cu | 4 +- cpp/tests/preprocessing/pq_serialization.cu | 2 +- 5 files changed, 95 insertions(+), 55 deletions(-) diff --git a/c/src/neighbors/cagra.cpp b/c/src/neighbors/cagra.cpp index 677968250d..6fd207bfc1 100644 --- a/c/src/neighbors/cagra.cpp +++ b/c/src/neighbors/cagra.cpp @@ -1052,7 +1052,7 @@ void dispatch_serialized_dataset_kind( // 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 PQ-compressed index to yet. - RAFT_FAIL("File holds a PQ-compressed (CAGRA-Q) dataset, which the C API has no dataset " + RAFT_FAIL("File holds a PQ-compressed dataset, which the C API has no dataset " "layout for; load it through the C++ API"); } } diff --git a/cpp/include/cuvs/neighbors/cagra.hpp b/cpp/include/cuvs/neighbors/cagra.hpp index ff7ee9e0d1..b1a046a2c2 100644 --- a/cpp/include/cuvs/neighbors/cagra.hpp +++ b/cpp/include/cuvs/neighbors/cagra.hpp @@ -921,10 +921,12 @@ using cagra_index_t = index>; /** - * The dataset type `deserialize` produces for `IndexT`: what that index views, made owning. + * @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 = @@ -1505,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. * @@ -1513,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] @@ -1539,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] @@ -1565,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] @@ -1591,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] @@ -1610,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. * @@ -1618,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] @@ -1644,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] @@ -1670,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] @@ -1696,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] @@ -1715,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. * @@ -1723,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] @@ -1751,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] @@ -1779,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] @@ -1807,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] @@ -1828,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. * @@ -1836,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] @@ -1864,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] @@ -1892,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] @@ -1920,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] @@ -2277,7 +2279,7 @@ enum class serialized_dataset_kind : std::uint32_t { /** Host-resident dataset using its standard row layout. */ host_standard = 4, /** - * Device-resident PQ-compressed dataset (CAGRA-Q). + * 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. @@ -2288,7 +2290,11 @@ enum class serialized_dataset_kind : std::uint32_t { /** Current experimental CAGRA serialization format version. */ inline constexpr int cagra_serialization_version = 6; -/** What a serialized CAGRA index says about itself. @see read_serialized_header */ +/** + * @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; @@ -2328,6 +2334,10 @@ auto read_serialized_header(raft::resources const& res, const std::string& filen * * 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; @@ -2888,102 +2898,132 @@ void deserialize(raft::resources const& handle, std::unique_ptr>* out_dataset = nullptr); -/* vpq_f16_index overloads (CAGRA-Q). +/** + * @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. As everywhere else the index holds a view: - * `deserialize` returns the owning dataset through `out_dataset`, which has to stay alive for as - * long as the index is used. + * and wants nothing from the file but the graph. * - * `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 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, diff --git a/cpp/include/cuvs/preprocessing/quantize/pq.hpp b/cpp/include/cuvs/preprocessing/quantize/pq.hpp index 44f120fffd..232ecc61aa 100644 --- a/cpp/include/cuvs/preprocessing/quantize/pq.hpp +++ b/cpp/include/cuvs/preprocessing/quantize/pq.hpp @@ -349,9 +349,9 @@ 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: 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. + * 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 @@ -373,7 +373,7 @@ inline constexpr int pq_serialization_version = 1; * 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-Q graph over them. + * // 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()); diff --git a/cpp/tests/neighbors/ann_cagra/test_pq_serialize.cu b/cpp/tests/neighbors/ann_cagra/test_pq_serialize.cu index 2e23b3196b..7152c0a0f7 100644 --- a/cpp/tests/neighbors/ann_cagra/test_pq_serialize.cu +++ b/cpp/tests/neighbors/ann_cagra/test_pq_serialize.cu @@ -4,7 +4,7 @@ */ /* - * Serializing a CAGRA index whose dataset is PQ-compressed (CAGRA-Q). + * 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 @@ -153,7 +153,7 @@ class CagraPqSerializeTest : public ::testing::Test { dataset_->data_handle(), std::min(n_queries, dataset_->extent(0)), dataset_->extent(1)); } - /** A CAGRA-Q index viewing `compressed` and the graph of `graph_source`; both must outlive it. */ + /** 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 { diff --git a/cpp/tests/preprocessing/pq_serialization.cu b/cpp/tests/preprocessing/pq_serialization.cu index 2fe38c432e..90dd237271 100644 --- a/cpp/tests/preprocessing/pq_serialization.cu +++ b/cpp/tests/preprocessing/pq_serialization.cu @@ -184,7 +184,7 @@ TEST_P(PqSerializationTest, RoundTrip) // 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 CAGRA-Q accepts. + // 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}, From b54113c2f1e2a06644fd33427a2e9530b58fe01c Mon Sep 17 00:00:00 2001 From: Irina Reshodko Date: Wed, 19 Aug 2026 08:14:30 -0700 Subject: [PATCH 8/8] removed C api redirection --- c/src/neighbors/cagra.cpp | 80 +++++++++++++++++++++------------------ 1 file changed, 44 insertions(+), 36 deletions(-) diff --git a/c/src/neighbors/cagra.cpp b/c/src/neighbors/cagra.cpp index 6fd207bfc1..f65f571713 100644 --- a/c/src/neighbors/cagra.cpp +++ b/c/src/neighbors/cagra.cpp @@ -15,8 +15,11 @@ #include #include #include +#include #include +#include #include +#include #include "../core/exceptions.hpp" #include "../core/interop.hpp" @@ -971,40 +974,47 @@ struct serialized_cagra_header { cuvs::neighbors::cagra::serialized_dataset_kind dataset_kind; }; -/** - * What the file holds, in the terms the C entry points dispatch on. - * - * The reading is the C++ API's: `cagra::read_serialized_header` parses the same preamble, validates - * the format version and the dataset kind, and leaves this with nothing to do but restate the - * element type in DLPack terms. - */ static auto read_serialized_header(cuvsResources_t res, const char *filename) -> serialized_cagra_header { auto res_ptr = reinterpret_cast(res); - auto const header = cuvs::neighbors::cagra::read_serialized_header( - *res_ptr, std::string(filename)); - DLDataType dtype{.code = 0, .bits = 0, .lanes = 1}; - switch (header.dtype) { - case CUDA_R_32F: - dtype.code = kDLFloat; - dtype.bits = 32; - break; - case CUDA_R_16F: - dtype.code = kDLFloat; - dtype.bits = 16; - break; - case CUDA_R_8I: - dtype.code = kDLInt; - dtype.bits = 8; - break; - case CUDA_R_8U: - dtype.code = kDLUInt; - dtype.bits = 8; - break; - default: - RAFT_FAIL("Unsupported dtype in file %s", filename); + std::ifstream is(filename, std::ios::in | std::ios::binary); + if (!is) { + RAFT_FAIL("Cannot open file %s", filename); } - return {dtype, header.dataset_kind}; + + char dtype_string[4]{}; + if (!is.read(dtype_string, sizeof(dtype_string))) { + RAFT_FAIL("Invalid or truncated index header in file %s", filename); + } + + auto const dtype = raft::numpy_serializer::parse_descr( + std::string(dtype_string, sizeof(dtype_string))); + DLDataType output_dtype{ + .code = 0, .bits = static_cast(dtype.itemsize * 8), .lanes = 1}; + if (dtype.kind == 'f' && dtype.itemsize == 4) { + output_dtype.code = kDLFloat; + } else if (dtype.kind == 'e' && dtype.itemsize == 2) { + output_dtype.code = kDLFloat; + } else if (dtype.kind == 'i' && dtype.itemsize == 1) { + output_dtype.code = kDLInt; + } else if (dtype.kind == 'u' && dtype.itemsize == 1) { + output_dtype.code = kDLUInt; + } else { + RAFT_FAIL("Unsupported dtype in file %s", filename); + } + + auto const version = raft::deserialize_scalar(*res_ptr, is); + auto const dataset_kind_raw = + raft::deserialize_scalar(*res_ptr, is); + RAFT_EXPECTS( + version == cuvs::neighbors::cagra::cagra_serialization_version, + "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), + "Invalid serialized dataset kind %u in file %s", + dataset_kind_raw, filename); + return {output_dtype, static_cast(dataset_kind_raw)}; } template @@ -1048,12 +1058,10 @@ void dispatch_serialized_dataset_kind( fn.template operator()< cuvs::neighbors::device_padded_dataset_view>(); break; - case serialized_kind::device_pq: - // 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 PQ-compressed index to yet. - RAFT_FAIL("File holds a PQ-compressed dataset, which the C API has no dataset " - "layout for; load it through the C++ API"); + // 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; } }