From 354f32e313bf4e79d22b296184bf76a2826bded7 Mon Sep 17 00:00:00 2001 From: achirkin Date: Mon, 31 Aug 2026 16:25:02 +0200 Subject: [PATCH 01/11] Implement dataset as a single type+specs --- cpp/include/cuvs/core/dataset.hpp | 489 ++++++++++++++++++++++++++++ examples/cpp/CMakeLists.txt | 2 + examples/cpp/src/dataset_example.cu | 134 ++++++++ 3 files changed, 625 insertions(+) create mode 100644 cpp/include/cuvs/core/dataset.hpp create mode 100644 examples/cpp/src/dataset_example.cu diff --git a/cpp/include/cuvs/core/dataset.hpp b/cpp/include/cuvs/core/dataset.hpp new file mode 100644 index 0000000000..e436f0502b --- /dev/null +++ b/cpp/include/cuvs/core/dataset.hpp @@ -0,0 +1,489 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include + +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +namespace CUVS_EXPORT cuvs { + +/************************************************************************************************ + * The core dataset type + ************************************************************************************************/ + +/** + * A spec is compressing if it defines a dictionary: the codebooks, the quantization range, or + * whatever else is needed to interpret the data. Plain specs declare `std::monostate` instead, so + * "there is no dictionary" and "the dictionary slot is empty" are the same vocabulary type. + */ +template +concept compressed_dataset_spec = requires { + typename SpecT::dictionary_type; + typename SpecT::dictionary_view_type; + typename SpecT::code_type; +} && !std::same_as; + +template +struct dataset { + using spec_type = typename SpecT::template apply; + using value_type = typename spec_type::value_type; + using index_type = typename spec_type::index_type; + + using data_type = typename spec_type::data_type; + using view_type = typename spec_type::view_type; + using const_view_type = typename spec_type::const_view_type; + + /* The second, optional storage slot; `std::monostate` for an uncompressed dataset. */ + using dictionary_type = typename spec_type::dictionary_type; + + static constexpr bool is_compressed = compressed_dataset_spec; + + explicit dataset(data_type&& data) + requires(!is_compressed) + : data_{std::make_shared(std::move(data))} + { + } + + dataset(data_type&& data, std::shared_ptr dictionary) + requires(is_compressed) + : data_{std::make_shared(std::move(data))}, dictionary_{std::move(dictionary)} + { + } + + /* + @return A non-owning view of the data, such as mdspan. + */ + [[nodiscard]] auto data_view() noexcept -> view_type { return spec_type::get_data_view(*data_); } + [[nodiscard]] auto data_const_view() const noexcept -> const_view_type + { + return spec_type::get_data_const_view(*data_); + } + [[nodiscard]] auto n_rows() const noexcept -> index_type { return spec_type::get_n_rows(*data_); } + // NOTE: a compressed dataset may not be able to tell its dimension from the encoded data alone + [[nodiscard]] auto dim() const noexcept -> uint32_t + { + return spec_type::get_dim(*data_, dict_ref()); + } + + /* + @return A non-owning view of the dictionary, such as the two codebooks of a VPQ dataset. + */ + [[nodiscard]] auto dictionary_view() const noexcept + requires(is_compressed) + { + return spec_type::get_dictionary_view(*dictionary_); + } + /* + A dictionary is immutable once trained, so several datasets may encode against the same one. + + @return A shared handle to the dictionary. + */ + [[nodiscard]] auto share_dictionary() const noexcept -> std::shared_ptr + requires(is_compressed) + { + return dictionary_; + } + + private: + std::shared_ptr data_; + /* Either a shared handle to the dictionary, or the empty dictionary itself. */ + [[no_unique_address]] std::conditional_t, + dictionary_type> dictionary_; + + [[nodiscard]] auto dict_ref() const noexcept -> const dictionary_type& + { + if constexpr (is_compressed) { + return *dictionary_; + } else { + return dictionary_; + } + } +}; + +/************************************************************************************************ + * Dataset specifications + ************************************************************************************************/ + +/** + * Empty dataset specification: contains no data, but keeps the dimension of the dataset implement + * the common interface. + */ +struct empty_spec { + struct empty_dataset_rep { + uint32_t dim; + }; + + template + struct apply { + using data_type = empty_dataset_rep; + using view_type = empty_dataset_rep; + using const_view_type = const empty_dataset_rep; + using value_type = std::remove_cv_t; + using index_type = std::remove_cv_t; + using dictionary_type = std::monostate; + + [[nodiscard]] static auto get_data_view(data_type& data) noexcept -> view_type { return data; } + [[nodiscard]] static auto get_data_const_view(const data_type& data) noexcept -> const_view_type + { + return data; + } + [[nodiscard]] static auto get_n_rows(const data_type& data) noexcept -> index_type { return 0; } + [[nodiscard]] static auto get_dim(const data_type& data, const dictionary_type&) noexcept + -> uint32_t + { + return data.dim; + } + }; +}; + +/** + * Dense plain/padded dataset implemented via raft::mdarray. + */ +template +struct mdarray_spec { + template + struct apply { + /* NOTE: index type != extents type + The index type can vary depending on the use case; we often may want to store indices in 32-bit + slots to save memory on large datasets. The extents are always fixed to 64-bit signed integers + to simplify the conversion and avoid integer overflow issues. + */ + using data_type = raft::mdarray, LayoutPolicy, ContainerPolicy>; + using view_type = typename data_type::view_type; + using const_view_type = typename data_type::const_view_type; + using value_type = std::remove_cv_t; + using index_type = std::remove_cv_t; + using dictionary_type = std::monostate; + + [[nodiscard]] static auto get_data_view(data_type& data) noexcept -> view_type + { + return data.view(); + } + + [[nodiscard]] static auto get_data_const_view(const data_type& data) noexcept -> const_view_type + { + return data.view(); + } + [[nodiscard]] static auto get_n_rows(const data_type& data) noexcept -> index_type + { + return static_cast(data.extent(0)); + } + [[nodiscard]] static auto get_dim(const data_type& data, const dictionary_type&) noexcept + -> uint32_t + { + return static_cast(data.extent(1)); + } + }; +}; + +/** + * Sparse dataset specification implemented via the raft sparse matrix hierarchy. `SparseLayoutT` + * selects the representation (`csr_layout`, `coo_layout`). + * + * NOTE: only the memory space of the container policy is used here; raft's sparse types bind the + * policy themselves, separately for the values, the offsets and the indices. + */ +template +struct sparse_spec { + template + struct apply { + using value_type = std::remove_cv_t; + using index_type = std::remove_cv_t; + using dictionary_type = std::monostate; + + static constexpr raft::memory_type mem_type = ContainerPolicy::mem_type; + static_assert(mem_type == raft::memory_type::host || mem_type == raft::memory_type::device, + "raft's sparse hierarchy only provides host and device containers"); + + using data_type = std::conditional_t< + mem_type == raft::memory_type::device, + typename SparseLayoutT:: + template matrix_type, + typename SparseLayoutT:: + template matrix_type>; + using view_type = typename data_type::view_type; + /* NOTE: raft's sparse hierarchy provides no read-only view yet (see the TODO in + raft/core/sparse_types.hpp); until then, even reading the structure requires a mutable matrix. + */ + using const_view_type = view_type; + + [[nodiscard]] static auto get_data_view(data_type& data) noexcept -> view_type + { + return data.view(); + } + + [[nodiscard]] static auto get_data_const_view(const data_type& data) noexcept -> const_view_type + { + return const_cast(data).view(); + } + [[nodiscard]] static auto get_n_rows(const data_type& data) noexcept -> index_type + { + return static_cast(const_cast(data).structure_view().get_n_rows()); + } + [[nodiscard]] static auto get_dim(const data_type& data, const dictionary_type&) noexcept + -> uint32_t + { + return static_cast(const_cast(data).structure_view().get_n_cols()); + } + }; +}; + +/** + * CSR representation: one compressed row per vector of the dataset. + * + * NOTE: the row offsets are 64-bit for the same reason the dense extents are: the number of + * non-zeros is not bounded by the index type. The column indices are feature ids, so they are + * bounded by the dataset dimension rather than by its size. + */ +struct csr_layout { + template typename ContainerPolicy> + using matrix_type = raft::csr_matrix; +}; + +/** + * COO representation: one (row, column, value) triple per non-zero. + * + * NOTE: unlike CSR, the row component is an index into the dataset rather than an offset into the + * non-zeros, hence the dataset index type. + */ +struct coo_layout { + template typename ContainerPolicy> + using matrix_type = raft::coo_matrix; +}; + +/** + * VQ+PQ compressed dataset specification, mirroring the VPQ dataset of cuvs/neighbors/common.hpp: + * the two codebooks form the dictionary, while the data slot keeps the encoded rows - the VQ label + * in the row prefix, followed by the packed PQ codes. + * + * `BookPolicy` is the container policy of the codebooks (elements of `MathT`) and `CodePolicy` the + * one of the encoded rows (elements of `uint8_t`); a bound policy cannot be rebound to another + * element type, hence the two parameters. + */ +template +struct vpq_spec { + /* The encoded rows are stored like the rows of any other dense dataset; only how a row is read + and where the dimension comes from are specific to the compression. + */ + template + using storage_spec = + typename mdarray_spec::template apply; + + template + struct apply : storage_spec { + /* The dataset stores codes, but it still represents vectors of `T`. */ + using value_type = std::remove_cv_t; + /* NOTE: the members of a dependent base are not visible to unqualified lookup, so `data_type` + has to be pulled into this scope to be usable in the signatures below. + */ + using typename storage_spec::data_type; + /* The type the codebooks are stored in; the rows themselves are always packed codes. */ + using math_type = MathT; + using code_type = uint8_t; + + /* [vq_n_centers, dim] */ + using vq_book_type = + raft::mdarray, raft::layout_c_contiguous, BookPolicy>; + /* [n_books, pq_n_centers, pq_len], where `n_books` is 1 for a codebook shared by all the + subspaces, or `pq_dim` for one codebook per subspace. A row-major 3d array with `n_books == 1` + is bit-identical to the flat 2d codebook of the current implementation. + */ + using pq_book_type = + raft::mdarray, raft::layout_c_contiguous, BookPolicy>; + + struct dictionary_type { + vq_book_type vq_code_book; + pq_book_type pq_code_book; + }; + + struct dictionary_view_type { + typename vq_book_type::const_view_type vq_code_book; + typename pq_book_type::const_view_type pq_code_book; + + /* NOTE: the dimension of the dataset is a property of the VQ codebook: the encoded rows tell + only how many bytes it took to compress a vector, which the padding makes ambiguous. + */ + [[nodiscard]] auto dim() const noexcept -> uint32_t + { + return static_cast(vq_code_book.extent(1)); + } + [[nodiscard]] auto vq_n_centers() const noexcept -> uint32_t + { + return static_cast(vq_code_book.extent(0)); + } + [[nodiscard]] auto n_books() const noexcept -> uint32_t + { + return static_cast(pq_code_book.extent(0)); + } + [[nodiscard]] auto pq_n_centers() const noexcept -> uint32_t + { + return static_cast(pq_code_book.extent(1)); + } + [[nodiscard]] auto pq_len() const noexcept -> uint32_t + { + return static_cast(pq_code_book.extent(2)); + } + [[nodiscard]] auto pq_bits() const noexcept -> uint32_t + { + return static_cast(std::countr_zero(pq_n_centers())); + } + [[nodiscard]] auto pq_dim() const noexcept -> uint32_t + { + return raft::div_rounding_up_safe(dim(), pq_len()); + } + /* Whether every subspace has a codebook of its own, as opposed to sharing a single one. */ + [[nodiscard]] auto per_subspace() const noexcept -> bool { return n_books() > 1; } + }; + + /* NOTE: the data accessors are inherited; only `get_dim` differs from a dense dataset. */ + [[nodiscard]] static auto get_dim(const data_type&, const dictionary_type& dict) noexcept + -> uint32_t + { + return static_cast(dict.vq_code_book.extent(1)); + } + [[nodiscard]] static auto get_dictionary_view(const dictionary_type& dict) noexcept + -> dictionary_view_type + { + return {dict.vq_code_book.view(), dict.pq_code_book.view()}; + } + /* The length of an encoded row in bytes, including the inlined VQ label. */ + [[nodiscard]] static auto get_encoded_row_length(const data_type& data) noexcept -> uint32_t + { + return static_cast(data.extent(1)); + } + }; +}; + +/** + * Scalar-quantized dataset specification: every component of a vector becomes one code, so the + * dictionary is only the range the codes are mapped back onto and the dimension is still a property + * of the data. + */ +template +struct sq_spec { + /* One code per component, so the codes are laid out exactly like a dense dataset. */ + template + using storage_spec = + typename mdarray_spec::template apply; + + template + struct apply : storage_spec { + /* The dataset stores codes, but it still represents vectors of `T`. */ + using value_type = std::remove_cv_t; + /* NOTE: the members of a dependent base are not visible to unqualified lookup, so `data_type` + has to be pulled into this scope to be usable in the signature below. + */ + using typename storage_spec::data_type; + using math_type = MathT; + using code_type = int8_t; + + /* Nothing to allocate: the whole dictionary is the interval the codes are dequantized into. + The members are named as in cuvs::preprocessing::quantize::scalar::quantizer. + */ + struct dictionary_type { + math_type min_; + math_type max_; + }; + using dictionary_view_type = dictionary_type; + + /* NOTE: only the signature differs from the inherited accessor, which cannot be reused because + it expects an empty dictionary. + */ + [[nodiscard]] static auto get_dim(const data_type& data, const dictionary_type&) noexcept + -> uint32_t + { + return static_cast(data.extent(1)); + } + [[nodiscard]] static auto get_dictionary_view(const dictionary_type& dict) noexcept + -> dictionary_view_type + { + return dict; + } + }; +}; + +/************************************************************************************************ + * Convenience aliases + ************************************************************************************************/ + +// Empty +template +using empty_dataset = dataset; + +// Dense +template +using contiguous_dataset = + dataset>; +template +using padded_dataset = dataset, ContainerPolicy>>; +template +using device_contiguous_dataset = + contiguous_dataset>>; +template +using device_padded_dataset = + padded_dataset>>; +template +using host_contiguous_dataset = + contiguous_dataset>>; +template +using host_padded_dataset = + padded_dataset>>; + +// Sparse +template +using csr_dataset = dataset>; +template +using coo_dataset = dataset>; +template +using device_csr_dataset = + csr_dataset>>; +template +using device_coo_dataset = + coo_dataset>>; +template +using host_csr_dataset = csr_dataset>>; +template +using host_coo_dataset = coo_dataset>>; + +// Compressed +template +using vpq_dataset = dataset>; +template +using sq_dataset = dataset>; +template +using device_vpq_dataset = + vpq_dataset>, + raft::device_accessor>>; +template +using host_vpq_dataset = vpq_dataset>, + raft::host_accessor>>; +template +using device_sq_dataset = + sq_dataset>>; +template +using host_sq_dataset = + sq_dataset>>; + +} // namespace CUVS_EXPORT cuvs diff --git a/examples/cpp/CMakeLists.txt b/examples/cpp/CMakeLists.txt index f4dffad3a9..7a4e0f504f 100644 --- a/examples/cpp/CMakeLists.txt +++ b/examples/cpp/CMakeLists.txt @@ -44,6 +44,7 @@ add_executable(CAGRA_BLOOM_FILTER_EXAMPLE src/cagra_bloom_filter_example.cu) add_executable(CAGRA_HNSW_ACE_BUILD_EXAMPLE src/cagra_hnsw_ace_build.cu) add_executable(CAGRA_HNSW_ACE_EXAMPLE src/cagra_hnsw_ace_example.cu) add_executable(CAGRA_PERSISTENT_EXAMPLE src/cagra_persistent_example.cu) +add_executable(DATASET_EXAMPLE src/dataset_example.cu) add_executable(DYNAMIC_BATCHING_EXAMPLE src/dynamic_batching_example.cu) add_executable(HNSW_ACE_EXAMPLE src/hnsw_ace_example.cu) add_executable(HNSW_OPENAI_EXAMPLE src/hnsw_openai_example.cu) @@ -73,6 +74,7 @@ target_link_libraries( target_link_libraries( DYNAMIC_BATCHING_EXAMPLE PRIVATE cuvs::cuvs $ Threads::Threads ) +target_link_libraries(DATASET_EXAMPLE PRIVATE cuvs::cuvs $) target_link_libraries(HNSW_ACE_EXAMPLE PRIVATE cuvs::cuvs $) target_link_libraries(HNSW_OPENAI_EXAMPLE PRIVATE cuvs::cuvs $) diff --git a/examples/cpp/src/dataset_example.cu b/examples/cpp/src/dataset_example.cu new file mode 100644 index 0000000000..39f4036baa --- /dev/null +++ b/examples/cpp/src/dataset_example.cu @@ -0,0 +1,134 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include + +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +/** The interface every spec provides, regardless of how the data is stored. */ +template +void print_dataset(const char* name, const DatasetT& dataset) +{ + std::cout << name << ": n_rows = " << dataset.n_rows() << ", dim = " << dataset.dim() << '\n'; +} + +/* The dictionary slot is free for the datasets that have no dictionary. */ +static_assert(!cuvs::device_contiguous_dataset::is_compressed); +static_assert( + std::same_as::dictionary_type, std::monostate>); +static_assert( + std::same_as::dictionary_type, std::monostate>); +static_assert(std::same_as::dictionary_type, std::monostate>); +static_assert(sizeof(cuvs::device_contiguous_dataset) == + sizeof(std::shared_ptr)); +static_assert(cuvs::device_vpq_dataset::is_compressed); + +int main() +{ + raft::resources res; + + int64_t n_rows = 8; + int64_t dim = 4; + + // An empty dataset carries the dimension only; there is no storage behind it. + cuvs::empty_dataset empty{{static_cast(dim)}}; + print_dataset("empty", empty); + + // A host dataset takes ownership of a raft::mdarray of a matching layout and storage policy. + auto host_data = raft::make_host_matrix(res, n_rows, dim); + for (int64_t i = 0; i < n_rows; i++) { + for (int64_t j = 0; j < dim; j++) { + host_data(i, j) = static_cast(i * dim + j); + } + } + cuvs::host_contiguous_dataset host_dataset{std::move(host_data)}; + print_dataset("host contiguous", host_dataset); + std::cout << " last element = " << host_dataset.data_view()(n_rows - 1, dim - 1) << '\n'; + + // The device variant differs only in the storage policy baked into the spec. + cuvs::device_contiguous_dataset device_dataset{ + raft::make_device_matrix(res, n_rows, dim)}; + print_dataset("device contiguous", device_dataset); + + // The same dataset over a padded layout: the spec, not the dataset, selects the layout. + using padded_dataset_type = cuvs::device_padded_dataset; + using padded_data_type = typename padded_dataset_type::data_type; + padded_dataset_type padded_dataset{padded_data_type{ + res, + typename padded_data_type::mapping_type{raft::make_extents(n_rows, dim)}, + typename padded_data_type::container_policy_type{}}}; + print_dataset("device padded", padded_dataset); + + // A sparse dataset: one CSR row per vector, so `dim` is the number of features. + uint64_t nnz = static_cast(n_rows) * 2; + cuvs::device_csr_dataset csr_dataset{ + raft::make_device_csr_matrix( + res, n_rows, static_cast(dim), nnz)}; + print_dataset("device csr", csr_dataset); + std::cout << " nnz = " << csr_dataset.data_view().structure_view().get_nnz() << '\n'; + + // The same dataset as a list of (row, column, value) triples. + cuvs::device_coo_dataset coo_dataset{ + raft::make_device_coo_matrix( + res, static_cast(n_rows), static_cast(dim), nnz)}; + print_dataset("device coo", coo_dataset); + + // A compressed dataset fills the second slot with a dictionary; for VPQ that is the pair of + // codebooks. Note the dimension of the dataset comes from the VQ codebook here, since the encoded + // rows know only how many bytes it took to compress a vector. + using vpq_dataset_type = cuvs::device_vpq_dataset; + using vpq_dictionary_type = typename vpq_dataset_type::dictionary_type; + + int64_t vq_n_centers = 4; + int64_t pq_n_centers = 256; + int64_t pq_len = 2; + // One byte per code at pq_bits == 8, prefixed with the inlined VQ label. + int64_t encoded_row_length = sizeof(uint32_t) + dim / pq_len; + + vpq_dictionary_type codebooks{ + raft::make_device_matrix(res, vq_n_centers, dim), + // A single codebook shared by all the subspaces; `pq_dim` of them would be per-subspace. + raft::make_device_mdarray( + res, raft::make_extents(1, pq_n_centers, pq_len))}; + auto dictionary = std::make_shared(std::move(codebooks)); + + vpq_dataset_type vpq_dataset{ + raft::make_device_matrix(res, n_rows, encoded_row_length), dictionary}; + print_dataset("device vpq", vpq_dataset); + auto books = vpq_dataset.dictionary_view(); + std::cout << " vq_n_centers = " << books.vq_n_centers() << ", pq_bits = " << books.pq_bits() + << ", pq_dim = " << books.pq_dim() << ", per_subspace = " << books.per_subspace() + << '\n'; + // The data slot is a plain dense matrix of bytes, hence the shared spec. + std::cout << " encoded row length = " << vpq_dataset.data_view().extent(1) << '\n'; + + // A dictionary is immutable, so another dataset may be encoded against the very same codebooks. + vpq_dataset_type shared_vpq_dataset{ + raft::make_device_matrix(res, n_rows * 2, encoded_row_length), + vpq_dataset.share_dictionary()}; + print_dataset("device vpq (shared dictionary)", shared_vpq_dataset); + + // Scalar quantization keeps one code per component, so the dimension is still a property of the + // data and the dictionary is only the interval to dequantize into. + using sq_dataset_type = cuvs::device_sq_dataset; + using sq_dictionary_type = typename sq_dataset_type::dictionary_type; + sq_dataset_type sq_dataset{ + raft::make_device_matrix(res, n_rows, dim), + std::make_shared(sq_dictionary_type{-1.0F, 1.0F})}; + print_dataset("device sq", sq_dataset); + std::cout << " range = [" << sq_dataset.dictionary_view().min_ << ", " + << sq_dataset.dictionary_view().max_ << "]\n"; + + return 0; +} From ea5c82911a1b00fc3468dc601e832898d3f43657 Mon Sep 17 00:00:00 2001 From: achirkin Date: Tue, 1 Sep 2026 11:52:01 +0200 Subject: [PATCH 02/11] Add an example of in-place dataset normalization exploiting shared ownership semantics --- cpp/include/cuvs/core/dataset.hpp | 23 +++++++++++++ examples/cpp/src/dataset_example.cu | 53 +++++++++++++++++++++++++++++ 2 files changed, 76 insertions(+) diff --git a/cpp/include/cuvs/core/dataset.hpp b/cpp/include/cuvs/core/dataset.hpp index e436f0502b..a5445db6f1 100644 --- a/cpp/include/cuvs/core/dataset.hpp +++ b/cpp/include/cuvs/core/dataset.hpp @@ -15,6 +15,7 @@ #include #include +#include #include #include #include @@ -80,6 +81,28 @@ struct dataset { return spec_type::get_dim(*data_, dict_ref()); } + /* + * Check if the caller is the only owner of the data. + * This is useful to determine if the data can be modified in place without copying. + */ + [[nodiscard]] auto is_data_unique() const noexcept -> bool + { + /* + The count cannot grow behind our back: only copying a `dataset` increments it and `data_` never + escapes the class, so no weak_ptr can resurrect a reference. Hence, if it reads one, the sole + reference is the one we hold and no other thread has anything to copy from. The converse does + not hold - a concurrent release may not be visible yet - but that only costs the optimization. + + The fence must follow the load: releasing a copy decrements with release ordering, yet + `use_count()` reads relaxed and shared_ptr pairs the decrement with an acquire only once the + counter reaches zero. Without the fence we could miss the writes of a copy that another thread + has already destroyed. + */ + if (data_.use_count() != 1) { return false; } + std::atomic_thread_fence(std::memory_order_acquire); + return true; + } + /* @return A non-owning view of the dictionary, such as the two codebooks of a VPQ dataset. */ diff --git a/examples/cpp/src/dataset_example.cu b/examples/cpp/src/dataset_example.cu index 39f4036baa..af51d6296a 100644 --- a/examples/cpp/src/dataset_example.cu +++ b/examples/cpp/src/dataset_example.cu @@ -9,11 +9,17 @@ #include #include #include +#include +#include #include +#include +#include +#include #include #include #include +#include #include /** The interface every spec provides, regardless of how the data is stored. */ @@ -23,6 +29,40 @@ void print_dataset(const char* name, const DatasetT& dataset) std::cout << name << ": n_rows = " << dataset.n_rows() << ", dim = " << dataset.dim() << '\n'; } +/** + * A transformation that takes its input by value and hands back the result. Written this way, it + * can skip the second buffer whenever the caller passes its only copy of the dataset: + * `d = normalize(res, std::move(d))` allocates only the column statistics. + */ +template +auto normalize(const raft::resources& res, cuvs::device_contiguous_dataset data) + -> cuvs::device_contiguous_dataset +{ + auto in = data.data_const_view(); + auto mean = raft::make_device_vector(res, in.extent(1)); + auto var = raft::make_device_vector(res, in.extent(1)); + // Both statistics come out of a single sweep over the data. + raft::stats::meanvar(res, in, mean.view(), var.view(), false); + auto mean_view = raft::make_const_mdspan(mean.view()); + auto var_view = raft::make_const_mdspan(var.view()); + // The linewise op broadcasts the per-column statistics for us; a constant column stays at zero. + auto standardize = [] __device__(T x, T mean, T var) -> T { + return var > T{0} ? (x - mean) / raft::sqrt(var) : T{0}; + }; + + if (data.is_data_unique()) { + // No other copy can observe the buffer, so the transform may write over its own input. + raft::linalg::matrix_vector_op( + res, in, mean_view, var_view, data.data_view(), standardize); + return data; + } + cuvs::device_contiguous_dataset out{ + raft::make_device_matrix(res, in.extent(0), in.extent(1))}; + raft::linalg::matrix_vector_op( + res, in, mean_view, var_view, out.data_view(), standardize); + return out; +} + /* The dictionary slot is free for the datasets that have no dictionary. */ static_assert(!cuvs::device_contiguous_dataset::is_compressed); static_assert( @@ -61,6 +101,19 @@ int main() raft::make_device_matrix(res, n_rows, dim)}; print_dataset("device contiguous", device_dataset); + // Moving the dataset in leaves `normalize` as its only owner, so the data is scaled in-place. + raft::linalg::map_offset(res, device_dataset.data_view(), raft::cast_op{}); + const auto* original_data = device_dataset.data_const_view().data_handle(); + device_dataset = normalize(res, std::move(device_dataset)); + std::cout << " reused the input buffer: " << std::boolalpha + << (device_dataset.data_const_view().data_handle() == original_data) << '\n'; + + // Passing a copy keeps a second owner alive, so the same call has to allocate its output. + auto shared_dataset = device_dataset; + auto rescaled = normalize(res, shared_dataset); + std::cout << " reused the input buffer: " + << (rescaled.data_const_view().data_handle() == original_data) << '\n'; + // The same dataset over a padded layout: the spec, not the dataset, selects the layout. using padded_dataset_type = cuvs::device_padded_dataset; using padded_data_type = typename padded_dataset_type::data_type; From 15b31c78c0cabcb4bcc36cd64592854e37a2a7ef Mon Sep 17 00:00:00 2001 From: achirkin Date: Wed, 14 Oct 2026 16:01:35 +0200 Subject: [PATCH 03/11] Dataset API without shared pointers --- cpp/include/cuvs/core/dataset.hpp | 585 ++++++++++++++------- examples/cpp/src/dataset_example.cu | 87 +-- fern/docs.yml | 2 + fern/pages/cpp_api/cpp-api-core-dataset.md | 208 ++++++++ fern/pages/cpp_api/index.md | 1 + 5 files changed, 673 insertions(+), 210 deletions(-) create mode 100644 fern/pages/cpp_api/cpp-api-core-dataset.md diff --git a/cpp/include/cuvs/core/dataset.hpp b/cpp/include/cuvs/core/dataset.hpp index a5445db6f1..18f839a569 100644 --- a/cpp/include/cuvs/core/dataset.hpp +++ b/cpp/include/cuvs/core/dataset.hpp @@ -15,11 +15,10 @@ #include #include -#include #include -#include #include #include +#include namespace CUVS_EXPORT cuvs { @@ -29,114 +28,270 @@ namespace CUVS_EXPORT cuvs { /** * A spec is compressing if it defines a dictionary: the codebooks, the quantization range, or - * whatever else is needed to interpret the data. Plain specs declare `std::monostate` instead, so - * "there is no dictionary" and "the dictionary slot is empty" are the same vocabulary type. + * whatever else is needed to interpret the data. Plain specs define none and the dataset + * substitutes `std::monostate`, so "there is no dictionary" and "the dictionary slot is empty" are + * the same vocabulary type. */ template -concept compressed_dataset_spec = requires { - typename SpecT::dictionary_type; - typename SpecT::dictionary_view_type; - typename SpecT::code_type; -} && !std::same_as; +concept compressed_dataset_spec = requires { typename SpecT::dictionary_type; }; +/** + * A non-owning view of a dataset: a view of the data and, for a compressed dataset, a view of the + * dictionary that gives it meaning. Cheap to copy and to pass by value, like an mdspan. + */ template -struct dataset { - using spec_type = typename SpecT::template apply; +struct dataset_view { + using spec_type = typename SpecT::template apply; + /* The read-only flavor, in which each view type collapses onto its const counterpart. */ + static constexpr bool is_const_view = std::is_const_v; + using value_type = typename spec_type::value_type; using index_type = typename spec_type::index_type; + /** + What one element of the data slot is + - for compressed datasets it is the raw data element type; often is uint8_t but may be + a larger type to enforce better alignment. + - for uncompressed datasets it is the value type of the dataset. + */ + using code_type = typename decltype([] { + if constexpr (requires { typename spec_type::code_type; }) { + return std::type_identity{}; + } else { + return std::type_identity{}; + } + }())::type; - using data_type = typename spec_type::data_type; - using view_type = typename spec_type::view_type; - using const_view_type = typename spec_type::const_view_type; + using data_type = typename spec_type::data_type; + using const_data_view_type = typename spec_type::const_data_view_type; + using data_view_type = + std::conditional_t; - /* The second, optional storage slot; `std::monostate` for an uncompressed dataset. */ - using dictionary_type = typename spec_type::dictionary_type; + /** + Optional slot that contains everything to give meaning to the encoded data; + `std::monostate` for an uncompressed dataset. + */ + using dictionary_type = typename decltype([] { + if constexpr (requires { typename spec_type::dictionary_type; }) { + return std::type_identity{}; + } else { + return std::type_identity{}; + } + }())::type; + // Optional, so that the specs don't have to declare it; a spec that declares neither view uses + // the dictionary itself, and one that declares only the mutable view uses it for both. + using const_dictionary_view_type = typename decltype([] { + if constexpr (requires { typename spec_type::const_dictionary_view_type; }) { + return std::type_identity{}; + } else if constexpr (requires { typename spec_type::dictionary_view_type; }) { + return std::type_identity{}; + } else { + return std::type_identity{}; + } + }())::type; + // Optional, so that the specs don't have to declare it + using dictionary_view_type = typename decltype([] { + if constexpr (is_const_view) { + return std::type_identity{}; + } else if constexpr (requires { typename spec_type::dictionary_view_type; }) { + return std::type_identity{}; + } else { + return std::type_identity{}; + } + }())::type; static constexpr bool is_compressed = compressed_dataset_spec; - explicit dataset(data_type&& data) + explicit dataset_view(data_view_type data_view) requires(!is_compressed) - : data_{std::make_shared(std::move(data))} + : data_view_{data_view} { } - dataset(data_type&& data, std::shared_ptr dictionary) + dataset_view(data_view_type data_view, dictionary_view_type dictionary_view) requires(is_compressed) - : data_{std::make_shared(std::move(data))}, dictionary_{std::move(dictionary)} + : data_view_{data_view}, dictionary_view_{dictionary_view} { } - /* - @return A non-owning view of the data, such as mdspan. + /* A mutable view decays to a read-only one, the way an mdspan does. + + NOTE: a constructor template, because in the mutable flavor the parameter would spell the very + same type; a non-template would be a copy constructor there, and declaring one - even a + constrained-away one - suppresses the implicitly declared copy constructor. + + NOTE: the dictionary is assigned in the body rather than in the init-list, because the accessor it + comes from is constrained on `is_compressed`. */ - [[nodiscard]] auto data_view() noexcept -> view_type { return spec_type::get_data_view(*data_); } - [[nodiscard]] auto data_const_view() const noexcept -> const_view_type + template + requires(is_const_view && std::same_as>) + dataset_view(const dataset_view& other) noexcept + : data_view_{other.data_const_view()} { - return spec_type::get_data_const_view(*data_); + if constexpr (is_compressed) { dictionary_view_ = other.const_dictionary_view(); } + } + + [[nodiscard]] auto n_rows() const noexcept -> index_type + { + return spec_type::get_n_rows(data_view_); } - [[nodiscard]] auto n_rows() const noexcept -> index_type { return spec_type::get_n_rows(*data_); } // NOTE: a compressed dataset may not be able to tell its dimension from the encoded data alone [[nodiscard]] auto dim() const noexcept -> uint32_t { - return spec_type::get_dim(*data_, dict_ref()); + return spec_type::get_dim(data_view_, dictionary_view_); } /* - * Check if the caller is the only owner of the data. - * This is useful to determine if the data can be modified in place without copying. - */ - [[nodiscard]] auto is_data_unique() const noexcept -> bool + @return A non-owning view of the data, such as mdspan. + */ + [[nodiscard]] auto data_view() const noexcept -> data_view_type { return data_view_; } + /* + NOTE: giving up write access is not a plain cast for most view types - a mutable mdspan, for + one, has to be rebuilt over a const accessor - hence the spec does the conversion. + + @return A read-only view of the data. + */ + [[nodiscard]] auto data_const_view() const noexcept -> const_data_view_type { - /* - The count cannot grow behind our back: only copying a `dataset` increments it and `data_` never - escapes the class, so no weak_ptr can resurrect a reference. Hence, if it reads one, the sole - reference is the one we hold and no other thread has anything to copy from. The converse does - not hold - a concurrent release may not be visible yet - but that only costs the optimization. - - The fence must follow the load: releasing a copy decrements with release ordering, yet - `use_count()` reads relaxed and shared_ptr pairs the decrement with an acquire only once the - counter reaches zero. Without the fence we could miss the writes of a copy that another thread - has already destroyed. - */ - if (data_.use_count() != 1) { return false; } - std::atomic_thread_fence(std::memory_order_acquire); - return true; + if constexpr (is_const_view) { + return data_view_; + } else { + return spec_type::get_data_const_view(data_view_); + } } /* @return A non-owning view of the dictionary, such as the two codebooks of a VPQ dataset. */ - [[nodiscard]] auto dictionary_view() const noexcept + [[nodiscard]] auto dictionary_view() const noexcept -> dictionary_view_type requires(is_compressed) { - return spec_type::get_dictionary_view(*dictionary_); + return dictionary_view_; } /* - A dictionary is immutable once trained, so several datasets may encode against the same one. - - @return A shared handle to the dictionary. + @return A read-only view of the dictionary. */ - [[nodiscard]] auto share_dictionary() const noexcept -> std::shared_ptr + [[nodiscard]] auto const_dictionary_view() const noexcept -> const_dictionary_view_type requires(is_compressed) { - return dictionary_; + if constexpr (is_const_view) { + return dictionary_view_; + } else { + return spec_type::get_dictionary_const_view(dictionary_view_); + } } private: - std::shared_ptr data_; - /* Either a shared handle to the dictionary, or the empty dictionary itself. */ - [[no_unique_address]] std::conditional_t, - dictionary_type> dictionary_; + data_view_type data_view_; + [[no_unique_address]] dictionary_view_type dictionary_view_; +}; + +/** + * An owning dataset: the data and, for a compressed dataset, the dictionary that gives it meaning. + * Ownership is unique, so the type is move-only whenever the underlying storage is. + */ +template +struct dataset { + /* NOTE: element type for an owning dataset cannot be const; + the constness of the dataset itself selects the flavor of the view it hands out. + */ + static_assert(!std::is_const_v, "Element type for an owning dataset must not be const."); + + using view_type = dataset_view; + using const_view_type = dataset_view; + + using spec_type = typename view_type::spec_type; + using value_type = typename view_type::value_type; + using index_type = typename view_type::index_type; + using code_type = typename view_type::code_type; + + using data_type = typename view_type::data_type; + using data_view_type = typename view_type::data_view_type; + using const_data_view_type = typename view_type::const_data_view_type; + + /** + Optional slot that contains everything to give meaning to the encoded data; + `std::monostate` for an uncompressed dataset. + */ + using dictionary_type = typename view_type::dictionary_type; + using dictionary_view_type = typename view_type::dictionary_view_type; + using const_dictionary_view_type = typename view_type::const_dictionary_view_type; + + static constexpr bool is_compressed = view_type::is_compressed; + + explicit dataset(data_type&& data) + requires(!is_compressed) + : data_{std::move(data)} + { + } + + dataset(data_type&& data, dictionary_type&& dictionary) + requires(is_compressed) + : data_{std::move(data)}, dictionary_{std::move(dictionary)} + { + } - [[nodiscard]] auto dict_ref() const noexcept -> const dictionary_type& + /* + @return A non-owning view of the whole dataset, data and dictionary alike. + */ + [[nodiscard]] auto view() noexcept -> view_type { if constexpr (is_compressed) { - return *dictionary_; + return view_type{data_view(), dictionary_view()}; } else { - return dictionary_; + return view_type{data_view()}; } } + /* + @return A read-only view of the whole dataset; the constness of the dataset selects this flavor. + */ + [[nodiscard]] auto view() const noexcept -> const_view_type + { + if constexpr (is_compressed) { + return const_view_type{data_const_view(), const_dictionary_view()}; + } else { + return const_view_type{data_const_view()}; + } + } + + [[nodiscard]] auto n_rows() const noexcept -> index_type { return spec_type::get_n_rows(data_); } + [[nodiscard]] auto dim() const noexcept -> uint32_t + { + // NOTE: a compressed dataset may not be able to tell its dimension from the encoded data alone + return spec_type::get_dim(data_, dictionary_); + } + + /* + @return A non-owning view of the data, such as mdspan. + */ + [[nodiscard]] auto data_view() noexcept -> data_view_type + { + return spec_type::get_data_view(data_); + } + [[nodiscard]] auto data_const_view() const noexcept -> const_data_view_type + { + return spec_type::get_data_const_view(data_); + } + + /* + @return A non-owning view of the dictionary, such as the two codebooks of a VPQ dataset. + */ + [[nodiscard]] auto dictionary_view() const noexcept -> dictionary_view_type + requires(is_compressed) + { + return spec_type::get_dictionary_view(dictionary_); + } + /* + @return A read-only view of the dictionary. + */ + [[nodiscard]] auto const_dictionary_view() const noexcept -> const_dictionary_view_type + requires(is_compressed) + { + return spec_type::get_dictionary_const_view(dictionary_); + } + + private: + data_type data_; + [[no_unique_address]] dictionary_type dictionary_; }; /************************************************************************************************ @@ -154,21 +309,33 @@ struct empty_spec { template struct apply { - using data_type = empty_dataset_rep; - using view_type = empty_dataset_rep; - using const_view_type = const empty_dataset_rep; - using value_type = std::remove_cv_t; - using index_type = std::remove_cv_t; - using dictionary_type = std::monostate; - - [[nodiscard]] static auto get_data_view(data_type& data) noexcept -> view_type { return data; } - [[nodiscard]] static auto get_data_const_view(const data_type& data) noexcept -> const_view_type + using data_type = empty_dataset_rep; + using data_view_type = empty_dataset_rep; + /* NOTE: not `const empty_dataset_rep`: a read-only view holds this as a member, and a const + member would delete its assignment operator. + */ + using const_data_view_type = empty_dataset_rep; + using value_type = std::remove_cv_t; + using index_type = std::remove_cv_t; + + [[nodiscard]] static auto get_data_view(data_type& data) noexcept -> data_view_type { return data; } - [[nodiscard]] static auto get_n_rows(const data_type& data) noexcept -> index_type { return 0; } - [[nodiscard]] static auto get_dim(const data_type& data, const dictionary_type&) noexcept - -> uint32_t + /* NOTE: the view of an empty dataset is the storage itself, so one overload serves both the + owning dataset and the view. + */ + [[nodiscard]] static auto get_data_const_view(const data_type& data) noexcept + -> const_data_view_type + { + return data; + } + /* NOTE: the shape accessors take the data in either form - the storage an owning dataset holds, + or the view a `dataset_view` holds. A plain spec defines no dictionary either, so it ignores + whatever the dataset puts in that slot instead. + */ + [[nodiscard]] static auto get_n_rows(const auto&) noexcept -> index_type { return 0; } + [[nodiscard]] static auto get_dim(const auto& data, const auto&) noexcept -> uint32_t { return data.dim; } @@ -187,28 +354,40 @@ struct mdarray_spec { slots to save memory on large datasets. The extents are always fixed to 64-bit signed integers to simplify the conversion and avoid integer overflow issues. */ - using data_type = raft::mdarray, LayoutPolicy, ContainerPolicy>; - using view_type = typename data_type::view_type; - using const_view_type = typename data_type::const_view_type; - using value_type = std::remove_cv_t; - using index_type = std::remove_cv_t; - using dictionary_type = std::monostate; - - [[nodiscard]] static auto get_data_view(data_type& data) noexcept -> view_type + using value_type = std::remove_cv_t; + using index_type = std::remove_cv_t; + using data_type = + raft::mdarray, LayoutPolicy, ContainerPolicy>; + using data_view_type = typename data_type::view_type; + using const_data_view_type = typename data_type::const_view_type; + + [[nodiscard]] static auto get_data_view(data_type& data) noexcept -> data_view_type { return data.view(); } - [[nodiscard]] static auto get_data_const_view(const data_type& data) noexcept -> const_view_type + [[nodiscard]] static auto get_data_const_view(const data_type& data) noexcept + -> const_data_view_type { return data.view(); } - [[nodiscard]] static auto get_n_rows(const data_type& data) noexcept -> index_type + /* A mutable mdspan does not decay to a read-only one on its own; it has to be rebuilt over an + accessor of `const T`. + */ + [[nodiscard]] static auto get_data_const_view(data_view_type data) noexcept + -> const_data_view_type + { + return raft::make_const_mdspan(data); + } + /* NOTE: the shape accessors take the data in either form - the storage an owning dataset holds, + or the view a `dataset_view` holds. A plain spec defines no dictionary either, so it ignores + whatever the dataset puts in that slot instead. + */ + [[nodiscard]] static auto get_n_rows(const auto& data) noexcept -> index_type { return static_cast(data.extent(0)); } - [[nodiscard]] static auto get_dim(const data_type& data, const dictionary_type&) noexcept - -> uint32_t + [[nodiscard]] static auto get_dim(const auto& data, const auto&) noexcept -> uint32_t { return static_cast(data.extent(1)); } @@ -226,9 +405,8 @@ template struct sparse_spec { template struct apply { - using value_type = std::remove_cv_t; - using index_type = std::remove_cv_t; - using dictionary_type = std::monostate; + using value_type = std::remove_cv_t; + using index_type = std::remove_cv_t; static constexpr raft::memory_type mem_type = ContainerPolicy::mem_type; static_assert(mem_type == raft::memory_type::host || mem_type == raft::memory_type::device, @@ -240,29 +418,50 @@ struct sparse_spec { template matrix_type, typename SparseLayoutT:: template matrix_type>; - using view_type = typename data_type::view_type; + using data_view_type = typename data_type::view_type; /* NOTE: raft's sparse hierarchy provides no read-only view yet (see the TODO in raft/core/sparse_types.hpp); until then, even reading the structure requires a mutable matrix. + + Consequently a read-only `dataset_view` over a sparse spec is read-only in name + only: the two view types coincide, so writing through it still compiles. That resolves itself + once raft grows a const sparse view. */ - using const_view_type = view_type; + using const_data_view_type = data_view_type; - [[nodiscard]] static auto get_data_view(data_type& data) noexcept -> view_type + [[nodiscard]] static auto get_data_view(data_type& data) noexcept -> data_view_type { return data.view(); } - [[nodiscard]] static auto get_data_const_view(const data_type& data) noexcept -> const_view_type + [[nodiscard]] static auto get_data_const_view(const data_type& data) noexcept + -> const_data_view_type { return const_cast(data).view(); } - [[nodiscard]] static auto get_n_rows(const data_type& data) noexcept -> index_type + /* NOTE: there is no read-only view to convert to (see above), so the view passes through. */ + [[nodiscard]] static auto get_data_const_view(data_view_type data) noexcept + -> const_data_view_type { - return static_cast(const_cast(data).structure_view().get_n_rows()); + return data; } - [[nodiscard]] static auto get_dim(const data_type& data, const dictionary_type&) noexcept - -> uint32_t + /* NOTE: raft's sparse hierarchy exposes the structure through a mutable accessor only, so + reading the shape of either the owning matrix or its view takes a cast. + */ + [[nodiscard]] static auto get_structure_view(const auto& data) noexcept { - return static_cast(const_cast(data).structure_view().get_n_cols()); + return const_cast&>(data).structure_view(); + } + /* NOTE: the shape accessors take the data in either form - the storage an owning dataset holds, + or the view a `dataset_view` holds. A plain spec defines no dictionary either, so it ignores + whatever the dataset puts in that slot instead. + */ + [[nodiscard]] static auto get_n_rows(const auto& data) noexcept -> index_type + { + return static_cast(get_structure_view(data).get_n_rows()); + } + [[nodiscard]] static auto get_dim(const auto& data, const auto&) noexcept -> uint32_t + { + return static_cast(get_structure_view(data).get_n_cols()); } }; }; @@ -308,75 +507,82 @@ struct vpq_spec { using storage_spec = typename mdarray_spec::template apply; + /* The type the codebooks are stored in; the rows themselves are always packed codes. */ + using math_type = MathT; + + /* [vq_n_centers, dim] */ + using vq_book_type = + raft::mdarray, raft::layout_c_contiguous, BookPolicy>; + /* [n_books, pq_n_centers, pq_len], where `n_books` is 1 for a codebook shared by all the + subspaces, or `pq_dim` for one codebook per subspace. A row-major 3d array with `n_books == 1` + is bit-identical to the flat 2d codebook of the current implementation. + */ + using pq_book_type = + raft::mdarray, raft::layout_c_contiguous, BookPolicy>; + + struct dictionary_type { + vq_book_type vq_code_book; + pq_book_type pq_code_book; + }; + + struct dictionary_view_type { + typename vq_book_type::const_view_type vq_code_book; + typename pq_book_type::const_view_type pq_code_book; + + /* NOTE: the dimension of the dataset is a property of the VQ codebook: the encoded rows tell + only how many bytes it took to compress a vector, which the padding makes ambiguous. + */ + [[nodiscard]] auto dim() const noexcept -> uint32_t + { + return static_cast(vq_code_book.extent(1)); + } + [[nodiscard]] auto vq_n_centers() const noexcept -> uint32_t + { + return static_cast(vq_code_book.extent(0)); + } + [[nodiscard]] auto n_books() const noexcept -> uint32_t + { + return static_cast(pq_code_book.extent(0)); + } + [[nodiscard]] auto pq_n_centers() const noexcept -> uint32_t + { + return static_cast(pq_code_book.extent(1)); + } + [[nodiscard]] auto pq_len() const noexcept -> uint32_t + { + return static_cast(pq_code_book.extent(2)); + } + [[nodiscard]] auto pq_bits() const noexcept -> uint32_t + { + return static_cast(std::countr_zero(pq_n_centers())); + } + [[nodiscard]] auto pq_dim() const noexcept -> uint32_t + { + return raft::div_rounding_up_safe(dim(), pq_len()); + } + /* Whether every subspace has a codebook of its own, as opposed to sharing a single one. */ + [[nodiscard]] auto per_subspace() const noexcept -> bool { return n_books() > 1; } + }; + template struct apply : storage_spec { /* The dataset stores codes, but it still represents vectors of `T`. */ using value_type = std::remove_cv_t; - /* NOTE: the members of a dependent base are not visible to unqualified lookup, so `data_type` - has to be pulled into this scope to be usable in the signatures below. + using code_type = uint8_t; + /* NOTE: re-exported rather than declared here, so that they do not belong to one instantiation + of `apply`; qualified lookup into `apply` does not reach the enclosing spec on its own. */ - using typename storage_spec::data_type; - /* The type the codebooks are stored in; the rows themselves are always packed codes. */ - using math_type = MathT; - using code_type = uint8_t; - - /* [vq_n_centers, dim] */ - using vq_book_type = - raft::mdarray, raft::layout_c_contiguous, BookPolicy>; - /* [n_books, pq_n_centers, pq_len], where `n_books` is 1 for a codebook shared by all the - subspaces, or `pq_dim` for one codebook per subspace. A row-major 3d array with `n_books == 1` - is bit-identical to the flat 2d codebook of the current implementation. + using math_type = vpq_spec::math_type; + using vq_book_type = vpq_spec::vq_book_type; + using pq_book_type = vpq_spec::pq_book_type; + using dictionary_type = vpq_spec::dictionary_type; + using dictionary_view_type = vpq_spec::dictionary_view_type; + + /* NOTE: the data accessors are inherited; only `get_dim` differs from a dense dataset. Like the + inherited ones, it reads the dictionary in either form: the codebooks an owning dataset holds, + or the views a `dataset_view` holds. */ - using pq_book_type = - raft::mdarray, raft::layout_c_contiguous, BookPolicy>; - - struct dictionary_type { - vq_book_type vq_code_book; - pq_book_type pq_code_book; - }; - - struct dictionary_view_type { - typename vq_book_type::const_view_type vq_code_book; - typename pq_book_type::const_view_type pq_code_book; - - /* NOTE: the dimension of the dataset is a property of the VQ codebook: the encoded rows tell - only how many bytes it took to compress a vector, which the padding makes ambiguous. - */ - [[nodiscard]] auto dim() const noexcept -> uint32_t - { - return static_cast(vq_code_book.extent(1)); - } - [[nodiscard]] auto vq_n_centers() const noexcept -> uint32_t - { - return static_cast(vq_code_book.extent(0)); - } - [[nodiscard]] auto n_books() const noexcept -> uint32_t - { - return static_cast(pq_code_book.extent(0)); - } - [[nodiscard]] auto pq_n_centers() const noexcept -> uint32_t - { - return static_cast(pq_code_book.extent(1)); - } - [[nodiscard]] auto pq_len() const noexcept -> uint32_t - { - return static_cast(pq_code_book.extent(2)); - } - [[nodiscard]] auto pq_bits() const noexcept -> uint32_t - { - return static_cast(std::countr_zero(pq_n_centers())); - } - [[nodiscard]] auto pq_dim() const noexcept -> uint32_t - { - return raft::div_rounding_up_safe(dim(), pq_len()); - } - /* Whether every subspace has a codebook of its own, as opposed to sharing a single one. */ - [[nodiscard]] auto per_subspace() const noexcept -> bool { return n_books() > 1; } - }; - - /* NOTE: the data accessors are inherited; only `get_dim` differs from a dense dataset. */ - [[nodiscard]] static auto get_dim(const data_type&, const dictionary_type& dict) noexcept - -> uint32_t + [[nodiscard]] static auto get_dim(const auto&, const auto& dict) noexcept -> uint32_t { return static_cast(dict.vq_code_book.extent(1)); } @@ -385,8 +591,21 @@ struct vpq_spec { { return {dict.vq_code_book.view(), dict.pq_code_book.view()}; } + /* NOTE: a codebook view is read-only to begin with, so the spec declares no separate const view + and these overloads exist only to spell the conversion out of either form. + */ + [[nodiscard]] static auto get_dictionary_const_view(const dictionary_type& dict) noexcept + -> dictionary_view_type + { + return get_dictionary_view(dict); + } + [[nodiscard]] static auto get_dictionary_const_view(dictionary_view_type dict) noexcept + -> dictionary_view_type + { + return dict; + } /* The length of an encoded row in bytes, including the inlined VQ label. */ - [[nodiscard]] static auto get_encoded_row_length(const data_type& data) noexcept -> uint32_t + [[nodiscard]] static auto get_encoded_row_length(const auto& data) noexcept -> uint32_t { return static_cast(data.extent(1)); } @@ -405,35 +624,41 @@ struct sq_spec { using storage_spec = typename mdarray_spec::template apply; + using math_type = MathT; + + /* Nothing to allocate: the whole dictionary is the interval the codes are dequantized into. + The members are named as in cuvs::preprocessing::quantize::scalar::quantizer. + */ + struct dictionary_type { + math_type min_; + math_type max_; + }; + using dictionary_view_type = dictionary_type; + template struct apply : storage_spec { /* The dataset stores codes, but it still represents vectors of `T`. */ using value_type = std::remove_cv_t; - /* NOTE: the members of a dependent base are not visible to unqualified lookup, so `data_type` - has to be pulled into this scope to be usable in the signature below. + using code_type = int8_t; + /* NOTE: re-exported rather than declared here, so that they do not belong to one instantiation + of `apply`; qualified lookup into `apply` does not reach the enclosing spec on its own. */ - using typename storage_spec::data_type; - using math_type = MathT; - using code_type = int8_t; + using math_type = sq_spec::math_type; + using dictionary_type = sq_spec::dictionary_type; + using dictionary_view_type = sq_spec::dictionary_view_type; - /* Nothing to allocate: the whole dictionary is the interval the codes are dequantized into. - The members are named as in cuvs::preprocessing::quantize::scalar::quantizer. - */ - struct dictionary_type { - math_type min_; - math_type max_; - }; - using dictionary_view_type = dictionary_type; - - /* NOTE: only the signature differs from the inherited accessor, which cannot be reused because - it expects an empty dictionary. + /* NOTE: every accessor is inherited; the dimension is still one code per component, so even + `get_dim` is the dense one. */ - [[nodiscard]] static auto get_dim(const data_type& data, const dictionary_type&) noexcept - -> uint32_t + [[nodiscard]] static auto get_dictionary_view(const dictionary_type& dict) noexcept + -> dictionary_view_type { - return static_cast(data.extent(1)); + return dict; } - [[nodiscard]] static auto get_dictionary_view(const dictionary_type& dict) noexcept + /* NOTE: the dictionary is its own view here, read-only or not, so one overload serves both the + owning dataset and the view. + */ + [[nodiscard]] static auto get_dictionary_const_view(const dictionary_type& dict) noexcept -> dictionary_view_type { return dict; diff --git a/examples/cpp/src/dataset_example.cu b/examples/cpp/src/dataset_example.cu index af51d6296a..a678f50786 100644 --- a/examples/cpp/src/dataset_example.cu +++ b/examples/cpp/src/dataset_example.cu @@ -18,7 +18,6 @@ #include #include -#include #include #include @@ -30,8 +29,8 @@ void print_dataset(const char* name, const DatasetT& dataset) } /** - * A transformation that takes its input by value and hands back the result. Written this way, it - * can skip the second buffer whenever the caller passes its only copy of the dataset: + * A transformation that takes its input by value and hands back the result. A dataset owns its data + * uniquely, so the caller has to move it in and the transform is free to write over its own input: * `d = normalize(res, std::move(d))` allocates only the column statistics. */ template @@ -50,17 +49,10 @@ auto normalize(const raft::resources& res, cuvs::device_contiguous_dataset T{0} ? (x - mean) / raft::sqrt(var) : T{0}; }; - if (data.is_data_unique()) { - // No other copy can observe the buffer, so the transform may write over its own input. - raft::linalg::matrix_vector_op( - res, in, mean_view, var_view, data.data_view(), standardize); - return data; - } - cuvs::device_contiguous_dataset out{ - raft::make_device_matrix(res, in.extent(0), in.extent(1))}; + // Nobody else can observe the buffer, so the transform writes over its own input. raft::linalg::matrix_vector_op( - res, in, mean_view, var_view, out.data_view(), standardize); - return out; + res, in, mean_view, var_view, data.data_view(), standardize); + return data; } /* The dictionary slot is free for the datasets that have no dictionary. */ @@ -71,8 +63,26 @@ static_assert( std::same_as::dictionary_type, std::monostate>); static_assert(std::same_as::dictionary_type, std::monostate>); static_assert(sizeof(cuvs::device_contiguous_dataset) == - sizeof(std::shared_ptr)); + sizeof(typename cuvs::device_contiguous_dataset::data_type)); static_assert(cuvs::device_vpq_dataset::is_compressed); +/* Ownership is unique, so an owning dataset moves rather than copies, while its view is a handle. + */ +static_assert(!std::is_copy_constructible_v>); +static_assert(std::is_copy_constructible_v< + typename cuvs::device_contiguous_dataset::view_type>); +/* The read-only flavor of the view is the same type with a const element type, in which the two + * data views coincide. A mutable view decays into it, but not the other way round. + */ +static_assert( + std::same_as< + typename cuvs::device_contiguous_dataset::const_view_type::data_view_type, + typename cuvs::device_contiguous_dataset::const_data_view_type>); +static_assert( + std::convertible_to::view_type, + typename cuvs::device_contiguous_dataset::const_view_type>); +static_assert( + !std::convertible_to::const_view_type, + typename cuvs::device_contiguous_dataset::view_type>); int main() { @@ -101,18 +111,33 @@ int main() raft::make_device_matrix(res, n_rows, dim)}; print_dataset("device contiguous", device_dataset); - // Moving the dataset in leaves `normalize` as its only owner, so the data is scaled in-place. + // The dataset has to be moved into `normalize`, which therefore scales the data in-place. raft::linalg::map_offset(res, device_dataset.data_view(), raft::cast_op{}); const auto* original_data = device_dataset.data_const_view().data_handle(); device_dataset = normalize(res, std::move(device_dataset)); std::cout << " reused the input buffer: " << std::boolalpha << (device_dataset.data_const_view().data_handle() == original_data) << '\n'; - // Passing a copy keeps a second owner alive, so the same call has to allocate its output. - auto shared_dataset = device_dataset; - auto rescaled = normalize(res, shared_dataset); - std::cout << " reused the input buffer: " - << (rescaled.data_const_view().data_handle() == original_data) << '\n'; + // A view is the non-owning counterpart of a dataset: it answers the same questions, but the data + // stays with the dataset it was taken from. + typename cuvs::device_contiguous_dataset::view_type device_view = + device_dataset.view(); + print_dataset("device contiguous (view)", device_view); + // The view hands out a read-only mdspan over the same buffer; rebuilding one by hand would mean + // spelling out the const accessor. + std::cout << " views the same buffer: " + << (device_view.data_const_view().data_handle() == original_data) << '\n'; + + // The constness of the dataset picks the flavor of the view: a const dataset hands out the + // read-only one, whose data view is already an mdspan of `const float`. + const auto& const_dataset = device_dataset; + typename cuvs::device_contiguous_dataset::const_view_type const_view = + const_dataset.view(); + print_dataset("device contiguous (const view)", const_view); + // ... and a mutable view decays into that same type on its own. + const_view = device_view; + std::cout << " views the same buffer: " + << (const_view.data_view().data_handle() == original_data) << '\n'; // The same dataset over a padded layout: the spec, not the dataset, selects the layout. using padded_dataset_type = cuvs::device_padded_dataset; @@ -154,10 +179,10 @@ int main() // A single codebook shared by all the subspaces; `pq_dim` of them would be per-subspace. raft::make_device_mdarray( res, raft::make_extents(1, pq_n_centers, pq_len))}; - auto dictionary = std::make_shared(std::move(codebooks)); vpq_dataset_type vpq_dataset{ - raft::make_device_matrix(res, n_rows, encoded_row_length), dictionary}; + raft::make_device_matrix(res, n_rows, encoded_row_length), + std::move(codebooks)}; print_dataset("device vpq", vpq_dataset); auto books = vpq_dataset.dictionary_view(); std::cout << " vq_n_centers = " << books.vq_n_centers() << ", pq_bits = " << books.pq_bits() @@ -166,19 +191,21 @@ int main() // The data slot is a plain dense matrix of bytes, hence the shared spec. std::cout << " encoded row length = " << vpq_dataset.data_view().extent(1) << '\n'; - // A dictionary is immutable, so another dataset may be encoded against the very same codebooks. - vpq_dataset_type shared_vpq_dataset{ - raft::make_device_matrix(res, n_rows * 2, encoded_row_length), - vpq_dataset.share_dictionary()}; - print_dataset("device vpq (shared dictionary)", shared_vpq_dataset); + // The view of a compressed dataset carries the dictionary along, so it can still tell the + // dimension that the encoded rows alone do not reveal. + auto vpq_view = vpq_dataset.view(); + print_dataset("device vpq (view)", vpq_view); + std::cout << " views the same codebook: " + << (vpq_view.dictionary_view().vq_code_book.data_handle() == + books.vq_code_book.data_handle()) + << '\n'; // Scalar quantization keeps one code per component, so the dimension is still a property of the // data and the dictionary is only the interval to dequantize into. using sq_dataset_type = cuvs::device_sq_dataset; using sq_dictionary_type = typename sq_dataset_type::dictionary_type; - sq_dataset_type sq_dataset{ - raft::make_device_matrix(res, n_rows, dim), - std::make_shared(sq_dictionary_type{-1.0F, 1.0F})}; + sq_dataset_type sq_dataset{raft::make_device_matrix(res, n_rows, dim), + sq_dictionary_type{-1.0F, 1.0F}}; print_dataset("device sq", sq_dataset); std::cout << " range = [" << sq_dataset.dictionary_view().min_ << ", " << sq_dataset.dictionary_view().max_ << "]\n"; diff --git a/fern/docs.yml b/fern/docs.yml index 77ee0dae2a..b055d844f1 100644 --- a/fern/docs.yml +++ b/fern/docs.yml @@ -358,6 +358,8 @@ navigation: path: "./pages/cpp_api/cpp-api-common-types-copy-serialization-and-utility-apis.md" - page: "Errors and Logging" path: "./pages/cpp_api/cpp-api-common-types-errors-and-logging.md" + - page: "Core Dataset" + path: "./pages/cpp_api/cpp-api-core-dataset.md" - page: "Core Roaring Allowlist" path: "./pages/cpp_api/cpp-api-core-roaring-allowlist.md" - page: "Distance Distance" diff --git a/fern/pages/cpp_api/cpp-api-core-dataset.md b/fern/pages/cpp_api/cpp-api-core-dataset.md new file mode 100644 index 0000000000..d7edff4151 --- /dev/null +++ b/fern/pages/cpp_api/cpp-api-core-dataset.md @@ -0,0 +1,208 @@ +--- +slug: api-reference/cpp-api-core-dataset +--- + +# Dataset + +_Source header: `cuvs/core/dataset.hpp`_ + +## Types + + +### dataset_view + +A non-owning view of a dataset: a view of the data and, for a compressed dataset, a view of the + +dictionary that gives it meaning. Cheap to copy and to pass by value, like an mdspan. + +```cpp +template +struct dataset_view { + /* The read-only flavor, in which each view type collapses onto its const counterpart. */ static constexpr bool is_const_view; + static constexpr bool is_compressed; +}; +``` + +**Fields** + +| Name | Type | Description | +| --- | --- | --- | +| `is_const_view` | `/* The read-only flavor, in which each view type collapses onto its const counterpart. */ static constexpr bool` | | +| `is_compressed` | `static constexpr bool` | | + + +### dataset_view::code_type + +What one element of the data slot is + +```cpp +using code_type = typename decltype([]; +``` + +- for compressed datasets it is the raw data element type; often is uint8_t but may be a larger type to enforce better alignment. +- for uncompressed datasets it is the value type of the dataset. + +**Returns** + +[`using code_type = typename`](/api-reference/cpp-api-core-dataset#dataset-view-code-type) + + +### dataset_view::dictionary_type + +Optional slot that contains everything to give meaning to the encoded data; + +```cpp +using dictionary_type = typename decltype([]; +``` + +`std::monostate` for an uncompressed dataset. + +**Returns** + +[`using dictionary_type = typename`](/api-reference/cpp-api-core-dataset#dictionary-type) + + +### dataset + +An owning dataset: the data and, for a compressed dataset, the dictionary that gives it meaning. + +Ownership is unique, so the type is move-only whenever the underlying storage is. + +```cpp +template +struct dataset { + using dictionary_type; + /* NOTE: element type for an owning dataset cannot be const; + static constexpr bool is_compressed; + /* @return A read-only view of the whole dataset; +}; +``` + +**Fields** + +| Name | Type | Description | +| --- | --- | --- | +| `dictionary_type` | `using` | Optional slot that contains everything to give meaning to the encoded data; `std::monostate` for an uncompressed dataset. | +| `const` | [`/* NOTE: element type for an owning dataset cannot be`](/api-reference/cpp-api-core-dataset#dataset) | | +| `is_compressed` | `static constexpr bool` | | +| `dataset` | `/* @return A read-only view of the whole` | | + + +### dictionary_type + +Optional slot that contains everything to give meaning to the encoded data; + +```cpp +using dictionary_type = typename view_type::dictionary_type; +``` + + +### empty_spec + +Empty dataset specification: contains no data, but keeps the dimension of the dataset implement + +the common interface. + +```cpp +struct empty_spec; +``` + + +### mdarray_spec + +Dense plain/padded dataset implemented via raft::mdarray. + +```cpp +template +struct mdarray_spec; +``` + + +### sparse_spec + +Sparse dataset specification implemented via the raft sparse matrix hierarchy. `SparseLayoutT` + +selects the representation (`csr_layout`, `coo_layout`). + +NOTE: only the memory space of the container policy is used here; raft's sparse types bind the policy themselves, separately for the values, the offsets and the indices. + +```cpp +template +struct sparse_spec; +``` + + +### csr_layout + +CSR representation: one compressed row per vector of the dataset. + +NOTE: the row offsets are 64-bit for the same reason the dense extents are: the number of non-zeros is not bounded by the index type. The column indices are feature ids, so they are bounded by the dataset dimension rather than by its size. + +```cpp +struct csr_layout; +``` + + +### coo_layout + +COO representation: one (row, column, value) triple per non-zero. + +NOTE: unlike CSR, the row component is an index into the dataset rather than an offset into the non-zeros, hence the dataset index type. + +```cpp +struct coo_layout; +``` + + +### vpq_spec + +VQ+PQ compressed dataset specification, mirroring the VPQ dataset of cuvs/neighbors/common.hpp: + +the two codebooks form the dictionary, while the data slot keeps the encoded rows - the VQ label in the row prefix, followed by the packed PQ codes. + +`BookPolicy` is the container policy of the codebooks (elements of `MathT`) and `CodePolicy` the one of the encoded rows (elements of `uint8_t`); a bound policy cannot be rebound to another element type, hence the two parameters. + +```cpp +template +struct vpq_spec { + /* The encoded rows are stored like the rows of any other dense dataset; + only how a row is read and where the dimension comes from are specific to the compression. */ template using storage_spec; + /* The type the codebooks are stored in; + the rows themselves are always packed codes. */ using math_type; + /* [vq_n_centers, dim] */ using vq_book_type; + /* [n_books, pq_n_centers, pq_len], where `n_books` is 1 for a codebook shared by all the subspaces, or `pq_dim` for one codebook per subspace. A row-major 3d array with ` n_books; +}; +``` + +**Fields** + +| Name | Type | Description | +| --- | --- | --- | +| `dataset` | `/* The encoded rows are stored like the rows of any other dense` | | +| `storage_spec` | `only how a row is read and where the dimension comes from are specific to the compression. */ template using` | | +| `in` | `/* The type the codebooks are stored` | | +| `math_type` | `the rows themselves are always packed codes. */ using` | | +| `vq_book_type` | `/* [vq_n_centers, dim] */ using` | | +| `n_books` | `/* [n_books, pq_n_centers, pq_len], where \`n_books\` is 1 for a codebook shared by all the subspaces, or \`pq_dim\` for one codebook per subspace. A row-major 3d array with \`` | | + + +### sq_spec + +Scalar-quantized dataset specification: every component of a vector becomes one code, so the + +dictionary is only the range the codes are mapped back onto and the dimension is still a property of the data. + +```cpp +template +struct sq_spec { + /* One code per component, so the codes are laid out exactly like a dense dataset. */ template using storage_spec; + /* Nothing to allocate: the whole dictionary is the interval the codes are dequantized into. The members are named as in cuvs::preprocessing::quantize::scalar::quantizer. */ struct dictionary_type; +}; +``` + +**Fields** + +| Name | Type | Description | +| --- | --- | --- | +| `storage_spec` | [`/* One code per component, so the codes are laid out exactly like a dense dataset. */ template using`](/api-reference/cpp-api-core-dataset#dataset) | | +| `dictionary_type` | [`/* Nothing to allocate: the whole dictionary is the interval the codes are dequantized into. The members are named as in cuvs::preprocessing::quantize::scalar::quantizer. */ struct`](/api-reference/cpp-api-preprocessing-quantize-scalar#preprocessing-quantize-scalar-quantizer) | | diff --git a/fern/pages/cpp_api/index.md b/fern/pages/cpp_api/index.md index 07115a5eca..67cb45e0f8 100644 --- a/fern/pages/cpp_api/index.md +++ b/fern/pages/cpp_api/index.md @@ -13,6 +13,7 @@ These pages are generated from the documented public headers in the cuVS source - [Bloom Filter](/api-reference/cpp-api-core-bloom-filter) - [Common Types](/api-reference/cpp-api-common-types) +- [Dataset](/api-reference/cpp-api-core-dataset) - [Roaring Allowlist](/api-reference/cpp-api-core-roaring-allowlist) ## Distance From fda265c187904a8806162d788fcaa4228da0be8b Mon Sep 17 00:00:00 2001 From: achirkin Date: Wed, 14 Oct 2026 16:20:33 +0200 Subject: [PATCH 04/11] Disaggregate Dataset API into multiple headers per spec --- cpp/include/cuvs/core/dataset.hpp | 734 +----------------- cpp/include/cuvs/core/dataset/dense.hpp | 91 +++ cpp/include/cuvs/core/dataset/empty.hpp | 63 ++ cpp/include/cuvs/core/dataset/sparse.hpp | 131 ++++ cpp/include/cuvs/core/dataset/sq.hpp | 84 ++ cpp/include/cuvs/core/dataset/type.hpp | 126 +++ cpp/include/cuvs/core/dataset/view.hpp | 182 +++++ cpp/include/cuvs/core/dataset/vpq.hpp | 162 ++++ examples/cpp/src/dataset_example.cu | 4 + fern/docs.yml | 16 +- .../cpp_api/cpp-api-core-dataset-dense.md | 19 + .../cpp_api/cpp-api-core-dataset-empty.md | 20 + .../cpp_api/cpp-api-core-dataset-sparse.md | 45 ++ fern/pages/cpp_api/cpp-api-core-dataset-sq.md | 31 + .../cpp_api/cpp-api-core-dataset-type.md | 44 ++ .../cpp_api/cpp-api-core-dataset-view.md | 62 ++ .../pages/cpp_api/cpp-api-core-dataset-vpq.md | 41 + fern/pages/cpp_api/cpp-api-core-dataset.md | 208 ----- fern/pages/cpp_api/index.md | 8 +- 19 files changed, 1133 insertions(+), 938 deletions(-) create mode 100644 cpp/include/cuvs/core/dataset/dense.hpp create mode 100644 cpp/include/cuvs/core/dataset/empty.hpp create mode 100644 cpp/include/cuvs/core/dataset/sparse.hpp create mode 100644 cpp/include/cuvs/core/dataset/sq.hpp create mode 100644 cpp/include/cuvs/core/dataset/type.hpp create mode 100644 cpp/include/cuvs/core/dataset/view.hpp create mode 100644 cpp/include/cuvs/core/dataset/vpq.hpp create mode 100644 fern/pages/cpp_api/cpp-api-core-dataset-dense.md create mode 100644 fern/pages/cpp_api/cpp-api-core-dataset-empty.md create mode 100644 fern/pages/cpp_api/cpp-api-core-dataset-sparse.md create mode 100644 fern/pages/cpp_api/cpp-api-core-dataset-sq.md create mode 100644 fern/pages/cpp_api/cpp-api-core-dataset-type.md create mode 100644 fern/pages/cpp_api/cpp-api-core-dataset-view.md create mode 100644 fern/pages/cpp_api/cpp-api-core-dataset-vpq.md delete mode 100644 fern/pages/cpp_api/cpp-api-core-dataset.md diff --git a/cpp/include/cuvs/core/dataset.hpp b/cpp/include/cuvs/core/dataset.hpp index 18f839a569..c3b5f2cfe2 100644 --- a/cpp/include/cuvs/core/dataset.hpp +++ b/cpp/include/cuvs/core/dataset.hpp @@ -5,733 +5,13 @@ #pragma once -#include - -#include -#include -#include -#include -#include -#include -#include - -#include -#include -#include -#include - -namespace CUVS_EXPORT cuvs { - -/************************************************************************************************ - * The core dataset type - ************************************************************************************************/ - -/** - * A spec is compressing if it defines a dictionary: the codebooks, the quantization range, or - * whatever else is needed to interpret the data. Plain specs define none and the dataset - * substitutes `std::monostate`, so "there is no dictionary" and "the dictionary slot is empty" are - * the same vocabulary type. - */ -template -concept compressed_dataset_spec = requires { typename SpecT::dictionary_type; }; - -/** - * A non-owning view of a dataset: a view of the data and, for a compressed dataset, a view of the - * dictionary that gives it meaning. Cheap to copy and to pass by value, like an mdspan. - */ -template -struct dataset_view { - using spec_type = typename SpecT::template apply; - /* The read-only flavor, in which each view type collapses onto its const counterpart. */ - static constexpr bool is_const_view = std::is_const_v; - - using value_type = typename spec_type::value_type; - using index_type = typename spec_type::index_type; - /** - What one element of the data slot is - - for compressed datasets it is the raw data element type; often is uint8_t but may be - a larger type to enforce better alignment. - - for uncompressed datasets it is the value type of the dataset. - */ - using code_type = typename decltype([] { - if constexpr (requires { typename spec_type::code_type; }) { - return std::type_identity{}; - } else { - return std::type_identity{}; - } - }())::type; - - using data_type = typename spec_type::data_type; - using const_data_view_type = typename spec_type::const_data_view_type; - using data_view_type = - std::conditional_t; - - /** - Optional slot that contains everything to give meaning to the encoded data; - `std::monostate` for an uncompressed dataset. - */ - using dictionary_type = typename decltype([] { - if constexpr (requires { typename spec_type::dictionary_type; }) { - return std::type_identity{}; - } else { - return std::type_identity{}; - } - }())::type; - // Optional, so that the specs don't have to declare it; a spec that declares neither view uses - // the dictionary itself, and one that declares only the mutable view uses it for both. - using const_dictionary_view_type = typename decltype([] { - if constexpr (requires { typename spec_type::const_dictionary_view_type; }) { - return std::type_identity{}; - } else if constexpr (requires { typename spec_type::dictionary_view_type; }) { - return std::type_identity{}; - } else { - return std::type_identity{}; - } - }())::type; - // Optional, so that the specs don't have to declare it - using dictionary_view_type = typename decltype([] { - if constexpr (is_const_view) { - return std::type_identity{}; - } else if constexpr (requires { typename spec_type::dictionary_view_type; }) { - return std::type_identity{}; - } else { - return std::type_identity{}; - } - }())::type; - - static constexpr bool is_compressed = compressed_dataset_spec; - - explicit dataset_view(data_view_type data_view) - requires(!is_compressed) - : data_view_{data_view} - { - } - - dataset_view(data_view_type data_view, dictionary_view_type dictionary_view) - requires(is_compressed) - : data_view_{data_view}, dictionary_view_{dictionary_view} - { - } - - /* A mutable view decays to a read-only one, the way an mdspan does. - - NOTE: a constructor template, because in the mutable flavor the parameter would spell the very - same type; a non-template would be a copy constructor there, and declaring one - even a - constrained-away one - suppresses the implicitly declared copy constructor. - - NOTE: the dictionary is assigned in the body rather than in the init-list, because the accessor it - comes from is constrained on `is_compressed`. - */ - template - requires(is_const_view && std::same_as>) - dataset_view(const dataset_view& other) noexcept - : data_view_{other.data_const_view()} - { - if constexpr (is_compressed) { dictionary_view_ = other.const_dictionary_view(); } - } - - [[nodiscard]] auto n_rows() const noexcept -> index_type - { - return spec_type::get_n_rows(data_view_); - } - // NOTE: a compressed dataset may not be able to tell its dimension from the encoded data alone - [[nodiscard]] auto dim() const noexcept -> uint32_t - { - return spec_type::get_dim(data_view_, dictionary_view_); - } - - /* - @return A non-owning view of the data, such as mdspan. - */ - [[nodiscard]] auto data_view() const noexcept -> data_view_type { return data_view_; } - /* - NOTE: giving up write access is not a plain cast for most view types - a mutable mdspan, for - one, has to be rebuilt over a const accessor - hence the spec does the conversion. - - @return A read-only view of the data. - */ - [[nodiscard]] auto data_const_view() const noexcept -> const_data_view_type - { - if constexpr (is_const_view) { - return data_view_; - } else { - return spec_type::get_data_const_view(data_view_); - } - } - - /* - @return A non-owning view of the dictionary, such as the two codebooks of a VPQ dataset. - */ - [[nodiscard]] auto dictionary_view() const noexcept -> dictionary_view_type - requires(is_compressed) - { - return dictionary_view_; - } - /* - @return A read-only view of the dictionary. - */ - [[nodiscard]] auto const_dictionary_view() const noexcept -> const_dictionary_view_type - requires(is_compressed) - { - if constexpr (is_const_view) { - return dictionary_view_; - } else { - return spec_type::get_dictionary_const_view(dictionary_view_); - } - } - - private: - data_view_type data_view_; - [[no_unique_address]] dictionary_view_type dictionary_view_; -}; - -/** - * An owning dataset: the data and, for a compressed dataset, the dictionary that gives it meaning. - * Ownership is unique, so the type is move-only whenever the underlying storage is. - */ -template -struct dataset { - /* NOTE: element type for an owning dataset cannot be const; - the constness of the dataset itself selects the flavor of the view it hands out. - */ - static_assert(!std::is_const_v, "Element type for an owning dataset must not be const."); - - using view_type = dataset_view; - using const_view_type = dataset_view; - - using spec_type = typename view_type::spec_type; - using value_type = typename view_type::value_type; - using index_type = typename view_type::index_type; - using code_type = typename view_type::code_type; - - using data_type = typename view_type::data_type; - using data_view_type = typename view_type::data_view_type; - using const_data_view_type = typename view_type::const_data_view_type; - - /** - Optional slot that contains everything to give meaning to the encoded data; - `std::monostate` for an uncompressed dataset. - */ - using dictionary_type = typename view_type::dictionary_type; - using dictionary_view_type = typename view_type::dictionary_view_type; - using const_dictionary_view_type = typename view_type::const_dictionary_view_type; - - static constexpr bool is_compressed = view_type::is_compressed; - - explicit dataset(data_type&& data) - requires(!is_compressed) - : data_{std::move(data)} - { - } - - dataset(data_type&& data, dictionary_type&& dictionary) - requires(is_compressed) - : data_{std::move(data)}, dictionary_{std::move(dictionary)} - { - } - - /* - @return A non-owning view of the whole dataset, data and dictionary alike. - */ - [[nodiscard]] auto view() noexcept -> view_type - { - if constexpr (is_compressed) { - return view_type{data_view(), dictionary_view()}; - } else { - return view_type{data_view()}; - } - } - /* - @return A read-only view of the whole dataset; the constness of the dataset selects this flavor. - */ - [[nodiscard]] auto view() const noexcept -> const_view_type - { - if constexpr (is_compressed) { - return const_view_type{data_const_view(), const_dictionary_view()}; - } else { - return const_view_type{data_const_view()}; - } - } - - [[nodiscard]] auto n_rows() const noexcept -> index_type { return spec_type::get_n_rows(data_); } - [[nodiscard]] auto dim() const noexcept -> uint32_t - { - // NOTE: a compressed dataset may not be able to tell its dimension from the encoded data alone - return spec_type::get_dim(data_, dictionary_); - } - - /* - @return A non-owning view of the data, such as mdspan. - */ - [[nodiscard]] auto data_view() noexcept -> data_view_type - { - return spec_type::get_data_view(data_); - } - [[nodiscard]] auto data_const_view() const noexcept -> const_data_view_type - { - return spec_type::get_data_const_view(data_); - } - - /* - @return A non-owning view of the dictionary, such as the two codebooks of a VPQ dataset. - */ - [[nodiscard]] auto dictionary_view() const noexcept -> dictionary_view_type - requires(is_compressed) - { - return spec_type::get_dictionary_view(dictionary_); - } - /* - @return A read-only view of the dictionary. - */ - [[nodiscard]] auto const_dictionary_view() const noexcept -> const_dictionary_view_type - requires(is_compressed) - { - return spec_type::get_dictionary_const_view(dictionary_); - } - - private: - data_type data_; - [[no_unique_address]] dictionary_type dictionary_; -}; - -/************************************************************************************************ - * Dataset specifications - ************************************************************************************************/ - -/** - * Empty dataset specification: contains no data, but keeps the dimension of the dataset implement - * the common interface. - */ -struct empty_spec { - struct empty_dataset_rep { - uint32_t dim; - }; - - template - struct apply { - using data_type = empty_dataset_rep; - using data_view_type = empty_dataset_rep; - /* NOTE: not `const empty_dataset_rep`: a read-only view holds this as a member, and a const - member would delete its assignment operator. - */ - using const_data_view_type = empty_dataset_rep; - using value_type = std::remove_cv_t; - using index_type = std::remove_cv_t; - - [[nodiscard]] static auto get_data_view(data_type& data) noexcept -> data_view_type - { - return data; - } - /* NOTE: the view of an empty dataset is the storage itself, so one overload serves both the - owning dataset and the view. - */ - [[nodiscard]] static auto get_data_const_view(const data_type& data) noexcept - -> const_data_view_type - { - return data; - } - /* NOTE: the shape accessors take the data in either form - the storage an owning dataset holds, - or the view a `dataset_view` holds. A plain spec defines no dictionary either, so it ignores - whatever the dataset puts in that slot instead. - */ - [[nodiscard]] static auto get_n_rows(const auto&) noexcept -> index_type { return 0; } - [[nodiscard]] static auto get_dim(const auto& data, const auto&) noexcept -> uint32_t - { - return data.dim; - } - }; -}; - -/** - * Dense plain/padded dataset implemented via raft::mdarray. - */ -template -struct mdarray_spec { - template - struct apply { - /* NOTE: index type != extents type - The index type can vary depending on the use case; we often may want to store indices in 32-bit - slots to save memory on large datasets. The extents are always fixed to 64-bit signed integers - to simplify the conversion and avoid integer overflow issues. - */ - using value_type = std::remove_cv_t; - using index_type = std::remove_cv_t; - using data_type = - raft::mdarray, LayoutPolicy, ContainerPolicy>; - using data_view_type = typename data_type::view_type; - using const_data_view_type = typename data_type::const_view_type; - - [[nodiscard]] static auto get_data_view(data_type& data) noexcept -> data_view_type - { - return data.view(); - } - - [[nodiscard]] static auto get_data_const_view(const data_type& data) noexcept - -> const_data_view_type - { - return data.view(); - } - /* A mutable mdspan does not decay to a read-only one on its own; it has to be rebuilt over an - accessor of `const T`. - */ - [[nodiscard]] static auto get_data_const_view(data_view_type data) noexcept - -> const_data_view_type - { - return raft::make_const_mdspan(data); - } - /* NOTE: the shape accessors take the data in either form - the storage an owning dataset holds, - or the view a `dataset_view` holds. A plain spec defines no dictionary either, so it ignores - whatever the dataset puts in that slot instead. - */ - [[nodiscard]] static auto get_n_rows(const auto& data) noexcept -> index_type - { - return static_cast(data.extent(0)); - } - [[nodiscard]] static auto get_dim(const auto& data, const auto&) noexcept -> uint32_t - { - return static_cast(data.extent(1)); - } - }; -}; - -/** - * Sparse dataset specification implemented via the raft sparse matrix hierarchy. `SparseLayoutT` - * selects the representation (`csr_layout`, `coo_layout`). - * - * NOTE: only the memory space of the container policy is used here; raft's sparse types bind the - * policy themselves, separately for the values, the offsets and the indices. - */ -template -struct sparse_spec { - template - struct apply { - using value_type = std::remove_cv_t; - using index_type = std::remove_cv_t; - - static constexpr raft::memory_type mem_type = ContainerPolicy::mem_type; - static_assert(mem_type == raft::memory_type::host || mem_type == raft::memory_type::device, - "raft's sparse hierarchy only provides host and device containers"); - - using data_type = std::conditional_t< - mem_type == raft::memory_type::device, - typename SparseLayoutT:: - template matrix_type, - typename SparseLayoutT:: - template matrix_type>; - using data_view_type = typename data_type::view_type; - /* NOTE: raft's sparse hierarchy provides no read-only view yet (see the TODO in - raft/core/sparse_types.hpp); until then, even reading the structure requires a mutable matrix. - - Consequently a read-only `dataset_view` over a sparse spec is read-only in name - only: the two view types coincide, so writing through it still compiles. That resolves itself - once raft grows a const sparse view. - */ - using const_data_view_type = data_view_type; - - [[nodiscard]] static auto get_data_view(data_type& data) noexcept -> data_view_type - { - return data.view(); - } - - [[nodiscard]] static auto get_data_const_view(const data_type& data) noexcept - -> const_data_view_type - { - return const_cast(data).view(); - } - /* NOTE: there is no read-only view to convert to (see above), so the view passes through. */ - [[nodiscard]] static auto get_data_const_view(data_view_type data) noexcept - -> const_data_view_type - { - return data; - } - /* NOTE: raft's sparse hierarchy exposes the structure through a mutable accessor only, so - reading the shape of either the owning matrix or its view takes a cast. - */ - [[nodiscard]] static auto get_structure_view(const auto& data) noexcept - { - return const_cast&>(data).structure_view(); - } - /* NOTE: the shape accessors take the data in either form - the storage an owning dataset holds, - or the view a `dataset_view` holds. A plain spec defines no dictionary either, so it ignores - whatever the dataset puts in that slot instead. - */ - [[nodiscard]] static auto get_n_rows(const auto& data) noexcept -> index_type - { - return static_cast(get_structure_view(data).get_n_rows()); - } - [[nodiscard]] static auto get_dim(const auto& data, const auto&) noexcept -> uint32_t - { - return static_cast(get_structure_view(data).get_n_cols()); - } - }; -}; - -/** - * CSR representation: one compressed row per vector of the dataset. - * - * NOTE: the row offsets are 64-bit for the same reason the dense extents are: the number of - * non-zeros is not bounded by the index type. The column indices are feature ids, so they are - * bounded by the dataset dimension rather than by its size. - */ -struct csr_layout { - template typename ContainerPolicy> - using matrix_type = raft::csr_matrix; -}; - -/** - * COO representation: one (row, column, value) triple per non-zero. - * - * NOTE: unlike CSR, the row component is an index into the dataset rather than an offset into the - * non-zeros, hence the dataset index type. - */ -struct coo_layout { - template typename ContainerPolicy> - using matrix_type = raft::coo_matrix; -}; - /** - * VQ+PQ compressed dataset specification, mirroring the VPQ dataset of cuvs/neighbors/common.hpp: - * the two codebooks form the dictionary, while the data slot keeps the encoded rows - the VQ label - * in the row prefix, followed by the packed PQ codes. - * - * `BookPolicy` is the container policy of the codebooks (elements of `MathT`) and `CodePolicy` the - * one of the encoded rows (elements of `uint8_t`); a bound policy cannot be rebound to another - * element type, hence the two parameters. + * The minimal viable dataset API: the owning dataset, its view, and the one spec that needs no + * storage backend at all. A dataset that does keep data around takes the spec that describes how - + * `cuvs/core/dataset/dense.hpp`, `sparse.hpp`, `vpq.hpp`, `sq.hpp` - so that including this header + * costs nothing beyond the types themselves. */ -template -struct vpq_spec { - /* The encoded rows are stored like the rows of any other dense dataset; only how a row is read - and where the dimension comes from are specific to the compression. - */ - template - using storage_spec = - typename mdarray_spec::template apply; - - /* The type the codebooks are stored in; the rows themselves are always packed codes. */ - using math_type = MathT; - - /* [vq_n_centers, dim] */ - using vq_book_type = - raft::mdarray, raft::layout_c_contiguous, BookPolicy>; - /* [n_books, pq_n_centers, pq_len], where `n_books` is 1 for a codebook shared by all the - subspaces, or `pq_dim` for one codebook per subspace. A row-major 3d array with `n_books == 1` - is bit-identical to the flat 2d codebook of the current implementation. - */ - using pq_book_type = - raft::mdarray, raft::layout_c_contiguous, BookPolicy>; - - struct dictionary_type { - vq_book_type vq_code_book; - pq_book_type pq_code_book; - }; - - struct dictionary_view_type { - typename vq_book_type::const_view_type vq_code_book; - typename pq_book_type::const_view_type pq_code_book; - - /* NOTE: the dimension of the dataset is a property of the VQ codebook: the encoded rows tell - only how many bytes it took to compress a vector, which the padding makes ambiguous. - */ - [[nodiscard]] auto dim() const noexcept -> uint32_t - { - return static_cast(vq_code_book.extent(1)); - } - [[nodiscard]] auto vq_n_centers() const noexcept -> uint32_t - { - return static_cast(vq_code_book.extent(0)); - } - [[nodiscard]] auto n_books() const noexcept -> uint32_t - { - return static_cast(pq_code_book.extent(0)); - } - [[nodiscard]] auto pq_n_centers() const noexcept -> uint32_t - { - return static_cast(pq_code_book.extent(1)); - } - [[nodiscard]] auto pq_len() const noexcept -> uint32_t - { - return static_cast(pq_code_book.extent(2)); - } - [[nodiscard]] auto pq_bits() const noexcept -> uint32_t - { - return static_cast(std::countr_zero(pq_n_centers())); - } - [[nodiscard]] auto pq_dim() const noexcept -> uint32_t - { - return raft::div_rounding_up_safe(dim(), pq_len()); - } - /* Whether every subspace has a codebook of its own, as opposed to sharing a single one. */ - [[nodiscard]] auto per_subspace() const noexcept -> bool { return n_books() > 1; } - }; - - template - struct apply : storage_spec { - /* The dataset stores codes, but it still represents vectors of `T`. */ - using value_type = std::remove_cv_t; - using code_type = uint8_t; - /* NOTE: re-exported rather than declared here, so that they do not belong to one instantiation - of `apply`; qualified lookup into `apply` does not reach the enclosing spec on its own. - */ - using math_type = vpq_spec::math_type; - using vq_book_type = vpq_spec::vq_book_type; - using pq_book_type = vpq_spec::pq_book_type; - using dictionary_type = vpq_spec::dictionary_type; - using dictionary_view_type = vpq_spec::dictionary_view_type; - - /* NOTE: the data accessors are inherited; only `get_dim` differs from a dense dataset. Like the - inherited ones, it reads the dictionary in either form: the codebooks an owning dataset holds, - or the views a `dataset_view` holds. - */ - [[nodiscard]] static auto get_dim(const auto&, const auto& dict) noexcept -> uint32_t - { - return static_cast(dict.vq_code_book.extent(1)); - } - [[nodiscard]] static auto get_dictionary_view(const dictionary_type& dict) noexcept - -> dictionary_view_type - { - return {dict.vq_code_book.view(), dict.pq_code_book.view()}; - } - /* NOTE: a codebook view is read-only to begin with, so the spec declares no separate const view - and these overloads exist only to spell the conversion out of either form. - */ - [[nodiscard]] static auto get_dictionary_const_view(const dictionary_type& dict) noexcept - -> dictionary_view_type - { - return get_dictionary_view(dict); - } - [[nodiscard]] static auto get_dictionary_const_view(dictionary_view_type dict) noexcept - -> dictionary_view_type - { - return dict; - } - /* The length of an encoded row in bytes, including the inlined VQ label. */ - [[nodiscard]] static auto get_encoded_row_length(const auto& data) noexcept -> uint32_t - { - return static_cast(data.extent(1)); - } - }; -}; - -/** - * Scalar-quantized dataset specification: every component of a vector becomes one code, so the - * dictionary is only the range the codes are mapped back onto and the dimension is still a property - * of the data. - */ -template -struct sq_spec { - /* One code per component, so the codes are laid out exactly like a dense dataset. */ - template - using storage_spec = - typename mdarray_spec::template apply; - - using math_type = MathT; - - /* Nothing to allocate: the whole dictionary is the interval the codes are dequantized into. - The members are named as in cuvs::preprocessing::quantize::scalar::quantizer. - */ - struct dictionary_type { - math_type min_; - math_type max_; - }; - using dictionary_view_type = dictionary_type; - - template - struct apply : storage_spec { - /* The dataset stores codes, but it still represents vectors of `T`. */ - using value_type = std::remove_cv_t; - using code_type = int8_t; - /* NOTE: re-exported rather than declared here, so that they do not belong to one instantiation - of `apply`; qualified lookup into `apply` does not reach the enclosing spec on its own. - */ - using math_type = sq_spec::math_type; - using dictionary_type = sq_spec::dictionary_type; - using dictionary_view_type = sq_spec::dictionary_view_type; - - /* NOTE: every accessor is inherited; the dimension is still one code per component, so even - `get_dim` is the dense one. - */ - [[nodiscard]] static auto get_dictionary_view(const dictionary_type& dict) noexcept - -> dictionary_view_type - { - return dict; - } - /* NOTE: the dictionary is its own view here, read-only or not, so one overload serves both the - owning dataset and the view. - */ - [[nodiscard]] static auto get_dictionary_const_view(const dictionary_type& dict) noexcept - -> dictionary_view_type - { - return dict; - } - }; -}; - -/************************************************************************************************ - * Convenience aliases - ************************************************************************************************/ - -// Empty -template -using empty_dataset = dataset; - -// Dense -template -using contiguous_dataset = - dataset>; -template -using padded_dataset = dataset, ContainerPolicy>>; -template -using device_contiguous_dataset = - contiguous_dataset>>; -template -using device_padded_dataset = - padded_dataset>>; -template -using host_contiguous_dataset = - contiguous_dataset>>; -template -using host_padded_dataset = - padded_dataset>>; - -// Sparse -template -using csr_dataset = dataset>; -template -using coo_dataset = dataset>; -template -using device_csr_dataset = - csr_dataset>>; -template -using device_coo_dataset = - coo_dataset>>; -template -using host_csr_dataset = csr_dataset>>; -template -using host_coo_dataset = coo_dataset>>; - -// Compressed -template -using vpq_dataset = dataset>; -template -using sq_dataset = dataset>; -template -using device_vpq_dataset = - vpq_dataset>, - raft::device_accessor>>; -template -using host_vpq_dataset = vpq_dataset>, - raft::host_accessor>>; -template -using device_sq_dataset = - sq_dataset>>; -template -using host_sq_dataset = - sq_dataset>>; -} // namespace CUVS_EXPORT cuvs +#include +#include +#include diff --git a/cpp/include/cuvs/core/dataset/dense.hpp b/cpp/include/cuvs/core/dataset/dense.hpp new file mode 100644 index 0000000000..e0e1775af3 --- /dev/null +++ b/cpp/include/cuvs/core/dataset/dense.hpp @@ -0,0 +1,91 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include + +#include +#include +#include +#include + +#include +#include + +namespace CUVS_EXPORT cuvs { + +/** + * Dense plain/padded dataset implemented via raft::mdarray. + */ +template +struct mdarray_spec { + template + struct apply { + /* NOTE: index type != extents type + The index type can vary depending on the use case; we often may want to store indices in 32-bit + slots to save memory on large datasets. The extents are always fixed to 64-bit signed integers + to simplify the conversion and avoid integer overflow issues. + */ + using value_type = std::remove_cv_t; + using index_type = std::remove_cv_t; + using data_type = + raft::mdarray, LayoutPolicy, ContainerPolicy>; + using data_view_type = typename data_type::view_type; + using const_data_view_type = typename data_type::const_view_type; + + [[nodiscard]] static auto get_data_view(data_type& data) noexcept -> data_view_type + { + return data.view(); + } + + [[nodiscard]] static auto get_data_const_view(const data_type& data) noexcept + -> const_data_view_type + { + return data.view(); + } + /* A mutable mdspan does not decay to a read-only one on its own; it has to be rebuilt over an + accessor of `const T`. + */ + [[nodiscard]] static auto get_data_const_view(data_view_type data) noexcept + -> const_data_view_type + { + return raft::make_const_mdspan(data); + } + /* NOTE: the shape accessors take the data in either form - the storage an owning dataset holds, + or the view a `dataset_view` holds. A plain spec defines no dictionary either, so it ignores + whatever the dataset puts in that slot instead. + */ + [[nodiscard]] static auto get_n_rows(const auto& data) noexcept -> index_type + { + return static_cast(data.extent(0)); + } + [[nodiscard]] static auto get_dim(const auto& data, const auto&) noexcept -> uint32_t + { + return static_cast(data.extent(1)); + } + }; +}; + +template +using contiguous_dataset = + dataset>; +template +using padded_dataset = dataset, ContainerPolicy>>; +template +using device_contiguous_dataset = + contiguous_dataset>>; +template +using device_padded_dataset = + padded_dataset>>; +template +using host_contiguous_dataset = + contiguous_dataset>>; +template +using host_padded_dataset = + padded_dataset>>; + +} // namespace CUVS_EXPORT cuvs diff --git a/cpp/include/cuvs/core/dataset/empty.hpp b/cpp/include/cuvs/core/dataset/empty.hpp new file mode 100644 index 0000000000..c2d68524f2 --- /dev/null +++ b/cpp/include/cuvs/core/dataset/empty.hpp @@ -0,0 +1,63 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include + +#include +#include + +namespace CUVS_EXPORT cuvs { + +/** + * Empty dataset specification: contains no data, but keeps the dimension of the dataset implement + * the common interface. + */ +struct empty_spec { + struct empty_dataset_rep { + uint32_t dim; + }; + + template + struct apply { + using data_type = empty_dataset_rep; + using data_view_type = empty_dataset_rep; + /* NOTE: not `const empty_dataset_rep`: a read-only view holds this as a member, and a const + member would delete its assignment operator. + */ + using const_data_view_type = empty_dataset_rep; + using value_type = std::remove_cv_t; + using index_type = std::remove_cv_t; + + [[nodiscard]] static auto get_data_view(data_type& data) noexcept -> data_view_type + { + return data; + } + /* NOTE: the view of an empty dataset is the storage itself, so one overload serves both the + owning dataset and the view. + */ + [[nodiscard]] static auto get_data_const_view(const data_type& data) noexcept + -> const_data_view_type + { + return data; + } + /* NOTE: the shape accessors take the data in either form - the storage an owning dataset holds, + or the view a `dataset_view` holds. A plain spec defines no dictionary either, so it ignores + whatever the dataset puts in that slot instead. + */ + [[nodiscard]] static auto get_n_rows(const auto&) noexcept -> index_type { return 0; } + [[nodiscard]] static auto get_dim(const auto& data, const auto&) noexcept -> uint32_t + { + return data.dim; + } + }; +}; + +template +using empty_dataset = dataset; + +} // namespace CUVS_EXPORT cuvs diff --git a/cpp/include/cuvs/core/dataset/sparse.hpp b/cpp/include/cuvs/core/dataset/sparse.hpp new file mode 100644 index 0000000000..617c42154b --- /dev/null +++ b/cpp/include/cuvs/core/dataset/sparse.hpp @@ -0,0 +1,131 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include + +#include +#include +#include +#include + +#include +#include + +namespace CUVS_EXPORT cuvs { + +/** + * Sparse dataset specification implemented via the raft sparse matrix hierarchy. `SparseLayoutT` + * selects the representation (`csr_layout`, `coo_layout`). + * + * NOTE: only the memory space of the container policy is used here; raft's sparse types bind the + * policy themselves, separately for the values, the offsets and the indices. + */ +template +struct sparse_spec { + template + struct apply { + using value_type = std::remove_cv_t; + using index_type = std::remove_cv_t; + + static constexpr raft::memory_type mem_type = ContainerPolicy::mem_type; + static_assert(mem_type == raft::memory_type::host || mem_type == raft::memory_type::device, + "raft's sparse hierarchy only provides host and device containers"); + + using data_type = std::conditional_t< + mem_type == raft::memory_type::device, + typename SparseLayoutT:: + template matrix_type, + typename SparseLayoutT:: + template matrix_type>; + using data_view_type = typename data_type::view_type; + /* NOTE: raft's sparse hierarchy provides no read-only view yet (see the TODO in + raft/core/sparse_types.hpp); until then, even reading the structure requires a mutable matrix. + + Consequently a read-only `dataset_view` over a sparse spec is read-only in name + only: the two view types coincide, so writing through it still compiles. That resolves itself + once raft grows a const sparse view. + */ + using const_data_view_type = data_view_type; + + [[nodiscard]] static auto get_data_view(data_type& data) noexcept -> data_view_type + { + return data.view(); + } + + [[nodiscard]] static auto get_data_const_view(const data_type& data) noexcept + -> const_data_view_type + { + return const_cast(data).view(); + } + /* NOTE: there is no read-only view to convert to (see above), so the view passes through. */ + [[nodiscard]] static auto get_data_const_view(data_view_type data) noexcept + -> const_data_view_type + { + return data; + } + /* NOTE: raft's sparse hierarchy exposes the structure through a mutable accessor only, so + reading the shape of either the owning matrix or its view takes a cast. + */ + [[nodiscard]] static auto get_structure_view(const auto& data) noexcept + { + return const_cast&>(data).structure_view(); + } + /* NOTE: the shape accessors take the data in either form - the storage an owning dataset holds, + or the view a `dataset_view` holds. A plain spec defines no dictionary either, so it ignores + whatever the dataset puts in that slot instead. + */ + [[nodiscard]] static auto get_n_rows(const auto& data) noexcept -> index_type + { + return static_cast(get_structure_view(data).get_n_rows()); + } + [[nodiscard]] static auto get_dim(const auto& data, const auto&) noexcept -> uint32_t + { + return static_cast(get_structure_view(data).get_n_cols()); + } + }; +}; + +/** + * CSR representation: one compressed row per vector of the dataset. + * + * NOTE: the row offsets are 64-bit for the same reason the dense extents are: the number of + * non-zeros is not bounded by the index type. The column indices are feature ids, so they are + * bounded by the dataset dimension rather than by its size. + */ +struct csr_layout { + template typename ContainerPolicy> + using matrix_type = raft::csr_matrix; +}; + +/** + * COO representation: one (row, column, value) triple per non-zero. + * + * NOTE: unlike CSR, the row component is an index into the dataset rather than an offset into the + * non-zeros, hence the dataset index type. + */ +struct coo_layout { + template typename ContainerPolicy> + using matrix_type = raft::coo_matrix; +}; + +template +using csr_dataset = dataset>; +template +using coo_dataset = dataset>; +template +using device_csr_dataset = + csr_dataset>>; +template +using device_coo_dataset = + coo_dataset>>; +template +using host_csr_dataset = csr_dataset>>; +template +using host_coo_dataset = coo_dataset>>; + +} // namespace CUVS_EXPORT cuvs diff --git a/cpp/include/cuvs/core/dataset/sq.hpp b/cpp/include/cuvs/core/dataset/sq.hpp new file mode 100644 index 0000000000..f1703f1913 --- /dev/null +++ b/cpp/include/cuvs/core/dataset/sq.hpp @@ -0,0 +1,84 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include +#include + +#include +#include +#include + +#include +#include + +namespace CUVS_EXPORT cuvs { + +/** + * Scalar-quantized dataset specification: every component of a vector becomes one code, so the + * dictionary is only the range the codes are mapped back onto and the dimension is still a property + * of the data. + */ +template +struct sq_spec { + /* One code per component, so the codes are laid out exactly like a dense dataset. */ + template + using storage_spec = + typename mdarray_spec::template apply; + + using math_type = MathT; + + /* Nothing to allocate: the whole dictionary is the interval the codes are dequantized into. + The members are named as in cuvs::preprocessing::quantize::scalar::quantizer. + */ + struct dictionary_type { + math_type min_; + math_type max_; + }; + using dictionary_view_type = dictionary_type; + + template + struct apply : storage_spec { + /* The dataset stores codes, but it still represents vectors of `T`. */ + using value_type = std::remove_cv_t; + using code_type = int8_t; + /* NOTE: re-exported rather than declared here, so that they do not belong to one instantiation + of `apply`; qualified lookup into `apply` does not reach the enclosing spec on its own. + */ + using math_type = sq_spec::math_type; + using dictionary_type = sq_spec::dictionary_type; + using dictionary_view_type = sq_spec::dictionary_view_type; + + /* NOTE: every accessor is inherited; the dimension is still one code per component, so even + `get_dim` is the dense one. + */ + [[nodiscard]] static auto get_dictionary_view(const dictionary_type& dict) noexcept + -> dictionary_view_type + { + return dict; + } + /* NOTE: the dictionary is its own view here, read-only or not, so one overload serves both the + owning dataset and the view. + */ + [[nodiscard]] static auto get_dictionary_const_view(const dictionary_type& dict) noexcept + -> dictionary_view_type + { + return dict; + } + }; +}; + +template +using sq_dataset = dataset>; +template +using device_sq_dataset = + sq_dataset>>; +template +using host_sq_dataset = + sq_dataset>>; + +} // namespace CUVS_EXPORT cuvs diff --git a/cpp/include/cuvs/core/dataset/type.hpp b/cpp/include/cuvs/core/dataset/type.hpp new file mode 100644 index 0000000000..2ad59b43ba --- /dev/null +++ b/cpp/include/cuvs/core/dataset/type.hpp @@ -0,0 +1,126 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include + +#include +#include +#include + +namespace CUVS_EXPORT cuvs { + +/** + * An owning dataset: the data and, for a compressed dataset, the dictionary that gives it meaning. + * Ownership is unique, so the type is move-only whenever the underlying storage is. + */ +template +struct dataset { + /* NOTE: element type for an owning dataset cannot be const; + the constness of the dataset itself selects the flavor of the view it hands out. + */ + static_assert(!std::is_const_v, "Element type for an owning dataset must not be const."); + + using view_type = dataset_view; + using const_view_type = dataset_view; + + using spec_type = typename view_type::spec_type; + using value_type = typename view_type::value_type; + using index_type = typename view_type::index_type; + using code_type = typename view_type::code_type; + + using data_type = typename view_type::data_type; + using data_view_type = typename view_type::data_view_type; + using const_data_view_type = typename view_type::const_data_view_type; + + /** + Optional slot that contains everything to give meaning to the encoded data; + `std::monostate` for an uncompressed dataset. + */ + using dictionary_type = typename view_type::dictionary_type; + using dictionary_view_type = typename view_type::dictionary_view_type; + using const_dictionary_view_type = typename view_type::const_dictionary_view_type; + + static constexpr bool is_compressed = view_type::is_compressed; + + explicit dataset(data_type&& data) + requires(!is_compressed) + : data_{std::move(data)} + { + } + + dataset(data_type&& data, dictionary_type&& dictionary) + requires(is_compressed) + : data_{std::move(data)}, dictionary_{std::move(dictionary)} + { + } + + /* + @return A non-owning view of the whole dataset, data and dictionary alike. + */ + [[nodiscard]] auto view() noexcept -> view_type + { + if constexpr (is_compressed) { + return view_type{data_view(), dictionary_view()}; + } else { + return view_type{data_view()}; + } + } + /* + @return A read-only view of the whole dataset; the constness of the dataset selects this flavor. + */ + [[nodiscard]] auto view() const noexcept -> const_view_type + { + if constexpr (is_compressed) { + return const_view_type{data_const_view(), const_dictionary_view()}; + } else { + return const_view_type{data_const_view()}; + } + } + + [[nodiscard]] auto n_rows() const noexcept -> index_type { return spec_type::get_n_rows(data_); } + [[nodiscard]] auto dim() const noexcept -> uint32_t + { + // NOTE: a compressed dataset may not be able to tell its dimension from the encoded data alone + return spec_type::get_dim(data_, dictionary_); + } + + /* + @return A non-owning view of the data, such as mdspan. + */ + [[nodiscard]] auto data_view() noexcept -> data_view_type + { + return spec_type::get_data_view(data_); + } + [[nodiscard]] auto data_const_view() const noexcept -> const_data_view_type + { + return spec_type::get_data_const_view(data_); + } + + /* + @return A non-owning view of the dictionary, such as the two codebooks of a VPQ dataset. + */ + [[nodiscard]] auto dictionary_view() const noexcept -> dictionary_view_type + requires(is_compressed) + { + return spec_type::get_dictionary_view(dictionary_); + } + /* + @return A read-only view of the dictionary. + */ + [[nodiscard]] auto const_dictionary_view() const noexcept -> const_dictionary_view_type + requires(is_compressed) + { + return spec_type::get_dictionary_const_view(dictionary_); + } + + private: + data_type data_; + [[no_unique_address]] dictionary_type dictionary_; +}; + +} // namespace CUVS_EXPORT cuvs diff --git a/cpp/include/cuvs/core/dataset/view.hpp b/cpp/include/cuvs/core/dataset/view.hpp new file mode 100644 index 0000000000..b2c7c4629f --- /dev/null +++ b/cpp/include/cuvs/core/dataset/view.hpp @@ -0,0 +1,182 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include + +#include +#include +#include +#include + +namespace CUVS_EXPORT cuvs { + +/** + * A spec is compressing if it defines a dictionary: the codebooks, the quantization range, or + * whatever else is needed to interpret the data. Plain specs define none and the dataset + * substitutes `std::monostate`, so "there is no dictionary" and "the dictionary slot is empty" are + * the same vocabulary type. + * + * NOTE: a spec declares its vocabulary types - the dictionary above, and any storage representation + * of its own - at spec scope and only re-exports them from `apply`. A type defined inside `apply` + * belongs to one instantiation of it, and the read-only and the mutable flavor of a view apply the + * spec to element types that differ in constness; the two would then disagree about a nested + * dictionary and could no longer be built from one another. + */ +template +concept compressed_dataset_spec = requires { typename SpecT::dictionary_type; }; + +/** + * A non-owning view of a dataset: a view of the data and, for a compressed dataset, a view of the + * dictionary that gives it meaning. Cheap to copy and to pass by value, like an mdspan. + */ +template +struct dataset_view { + using spec_type = typename SpecT::template apply; + /* The read-only flavor, in which each view type collapses onto its const counterpart. */ + static constexpr bool is_const_view = std::is_const_v; + + using value_type = typename spec_type::value_type; + using index_type = typename spec_type::index_type; + /** + What one element of the data slot is + - for compressed datasets it is the raw data element type; often is uint8_t but may be + a larger type to enforce better alignment. + - for uncompressed datasets it is the value type of the dataset. + */ + using code_type = typename decltype([] { + if constexpr (requires { typename spec_type::code_type; }) { + return std::type_identity{}; + } else { + return std::type_identity{}; + } + }())::type; + + using data_type = typename spec_type::data_type; + using const_data_view_type = typename spec_type::const_data_view_type; + using data_view_type = + std::conditional_t; + + /** + Optional slot that contains everything to give meaning to the encoded data; + `std::monostate` for an uncompressed dataset. + */ + using dictionary_type = typename decltype([] { + if constexpr (requires { typename spec_type::dictionary_type; }) { + return std::type_identity{}; + } else { + return std::type_identity{}; + } + }())::type; + // Optional, so that the specs don't have to declare it; a spec that declares neither view uses + // the dictionary itself, and one that declares only the mutable view uses it for both. + using const_dictionary_view_type = typename decltype([] { + if constexpr (requires { typename spec_type::const_dictionary_view_type; }) { + return std::type_identity{}; + } else if constexpr (requires { typename spec_type::dictionary_view_type; }) { + return std::type_identity{}; + } else { + return std::type_identity{}; + } + }())::type; + // Optional, so that the specs don't have to declare it + using dictionary_view_type = typename decltype([] { + if constexpr (is_const_view) { + return std::type_identity{}; + } else if constexpr (requires { typename spec_type::dictionary_view_type; }) { + return std::type_identity{}; + } else { + return std::type_identity{}; + } + }())::type; + + static constexpr bool is_compressed = compressed_dataset_spec; + + explicit dataset_view(data_view_type data_view) + requires(!is_compressed) + : data_view_{data_view} + { + } + + dataset_view(data_view_type data_view, dictionary_view_type dictionary_view) + requires(is_compressed) + : data_view_{data_view}, dictionary_view_{dictionary_view} + { + } + + /* A mutable view decays to a read-only one, the way an mdspan does. + + NOTE: a constructor template, because in the mutable flavor the parameter would spell the very + same type; a non-template would be a copy constructor there, and declaring one - even a + constrained-away one - suppresses the implicitly declared copy constructor. + + NOTE: the dictionary is assigned in the body rather than in the init-list, because the accessor it + comes from is constrained on `is_compressed`. + */ + template + requires(is_const_view && std::same_as>) + dataset_view(const dataset_view& other) noexcept + : data_view_{other.data_const_view()} + { + if constexpr (is_compressed) { dictionary_view_ = other.const_dictionary_view(); } + } + + [[nodiscard]] auto n_rows() const noexcept -> index_type + { + return spec_type::get_n_rows(data_view_); + } + // NOTE: a compressed dataset may not be able to tell its dimension from the encoded data alone + [[nodiscard]] auto dim() const noexcept -> uint32_t + { + return spec_type::get_dim(data_view_, dictionary_view_); + } + + /* + @return A non-owning view of the data, such as mdspan. + */ + [[nodiscard]] auto data_view() const noexcept -> data_view_type { return data_view_; } + /* + NOTE: giving up write access is not a plain cast for most view types - a mutable mdspan, for + one, has to be rebuilt over a const accessor - hence the spec does the conversion. + + @return A read-only view of the data. + */ + [[nodiscard]] auto data_const_view() const noexcept -> const_data_view_type + { + if constexpr (is_const_view) { + return data_view_; + } else { + return spec_type::get_data_const_view(data_view_); + } + } + + /* + @return A non-owning view of the dictionary, such as the two codebooks of a VPQ dataset. + */ + [[nodiscard]] auto dictionary_view() const noexcept -> dictionary_view_type + requires(is_compressed) + { + return dictionary_view_; + } + /* + @return A read-only view of the dictionary. + */ + [[nodiscard]] auto const_dictionary_view() const noexcept -> const_dictionary_view_type + requires(is_compressed) + { + if constexpr (is_const_view) { + return dictionary_view_; + } else { + return spec_type::get_dictionary_const_view(dictionary_view_); + } + } + + private: + data_view_type data_view_; + [[no_unique_address]] dictionary_view_type dictionary_view_; +}; + +} // namespace CUVS_EXPORT cuvs diff --git a/cpp/include/cuvs/core/dataset/vpq.hpp b/cpp/include/cuvs/core/dataset/vpq.hpp new file mode 100644 index 0000000000..bc47b9fc23 --- /dev/null +++ b/cpp/include/cuvs/core/dataset/vpq.hpp @@ -0,0 +1,162 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include +#include + +#include +#include +#include +#include + +#include +#include +#include + +namespace CUVS_EXPORT cuvs { + +/** + * VQ+PQ compressed dataset specification, mirroring the VPQ dataset of cuvs/neighbors/common.hpp: + * the two codebooks form the dictionary, while the data slot keeps the encoded rows - the VQ label + * in the row prefix, followed by the packed PQ codes. + * + * `BookPolicy` is the container policy of the codebooks (elements of `MathT`) and `CodePolicy` the + * one of the encoded rows (elements of `uint8_t`); a bound policy cannot be rebound to another + * element type, hence the two parameters. + */ +template +struct vpq_spec { + /* The encoded rows are stored like the rows of any other dense dataset; only how a row is read + and where the dimension comes from are specific to the compression. + */ + template + using storage_spec = + typename mdarray_spec::template apply; + + /* The type the codebooks are stored in; the rows themselves are always packed codes. */ + using math_type = MathT; + + /* [vq_n_centers, dim] */ + using vq_book_type = + raft::mdarray, raft::layout_c_contiguous, BookPolicy>; + /* [n_books, pq_n_centers, pq_len], where `n_books` is 1 for a codebook shared by all the + subspaces, or `pq_dim` for one codebook per subspace. A row-major 3d array with `n_books == 1` + is bit-identical to the flat 2d codebook of the current implementation. + */ + using pq_book_type = + raft::mdarray, raft::layout_c_contiguous, BookPolicy>; + + struct dictionary_type { + vq_book_type vq_code_book; + pq_book_type pq_code_book; + }; + + struct dictionary_view_type { + typename vq_book_type::const_view_type vq_code_book; + typename pq_book_type::const_view_type pq_code_book; + + /* NOTE: the dimension of the dataset is a property of the VQ codebook: the encoded rows tell + only how many bytes it took to compress a vector, which the padding makes ambiguous. + */ + [[nodiscard]] auto dim() const noexcept -> uint32_t + { + return static_cast(vq_code_book.extent(1)); + } + [[nodiscard]] auto vq_n_centers() const noexcept -> uint32_t + { + return static_cast(vq_code_book.extent(0)); + } + [[nodiscard]] auto n_books() const noexcept -> uint32_t + { + return static_cast(pq_code_book.extent(0)); + } + [[nodiscard]] auto pq_n_centers() const noexcept -> uint32_t + { + return static_cast(pq_code_book.extent(1)); + } + [[nodiscard]] auto pq_len() const noexcept -> uint32_t + { + return static_cast(pq_code_book.extent(2)); + } + [[nodiscard]] auto pq_bits() const noexcept -> uint32_t + { + return static_cast(std::countr_zero(pq_n_centers())); + } + [[nodiscard]] auto pq_dim() const noexcept -> uint32_t + { + return raft::div_rounding_up_safe(dim(), pq_len()); + } + /* Whether every subspace has a codebook of its own, as opposed to sharing a single one. */ + [[nodiscard]] auto per_subspace() const noexcept -> bool { return n_books() > 1; } + }; + + template + struct apply : storage_spec { + /* The dataset stores codes, but it still represents vectors of `T`. */ + using value_type = std::remove_cv_t; + using code_type = uint8_t; + /* NOTE: re-exported rather than declared here, so that they do not belong to one instantiation + of `apply`; qualified lookup into `apply` does not reach the enclosing spec on its own. + */ + using math_type = vpq_spec::math_type; + using vq_book_type = vpq_spec::vq_book_type; + using pq_book_type = vpq_spec::pq_book_type; + using dictionary_type = vpq_spec::dictionary_type; + using dictionary_view_type = vpq_spec::dictionary_view_type; + + /* NOTE: the data accessors are inherited; only `get_dim` differs from a dense dataset. Like the + inherited ones, it reads the dictionary in either form: the codebooks an owning dataset holds, + or the views a `dataset_view` holds. + */ + [[nodiscard]] static auto get_dim(const auto&, const auto& dict) noexcept -> uint32_t + { + return static_cast(dict.vq_code_book.extent(1)); + } + [[nodiscard]] static auto get_dictionary_view(const dictionary_type& dict) noexcept + -> dictionary_view_type + { + return {dict.vq_code_book.view(), dict.pq_code_book.view()}; + } + /* NOTE: a codebook view is read-only to begin with, so the spec declares no separate const view + and these overloads exist only to spell the conversion out of either form. + */ + [[nodiscard]] static auto get_dictionary_const_view(const dictionary_type& dict) noexcept + -> dictionary_view_type + { + return get_dictionary_view(dict); + } + [[nodiscard]] static auto get_dictionary_const_view(dictionary_view_type dict) noexcept + -> dictionary_view_type + { + return dict; + } + /* The length of an encoded row in bytes, including the inlined VQ label. */ + [[nodiscard]] static auto get_encoded_row_length(const auto& data) noexcept -> uint32_t + { + return static_cast(data.extent(1)); + } + }; +}; + +template +using vpq_dataset = dataset>; +template +using device_vpq_dataset = + vpq_dataset>, + raft::device_accessor>>; +template +using host_vpq_dataset = vpq_dataset>, + raft::host_accessor>>; + +} // namespace CUVS_EXPORT cuvs diff --git a/examples/cpp/src/dataset_example.cu b/examples/cpp/src/dataset_example.cu index a678f50786..982ff2ab94 100644 --- a/examples/cpp/src/dataset_example.cu +++ b/examples/cpp/src/dataset_example.cu @@ -4,6 +4,10 @@ */ #include +#include +#include +#include +#include #include #include diff --git a/fern/docs.yml b/fern/docs.yml index b055d844f1..8b396c384a 100644 --- a/fern/docs.yml +++ b/fern/docs.yml @@ -358,10 +358,22 @@ navigation: path: "./pages/cpp_api/cpp-api-common-types-copy-serialization-and-utility-apis.md" - page: "Errors and Logging" path: "./pages/cpp_api/cpp-api-common-types-errors-and-logging.md" - - page: "Core Dataset" - path: "./pages/cpp_api/cpp-api-core-dataset.md" + - page: "Core Dataset Dense" + path: "./pages/cpp_api/cpp-api-core-dataset-dense.md" + - page: "Core Dataset Empty" + path: "./pages/cpp_api/cpp-api-core-dataset-empty.md" - page: "Core Roaring Allowlist" path: "./pages/cpp_api/cpp-api-core-roaring-allowlist.md" + - page: "Core Dataset Sparse" + path: "./pages/cpp_api/cpp-api-core-dataset-sparse.md" + - page: "Core Dataset SQ" + path: "./pages/cpp_api/cpp-api-core-dataset-sq.md" + - page: "Core Dataset Type" + path: "./pages/cpp_api/cpp-api-core-dataset-type.md" + - page: "Core Dataset View" + path: "./pages/cpp_api/cpp-api-core-dataset-view.md" + - page: "Core Dataset Vpq" + path: "./pages/cpp_api/cpp-api-core-dataset-vpq.md" - page: "Distance Distance" path: "./pages/cpp_api/cpp-api-distance-distance.md" - page: "Distance Grammian" diff --git a/fern/pages/cpp_api/cpp-api-core-dataset-dense.md b/fern/pages/cpp_api/cpp-api-core-dataset-dense.md new file mode 100644 index 0000000000..891a14ea43 --- /dev/null +++ b/fern/pages/cpp_api/cpp-api-core-dataset-dense.md @@ -0,0 +1,19 @@ +--- +slug: api-reference/cpp-api-core-dataset-dense +--- + +# Dense + +_Source header: `cuvs/core/dataset/dense.hpp`_ + +## Types + + +### mdarray_spec + +Dense plain/padded dataset implemented via raft::mdarray. + +```cpp +template +struct mdarray_spec; +``` diff --git a/fern/pages/cpp_api/cpp-api-core-dataset-empty.md b/fern/pages/cpp_api/cpp-api-core-dataset-empty.md new file mode 100644 index 0000000000..a204a03382 --- /dev/null +++ b/fern/pages/cpp_api/cpp-api-core-dataset-empty.md @@ -0,0 +1,20 @@ +--- +slug: api-reference/cpp-api-core-dataset-empty +--- + +# Empty + +_Source header: `cuvs/core/dataset/empty.hpp`_ + +## Types + + +### empty_spec + +Empty dataset specification: contains no data, but keeps the dimension of the dataset implement + +the common interface. + +```cpp +struct empty_spec; +``` diff --git a/fern/pages/cpp_api/cpp-api-core-dataset-sparse.md b/fern/pages/cpp_api/cpp-api-core-dataset-sparse.md new file mode 100644 index 0000000000..0e3788c6b9 --- /dev/null +++ b/fern/pages/cpp_api/cpp-api-core-dataset-sparse.md @@ -0,0 +1,45 @@ +--- +slug: api-reference/cpp-api-core-dataset-sparse +--- + +# Sparse + +_Source header: `cuvs/core/dataset/sparse.hpp`_ + +## Types + + +### sparse_spec + +Sparse dataset specification implemented via the raft sparse matrix hierarchy. `SparseLayoutT` + +selects the representation (`csr_layout`, `coo_layout`). + +NOTE: only the memory space of the container policy is used here; raft's sparse types bind the policy themselves, separately for the values, the offsets and the indices. + +```cpp +template +struct sparse_spec; +``` + + +### csr_layout + +CSR representation: one compressed row per vector of the dataset. + +NOTE: the row offsets are 64-bit for the same reason the dense extents are: the number of non-zeros is not bounded by the index type. The column indices are feature ids, so they are bounded by the dataset dimension rather than by its size. + +```cpp +struct csr_layout; +``` + + +### coo_layout + +COO representation: one (row, column, value) triple per non-zero. + +NOTE: unlike CSR, the row component is an index into the dataset rather than an offset into the non-zeros, hence the dataset index type. + +```cpp +struct coo_layout; +``` diff --git a/fern/pages/cpp_api/cpp-api-core-dataset-sq.md b/fern/pages/cpp_api/cpp-api-core-dataset-sq.md new file mode 100644 index 0000000000..04a3a8021d --- /dev/null +++ b/fern/pages/cpp_api/cpp-api-core-dataset-sq.md @@ -0,0 +1,31 @@ +--- +slug: api-reference/cpp-api-core-dataset-sq +--- + +# SQ + +_Source header: `cuvs/core/dataset/sq.hpp`_ + +## Types + + +### sq_spec + +Scalar-quantized dataset specification: every component of a vector becomes one code, so the + +dictionary is only the range the codes are mapped back onto and the dimension is still a property of the data. + +```cpp +template +struct sq_spec { + /* One code per component, so the codes are laid out exactly like a dense dataset. */ template using storage_spec; + /* Nothing to allocate: the whole dictionary is the interval the codes are dequantized into. The members are named as in cuvs::preprocessing::quantize::scalar::quantizer. */ struct dictionary_type; +}; +``` + +**Fields** + +| Name | Type | Description | +| --- | --- | --- | +| `storage_spec` | [`/* One code per component, so the codes are laid out exactly like a dense dataset. */ template using`](/api-reference/cpp-api-core-dataset-type#dataset) | | +| `dictionary_type` | [`/* Nothing to allocate: the whole dictionary is the interval the codes are dequantized into. The members are named as in cuvs::preprocessing::quantize::scalar::quantizer. */ struct`](/api-reference/cpp-api-preprocessing-quantize-scalar#preprocessing-quantize-scalar-quantizer) | | diff --git a/fern/pages/cpp_api/cpp-api-core-dataset-type.md b/fern/pages/cpp_api/cpp-api-core-dataset-type.md new file mode 100644 index 0000000000..15f994c658 --- /dev/null +++ b/fern/pages/cpp_api/cpp-api-core-dataset-type.md @@ -0,0 +1,44 @@ +--- +slug: api-reference/cpp-api-core-dataset-type +--- + +# Type + +_Source header: `cuvs/core/dataset/type.hpp`_ + +## Types + + +### dataset + +An owning dataset: the data and, for a compressed dataset, the dictionary that gives it meaning. + +Ownership is unique, so the type is move-only whenever the underlying storage is. + +```cpp +template +struct dataset { + using dictionary_type; + /* NOTE: element type for an owning dataset cannot be const; + static constexpr bool is_compressed; + /* @return A read-only view of the whole dataset; +}; +``` + +**Fields** + +| Name | Type | Description | +| --- | --- | --- | +| `dictionary_type` | `using` | Optional slot that contains everything to give meaning to the encoded data; `std::monostate` for an uncompressed dataset. | +| `const` | [`/* NOTE: element type for an owning dataset cannot be`](/api-reference/cpp-api-core-dataset-type#dataset) | | +| `is_compressed` | `static constexpr bool` | | +| `dataset` | `/* @return A read-only view of the whole` | | + + +### dictionary_type + +Optional slot that contains everything to give meaning to the encoded data; + +```cpp +using dictionary_type = typename view_type::dictionary_type; +``` diff --git a/fern/pages/cpp_api/cpp-api-core-dataset-view.md b/fern/pages/cpp_api/cpp-api-core-dataset-view.md new file mode 100644 index 0000000000..202bf82eaa --- /dev/null +++ b/fern/pages/cpp_api/cpp-api-core-dataset-view.md @@ -0,0 +1,62 @@ +--- +slug: api-reference/cpp-api-core-dataset-view +--- + +# View + +_Source header: `cuvs/core/dataset/view.hpp`_ + +## Types + + +### dataset_view + +A non-owning view of a dataset: a view of the data and, for a compressed dataset, a view of the + +dictionary that gives it meaning. Cheap to copy and to pass by value, like an mdspan. + +```cpp +template +struct dataset_view { + /* The read-only flavor, in which each view type collapses onto its const counterpart. */ static constexpr bool is_const_view; + static constexpr bool is_compressed; +}; +``` + +**Fields** + +| Name | Type | Description | +| --- | --- | --- | +| `is_const_view` | `/* The read-only flavor, in which each view type collapses onto its const counterpart. */ static constexpr bool` | | +| `is_compressed` | `static constexpr bool` | | + + +### dataset_view::code_type + +What one element of the data slot is + +```cpp +using code_type = typename decltype([]; +``` + +- for compressed datasets it is the raw data element type; often is uint8_t but may be a larger type to enforce better alignment. +- for uncompressed datasets it is the value type of the dataset. + +**Returns** + +[`using code_type = typename`](/api-reference/cpp-api-core-dataset-view#dataset-view-code-type) + + +### dataset_view::dictionary_type + +Optional slot that contains everything to give meaning to the encoded data; + +```cpp +using dictionary_type = typename decltype([]; +``` + +`std::monostate` for an uncompressed dataset. + +**Returns** + +[`using dictionary_type = typename`](/api-reference/cpp-api-core-dataset-view#dataset-view-dictionary-type) diff --git a/fern/pages/cpp_api/cpp-api-core-dataset-vpq.md b/fern/pages/cpp_api/cpp-api-core-dataset-vpq.md new file mode 100644 index 0000000000..f252c59dac --- /dev/null +++ b/fern/pages/cpp_api/cpp-api-core-dataset-vpq.md @@ -0,0 +1,41 @@ +--- +slug: api-reference/cpp-api-core-dataset-vpq +--- + +# Vpq + +_Source header: `cuvs/core/dataset/vpq.hpp`_ + +## Types + + +### vpq_spec + +VQ+PQ compressed dataset specification, mirroring the VPQ dataset of cuvs/neighbors/common.hpp: + +the two codebooks form the dictionary, while the data slot keeps the encoded rows - the VQ label in the row prefix, followed by the packed PQ codes. + +`BookPolicy` is the container policy of the codebooks (elements of `MathT`) and `CodePolicy` the one of the encoded rows (elements of `uint8_t`); a bound policy cannot be rebound to another element type, hence the two parameters. + +```cpp +template +struct vpq_spec { + /* The encoded rows are stored like the rows of any other dense dataset; + only how a row is read and where the dimension comes from are specific to the compression. */ template using storage_spec; + /* The type the codebooks are stored in; + the rows themselves are always packed codes. */ using math_type; + /* [vq_n_centers, dim] */ using vq_book_type; + /* [n_books, pq_n_centers, pq_len], where `n_books` is 1 for a codebook shared by all the subspaces, or `pq_dim` for one codebook per subspace. A row-major 3d array with ` n_books; +}; +``` + +**Fields** + +| Name | Type | Description | +| --- | --- | --- | +| `dataset` | `/* The encoded rows are stored like the rows of any other dense` | | +| `storage_spec` | `only how a row is read and where the dimension comes from are specific to the compression. */ template using` | | +| `in` | `/* The type the codebooks are stored` | | +| `math_type` | `the rows themselves are always packed codes. */ using` | | +| `vq_book_type` | `/* [vq_n_centers, dim] */ using` | | +| `n_books` | `/* [n_books, pq_n_centers, pq_len], where \`n_books\` is 1 for a codebook shared by all the subspaces, or \`pq_dim\` for one codebook per subspace. A row-major 3d array with \`` | | diff --git a/fern/pages/cpp_api/cpp-api-core-dataset.md b/fern/pages/cpp_api/cpp-api-core-dataset.md deleted file mode 100644 index d7edff4151..0000000000 --- a/fern/pages/cpp_api/cpp-api-core-dataset.md +++ /dev/null @@ -1,208 +0,0 @@ ---- -slug: api-reference/cpp-api-core-dataset ---- - -# Dataset - -_Source header: `cuvs/core/dataset.hpp`_ - -## Types - - -### dataset_view - -A non-owning view of a dataset: a view of the data and, for a compressed dataset, a view of the - -dictionary that gives it meaning. Cheap to copy and to pass by value, like an mdspan. - -```cpp -template -struct dataset_view { - /* The read-only flavor, in which each view type collapses onto its const counterpart. */ static constexpr bool is_const_view; - static constexpr bool is_compressed; -}; -``` - -**Fields** - -| Name | Type | Description | -| --- | --- | --- | -| `is_const_view` | `/* The read-only flavor, in which each view type collapses onto its const counterpart. */ static constexpr bool` | | -| `is_compressed` | `static constexpr bool` | | - - -### dataset_view::code_type - -What one element of the data slot is - -```cpp -using code_type = typename decltype([]; -``` - -- for compressed datasets it is the raw data element type; often is uint8_t but may be a larger type to enforce better alignment. -- for uncompressed datasets it is the value type of the dataset. - -**Returns** - -[`using code_type = typename`](/api-reference/cpp-api-core-dataset#dataset-view-code-type) - - -### dataset_view::dictionary_type - -Optional slot that contains everything to give meaning to the encoded data; - -```cpp -using dictionary_type = typename decltype([]; -``` - -`std::monostate` for an uncompressed dataset. - -**Returns** - -[`using dictionary_type = typename`](/api-reference/cpp-api-core-dataset#dictionary-type) - - -### dataset - -An owning dataset: the data and, for a compressed dataset, the dictionary that gives it meaning. - -Ownership is unique, so the type is move-only whenever the underlying storage is. - -```cpp -template -struct dataset { - using dictionary_type; - /* NOTE: element type for an owning dataset cannot be const; - static constexpr bool is_compressed; - /* @return A read-only view of the whole dataset; -}; -``` - -**Fields** - -| Name | Type | Description | -| --- | --- | --- | -| `dictionary_type` | `using` | Optional slot that contains everything to give meaning to the encoded data; `std::monostate` for an uncompressed dataset. | -| `const` | [`/* NOTE: element type for an owning dataset cannot be`](/api-reference/cpp-api-core-dataset#dataset) | | -| `is_compressed` | `static constexpr bool` | | -| `dataset` | `/* @return A read-only view of the whole` | | - - -### dictionary_type - -Optional slot that contains everything to give meaning to the encoded data; - -```cpp -using dictionary_type = typename view_type::dictionary_type; -``` - - -### empty_spec - -Empty dataset specification: contains no data, but keeps the dimension of the dataset implement - -the common interface. - -```cpp -struct empty_spec; -``` - - -### mdarray_spec - -Dense plain/padded dataset implemented via raft::mdarray. - -```cpp -template -struct mdarray_spec; -``` - - -### sparse_spec - -Sparse dataset specification implemented via the raft sparse matrix hierarchy. `SparseLayoutT` - -selects the representation (`csr_layout`, `coo_layout`). - -NOTE: only the memory space of the container policy is used here; raft's sparse types bind the policy themselves, separately for the values, the offsets and the indices. - -```cpp -template -struct sparse_spec; -``` - - -### csr_layout - -CSR representation: one compressed row per vector of the dataset. - -NOTE: the row offsets are 64-bit for the same reason the dense extents are: the number of non-zeros is not bounded by the index type. The column indices are feature ids, so they are bounded by the dataset dimension rather than by its size. - -```cpp -struct csr_layout; -``` - - -### coo_layout - -COO representation: one (row, column, value) triple per non-zero. - -NOTE: unlike CSR, the row component is an index into the dataset rather than an offset into the non-zeros, hence the dataset index type. - -```cpp -struct coo_layout; -``` - - -### vpq_spec - -VQ+PQ compressed dataset specification, mirroring the VPQ dataset of cuvs/neighbors/common.hpp: - -the two codebooks form the dictionary, while the data slot keeps the encoded rows - the VQ label in the row prefix, followed by the packed PQ codes. - -`BookPolicy` is the container policy of the codebooks (elements of `MathT`) and `CodePolicy` the one of the encoded rows (elements of `uint8_t`); a bound policy cannot be rebound to another element type, hence the two parameters. - -```cpp -template -struct vpq_spec { - /* The encoded rows are stored like the rows of any other dense dataset; - only how a row is read and where the dimension comes from are specific to the compression. */ template using storage_spec; - /* The type the codebooks are stored in; - the rows themselves are always packed codes. */ using math_type; - /* [vq_n_centers, dim] */ using vq_book_type; - /* [n_books, pq_n_centers, pq_len], where `n_books` is 1 for a codebook shared by all the subspaces, or `pq_dim` for one codebook per subspace. A row-major 3d array with ` n_books; -}; -``` - -**Fields** - -| Name | Type | Description | -| --- | --- | --- | -| `dataset` | `/* The encoded rows are stored like the rows of any other dense` | | -| `storage_spec` | `only how a row is read and where the dimension comes from are specific to the compression. */ template using` | | -| `in` | `/* The type the codebooks are stored` | | -| `math_type` | `the rows themselves are always packed codes. */ using` | | -| `vq_book_type` | `/* [vq_n_centers, dim] */ using` | | -| `n_books` | `/* [n_books, pq_n_centers, pq_len], where \`n_books\` is 1 for a codebook shared by all the subspaces, or \`pq_dim\` for one codebook per subspace. A row-major 3d array with \`` | | - - -### sq_spec - -Scalar-quantized dataset specification: every component of a vector becomes one code, so the - -dictionary is only the range the codes are mapped back onto and the dimension is still a property of the data. - -```cpp -template -struct sq_spec { - /* One code per component, so the codes are laid out exactly like a dense dataset. */ template using storage_spec; - /* Nothing to allocate: the whole dictionary is the interval the codes are dequantized into. The members are named as in cuvs::preprocessing::quantize::scalar::quantizer. */ struct dictionary_type; -}; -``` - -**Fields** - -| Name | Type | Description | -| --- | --- | --- | -| `storage_spec` | [`/* One code per component, so the codes are laid out exactly like a dense dataset. */ template using`](/api-reference/cpp-api-core-dataset#dataset) | | -| `dictionary_type` | [`/* Nothing to allocate: the whole dictionary is the interval the codes are dequantized into. The members are named as in cuvs::preprocessing::quantize::scalar::quantizer. */ struct`](/api-reference/cpp-api-preprocessing-quantize-scalar#preprocessing-quantize-scalar-quantizer) | | diff --git a/fern/pages/cpp_api/index.md b/fern/pages/cpp_api/index.md index 67cb45e0f8..614a831b1e 100644 --- a/fern/pages/cpp_api/index.md +++ b/fern/pages/cpp_api/index.md @@ -13,8 +13,14 @@ These pages are generated from the documented public headers in the cuVS source - [Bloom Filter](/api-reference/cpp-api-core-bloom-filter) - [Common Types](/api-reference/cpp-api-common-types) -- [Dataset](/api-reference/cpp-api-core-dataset) +- [Dense](/api-reference/cpp-api-core-dataset-dense) +- [Empty](/api-reference/cpp-api-core-dataset-empty) - [Roaring Allowlist](/api-reference/cpp-api-core-roaring-allowlist) +- [Sparse](/api-reference/cpp-api-core-dataset-sparse) +- [SQ](/api-reference/cpp-api-core-dataset-sq) +- [Type](/api-reference/cpp-api-core-dataset-type) +- [View](/api-reference/cpp-api-core-dataset-view) +- [Vpq](/api-reference/cpp-api-core-dataset-vpq) ## Distance From 57544ad05189563f1e94541797c03b02668e1221 Mon Sep 17 00:00:00 2001 From: achirkin Date: Wed, 14 Oct 2026 18:33:26 +0200 Subject: [PATCH 05/11] Minimum changeset to use the new api --- c/src/neighbors/cagra.cpp | 15 +- c/src/preprocessing/quantize/pq.cpp | 6 +- cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h | 29 +- .../ann/src/cuvs/cuvs_mg_cagra_wrapper.h | 5 +- cpp/include/cuvs/core/dataset/dense.hpp | 13 +- cpp/include/cuvs/core/dataset/type.hpp | 12 +- cpp/include/cuvs/core/dataset/view.hpp | 10 + cpp/include/cuvs/core/dataset/vpq.hpp | 76 +- cpp/include/cuvs/neighbors/cagra.hpp | 23 +- cpp/include/cuvs/neighbors/common.hpp | 776 ++++-------------- .../cuvs/preprocessing/quantize/pq.hpp | 23 +- cpp/src/neighbors/cagra.cuh | 8 +- cpp/src/neighbors/detail/cagra/add_nodes.cuh | 7 +- .../neighbors/detail/cagra/cagra_build.cuh | 59 +- .../neighbors/detail/cagra/cagra_merge.cuh | 8 +- .../neighbors/detail/cagra/cagra_search.cuh | 9 +- .../detail/cagra/cagra_serialize.cuh | 4 +- .../cagra/compute_distance_standard.hpp | 2 +- .../detail/cagra/compute_distance_vpq.hpp | 14 +- cpp/src/neighbors/detail/cagra/factory.cuh | 21 +- .../neighbors/detail/dataset_serialize.hpp | 46 +- cpp/src/neighbors/detail/hnsw.hpp | 19 +- .../neighbors/detail/vamana/vamana_build.cuh | 6 +- .../detail/vamana/vamana_serialize.cuh | 4 +- cpp/src/neighbors/iface/iface.hpp | 8 +- cpp/src/neighbors/mg/mg_cagra_inst.cu.in | 25 +- .../neighbors/scann/detail/scann_build.cuh | 2 +- cpp/src/neighbors/tiered_index.cu | 15 +- cpp/src/preprocessing/quantize/detail/pq.cuh | 57 +- cpp/tests/neighbors/ann_cagra.cuh | 25 +- .../ann_cagra/test_iterative_cagra_q.cu | 9 +- .../ann_cagra/test_merge_fastener.cu | 18 +- cpp/tests/neighbors/ann_scann.cuh | 2 +- cpp/tests/neighbors/vpq_utils.cuh | 24 +- .../preprocessing/product_quantization.cu | 10 +- examples/cpp/src/cagra_hnsw_ace_example.cu | 8 +- examples/cpp/src/dataset_example.cu | 9 +- .../cpp_api/cpp-api-core-dataset-type.md | 2 + .../cpp_api/cpp-api-core-dataset-view.md | 2 + .../pages/cpp_api/cpp-api-core-dataset-vpq.md | 6 +- .../pages/cpp_api/cpp-api-neighbors-common.md | 59 +- .../cpp-api-preprocessing-quantize-pq.md | 2 +- 42 files changed, 533 insertions(+), 945 deletions(-) diff --git a/c/src/neighbors/cagra.cpp b/c/src/neighbors/cagra.cpp index 97cc9d8bac..a8392af245 100644 --- a/c/src/neighbors/cagra.cpp +++ b/c/src/neighbors/cagra.cpp @@ -182,9 +182,9 @@ static void merge_indices_for_layout( auto const stride = static_cast(index_ptrs.front()->dataset().stride()); try { - auto matrix = raft::make_device_matrix(*res_ptr, final_row_count, stride); using owner_t = cuvs::neighbors::owning_dataset_for_view_t; - auto owner = std::make_unique(std::move(matrix), dim); + auto owner = std::make_unique( + cuvs::neighbors::detail::make_dense_array(*res_ptr, final_row_count, dim, stride)); auto view = owner->as_dataset_view(); auto merged_idx = cuvs::neighbors::cagra::merge( @@ -218,8 +218,9 @@ static void merge_indices_for_layout( cuvs::neighbors::host_standard_dataset_view>; using host_owner_t = cuvs::neighbors::owning_dataset_for_view_t; - auto matrix = raft::make_host_matrix(final_row_count, stride); - std::fill_n(matrix.data_handle(), static_cast(matrix.size()), T{}); + auto matrix = cuvs::neighbors::detail::make_dense_array( + *res_ptr, final_row_count, dim, stride); + std::fill_n(matrix.data_handle(), matrix.mapping().required_span_size(), T{}); std::size_t row_offset = 0; auto stream = raft::resource::get_cuda_stream(*res_ptr); @@ -227,7 +228,7 @@ static void merge_indices_for_layout( auto const& input = index->dataset(); raft::copy_matrix(matrix.data_handle() + row_offset * static_cast(stride), static_cast(stride), - input.view().data_handle(), + input.data_const_view().data_handle(), static_cast(input.stride()), static_cast(dim), static_cast(input.n_rows()), @@ -236,7 +237,7 @@ static void merge_indices_for_layout( } raft::resource::sync_stream(*res_ptr); - auto owner = std::make_unique(std::move(matrix), dim); + auto owner = std::make_unique(std::move(matrix)); auto view = owner->as_dataset_view(); auto merged_idx = cuvs::neighbors::cagra::build(*res_ptr, params_cpp, view); auto* holder = @@ -1383,7 +1384,7 @@ void get_dataset_view(cuvsCagraIndex_t index, DLManagedTensor* dataset) RAFT_FAIL( "cuvsCagraIndexGetDataset does not expose quantized datasets as dense DLPack tensors"); } else { - cuvs::core::to_dlpack(idx.dataset().view(), dataset); + cuvs::core::to_dlpack(cuvs::neighbors::dense_row_major_view(idx.dataset()), dataset); } }); } diff --git a/c/src/preprocessing/quantize/pq.cpp b/c/src/preprocessing/quantize/pq.cpp index 1e3a48694a..169c24a001 100644 --- a/c/src/preprocessing/quantize/pq.cpp +++ b/c/src/preprocessing/quantize/pq.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -244,7 +244,7 @@ extern "C" cuvsError_t cuvsProductQuantizerGetPqCodebook(cuvsProductQuantizer_t if (quantizer->dtype.code == kDLFloat && quantizer->dtype.bits == 32) { auto pq_mdspan = (reinterpret_cast*>(quant_addr)) - ->vpq_codebooks.pq_code_book.view(); + ->vpq_codebooks.dictionary_view().pq_code_book; cuvs::core::to_dlpack(pq_mdspan, pq_codebook); } else { RAFT_FAIL("Unsupported quantizer dtype: %d and bits: %d", @@ -266,7 +266,7 @@ extern "C" cuvsError_t cuvsProductQuantizerGetVqCodebook(cuvsProductQuantizer_t if (quantizer->dtype.code == kDLFloat && quantizer->dtype.bits == 32) { auto pq_mdspan = (reinterpret_cast*>(quant_addr)) - ->vpq_codebooks.vq_code_book.view(); + ->vpq_codebooks.dictionary_view().vq_code_book; cuvs::core::to_dlpack(pq_mdspan, vq_codebook); } else { RAFT_FAIL("Unsupported quantizer dtype: %d and bits: %d", diff --git a/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h index da5e5baa9d..0842d011fb 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h +++ b/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h @@ -93,7 +93,12 @@ auto make_padded_view(const raft::resources& res, } } cuvs::neighbors::cagra::detail::copy_with_padding(res, buffer, src); - return {raft::make_const_mdspan(buffer.view()), static_cast(src.extent(1))}; + return cuvs::neighbors::detail::make_dense_view< + cuvs::neighbors::device_padded_dataset_view>( + buffer.data_handle(), + buffer.extent(0), + static_cast(src.extent(1)), + static_cast(buffer.extent(1))); } /** Grow `buffers` to `size` empty matrices (device matrices are not default-constructible). */ @@ -373,8 +378,9 @@ void cuvs_cagra::build(const T* dataset, size_t nrow) auto const stride = static_cast( cuvs::neighbors::cagra_required_row_width(static_cast(dim_))); *dataset_ = raft::make_device_matrix(handle_, merged_rows, stride); - auto merged_dataset_view = cuvs::neighbors::device_padded_dataset_view( - raft::make_const_mdspan(dataset_->view()), static_cast(dim_)); + auto merged_dataset_view = cuvs::neighbors::detail::make_dense_view< + cuvs::neighbors::device_padded_dataset_view>( + dataset_->data_handle(), merged_rows, static_cast(dim_), stride); index_ = std::make_shared(cuvs::neighbors::cagra::merge(handle_, params, @@ -416,8 +422,8 @@ void cuvs_cagra::compress_dataset(const T* dataset, size_t nrow) handle_, parse_metric_type(metric_), vpq_dataset_->as_dataset_view(), index_->graph()); // Search runs on the compressed rows and the graph, so release the dense copy of the dataset. - cuvs::neighbors::device_padded_dataset_view empty_dv( - raft::make_device_matrix_view(static_cast(nullptr), 0, this->dim_), this->dim_); + auto empty_dv = cuvs::neighbors::detail::make_empty_dataset_view< + cuvs::neighbors::device_padded_dataset_view>(this->dim_); *index_ = cuvs::neighbors::cagra::update_dataset(handle_, std::move(*index_), empty_dv); *dataset_ = raft::make_device_matrix(handle_, 0, 0); need_dataset_update_ = false; @@ -482,9 +488,9 @@ void cuvs_cagra::set_search_param(const search_param_base& param, dataset_mem_ = sp.dataset_mem; // First free up existing memory - *dataset_ = raft::make_device_matrix(handle_, 0, 0); - cuvs::neighbors::device_padded_dataset_view empty_dv( - raft::make_device_matrix_view(static_cast(nullptr), 0, this->dim_), this->dim_); + *dataset_ = raft::make_device_matrix(handle_, 0, 0); + auto empty_dv = cuvs::neighbors::detail::make_empty_dataset_view< + cuvs::neighbors::device_padded_dataset_view>(this->dim_); *index_ = cuvs::neighbors::cagra::update_dataset(handle_, std::move(*index_), empty_dv); // Allocate space using the correct memory resource. @@ -494,10 +500,9 @@ void cuvs_cagra::set_search_param(const search_param_base& param, auto mr = get_mr(dataset_mem_); cuvs::neighbors::cagra::detail::copy_with_padding(handle_, *dataset_, *input_dataset_v_, mr); - cuvs::neighbors::device_padded_dataset_view dv( - raft::make_device_matrix_view( - dataset_->data_handle(), dataset_->extent(0), dataset_->extent(1)), - this->dim_); + auto dv = cuvs::neighbors::detail::make_dense_view< + cuvs::neighbors::device_padded_dataset_view>( + dataset_->data_handle(), dataset_->extent(0), this->dim_, dataset_->extent(1)); *index_ = cuvs::neighbors::cagra::update_dataset(handle_, std::move(*index_), dv); need_dataset_update_ = false; diff --git a/cpp/bench/ann/src/cuvs/cuvs_mg_cagra_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_mg_cagra_wrapper.h index 7a89b68394..a01c047b59 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_mg_cagra_wrapper.h +++ b/cpp/bench/ann/src/cuvs/cuvs_mg_cagra_wrapper.h @@ -95,8 +95,9 @@ void cuvs_mg_cagra::build(const T* dataset, size_t nrow) raft::make_host_matrix_view(dataset, nrow, dim_); // The row alignment of the host view is irrelevant: every per-rank device shard is padded // individually during the multi-GPU build. - cuvs::neighbors::host_padded_dataset_view dataset_view(dataset_mds, - static_cast(dim_)); + auto dataset_view = + cuvs::neighbors::detail::make_dense_view>( + dataset_mds.data_handle(), nrow, static_cast(dim_), static_cast(dim_)); auto idx = cuvs::neighbors::cagra::build(clique_, build_params, dataset_view); index_ = std::make_shared< cuvs::neighbors::mg_index, T, IdxT>>( diff --git a/cpp/include/cuvs/core/dataset/dense.hpp b/cpp/include/cuvs/core/dataset/dense.hpp index e0e1775af3..69bafe1c10 100644 --- a/cpp/include/cuvs/core/dataset/dense.hpp +++ b/cpp/include/cuvs/core/dataset/dense.hpp @@ -67,14 +67,25 @@ struct mdarray_spec { { return static_cast(data.extent(1)); } + [[nodiscard]] static auto get_stride(const auto& data) noexcept -> uint32_t + { + return static_cast(data.stride(0)); + } }; }; +/** + * Row-major rows padded to 16 bytes, the alignment the CAGRA kernels require. A padded matrix keeps + * the logical dimension as its second extent; the row stride follows from the layout. + */ +template +using padded_layout = raft::layout_right_padded_impl<16 / sizeof(T)>; + template using contiguous_dataset = dataset>; template -using padded_dataset = dataset, ContainerPolicy>>; +using padded_dataset = dataset, ContainerPolicy>>; template using device_contiguous_dataset = contiguous_dataset>>; diff --git a/cpp/include/cuvs/core/dataset/type.hpp b/cpp/include/cuvs/core/dataset/type.hpp index 2ad59b43ba..c9d48f002d 100644 --- a/cpp/include/cuvs/core/dataset/type.hpp +++ b/cpp/include/cuvs/core/dataset/type.hpp @@ -82,12 +82,22 @@ struct dataset { } } + /* The former name of the read-only `view()`, kept so that the existing callers keep working. */ + [[nodiscard]] auto as_dataset_view() const noexcept -> const_view_type { return view(); } + [[nodiscard]] auto n_rows() const noexcept -> index_type { return spec_type::get_n_rows(data_); } [[nodiscard]] auto dim() const noexcept -> uint32_t { // NOTE: a compressed dataset may not be able to tell its dimension from the encoded data alone return spec_type::get_dim(data_, dictionary_); } + /* The distance between the starts of two rows, in elements; defined if the spec has strided rows. + */ + [[nodiscard]] auto stride() const noexcept -> uint32_t + requires requires(const data_type& d) { spec_type::get_stride(d); } + { + return spec_type::get_stride(data_); + } /* @return A non-owning view of the data, such as mdspan. @@ -104,7 +114,7 @@ struct dataset { /* @return A non-owning view of the dictionary, such as the two codebooks of a VPQ dataset. */ - [[nodiscard]] auto dictionary_view() const noexcept -> dictionary_view_type + [[nodiscard]] auto dictionary_view() noexcept -> dictionary_view_type requires(is_compressed) { return spec_type::get_dictionary_view(dictionary_); diff --git a/cpp/include/cuvs/core/dataset/view.hpp b/cpp/include/cuvs/core/dataset/view.hpp index b2c7c4629f..89bd5cd614 100644 --- a/cpp/include/cuvs/core/dataset/view.hpp +++ b/cpp/include/cuvs/core/dataset/view.hpp @@ -95,6 +95,9 @@ struct dataset_view { static constexpr bool is_compressed = compressed_dataset_spec; + /* An empty view, like a default-constructed mdspan: no rows, no data. */ + dataset_view() noexcept = default; + explicit dataset_view(data_view_type data_view) requires(!is_compressed) : data_view_{data_view} @@ -133,6 +136,13 @@ struct dataset_view { { return spec_type::get_dim(data_view_, dictionary_view_); } + /* The distance between the starts of two rows, in elements; defined if the spec has strided rows. + */ + [[nodiscard]] auto stride() const noexcept -> uint32_t + requires requires(const data_view_type& d) { spec_type::get_stride(d); } + { + return spec_type::get_stride(data_view_); + } /* @return A non-owning view of the data, such as mdspan. diff --git a/cpp/include/cuvs/core/dataset/vpq.hpp b/cpp/include/cuvs/core/dataset/vpq.hpp index bc47b9fc23..97539b03fc 100644 --- a/cpp/include/cuvs/core/dataset/vpq.hpp +++ b/cpp/include/cuvs/core/dataset/vpq.hpp @@ -41,48 +41,28 @@ struct vpq_spec { /* The type the codebooks are stored in; the rows themselves are always packed codes. */ using math_type = MathT; - /* [vq_n_centers, dim] */ - using vq_book_type = - raft::mdarray, raft::layout_c_contiguous, BookPolicy>; - /* [n_books, pq_n_centers, pq_len], where `n_books` is 1 for a codebook shared by all the - subspaces, or `pq_dim` for one codebook per subspace. A row-major 3d array with `n_books == 1` - is bit-identical to the flat 2d codebook of the current implementation. - */ - using pq_book_type = - raft::mdarray, raft::layout_c_contiguous, BookPolicy>; + /* [vq_n_centers, dim] for the VQ codebook, [pq_n_centers, pq_len] for the PQ one. */ + using book_type = + raft::mdarray, raft::layout_c_contiguous, BookPolicy>; struct dictionary_type { - vq_book_type vq_code_book; - pq_book_type pq_code_book; + book_type vq_code_book; + book_type pq_code_book; }; - struct dictionary_view_type { - typename vq_book_type::const_view_type vq_code_book; - typename pq_book_type::const_view_type pq_code_book; + /* The views of the codebooks, writable or read-only alike. */ + template + struct dictionary_view_of { + BookViewT vq_code_book; + BookViewT pq_code_book; /* NOTE: the dimension of the dataset is a property of the VQ codebook: the encoded rows tell only how many bytes it took to compress a vector, which the padding makes ambiguous. */ - [[nodiscard]] auto dim() const noexcept -> uint32_t - { - return static_cast(vq_code_book.extent(1)); - } - [[nodiscard]] auto vq_n_centers() const noexcept -> uint32_t - { - return static_cast(vq_code_book.extent(0)); - } - [[nodiscard]] auto n_books() const noexcept -> uint32_t - { - return static_cast(pq_code_book.extent(0)); - } - [[nodiscard]] auto pq_n_centers() const noexcept -> uint32_t - { - return static_cast(pq_code_book.extent(1)); - } - [[nodiscard]] auto pq_len() const noexcept -> uint32_t - { - return static_cast(pq_code_book.extent(2)); - } + [[nodiscard]] auto dim() const noexcept -> uint32_t { return vq_code_book.extent(1); } + [[nodiscard]] auto vq_n_centers() const noexcept -> uint32_t { return vq_code_book.extent(0); } + [[nodiscard]] auto pq_n_centers() const noexcept -> uint32_t { return pq_code_book.extent(0); } + [[nodiscard]] auto pq_len() const noexcept -> uint32_t { return pq_code_book.extent(1); } [[nodiscard]] auto pq_bits() const noexcept -> uint32_t { return static_cast(std::countr_zero(pq_n_centers())); @@ -91,9 +71,9 @@ struct vpq_spec { { return raft::div_rounding_up_safe(dim(), pq_len()); } - /* Whether every subspace has a codebook of its own, as opposed to sharing a single one. */ - [[nodiscard]] auto per_subspace() const noexcept -> bool { return n_books() > 1; } }; + using dictionary_view_type = dictionary_view_of; + using const_dictionary_view_type = dictionary_view_of; template struct apply : storage_spec { @@ -103,11 +83,11 @@ struct vpq_spec { /* NOTE: re-exported rather than declared here, so that they do not belong to one instantiation of `apply`; qualified lookup into `apply` does not reach the enclosing spec on its own. */ - using math_type = vpq_spec::math_type; - using vq_book_type = vpq_spec::vq_book_type; - using pq_book_type = vpq_spec::pq_book_type; - using dictionary_type = vpq_spec::dictionary_type; - using dictionary_view_type = vpq_spec::dictionary_view_type; + using math_type = vpq_spec::math_type; + using book_type = vpq_spec::book_type; + using dictionary_type = vpq_spec::dictionary_type; + using dictionary_view_type = vpq_spec::dictionary_view_type; + using const_dictionary_view_type = vpq_spec::const_dictionary_view_type; /* NOTE: the data accessors are inherited; only `get_dim` differs from a dense dataset. Like the inherited ones, it reads the dictionary in either form: the codebooks an owning dataset holds, @@ -117,23 +97,21 @@ struct vpq_spec { { return static_cast(dict.vq_code_book.extent(1)); } - [[nodiscard]] static auto get_dictionary_view(const dictionary_type& dict) noexcept + [[nodiscard]] static auto get_dictionary_view(dictionary_type& dict) noexcept -> dictionary_view_type { return {dict.vq_code_book.view(), dict.pq_code_book.view()}; } - /* NOTE: a codebook view is read-only to begin with, so the spec declares no separate const view - and these overloads exist only to spell the conversion out of either form. - */ [[nodiscard]] static auto get_dictionary_const_view(const dictionary_type& dict) noexcept - -> dictionary_view_type + -> const_dictionary_view_type { - return get_dictionary_view(dict); + return {dict.vq_code_book.view(), dict.pq_code_book.view()}; } [[nodiscard]] static auto get_dictionary_const_view(dictionary_view_type dict) noexcept - -> dictionary_view_type + -> const_dictionary_view_type { - return dict; + return {raft::make_const_mdspan(dict.vq_code_book), + raft::make_const_mdspan(dict.pq_code_book)}; } /* The length of an encoded row in bytes, including the inlined VQ label. */ [[nodiscard]] static auto get_encoded_row_length(const auto& data) noexcept -> uint32_t diff --git a/cpp/include/cuvs/neighbors/cagra.hpp b/cpp/include/cuvs/neighbors/cagra.hpp index 418db9d88f..3f22a4ac3e 100644 --- a/cpp/include/cuvs/neighbors/cagra.hpp +++ b/cpp/include/cuvs/neighbors/cagra.hpp @@ -592,7 +592,7 @@ struct CUVS_EXPORT index : cuvs::neighbors::index { : cuvs::neighbors::index(), metric_(metric), graph_(raft::make_device_matrix(res, 0, 0)), - dataset_{}, + dataset_{cuvs::neighbors::detail::make_empty_dataset_view()}, dataset_norms_(std::nullopt) { } @@ -793,24 +793,9 @@ struct CUVS_EXPORT index : cuvs::neighbors::index { // Re-open the file descriptor in read-only mode for subsequent operations dataset_fd_.emplace(std::move(fd)); - if constexpr (cuvs::neighbors::is_device_padded_dataset_view_v) { - auto v = raft::make_device_matrix_view( - static_cast(nullptr), int64_t{0}, dim_); - dataset_ = DatasetViewT(v, dim_); - } else if constexpr (cuvs::neighbors::is_device_standard_dataset_view_v) { - auto v = raft::make_device_matrix_view( - static_cast(nullptr), int64_t{0}, dim_); - dataset_ = DatasetViewT(v); - } else if constexpr (cuvs::neighbors::is_host_padded_dataset_view_v) { - auto v = raft::make_host_matrix_view( - static_cast(nullptr), int64_t{0}, dim_); - dataset_ = DatasetViewT(v, dim_); - } else if constexpr (cuvs::neighbors::is_host_standard_dataset_view_v) { - auto v = raft::make_host_matrix_view( - static_cast(nullptr), int64_t{0}, dim_); - dataset_ = DatasetViewT(v); - } else if constexpr (cuvs::neighbors::is_empty_dataset_view_v) { - dataset_ = DatasetViewT{dim_}; + if constexpr (cuvs::neighbors::is_dense_row_major_dataset_view_v || + cuvs::neighbors::is_empty_dataset_view_v) { + dataset_ = cuvs::neighbors::detail::make_empty_dataset_view(dim_); } else { RAFT_FAIL("update_dataset(fd): unsupported DatasetViewT for disk-backed dataset"); } diff --git a/cpp/include/cuvs/neighbors/common.hpp b/cpp/include/cuvs/neighbors/common.hpp index baa5b44da2..cced5911b8 100644 --- a/cpp/include/cuvs/neighbors/common.hpp +++ b/cpp/include/cuvs/neighbors/common.hpp @@ -24,6 +24,8 @@ #include #include +#include +#include #include #include @@ -192,27 +194,6 @@ using dataset_owning_accessor_for_view = std::conditional_t, host_owning_accessor>; -template -using dense_owning_matrix = std::conditional_t, - raft::host_matrix>; - -template -using dense_view_matrix = - std::conditional_t, - raft::host_matrix_view>; - -template -using vpq_vq_book_matrix = std::conditional_t, - raft::host_matrix>; - -template -using vpq_data_matrix = std::conditional_t, - raft::host_matrix>; - // ----------------------------------------------------------------------------- // empty // ----------------------------------------------------------------------------- @@ -232,186 +213,6 @@ using empty_dataset_owning_storage = empty_dataset_storage; template using empty_dataset_view_storage = empty_dataset_storage; -// ----------------------------------------------------------------------------- -// dense row-major (logical dim may differ from row pitch; shared by padded & standard) -// ----------------------------------------------------------------------------- - -/** - * Dense row-major owning storage shared by padded and standard dataset containers. - * - * Template parameters: - * - MatrixT: owning matrix type that stores the payload (host/device matrix). - * - ViewT: non-owning row-major view type returned by `view()`. - * - DataT: scalar element type of the dataset payload. - * - IdxT: index type used for row counts (`n_rows()` return type). - */ -template -struct dense_row_major_dataset_owning_storage { - MatrixT data_; - uint32_t logical_dim_; - - dense_row_major_dataset_owning_storage(MatrixT&& data, uint32_t logical_dim) noexcept - : data_{std::move(data)}, logical_dim_{logical_dim} - { - } - - [[nodiscard]] auto n_rows() const noexcept -> IdxT { return data_.extent(0); } - [[nodiscard]] auto dim() const noexcept -> uint32_t { return logical_dim_; } - [[nodiscard]] auto stride() const noexcept -> uint32_t - { - return static_cast(data_.extent(1)); - } - [[nodiscard]] auto view() const noexcept -> ViewT { return data_.view(); } - [[nodiscard]] auto data_handle() noexcept -> DataT* { return data_.data_handle(); } - [[nodiscard]] auto data_handle() const noexcept -> const DataT* { return data_.data_handle(); } -}; - -template -struct dense_row_major_dataset_view_storage { - ViewT data_; - uint32_t logical_dim_; - - dense_row_major_dataset_view_storage() noexcept = default; - - explicit dense_row_major_dataset_view_storage(ViewT v) noexcept - : data_(v), logical_dim_(static_cast(v.extent(1))) - { - } - - dense_row_major_dataset_view_storage(ViewT v, uint32_t logical_dim) noexcept - : data_(v), logical_dim_(logical_dim) - { - } - - dense_row_major_dataset_view_storage(dense_row_major_dataset_view_storage const& other) noexcept - : data_(other.data_), logical_dim_(other.logical_dim_) - { - } - - [[nodiscard]] auto n_rows() const noexcept -> IdxT { return data_.extent(0); } - [[nodiscard]] auto dim() const noexcept -> uint32_t { return logical_dim_; } - [[nodiscard]] auto stride() const noexcept -> uint32_t - { - return static_cast(data_.stride(0) > 0 ? data_.stride(0) : data_.extent(1)); - } - [[nodiscard]] auto view() const noexcept -> ViewT { return data_; } -}; - -template -using padded_dataset_owning_storage = - dense_row_major_dataset_owning_storage; - -template -using padded_dataset_view_storage = dense_row_major_dataset_view_storage; - -template -using standard_dataset_owning_storage = - dense_row_major_dataset_owning_storage; - -template -using standard_dataset_view_storage = dense_row_major_dataset_view_storage; - -// ----------------------------------------------------------------------------- -// VPQ compressed -// ----------------------------------------------------------------------------- - -/** - * Owning storage for VPQ-compressed datasets. - * - * Template parameters: - * - VqBookMatrixT: owning matrix type for the VQ codebook. - * - PqBookMatrixT: owning matrix type for the PQ codebook. - * - DataMatrixT: owning matrix type for encoded row data (uint8 codes). - * - MathT: floating-point type used by VQ/PQ codebooks. - * - IdxT: index type used for row counts (`n_rows()` return type). - */ -template -struct vpq_dataset_owning_storage { - /** Floating-point type used for VQ/PQ codebooks (rows are still uint8 codes). */ - using math_type = MathT; - - VqBookMatrixT vq_code_book; - PqBookMatrixT pq_code_book; - DataMatrixT data; - - vpq_dataset_owning_storage(VqBookMatrixT&& vq_code_book, - PqBookMatrixT&& pq_code_book, - DataMatrixT&& data) noexcept - : vq_code_book{std::move(vq_code_book)}, - pq_code_book{std::move(pq_code_book)}, - data{std::move(data)} - { - } - - [[nodiscard]] auto n_rows() const noexcept -> IdxT { return data.extent(0); } - [[nodiscard]] auto dim() const noexcept -> uint32_t { return vq_code_book.extent(1); } - - [[nodiscard]] constexpr inline auto encoded_row_length() const noexcept -> uint32_t - { - return data.extent(1); - } - [[nodiscard]] constexpr inline auto vq_n_centers() const noexcept -> uint32_t - { - return vq_code_book.extent(0); - } - [[nodiscard]] constexpr inline auto pq_bits() const noexcept -> uint32_t - { - auto pq_width = pq_n_centers(); -#ifdef __cpp_lib_bitops - return std::countr_zero(pq_width); -#else - uint32_t pq_bits = 0; - while (pq_width > 1) { - pq_bits++; - pq_width >>= 1; - } - return pq_bits; -#endif - } - [[nodiscard]] constexpr inline auto pq_dim() const noexcept -> uint32_t - { - return raft::div_rounding_up_unsafe(dim(), pq_len()); - } - [[nodiscard]] constexpr inline auto pq_len() const noexcept -> uint32_t - { - return pq_code_book.extent(1); - } - [[nodiscard]] constexpr inline auto pq_n_centers() const noexcept -> uint32_t - { - return pq_code_book.extent(0); - } -}; - -template -struct vpq_dataset_view_storage { - using owning_dataset_type = - dataset>; - - owning_dataset_type const* dataset_{nullptr}; - - vpq_dataset_view_storage() = default; - - explicit vpq_dataset_view_storage(owning_dataset_type const* ptr) : dataset_(ptr) - { - RAFT_EXPECTS(ptr != nullptr, "vpq_dataset_view: null dataset pointer"); - } - - [[nodiscard]] auto n_rows() const noexcept - { - using idx_type = decltype(std::declval().n_rows()); - return dataset_ != nullptr ? dataset_->n_rows() : idx_type{0}; - } - [[nodiscard]] auto dim() const noexcept -> uint32_t - { - return dataset_ != nullptr ? dataset_->dim() : uint32_t{0}; - } - [[nodiscard]] owning_dataset_type const& dset() const noexcept { return *dataset_; } -}; - } // namespace detail // ----------------------------------------------------------------------------- @@ -425,55 +226,6 @@ struct empty_dataset_container { using view_storage = detail::empty_dataset_view_storage; }; -// ----------------------------------------------------------------------------- -// padded (row-major with logical dim vs stride) -// ----------------------------------------------------------------------------- - -struct padded_dataset_container { - template - using owning_storage = - detail::padded_dataset_owning_storage, - detail::dense_view_matrix, - DataT, - IdxT>; - template - using view_storage = detail:: - padded_dataset_view_storage, DataT, IdxT>; -}; - -// ----------------------------------------------------------------------------- -// standard (row-major with arbitrary stride; no CAGRA alignment requirement) -// ----------------------------------------------------------------------------- - -struct standard_dataset_container { - template - using owning_storage = - detail::standard_dataset_owning_storage, - detail::dense_view_matrix, - DataT, - IdxT>; - template - using view_storage = detail:: - standard_dataset_view_storage, DataT, IdxT>; -}; - -// ----------------------------------------------------------------------------- -// VPQ compressed -// ----------------------------------------------------------------------------- - -struct vpq_dataset_container { - template - using owning_storage = - detail::vpq_dataset_owning_storage, - detail::vpq_vq_book_matrix, - detail::vpq_data_matrix, - MathT, - IdxT>; - template - using view_storage = - detail::vpq_dataset_view_storage; -}; - template struct dataset { static_assert(!std::is_same_v, @@ -518,107 +270,6 @@ struct dataset_view using view_storage_type::view_storage_type; }; -// ----------------------------------------------------------------------------- -// standard (row-major with arbitrary stride) -// ----------------------------------------------------------------------------- - -template -struct dataset - : standard_dataset_container::template owning_storage { - using container_type = standard_dataset_container; - using owning_storage_type = - typename container_type::template owning_storage; - using owning_storage_type::owning_storage_type; - - [[nodiscard]] auto as_dataset_view() const noexcept - -> dataset_view> - { - return dataset_view>(this->view(), - this->dim()); - } -}; - -template -struct dataset_view - : standard_dataset_container::template view_storage { - using container_type = standard_dataset_container; - using view_storage_type = typename container_type::template view_storage; - using view_storage_type::view_storage_type; -}; - -// ----------------------------------------------------------------------------- -// padded (row-major with logical dim vs stride) -// ----------------------------------------------------------------------------- - -template -struct dataset - : padded_dataset_container::template owning_storage { - using container_type = padded_dataset_container; - using owning_storage_type = - typename container_type::template owning_storage; - using owning_storage_type::owning_storage_type; - - [[nodiscard]] auto as_dataset_view() const noexcept - -> dataset_view> - { - return dataset_view>(this->view(), - this->dim()); - } -}; - -template -struct dataset_view - : padded_dataset_container::template view_storage { - using container_type = padded_dataset_container; - using view_storage_type = typename container_type::template view_storage; - using view_storage_type::view_storage_type; -}; - -// ----------------------------------------------------------------------------- -// VPQ compressed (view holds non-owning pointer to owning dataset) -// ----------------------------------------------------------------------------- - -template -struct dataset - : vpq_dataset_container::template owning_storage { - using container_type = vpq_dataset_container; - using owning_storage_type = - typename container_type::template owning_storage; - using owning_storage_type::owning_storage_type; - - [[nodiscard]] auto as_dataset_view() const - -> dataset_view> - { - return dataset_view>{this}; - } -}; - -template -struct dataset_view - : vpq_dataset_container::template view_storage { - using container_type = vpq_dataset_container; - using view_storage_type = typename container_type::template view_storage; - using view_storage_type::view_storage_type; -}; - /** * @brief Aliases for concrete `dataset` / `dataset_view` layouts. */ @@ -638,53 +289,47 @@ template using host_empty_dataset_view = dataset_view>; -template -using device_padded_dataset = - dataset>; - -template -using device_padded_dataset_view = - dataset_view>; - -template -using host_padded_dataset = - dataset>; - -template -using host_padded_dataset_view = - dataset_view>; - -template -using device_standard_dataset = - dataset>; - -template -using device_standard_dataset_view = - dataset_view>; - -template -using host_standard_dataset = - dataset>; - -template -using host_standard_dataset_view = - dataset_view>; - -template -using device_vpq_dataset = - dataset>; - -template -using device_vpq_dataset_view = - dataset_view>; - -template -using host_vpq_dataset = - dataset>; +// ----------------------------------------------------------------------------- +// Compatibility spellings of the dataset types, so that the code written against the former +// `cuvs::neighbors` dataset API keeps compiling on top of `cuvs::dataset` and `cuvs::dataset_view` +// (the headers under cuvs/core/dataset). It is all names, with the `standard` dataset as the only +// new kind. Once the callers spell the types `cuvs::...` directly, this section can be deleted. +// ----------------------------------------------------------------------------- -template -using host_vpq_dataset_view = - dataset_view>; +using cuvs::device_padded_dataset; +using cuvs::host_padded_dataset; + +/* The `standard` dataset: row-major rows of any stride (at least the dimension), no alignment. */ +template +using device_standard_dataset = cuvs::dataset< + T, + IdxT, + cuvs::mdarray_spec>>>; +template +using host_standard_dataset = cuvs::dataset< + T, + IdxT, + cuvs::mdarray_spec>>>; + +/* The former spelling carried the math type as an optional third argument. */ +template +using device_vpq_dataset = cuvs::device_vpq_dataset; +template +using host_vpq_dataset = cuvs::host_vpq_dataset; + +/* The read-only views are the const views of the owning types. */ +template +using device_padded_dataset_view = typename device_padded_dataset::const_view_type; +template +using host_padded_dataset_view = typename host_padded_dataset::const_view_type; +template +using device_standard_dataset_view = typename device_standard_dataset::const_view_type; +template +using host_standard_dataset_view = typename host_standard_dataset::const_view_type; +template +using device_vpq_dataset_view = typename device_vpq_dataset::const_view_type; +template +using host_vpq_dataset_view = typename host_vpq_dataset::const_view_type; // Maps a dataset view type to its owning (allocating) dataset counterpart. // Used by serialize/deserialize to type the out_dataset output parameter; @@ -692,71 +337,14 @@ using host_vpq_dataset_view = template struct owning_dataset_for_view; -template -struct owning_dataset_for_view> { - using type = device_padded_dataset; -}; - -template -struct owning_dataset_for_view> { - using type = device_standard_dataset; -}; - -template -struct owning_dataset_for_view> { - using type = host_padded_dataset; -}; - -template -struct owning_dataset_for_view> { - using type = host_standard_dataset; -}; - -template -struct owning_dataset_for_view> { - using type = device_vpq_dataset; +template +struct owning_dataset_for_view> { + using type = cuvs::dataset; }; template using owning_dataset_for_view_t = typename owning_dataset_for_view::type; -template -struct is_padded_dataset : std::false_type {}; - -template -struct is_padded_dataset> - : std::true_type {}; - -template -struct is_padded_dataset> - : std::true_type {}; - -template -inline constexpr bool is_padded_dataset_v = is_padded_dataset::value; - -template -struct is_standard_dataset : std::false_type {}; - -template -struct is_standard_dataset> - : std::true_type {}; - -template -struct is_standard_dataset> - : std::true_type {}; - -template -inline constexpr bool is_standard_dataset_v = is_standard_dataset::value; - -template -struct is_vpq_dataset : std::false_type {}; - -template -struct is_vpq_dataset> : std::true_type {}; - -template -inline constexpr bool is_vpq_dataset_v = is_vpq_dataset::value; - // ----------------------------------------------------------------------------- // Dataset view compile-time classification (replaces runtime std::variant dispatch). // ----------------------------------------------------------------------------- @@ -794,18 +382,32 @@ struct dataset_view_kind_of -struct dataset_view_kind_of> { +/** The migrated kinds are told apart by their spec, for a dataset as well as for its view. */ +template