From 2f07757c0ea61afa3c6813c9c2bc095706d4bf6b Mon Sep 17 00:00:00 2001 From: Howard Huang Date: Sun, 30 Aug 2026 16:50:28 -0700 Subject: [PATCH 01/15] Reuse raft::mdarray/mdspan directly in dense dataset storage (#2395) dense_owning_matrix/dense_view_matrix (and the VPQ codebook/code matrix aliases) picked between raft::device_matrix and raft::host_matrix via std::conditional_t, even though those are themselves just aliases for raft::mdarray/raft::mdspan with the exact accessor already computed as Accessor. Point the aliases at raft::mdarray/raft::mdspan directly instead. dense_row_major_dataset_owning_storage/_view_storage also wrapped their matrix/view as a field and hand-forwarded view()/data_handle(), which raft::mdarray/raft::mdspan already provide natively. They now inherit from the matrix/view type instead, so those forwards are no longer needed. Co-Authored-By: Claude Sonnet 5 --- cpp/include/cuvs/neighbors/common.hpp | 83 ++++++++++++++++----------- 1 file changed, 49 insertions(+), 34 deletions(-) diff --git a/cpp/include/cuvs/neighbors/common.hpp b/cpp/include/cuvs/neighbors/common.hpp index 3be6cf1fa1..46b9e7b599 100644 --- a/cpp/include/cuvs/neighbors/common.hpp +++ b/cpp/include/cuvs/neighbors/common.hpp @@ -189,26 +189,36 @@ using dataset_owning_accessor_for_view = std::conditional_t, host_owning_accessor>; +// Accessor here is already device_owning_accessor / host_owning_accessor at every +// call site -- exactly the container policy raft::device_mdarray/host_mdarray default to for +// element type DataT -- so pass it straight through instead of re-deriving a +// raft::device_matrix/host_matrix from scratch. template -using dense_owning_matrix = std::conditional_t, - raft::host_matrix>; +using dense_owning_matrix = + raft::mdarray, raft::row_major, Accessor>; template -using dense_view_matrix = - std::conditional_t, - raft::host_matrix_view>; +using dense_view_matrix = raft::mdspan, + raft::row_major, + dataset_view_accessor_for_owning>; template -using vpq_vq_book_matrix = std::conditional_t, - raft::host_matrix>; +using vpq_vq_book_matrix = + raft::mdarray, raft::row_major, Accessor>; + +// VPQ codes are always uint8_t regardless of MathT, so retarget the owning accessor's element +// type instead of re-deriving a device/host matrix; residency is still driven by Accessor. +template +using owning_accessor_with_value_type = std::conditional_t, + host_owning_accessor>; template -using vpq_data_matrix = std::conditional_t, - raft::host_matrix>; +using vpq_data_matrix = raft::mdarray, + raft::row_major, + owning_accessor_with_value_type>; // ----------------------------------------------------------------------------- // empty @@ -234,7 +244,9 @@ using empty_dataset_view_storage = empty_dataset_storage; // ----------------------------------------------------------------------------- /** - * Dense row-major owning storage shared by padded and standard dataset containers. + * Dense row-major owning storage shared by padded and standard dataset containers. Publicly + * inherits from MatrixT (an `raft::mdarray`) so `view()`/`data_handle()`/`extent()` etc. are + * reused as-is rather than hand-forwarded; `logical_dim_` is the only state this struct adds. * * Template parameters: * - MatrixT: owning matrix type that stores the payload (host/device matrix). @@ -243,55 +255,58 @@ using empty_dataset_view_storage = empty_dataset_storage; * - IdxT: index type used for row counts (`n_rows()` return type). */ template -struct dense_row_major_dataset_owning_storage { - MatrixT data_; +struct dense_row_major_dataset_owning_storage : public MatrixT { uint32_t logical_dim_; + // MatrixT (mdarray) also has its own stride(size_t); pull it back into scope since declaring + // our own no-arg stride() below would otherwise hide it entirely (C++ name hiding). + using MatrixT::stride; + dense_row_major_dataset_owning_storage(MatrixT&& data, uint32_t logical_dim) noexcept - : data_{std::move(data)}, logical_dim_{logical_dim} + : MatrixT{std::move(data)}, logical_dim_{logical_dim} { } - [[nodiscard]] auto n_rows() const noexcept -> IdxT { return data_.extent(0); } + [[nodiscard]] auto n_rows() const noexcept -> IdxT { return this->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)); + return static_cast(this->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(); } + // view() and data_handle() are inherited directly from MatrixT (raft::mdarray); no hand-written + // forwarding needed since MatrixT::view() const already returns exactly ViewT. }; template -struct dense_row_major_dataset_view_storage { - ViewT data_; +struct dense_row_major_dataset_view_storage : public ViewT { uint32_t logical_dim_; + // ViewT (mdspan) also has its own stride(size_t); pull it back into scope since declaring our + // own no-arg stride() below would otherwise hide it entirely (C++ name hiding), and the body of + // that stride() itself needs to call the inherited one. + using ViewT::stride; + 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))) + : ViewT(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_) + : ViewT(v), logical_dim_(logical_dim) { } - [[nodiscard]] auto n_rows() const noexcept -> IdxT { return data_.extent(0); } + [[nodiscard]] auto n_rows() const noexcept -> IdxT { return this->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)); + return static_cast(ViewT::stride(0) > 0 ? ViewT::stride(0) : this->extent(1)); } - [[nodiscard]] auto view() const noexcept -> ViewT { return data_; } + // ViewT (mdspan) has no view() of its own -- it already *is* the view -- so this shrinks to a + // plain upcast instead of reaching into a wrapped field. + [[nodiscard]] auto view() const noexcept -> ViewT { return *this; } }; template From ba608e5f8e1a01618296f9d71cd743b5df77ea4d Mon Sep 17 00:00:00 2001 From: Howard Huang Date: Wed, 2 Sep 2026 08:29:43 -0700 Subject: [PATCH 02/15] Add experimental Spec-based dataset/dataset_view prototype (#2395) as_dataset_view()/n_rows()/dim() can be shared across all dataset kinds with zero per-kind dispatch inside dataset/dataset_view themselves. Adds cuvs::neighbors::experimental::{dataset,dataset_view} as a Spec/policy-based prototype: every member is a one-line forward to spec_type::get_*(), with all kind-specific logic (empty_spec, mdarray_spec, vpq_spec) living outside dataset/dataset_view, which never name or branch on a concrete kind. dataset and dataset_view stay two independent, non-inheriting types (no shared_ptr, no "sometimes owning" object). Purely additive: lives in its own namespace, not referenced by any existing type, alias, or call site. Verified via a standalone sandbox (all three specs, both dataset and dataset_view, asserting n_rows/dim/ as_dataset_view) and a full rebuild + DATASET_C_TEST (7/7) + CAGRA_C_TEST (14/14) + PREPROCESSING_TEST (226/226) + NEIGHBORS_TEST (339/339), all passing unchanged from before this commit. Limitations / not yet done: - Not wired up: the real dataset/dataset_view types (Container-tagged) and every downstream consumer (CAGRA build/search, serialization, the C API, compile-time classification traits like dataset_view_kind_of/is_padded_dataset_view_v) are untouched. This prototype does not replace them yet. - mdarray_spec's layout choice for the padded case is unverified against CAGRA's actual alignment requirements -- needs confirmation before real use. - No sparse (CSR/COO) or scalar-quantized specs; only empty/dense/vpq, matching today's four kinds minus the empty/dense split. - This design is still an open thread with the team, not finalized; land as prototype only, pending further review. Co-Authored-By: Claude Sonnet 5 --- cpp/include/cuvs/neighbors/common.hpp | 237 ++++++++++++++++++++++++++ 1 file changed, 237 insertions(+) diff --git a/cpp/include/cuvs/neighbors/common.hpp b/cpp/include/cuvs/neighbors/common.hpp index e8654e6065..8f948b9bdc 100644 --- a/cpp/include/cuvs/neighbors/common.hpp +++ b/cpp/include/cuvs/neighbors/common.hpp @@ -36,6 +36,7 @@ #include #include #include +#include #ifdef __cpp_lib_bitops #include #endif @@ -1354,6 +1355,242 @@ auto make_host_standard_dataset_view(SrcT const& src) host_standard_dataset_view>(src, static_cast(src.extent(1))); } +// ===================================================================================== +// Experimental: Spec-based dataset/dataset_view prototype (#2395 follow-up). +// +// Not wired up to any public alias, trait, or downstream call site yet -- exists to validate the +// design in isolation. `dataset` and `dataset_view` are single generic +// templates with zero per-kind dispatch inside them: every member is a one-line forward to +// `spec_type::get_*(...)`, and all kind-specific logic lives in the per-kind Spec structs below +// (`empty_spec`, `mdarray_spec`, `vpq_spec`), which dataset/dataset_view never name or branch on. +// `dataset` and `dataset_view` are deliberately two independent, non-inheriting types (no +// shared_ptr, no "sometimes owning" object): `dataset` holds owning storage (mdarray-shaped), +// `dataset_view` holds the corresponding view storage (mdspan-shaped); the same +// `get_n_rows`/`get_dim` spec functions serve both, since `raft::mdarray`/`raft::mdspan` both +// expose `.extent(r)`. +// ===================================================================================== +namespace experimental { + +/** + * A spec defines a dictionary iff it needs a second storage slot to interpret the data (e.g. PQ + * codebooks). Non-compressed specs declare `dictionary_type = std::monostate` -- the same + * vocabulary type for "no dictionary," not just an omitted member -- so `dataset`/`dataset_view` + * never need to branch on whether the slot exists; they just always have one, sometimes empty. + */ +template +concept compressed_dataset_spec = requires { + typename SpecT::dictionary_type; + typename SpecT::dictionary_view_type; +} && !std::is_same_v; + +// ----------------------------------------------------------------------------- +// empty +// ----------------------------------------------------------------------------- + +struct empty_spec { + struct rep { + uint32_t dim; + }; + + template + struct apply { + using value_type = std::remove_cv_t; + using index_type = std::remove_cv_t; + using data_type = rep; + using view_type = rep; + using dictionary_type = std::monostate; + using dictionary_view_type = std::monostate; + + [[nodiscard]] static auto get_data_view(data_type const& data) noexcept -> view_type + { + return data; + } + [[nodiscard]] static auto get_n_rows(rep const&) noexcept -> index_type { return 0; } + [[nodiscard]] static auto get_dim(rep const& data, dictionary_type const&) noexcept -> uint32_t + { + return data.dim; + } + [[nodiscard]] static auto get_dictionary_view(dictionary_type const&) noexcept + -> dictionary_view_type + { + return {}; + } + }; +}; + +// ----------------------------------------------------------------------------- +// dense (plain or padded), implemented via raft::mdarray +// ----------------------------------------------------------------------------- + +template +struct mdarray_spec { + template + struct apply { + using value_type = std::remove_cv_t; + using index_type = std::remove_cv_t; + using data_type = raft::mdarray, LayoutPolicy, ContainerPolicy>; + // `get_data_view` takes `data_type const&`, so `.view()` resolves to the const overload, + // returning `const_view_type` (const element type) -- match that here, not the mutable + // `view_type`. + using view_type = typename data_type::const_view_type; + using dictionary_type = std::monostate; + using dictionary_view_type = std::monostate; + + [[nodiscard]] static auto get_data_view(data_type const& data) noexcept -> view_type + { + return data.view(); + } + template + [[nodiscard]] static auto get_n_rows(AnyExtentShaped const& data) noexcept -> index_type + { + return static_cast(data.extent(0)); + } + template + [[nodiscard]] static auto get_dim(AnyExtentShaped const& data, dictionary_type const&) noexcept + -> uint32_t + { + return static_cast(data.extent(1)); + } + [[nodiscard]] static auto get_dictionary_view(dictionary_type const&) noexcept + -> dictionary_view_type + { + return {}; + } + }; +}; + +// ----------------------------------------------------------------------------- +// VPQ compressed: data = encoded rows (uint8_t codes); dictionary = {vq_code_book, pq_code_book} +// ----------------------------------------------------------------------------- + +template +struct vpq_spec { + template + using storage_spec = + typename mdarray_spec::template apply; + + template + struct apply : storage_spec { + using value_type = std::remove_cv_t; + /* Members of a dependent base aren't visible to unqualified lookup, so pull these in. */ + using typename storage_spec::data_type; + using typename storage_spec::view_type; + using math_type = MathT; + + using vq_book_type = + raft::mdarray, raft::row_major, BookPolicy>; + using pq_book_type = + raft::mdarray, raft::row_major, 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; + }; + + /* get_data_view/get_n_rows are inherited from storage_spec unchanged; only get_dim and + get_dictionary_view differ from a plain dense dataset, since the dimension comes from the VQ + codebook, not the encoded rows. */ + template + [[nodiscard]] static auto get_dim(AnyData const&, AnyDict const& dict) noexcept -> uint32_t + { + return static_cast(dict.vq_code_book.extent(1)); + } + [[nodiscard]] static auto get_dictionary_view(dictionary_type const& dict) noexcept + -> dictionary_view_type + { + return {dict.vq_code_book.view(), dict.pq_code_book.view()}; + } + }; +}; + +// ----------------------------------------------------------------------------- +// dataset / dataset_view +// ----------------------------------------------------------------------------- + +template +struct dataset_view; + +/** Owning dataset: value-held storage (no shared_ptr -- exclusive ownership, like today's + * `dataset`). Every member is a one-line forward to `spec_type::get_*`; all + * per-kind logic lives in `SpecT`, never inside this struct. */ +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 dictionary_type = typename spec_type::dictionary_type; + + explicit dataset(data_type&& data, dictionary_type&& dictionary = dictionary_type{}) + : data_{std::move(data)}, dictionary_{std::move(dictionary)} + { + } + + [[nodiscard]] auto n_rows() const noexcept -> index_type { return spec_type::get_n_rows(data_); } + [[nodiscard]] auto dim() const noexcept -> uint32_t + { + return spec_type::get_dim(data_, dictionary_); + } + [[nodiscard]] auto data_view() const noexcept { return spec_type::get_data_view(data_); } + [[nodiscard]] auto dictionary_view() const noexcept + { + return spec_type::get_dictionary_view(dictionary_); + } + + [[nodiscard]] auto as_dataset_view() const noexcept -> dataset_view + { + return dataset_view(data_view(), dictionary_view()); + } + + private: + data_type data_; + [[no_unique_address]] dictionary_type dictionary_; +}; + +/** Non-owning dataset view: holds only view-shaped storage (mdspan, not mdarray). Deliberately not + * derived from `dataset` -- a view type should hold "all view state" with no inheritance and no + * shared ownership tying it to the owning type. Reuses the same `get_n_rows`/`get_dim` spec + * functions as `dataset`, fed view-shaped arguments instead of owning ones, since `raft::mdspan` + * exposes the same `.extent(r)` shape as `raft::mdarray`. */ +template +struct dataset_view { + using spec_type = typename SpecT::template apply; + using value_type = typename spec_type::value_type; + using index_type = typename spec_type::index_type; + using view_type = typename spec_type::view_type; + using dictionary_view_type = typename spec_type::dictionary_view_type; + + explicit dataset_view(view_type data_view, + dictionary_view_type dictionary_view = dictionary_view_type{}) noexcept + : data_view_{data_view}, dictionary_view_{dictionary_view} + { + } + + [[nodiscard]] auto n_rows() const noexcept -> index_type + { + return spec_type::get_n_rows(data_view_); + } + [[nodiscard]] auto dim() const noexcept -> uint32_t + { + return spec_type::get_dim(data_view_, dictionary_view_); + } + [[nodiscard]] auto data_view() const noexcept -> view_type { return data_view_; } + [[nodiscard]] auto dictionary_view() const noexcept -> dictionary_view_type + { + return dictionary_view_; + } + + private: + view_type data_view_; + [[no_unique_address]] dictionary_view_type dictionary_view_; +}; + +} // namespace experimental + namespace filtering { /** From f56f6970948af7e6e308a90bf991081540806e0d Mon Sep 17 00:00:00 2001 From: Howard Huang Date: Fri, 4 Sep 2026 16:23:30 -0700 Subject: [PATCH 03/15] Replace Container-tagged dataset/dataset_view with Spec-based design (#2395) common.hpp's dataset/dataset_view<...> (four hand-specialized structs tagged by ContainerType) are replaced by dataset/dataset_view: two independent, non-inheriting types where every member is a one-line forward to a per-kind Spec (empty_dataset_spec, padded_dataset_spec, standard_dataset_spec, vpq_dataset_spec) -- no dispatch inside dataset/dataset_view themselves. VPQ's codebooks move from three flat, ungrouped matrices (vq_code_book, pq_code_book, data as direct members) into a dictionary_type{vq_code_book, pq_code_book} bundle alongside data, giving every kind the same shape (one data slot + one optional dictionary slot). The old dataset_view -> owner back-pointer (.dset()) is removed; a VPQ view now holds its own dictionary_view() directly, matching padded/standard views holding their own state. Public aliases (device_padded_dataset etc.) keep their names and 2-arg construction signatures, so most call sites are unaffected. Call sites that reached into dataset internals directly needed updating: .view()/.stride() on a dataset or dataset_view -> .data_view(); raw .vq_code_book/.pq_code_book/.data member access -> .dictionary_view()/ .data_view(); old 3-arg VPQ construction -> 2-arg (codes, dictionary_type{vq,pq}). Touches CAGRA build/search/serialize/merge, VPQ training (pq.cuh), SCANN, Vamana, HNSW export, multi-GPU CAGRA, and the C API's product-quantizer codebook accessors (pq.cpp) -- the one place the C API reaches into dataset internals rather than going through the opaque cuvsDataset handle. Also: removes detail::vpq_dataset_spec_impl (only empty/padded/standard/ vpq_dataset_spec exist as public specs; padded and standard share dense_dataset_spec_impl since two tags need the same body, but VPQ had no second tag to share with, so its body is now inlined directly into vpq_dataset_spec, consistent with how empty_dataset_spec is already written). Verified: full rebuild clean; DATASET_C_TEST (7/7), CAGRA_C_TEST (14/14), PREPROCESSING_TEST (226/226), NEIGHBORS_TEST (339/339) all pass. Co-Authored-By: Claude Sonnet 5 --- c/src/neighbors/cagra.cpp | 8 +- c/src/preprocessing/quantize/pq.cpp | 8 +- cpp/include/cuvs/neighbors/common.hpp | 1180 +++++++---------- .../cuvs/preprocessing/quantize/pq.hpp | 10 +- cpp/src/neighbors/cagra.cuh | 2 +- cpp/src/neighbors/detail/cagra/add_nodes.cuh | 2 +- .../neighbors/detail/cagra/cagra_build.cuh | 10 +- .../neighbors/detail/cagra/cagra_merge.cuh | 24 +- .../neighbors/detail/cagra/cagra_search.cuh | 10 +- .../detail/cagra/cagra_serialize.cuh | 24 +- .../cagra/compute_distance_standard.hpp | 7 +- .../detail/cagra/compute_distance_vpq.hpp | 25 +- cpp/src/neighbors/detail/cagra/factory.cuh | 17 +- .../neighbors/detail/dataset_serialize.hpp | 14 +- cpp/src/neighbors/detail/hnsw.hpp | 17 +- .../neighbors/detail/vamana/vamana_build.cuh | 10 +- .../detail/vamana/vamana_serialize.cuh | 6 +- cpp/src/neighbors/mg/mg_cagra_inst.cu.in | 10 +- .../neighbors/scann/detail/scann_build.cuh | 4 +- cpp/src/neighbors/tiered_index.cu | 4 +- cpp/src/preprocessing/quantize/detail/pq.cuh | 66 +- cpp/tests/neighbors/ann_cagra.cuh | 4 +- .../ann_cagra/test_merge_fastener.cu | 12 +- cpp/tests/neighbors/ann_scann.cuh | 7 +- cpp/tests/neighbors/vpq_utils.cuh | 14 +- .../preprocessing/product_quantization.cu | 4 +- 26 files changed, 685 insertions(+), 814 deletions(-) diff --git a/c/src/neighbors/cagra.cpp b/c/src/neighbors/cagra.cpp index 82dba5424d..5226b52737 100644 --- a/c/src/neighbors/cagra.cpp +++ b/c/src/neighbors/cagra.cpp @@ -144,7 +144,7 @@ static void merge_indices_for_layout( cuvs::neighbors::cagra::detail::merged_dataset_size( *res_ptr, index_ptrs, row_filter); auto const dim = static_cast(index_ptrs.front()->dim()); - auto const stride = static_cast(index_ptrs.front()->dataset().stride()); + auto const stride = static_cast(index_ptrs.front()->dataset().data_view().stride()); try { auto matrix = raft::make_device_matrix(*res_ptr, final_row_count, stride); @@ -192,8 +192,8 @@ 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(), - static_cast(input.stride()), + input.data_view().data_handle(), + static_cast(input.data_view().stride()), static_cast(dim), static_cast(input.n_rows()), stream); @@ -1159,7 +1159,7 @@ void get_dataset_view(cuvsCagraIndex_t index, DLManagedTensor* dataset) box, "cuvsCagraIndexGetDataset: null index handle", "cuvsCagraIndexGetDataset: host indices are allowed", - [&](auto& idx) { cuvs::core::to_dlpack(idx.dataset().view(), dataset); }); + [&](auto& idx) { cuvs::core::to_dlpack(idx.dataset().data_view(), dataset); }); } template diff --git a/c/src/preprocessing/quantize/pq.cpp b/c/src/preprocessing/quantize/pq.cpp index 1e3a48694a..9b6a616e0f 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,8 @@ 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 +267,8 @@ 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/include/cuvs/neighbors/common.hpp b/cpp/include/cuvs/neighbors/common.hpp index 8f948b9bdc..46ef145412 100644 --- a/cpp/include/cuvs/neighbors/common.hpp +++ b/cpp/include/cuvs/neighbors/common.hpp @@ -148,21 +148,37 @@ enum class MergeStrategy { /** @} */ // end group neighbors_index /** - * @brief Tags selecting dataset representation for `dataset` / `dataset_view`. + * @brief Spec-based `dataset` / `dataset_view`. * - * Each container defines nested `owning_storage` then `view_storage` (aliases into `detail::*` - * storage types shared by device/host). Accessibility (device vs host) is selected by the - * `Accessor` template parameter on `dataset` / `dataset_view`, not by duplicating containers. - * Layout kinds: empty, padded, standard, VPQ. `dataset` / `dataset_view` only express ownership - * vs view. + * `dataset` and `dataset_view` are single generic templates with zero + * per-kind dispatch inside them: every member is a one-line forward to `spec_type::get_*(...)`, + * and all kind-specific logic lives in the per-kind Spec structs below (`empty_dataset_spec`, + * `padded_dataset_spec`, `standard_dataset_spec`, `vpq_dataset_spec`), which `dataset`/ + * `dataset_view` never name or branch on. `dataset` and `dataset_view` are deliberately two + * independent, non-inheriting types (no shared_ptr, no "sometimes owning" object): `dataset` holds + * owning storage (mdarray-shaped), `dataset_view` holds the corresponding view storage + * (mdspan-shaped). The same `get_n_rows`/`get_dim` spec functions serve both, since + * `raft::mdarray`/`raft::mdspan` both expose `.extent(r)`. */ -template +template struct dataset; -template +template struct dataset_view; +/** + * A spec defines a dictionary iff it needs a second storage slot to interpret the data (e.g. PQ + * codebooks). Non-compressed specs declare `dictionary_type = std::monostate` -- the same + * vocabulary type for "no dictionary," not just an omitted member -- so `dataset`/`dataset_view` + * never need to branch on whether the slot exists; they just always have one, sometimes empty. + */ +template +concept compressed_dataset_spec = requires { + typename SpecT::dictionary_type; + typename SpecT::dictionary_view_type; +} && !std::is_same_v; + namespace detail { // Default owning/view accessors for public dataset aliases. @@ -234,20 +250,14 @@ struct empty_dataset_storage { [[nodiscard]] auto dim() const noexcept -> uint32_t { return suggested_dim; } }; -template -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. Publicly - * inherits from MatrixT (an `raft::mdarray`) so `view()`/`data_handle()`/`extent()` etc. are - * reused as-is rather than hand-forwarded; `logical_dim_` is the only state this struct adds. + * Dense row-major owning storage shared by padded and standard dataset specs. Publicly inherits + * from MatrixT (a `raft::mdarray`) so `view()`/`data_handle()`/`extent()` etc. are reused as-is + * rather than hand-forwarded; `logical_dim_` is the only state this struct adds. * * Template parameters: * - MatrixT: owning matrix type that stores the payload (host/device matrix). @@ -310,326 +320,388 @@ struct dense_row_major_dataset_view_storage : public ViewT { [[nodiscard]] auto view() const noexcept -> ViewT { return *this; } }; -template -using padded_dataset_owning_storage = - dense_row_major_dataset_owning_storage; - -template -using padded_dataset_view_storage = dense_row_major_dataset_view_storage; +/** Spec-side implementation shared by `padded_dataset_spec`/`standard_dataset_spec`; those two + * stay distinct top-level types (identical bodies) purely so classification traits can tell them + * apart -- exactly mirroring today's `padded_dataset_container`/`standard_dataset_container`, + * which are likewise two differently-named tags over one shared storage implementation. */ +template +struct dense_dataset_spec_impl { + template + struct apply { + using value_type = std::remove_cv_t; + using index_type = std::remove_cv_t; + using MatrixT = dense_owning_matrix; + using ViewT = dense_view_matrix; + using data_type = dense_row_major_dataset_owning_storage; + using view_type = dense_row_major_dataset_view_storage; + using dictionary_type = std::monostate; + using dictionary_view_type = std::monostate; -template -using standard_dataset_owning_storage = - dense_row_major_dataset_owning_storage; + [[nodiscard]] static auto get_data_view(data_type const& data) noexcept -> view_type + { + return view_type(data.view(), data.dim()); + } + template + [[nodiscard]] static auto get_n_rows(AnyDatasetOrView const& data) noexcept -> index_type + { + return data.n_rows(); + } + template + [[nodiscard]] static auto get_dim(AnyDatasetOrView const& data, dictionary_type const&) noexcept + -> uint32_t + { + return data.dim(); + } + [[nodiscard]] static auto get_dictionary_view(dictionary_type const&) noexcept + -> dictionary_view_type + { + return {}; + } + }; +}; -template -using standard_dataset_view_storage = dense_row_major_dataset_view_storage; +} // namespace detail // ----------------------------------------------------------------------------- -// VPQ compressed +// Public specs -- the only place per-kind logic lives. // ----------------------------------------------------------------------------- -/** - * 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)} - { - } +template +struct empty_dataset_spec { + using accessor_type = Accessor; - [[nodiscard]] auto n_rows() const noexcept -> IdxT { return data.extent(0); } - [[nodiscard]] auto dim() const noexcept -> uint32_t { return vq_code_book.extent(1); } + template + struct apply { + using value_type = std::remove_cv_t; + using index_type = std::remove_cv_t; + using data_type = detail::empty_dataset_storage; + using view_type = detail::empty_dataset_storage; + using dictionary_type = std::monostate; + using dictionary_view_type = std::monostate; - [[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; + [[nodiscard]] static auto get_data_view(data_type const& data) noexcept -> view_type + { + return data; } - 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); - } + [[nodiscard]] static auto get_n_rows(data_type const& data) noexcept -> index_type + { + return static_cast(data.n_rows()); + } + [[nodiscard]] static auto get_dim(data_type const& data, dictionary_type const&) noexcept + -> uint32_t + { + return data.dim(); + } + [[nodiscard]] static auto get_dictionary_view(dictionary_type const&) noexcept + -> dictionary_view_type + { + return {}; + } + }; }; -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"); - } +template +struct padded_dataset_spec { + using accessor_type = ContainerPolicy; + template + struct apply : detail::dense_dataset_spec_impl::template apply {}; +}; - [[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_; } +template +struct standard_dataset_spec { + using accessor_type = ContainerPolicy; + template + struct apply : detail::dense_dataset_spec_impl::template apply {}; }; -} // namespace detail +/** `Accessor` drives both codebook and code residency, mirroring today's + * single-`Accessor`-per-VPQ-dataset design (`vpq_vq_book_matrix`/`vpq_data_matrix` are both keyed + * off one `Accessor`). Data = encoded rows (uint8_t codes); dictionary = {vq_code_book, + * pq_code_book}. Inlined directly (unlike padded/standard) since no second tag shares this body. */ +template +struct vpq_dataset_spec { + using accessor_type = Accessor; -// ----------------------------------------------------------------------------- -// empty -// ----------------------------------------------------------------------------- + template + struct apply { + using value_type = std::remove_cv_t; + using index_type = std::remove_cv_t; + using math_type = MathT; -struct empty_dataset_container { - template - using owning_storage = detail::empty_dataset_owning_storage; - template - using view_storage = detail::empty_dataset_view_storage; -}; + using data_type = detail::vpq_data_matrix; + using view_type = raft::mdspan, + raft::row_major, + detail::dataset_view_accessor_for_owning>; -// ----------------------------------------------------------------------------- -// padded (row-major with logical dim vs stride) -// ----------------------------------------------------------------------------- + using vq_book_type = detail::vpq_vq_book_matrix; + using pq_book_type = detail::vpq_vq_book_matrix; -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>; -}; + 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; -// ----------------------------------------------------------------------------- -// standard (row-major with arbitrary stride; no CAGRA alignment requirement) -// ----------------------------------------------------------------------------- + [[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 pq_n_centers() const noexcept -> uint32_t + { + return static_cast(pq_code_book.extent(0)); + } + [[nodiscard]] auto pq_len() const noexcept -> uint32_t + { + return static_cast(pq_code_book.extent(1)); + } + [[nodiscard]] 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 bits = 0; + while (pq_width > 1) { + bits++; + pq_width >>= 1; + } + return bits; +#endif + } + [[nodiscard]] auto pq_dim() const noexcept -> uint32_t + { + return raft::div_rounding_up_unsafe(dim(), pq_len()); + } + }; -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>; + [[nodiscard]] static auto get_data_view(data_type const& data) noexcept -> view_type + { + return data.view(); + } + template + [[nodiscard]] static auto get_n_rows(AnyExtentShaped const& data) noexcept -> index_type + { + return static_cast(data.extent(0)); + } + /* get_dim differs from a plain dense dataset: the dimension comes from the VQ codebook, not + the encoded rows (row padding makes the encoded-row width ambiguous as a dimension). */ + template + [[nodiscard]] static auto get_dim(AnyData const&, dictionary_type const& dict) noexcept + -> uint32_t + { + return static_cast(dict.vq_code_book.extent(1)); + } + template + [[nodiscard]] static auto get_dim(AnyData const&, dictionary_view_type const& dict) noexcept + -> uint32_t + { + return dict.dim(); + } + [[nodiscard]] static auto get_dictionary_view(dictionary_type const& dict) noexcept + -> dictionary_view_type + { + return {dict.vq_code_book.view(), dict.pq_code_book.view()}; + } + [[nodiscard]] static auto get_encoded_row_length(data_type const& data) noexcept -> uint32_t + { + return static_cast(data.extent(1)); + } + [[nodiscard]] static auto get_encoded_row_length(view_type const& data) noexcept -> uint32_t + { + return static_cast(data.extent(1)); + } + }; }; // ----------------------------------------------------------------------------- -// VPQ compressed +// dataset / dataset_view // ----------------------------------------------------------------------------- -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 +/** Owning dataset: value-held storage (no shared_ptr -- exclusive ownership). Every member is a + * one-line forward to `spec_type::get_*`; all per-kind logic lives in `SpecT`, never inside this + * struct. */ +template struct dataset { - static_assert(!std::is_same_v, - "dataset: unsupported ContainerType / type-parameter combination"); -}; + 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 dictionary_type = typename spec_type::dictionary_type; -template -struct dataset_view { - static_assert(!std::is_same_v, - "dataset_view: unsupported ContainerType / type-parameter combination"); -}; + // Non-compressed: forward constructor args straight to data_type's own constructor (e.g. + // (MatrixT&&, uint32_t logical_dim) for dense, (uint32_t dim) for empty) -- preserves today's + // construction call sites unchanged. + template + explicit dataset(Args&&... args) + requires(!compressed_dataset_spec && std::is_constructible_v) + : data_(std::forward(args)...), dictionary_{} + { + } -// ----------------------------------------------------------------------------- -// empty -// ----------------------------------------------------------------------------- + // Compressed: data (codes) and dictionary (codebooks) constructed independently. + dataset(data_type&& data, dictionary_type&& dictionary) + requires(compressed_dataset_spec) + : data_(std::move(data)), dictionary_(std::move(dictionary)) + { + } -template -struct dataset - : empty_dataset_container::template owning_storage { - using container_type = empty_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> + [[nodiscard]] auto n_rows() const noexcept -> index_type { return spec_type::get_n_rows(data_); } + [[nodiscard]] auto dim() const noexcept -> uint32_t { - return dataset_view>{this->dim()}; + return spec_type::get_dim(data_, dictionary_); + } + [[nodiscard]] auto data_view() const noexcept { return spec_type::get_data_view(data_); } + [[nodiscard]] auto dictionary_view() const noexcept + { + return spec_type::get_dictionary_view(dictionary_); } -}; -template -struct dataset_view - : empty_dataset_container::template view_storage { - using container_type = empty_dataset_container; - using view_storage_type = typename container_type::template view_storage; - using view_storage_type::view_storage_type; -}; + [[nodiscard]] auto as_dataset_view() const noexcept -> dataset_view + { + return dataset_view(data_view(), dictionary_view()); + } -// ----------------------------------------------------------------------------- -// standard (row-major with arbitrary stride) -// ----------------------------------------------------------------------------- + // Move the owning storage out (e.g. to reuse an already-encoded codes matrix while rebuilding + // only the dictionary at a different math_type, as in VPQ's f32->f16 conversion path). + [[nodiscard]] auto release_data() noexcept -> data_type&& { return std::move(data_); } + [[nodiscard]] auto release_dictionary() noexcept -> dictionary_type&& + { + return std::move(dictionary_); + } -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> + // Dictionary-derived helpers (VPQ: encoded_row_length/vq_n_centers/pq_bits/pq_dim/pq_len/ + // pq_n_centers) forward through dictionary_view() when the dictionary provides them; SFINAE'd + // away for kinds without a dictionary, matching today's VPQ-only surface without dataset<> + // itself branching on which kind it is. + [[nodiscard]] auto encoded_row_length() const noexcept + requires requires(data_type const& d) { spec_type::get_encoded_row_length(d); } { - return dataset_view>(this->view(), - this->dim()); + return spec_type::get_encoded_row_length(data_); + } + [[nodiscard]] auto vq_n_centers() const noexcept + requires requires(decltype(dictionary_view()) const& d) { d.vq_n_centers(); } + { + return dictionary_view().vq_n_centers(); + } + [[nodiscard]] auto pq_n_centers() const noexcept + requires requires(decltype(dictionary_view()) const& d) { d.pq_n_centers(); } + { + return dictionary_view().pq_n_centers(); + } + [[nodiscard]] auto pq_len() const noexcept + requires requires(decltype(dictionary_view()) const& d) { d.pq_len(); } + { + return dictionary_view().pq_len(); + } + [[nodiscard]] auto pq_bits() const noexcept + requires requires(decltype(dictionary_view()) const& d) { d.pq_bits(); } + { + return dictionary_view().pq_bits(); + } + [[nodiscard]] auto pq_dim() const noexcept + requires requires(decltype(dictionary_view()) const& d) { d.pq_dim(); } + { + return dictionary_view().pq_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; + private: + data_type data_; + [[no_unique_address]] dictionary_type dictionary_; }; -// ----------------------------------------------------------------------------- -// padded (row-major with logical dim vs stride) -// ----------------------------------------------------------------------------- +/** Non-owning dataset view: holds only view-shaped storage (mdspan, not mdarray). Deliberately not + * derived from `dataset` -- a view type holds "all view state" with no inheritance and no shared + * ownership tying it to the owning type. Reuses the same `get_n_rows`/`get_dim` spec functions as + * `dataset`, fed view-shaped arguments instead of owning ones. */ +template +struct dataset_view { + using spec_type = typename SpecT::template apply; + using value_type = typename spec_type::value_type; + using index_type = typename spec_type::index_type; + using view_type = typename spec_type::view_type; + using dictionary_view_type = typename spec_type::dictionary_view_type; -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> + dataset_view() noexcept = default; + + // Already-constructed (view_type, dictionary_view_type) pair -- the shape `as_dataset_view()` + // always constructs with, for every kind (dictionary_view_type is std::monostate and + // defaults away when there's no dictionary). Not a template, so it's preferred over the + // forwarding constructor below whenever both could apply. + dataset_view(view_type data_view, dictionary_view_type dictionary_view = {}) noexcept + : data_view_{data_view}, dictionary_view_{dictionary_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; -}; + // Forward raw constructor args straight to view_type's own constructor (e.g. (ViewT, uint32_t + // logical_dim) for dense, (uint32_t dim) for empty) -- preserves today's direct-construction + // call sites (e.g. `device_padded_dataset_view(raw_mdspan, dim)`) unchanged. `view_type` + // is never itself constructible from `(view_type, dictionary_view_type)` (its own constructors + // only take mdspan-shaped args), so this and the plain constructor above never both match the + // same call -- no ambiguity. + template + explicit dataset_view(Args&&... args) + requires(std::is_constructible_v) + : data_view_(std::forward(args)...), dictionary_view_{} + { + } -// ----------------------------------------------------------------------------- -// VPQ compressed (view holds non-owning pointer to owning dataset) -// ----------------------------------------------------------------------------- + [[nodiscard]] auto n_rows() const noexcept -> index_type + { + return spec_type::get_n_rows(data_view_); + } + [[nodiscard]] auto dim() const noexcept -> uint32_t + { + return spec_type::get_dim(data_view_, dictionary_view_); + } + [[nodiscard]] auto data_view() const noexcept -> view_type { return data_view_; } + [[nodiscard]] auto dictionary_view() const noexcept -> dictionary_view_type + { + return dictionary_view_; + } -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> + // See dataset<>'s equivalent block: VPQ-only helpers, SFINAE'd away for kinds without a + // dictionary. + [[nodiscard]] auto encoded_row_length() const noexcept + requires requires(view_type const& d) { spec_type::get_encoded_row_length(d); } { - return dataset_view>{this}; + return spec_type::get_encoded_row_length(data_view_); + } + [[nodiscard]] auto vq_n_centers() const noexcept + requires requires(dictionary_view_type const& d) { d.vq_n_centers(); } + { + return dictionary_view_.vq_n_centers(); + } + [[nodiscard]] auto pq_n_centers() const noexcept + requires requires(dictionary_view_type const& d) { d.pq_n_centers(); } + { + return dictionary_view_.pq_n_centers(); + } + [[nodiscard]] auto pq_len() const noexcept + requires requires(dictionary_view_type const& d) { d.pq_len(); } + { + return dictionary_view_.pq_len(); + } + [[nodiscard]] auto pq_bits() const noexcept + requires requires(dictionary_view_type const& d) { d.pq_bits(); } + { + return dictionary_view_.pq_bits(); + } + [[nodiscard]] auto pq_dim() const noexcept + requires requires(dictionary_view_type const& d) { d.pq_dim(); } + { + return dictionary_view_.pq_dim(); } -}; -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; + private: + view_type data_view_{}; + [[no_unique_address]] dictionary_view_type dictionary_view_{}; }; /** @@ -637,136 +709,152 @@ struct dataset_view */ template using device_empty_dataset = - dataset>; + dataset>>; template using device_empty_dataset_view = - dataset_view>; + dataset_view>>; template using host_empty_dataset = - dataset>; + dataset>>; template using host_empty_dataset_view = - dataset_view>; + dataset_view>>; template using device_padded_dataset = - dataset>; + dataset>>; template using device_padded_dataset_view = - dataset_view>; + dataset_view>>; template using host_padded_dataset = - dataset>; + dataset>>; template using host_padded_dataset_view = - dataset_view>; + dataset_view>>; template using device_standard_dataset = - dataset>; + dataset>>; template using device_standard_dataset_view = - dataset_view>; + dataset_view>>; template using host_standard_dataset = - dataset>; + dataset>>; template using host_standard_dataset_view = - dataset_view>; + dataset_view>>; template using device_vpq_dataset = - dataset>; + dataset>>; template using device_vpq_dataset_view = - dataset_view>; + dataset_view>>; template using host_vpq_dataset = - dataset>; + dataset>>; template using host_vpq_dataset_view = - dataset_view>; + dataset_view>>; -// Maps a dataset view type to its owning (allocating) dataset counterpart. -// Used by serialize/deserialize to type the out_dataset output parameter; -// adding a new dataset type only requires adding a new specialization here. +// Maps a dataset view type to its owning (allocating) dataset counterpart. Trivial and total under +// the Spec design: the owning type for `dataset_view` is always +// `dataset` +// -- no per-kind specialization table needed (unlike the old Container-tagged design). 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 = dataset; }; -template -struct owning_dataset_for_view> { - using type = device_standard_dataset; -}; +template +using owning_dataset_for_view_t = typename owning_dataset_for_view::type; + +// ----------------------------------------------------------------------------- +// Spec-kind classification (all derived from SpecT; dataset/dataset_view never branch on kind). +// ----------------------------------------------------------------------------- + +template +struct is_empty_spec : std::false_type {}; +template +struct is_empty_spec> : std::true_type {}; +template +inline constexpr bool is_empty_spec_v = is_empty_spec::value; + +template +struct is_padded_spec : std::false_type {}; +template +struct is_padded_spec> : std::true_type {}; +template +inline constexpr bool is_padded_spec_v = is_padded_spec::value; -template -struct owning_dataset_for_view> { - using type = host_padded_dataset; -}; +template +struct is_standard_spec : std::false_type {}; +template +struct is_standard_spec> : std::true_type {}; +template +inline constexpr bool is_standard_spec_v = is_standard_spec::value; -template -struct owning_dataset_for_view> { - using type = host_standard_dataset; -}; +template +struct is_vpq_spec : std::false_type {}; +template +struct is_vpq_spec> : std::true_type {}; +template +inline constexpr bool is_vpq_spec_v = is_vpq_spec::value; -template -struct owning_dataset_for_view> { - using type = device_vpq_dataset; +template +struct vpq_spec_math_type {}; +template +struct vpq_spec_math_type> { + using type = MathT; }; +template +using vpq_spec_math_type_t = typename vpq_spec_math_type::type; -template -using owning_dataset_for_view_t = typename owning_dataset_for_view::type; - +/** Owning-side kind traits (mirror today's `is_padded_dataset_v`/`is_standard_dataset_v`/ + * `is_vpq_dataset_v`, used for SFINAE overload selection in factory.cuh/compute_distance_vpq.hpp). + */ 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 +struct is_padded_dataset> : std::bool_constant> {}; +template +struct is_padded_dataset> + : std::bool_constant> {}; 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 +struct is_standard_dataset> + : std::bool_constant> {}; +template +struct is_standard_dataset> + : std::bool_constant> {}; 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 +struct is_vpq_dataset> : std::bool_constant> {}; template inline constexpr bool is_vpq_dataset_v = is_vpq_dataset::value; @@ -794,6 +882,9 @@ enum class dataset_view_kind { vpq_f32, }; +template +using dataset_view_type_t = std::remove_cvref_t; + /** Primary template returns `unknown` so traits safely return `false` for non-dataset-view types. */ template @@ -801,39 +892,34 @@ struct dataset_view_kind_of { static constexpr dataset_view_kind value = dataset_view_kind::unknown; }; -template -struct dataset_view_kind_of> { - static constexpr dataset_view_kind value = dataset_view_kind::empty; -}; - -template -struct dataset_view_kind_of> { - static constexpr dataset_view_kind value = dataset_view_kind::padded; -}; - -template -struct dataset_view_kind_of> { - static constexpr dataset_view_kind value = dataset_view_kind::standard; -}; - -template -struct dataset_view_kind_of> { - static_assert(std::is_same_v || std::is_same_v, - "VPQ dataset_view_kind_of expects MathT to be half or float"); - static constexpr dataset_view_kind value = - std::is_same_v ? dataset_view_kind::vpq_f16 : dataset_view_kind::vpq_f32; +template +struct dataset_view_kind_of> { + static constexpr dataset_view_kind value = []() constexpr { + if constexpr (is_empty_spec_v) { + return dataset_view_kind::empty; + } else if constexpr (is_padded_spec_v) { + return dataset_view_kind::padded; + } else if constexpr (is_standard_spec_v) { + return dataset_view_kind::standard; + } else if constexpr (is_vpq_spec_v) { + static_assert(std::is_same_v, half> || + std::is_same_v, float>, + "VPQ dataset_view_kind_of expects MathT to be half or float"); + return std::is_same_v, half> ? dataset_view_kind::vpq_f16 + : dataset_view_kind::vpq_f32; + } else { + return dataset_view_kind::unknown; + } + }(); }; -template -using dataset_view_type_t = std::remove_cvref_t; - /** True when the dataset view accessor is device-accessible. */ template struct dataset_view_is_device_accessible : std::false_type {}; -template -struct dataset_view_is_device_accessible> - : std::bool_constant {}; +template +struct dataset_view_is_device_accessible> + : std::bool_constant {}; template inline constexpr bool dataset_view_is_device_accessible_v = @@ -939,28 +1025,51 @@ inline constexpr bool compatible_host_device_dataset_views_v = /** * Generic accessor retargeting while preserving the dataset tag/layout and value/index types: - * `dataset -> dataset` - * `dataset_view -> dataset_view` + * `dataset> -> dataset>` + * `dataset_view> -> dataset_view>` */ template struct with_accessor; -template -struct with_accessor, NewAccessor> { - using type = dataset; +template +struct with_accessor>, NewAccessor> { + using type = dataset>; +}; + +template +struct with_accessor>, NewAccessor> { + using type = dataset>; +}; + +template +struct with_accessor>, NewAccessor> { + using type = dataset>; +}; + +template +struct with_accessor>, NewAccessor> { + using type = dataset>; +}; + +template +struct with_accessor>, NewAccessor> { + using type = dataset_view>; +}; + +template +struct with_accessor>, NewAccessor> { + using type = dataset_view>; +}; + +template +struct with_accessor>, NewAccessor> { + using type = dataset_view>; }; -template -struct with_accessor, NewAccessor> { - using type = dataset_view; +template +struct with_accessor>, NewAccessor> { + using type = dataset_view>; }; template @@ -990,10 +1099,10 @@ using to_device_accessor_t = typename to_device_accessor::type; template struct device_counterpart; -template -struct device_counterpart> { - using type = with_accessor_t, - to_device_accessor_t>; +template +struct device_counterpart> { + using type = with_accessor_t, + to_device_accessor_t>; }; template @@ -1009,37 +1118,11 @@ template inline constexpr bool is_dense_row_major_dataset_view_v = is_padded_dataset_view_v || is_standard_dataset_view_v; -/** Element type `T` for `cagra::build(res, params, dataset_view)` (deduced, not a template arg). */ -template -struct cagra_view_element_type; - -template -struct cagra_view_element_type> { - using type = DataT; -}; - -template -struct cagra_view_element_type> { - using type = DataT; -}; - -template -struct cagra_view_element_type> { - using type = DataT; -}; - -template -struct cagra_view_element_type> { - using type = DataT; -}; - -template -struct cagra_view_element_type> { - using type = MathT; -}; - +/** Element type `T` for `cagra::build(res, params, dataset_view)` (deduced, not a template arg). + * Trivial under the Spec design: every `dataset_view` already carries `T` directly. + */ template -using cagra_view_element_type_t = typename cagra_view_element_type>::type; +using cagra_view_element_type_t = typename dataset_view_type_t::value_type; // ----------------------------------------------------------------------------- // CAGRA row width in elements (same for make_device_padded_dataset* and index layout checks). @@ -1322,11 +1405,6 @@ auto make_device_standard_dataset_view(SrcT const& src) * wire-format `(logical_dim, stride)` because the deserialized host buffer is tight `[n_rows x * dim]` while the on-disk stride may be larger. Do not call from user code; prefer * `make_device_standard_dataset_view()` when wrapping existing correctly-strided storage. - * - * Potential future call sites if an owning copy with explicit stride is needed: - * - C API dataset upload (mirroring `make_device_padded_dataset` in `c/src/neighbors/cagra.cpp`) - * - `tiered_index` / composite index paths that materialize standard-layout device storage - * - Multigpu (MG) index build or merge when rehydrating a strided dataset from host fragments */ template auto make_device_standard_dataset(const raft::resources& res, @@ -1355,242 +1433,6 @@ auto make_host_standard_dataset_view(SrcT const& src) host_standard_dataset_view>(src, static_cast(src.extent(1))); } -// ===================================================================================== -// Experimental: Spec-based dataset/dataset_view prototype (#2395 follow-up). -// -// Not wired up to any public alias, trait, or downstream call site yet -- exists to validate the -// design in isolation. `dataset` and `dataset_view` are single generic -// templates with zero per-kind dispatch inside them: every member is a one-line forward to -// `spec_type::get_*(...)`, and all kind-specific logic lives in the per-kind Spec structs below -// (`empty_spec`, `mdarray_spec`, `vpq_spec`), which dataset/dataset_view never name or branch on. -// `dataset` and `dataset_view` are deliberately two independent, non-inheriting types (no -// shared_ptr, no "sometimes owning" object): `dataset` holds owning storage (mdarray-shaped), -// `dataset_view` holds the corresponding view storage (mdspan-shaped); the same -// `get_n_rows`/`get_dim` spec functions serve both, since `raft::mdarray`/`raft::mdspan` both -// expose `.extent(r)`. -// ===================================================================================== -namespace experimental { - -/** - * A spec defines a dictionary iff it needs a second storage slot to interpret the data (e.g. PQ - * codebooks). Non-compressed specs declare `dictionary_type = std::monostate` -- the same - * vocabulary type for "no dictionary," not just an omitted member -- so `dataset`/`dataset_view` - * never need to branch on whether the slot exists; they just always have one, sometimes empty. - */ -template -concept compressed_dataset_spec = requires { - typename SpecT::dictionary_type; - typename SpecT::dictionary_view_type; -} && !std::is_same_v; - -// ----------------------------------------------------------------------------- -// empty -// ----------------------------------------------------------------------------- - -struct empty_spec { - struct rep { - uint32_t dim; - }; - - template - struct apply { - using value_type = std::remove_cv_t; - using index_type = std::remove_cv_t; - using data_type = rep; - using view_type = rep; - using dictionary_type = std::monostate; - using dictionary_view_type = std::monostate; - - [[nodiscard]] static auto get_data_view(data_type const& data) noexcept -> view_type - { - return data; - } - [[nodiscard]] static auto get_n_rows(rep const&) noexcept -> index_type { return 0; } - [[nodiscard]] static auto get_dim(rep const& data, dictionary_type const&) noexcept -> uint32_t - { - return data.dim; - } - [[nodiscard]] static auto get_dictionary_view(dictionary_type const&) noexcept - -> dictionary_view_type - { - return {}; - } - }; -}; - -// ----------------------------------------------------------------------------- -// dense (plain or padded), implemented via raft::mdarray -// ----------------------------------------------------------------------------- - -template -struct mdarray_spec { - template - struct apply { - using value_type = std::remove_cv_t; - using index_type = std::remove_cv_t; - using data_type = raft::mdarray, LayoutPolicy, ContainerPolicy>; - // `get_data_view` takes `data_type const&`, so `.view()` resolves to the const overload, - // returning `const_view_type` (const element type) -- match that here, not the mutable - // `view_type`. - using view_type = typename data_type::const_view_type; - using dictionary_type = std::monostate; - using dictionary_view_type = std::monostate; - - [[nodiscard]] static auto get_data_view(data_type const& data) noexcept -> view_type - { - return data.view(); - } - template - [[nodiscard]] static auto get_n_rows(AnyExtentShaped const& data) noexcept -> index_type - { - return static_cast(data.extent(0)); - } - template - [[nodiscard]] static auto get_dim(AnyExtentShaped const& data, dictionary_type const&) noexcept - -> uint32_t - { - return static_cast(data.extent(1)); - } - [[nodiscard]] static auto get_dictionary_view(dictionary_type const&) noexcept - -> dictionary_view_type - { - return {}; - } - }; -}; - -// ----------------------------------------------------------------------------- -// VPQ compressed: data = encoded rows (uint8_t codes); dictionary = {vq_code_book, pq_code_book} -// ----------------------------------------------------------------------------- - -template -struct vpq_spec { - template - using storage_spec = - typename mdarray_spec::template apply; - - template - struct apply : storage_spec { - using value_type = std::remove_cv_t; - /* Members of a dependent base aren't visible to unqualified lookup, so pull these in. */ - using typename storage_spec::data_type; - using typename storage_spec::view_type; - using math_type = MathT; - - using vq_book_type = - raft::mdarray, raft::row_major, BookPolicy>; - using pq_book_type = - raft::mdarray, raft::row_major, 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; - }; - - /* get_data_view/get_n_rows are inherited from storage_spec unchanged; only get_dim and - get_dictionary_view differ from a plain dense dataset, since the dimension comes from the VQ - codebook, not the encoded rows. */ - template - [[nodiscard]] static auto get_dim(AnyData const&, AnyDict const& dict) noexcept -> uint32_t - { - return static_cast(dict.vq_code_book.extent(1)); - } - [[nodiscard]] static auto get_dictionary_view(dictionary_type const& dict) noexcept - -> dictionary_view_type - { - return {dict.vq_code_book.view(), dict.pq_code_book.view()}; - } - }; -}; - -// ----------------------------------------------------------------------------- -// dataset / dataset_view -// ----------------------------------------------------------------------------- - -template -struct dataset_view; - -/** Owning dataset: value-held storage (no shared_ptr -- exclusive ownership, like today's - * `dataset`). Every member is a one-line forward to `spec_type::get_*`; all - * per-kind logic lives in `SpecT`, never inside this struct. */ -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 dictionary_type = typename spec_type::dictionary_type; - - explicit dataset(data_type&& data, dictionary_type&& dictionary = dictionary_type{}) - : data_{std::move(data)}, dictionary_{std::move(dictionary)} - { - } - - [[nodiscard]] auto n_rows() const noexcept -> index_type { return spec_type::get_n_rows(data_); } - [[nodiscard]] auto dim() const noexcept -> uint32_t - { - return spec_type::get_dim(data_, dictionary_); - } - [[nodiscard]] auto data_view() const noexcept { return spec_type::get_data_view(data_); } - [[nodiscard]] auto dictionary_view() const noexcept - { - return spec_type::get_dictionary_view(dictionary_); - } - - [[nodiscard]] auto as_dataset_view() const noexcept -> dataset_view - { - return dataset_view(data_view(), dictionary_view()); - } - - private: - data_type data_; - [[no_unique_address]] dictionary_type dictionary_; -}; - -/** Non-owning dataset view: holds only view-shaped storage (mdspan, not mdarray). Deliberately not - * derived from `dataset` -- a view type should hold "all view state" with no inheritance and no - * shared ownership tying it to the owning type. Reuses the same `get_n_rows`/`get_dim` spec - * functions as `dataset`, fed view-shaped arguments instead of owning ones, since `raft::mdspan` - * exposes the same `.extent(r)` shape as `raft::mdarray`. */ -template -struct dataset_view { - using spec_type = typename SpecT::template apply; - using value_type = typename spec_type::value_type; - using index_type = typename spec_type::index_type; - using view_type = typename spec_type::view_type; - using dictionary_view_type = typename spec_type::dictionary_view_type; - - explicit dataset_view(view_type data_view, - dictionary_view_type dictionary_view = dictionary_view_type{}) noexcept - : data_view_{data_view}, dictionary_view_{dictionary_view} - { - } - - [[nodiscard]] auto n_rows() const noexcept -> index_type - { - return spec_type::get_n_rows(data_view_); - } - [[nodiscard]] auto dim() const noexcept -> uint32_t - { - return spec_type::get_dim(data_view_, dictionary_view_); - } - [[nodiscard]] auto data_view() const noexcept -> view_type { return data_view_; } - [[nodiscard]] auto dictionary_view() const noexcept -> dictionary_view_type - { - return dictionary_view_; - } - - private: - view_type data_view_; - [[no_unique_address]] dictionary_view_type dictionary_view_; -}; - -} // namespace experimental - namespace filtering { /** diff --git a/cpp/include/cuvs/preprocessing/quantize/pq.hpp b/cpp/include/cuvs/preprocessing/quantize/pq.hpp index a633ac6672..b5617767a9 100644 --- a/cpp/include/cuvs/preprocessing/quantize/pq.hpp +++ b/cpp/include/cuvs/preprocessing/quantize/pq.hpp @@ -302,13 +302,13 @@ template SrcT const& src) -> cuvs::neighbors::device_vpq_dataset { - // A cuVS dataset keeps its logical width in `dim()` while `view()` spans the full row pitch. + // A cuVS dataset keeps its logical width in `dim()` while `data_view()` spans the full row + // pitch. if constexpr (requires { - src.view(); + src.data_view(); src.dim(); - src.stride(); }) { - auto const rows = src.view(); + auto const rows = src.data_view(); using value_type = typename decltype(rows)::value_type; using extents_type = raft::matrix_extent; return make_vpq_dataset( @@ -317,7 +317,7 @@ template raft::mdspan{ rows.data_handle(), raft::make_strided_layout(extents_type{rows.extent(0), int64_t{src.dim()}}, - cuda::std::array{int64_t{src.stride()}, 1})}); + cuda::std::array{int64_t{rows.stride()}, 1})}); } else { using value_type = typename SrcT::value_type; static_assert(std::is_same_v || std::is_same_v || diff --git a/cpp/src/neighbors/cagra.cuh b/cpp/src/neighbors/cagra.cuh index 7af1251ba0..f0dfe13bec 100644 --- a/cpp/src/neighbors/cagra.cuh +++ b/cpp/src/neighbors/cagra.cuh @@ -45,7 +45,7 @@ CUVS_EXPORT void index::compute_dataset_norms_(raft::reso if constexpr (nb::is_padded_dataset_view_v || nb::is_standard_dataset_view_v) { - rm_dataset = dataset_.view(); + rm_dataset = dataset_.data_view().view(); } else if constexpr (nb::is_vpq_dataset_view_v) { skip_norms = true; } diff --git a/cpp/src/neighbors/detail/cagra/add_nodes.cuh b/cpp/src/neighbors/detail/cagra/add_nodes.cuh index c19314a835..e5030ede5b 100644 --- a/cpp/src/neighbors/detail/cagra/add_nodes.cuh +++ b/cpp/src/neighbors/detail/cagra/add_nodes.cuh @@ -363,7 +363,7 @@ void extend_core(raft::resources const& handle, "cuvs::neighbors::hnsw::deserialize() before calling extend()."); const std::size_t initial_dataset_size = index.size(); - const auto extended_view = extended_dataset.view(); + const auto extended_view = extended_dataset.data_view(); const std::size_t new_dataset_size = static_cast(extended_view.extent(0)); const std::size_t degree = index.graph_degree(); const std::size_t dim = index.dim(); diff --git a/cpp/src/neighbors/detail/cagra/cagra_build.cuh b/cpp/src/neighbors/detail/cagra/cagra_build.cuh index 44770fda8f..900f62154f 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_build.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_build.cuh @@ -1341,7 +1341,7 @@ auto build_ace(raft::resources const& res, const index_params& params, DatasetVi params.graph_degree, npartitions); - auto dataset_view = dataset.view(); + auto dataset_view = dataset.data_view(); size_t dataset_size = dataset.n_rows(); size_t dataset_dim = dataset.dim(); @@ -2290,7 +2290,7 @@ auto ensure_device_padded_for_iterative_search( if constexpr (cuvs::neighbors::is_device_padded_dataset_view_v) { return dataset; } else { - padded_own = cuvs::neighbors::make_device_padded_dataset(res, dataset.view()); + padded_own = cuvs::neighbors::make_device_padded_dataset(res, dataset.data_view()); return padded_own->as_dataset_view(); } } @@ -2315,7 +2315,7 @@ auto iterative_build_graph(raft::resources const& res, std::unique_ptr> padded_own; auto search_dataset = ensure_device_padded_for_iterative_search(res, dataset, padded_own); - auto dev_dataset = search_dataset.view(); + auto dev_dataset = search_dataset.data_view(); uint32_t logical_dim = search_dataset.dim(); // Determine initial graph size. @@ -2636,7 +2636,7 @@ auto build_from_host_matrix(raft::resources const& res, static_cast(n_rows), intermediate_degree, graph_degree, - dataset.view()); + dataset.data_view()); }(); RAFT_LOG_TRACE("Graph optimized, creating index"); @@ -2688,7 +2688,7 @@ auto build_from_device_matrix(raft::resources const& res, device_dataset.n_rows(), intermediate_degree, graph_degree, - device_dataset.view()); + device_dataset.data_view()); }(); RAFT_LOG_TRACE("Graph optimized, creating index"); diff --git a/cpp/src/neighbors/detail/cagra/cagra_merge.cuh b/cpp/src/neighbors/detail/cagra/cagra_merge.cuh index aa00f73cfe..4eb65fb391 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_merge.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_merge.cuh @@ -91,10 +91,10 @@ cuvs::neighbors::cagra::index merge_rebuild( "before merge."); if (dim == 0) { dim = index->dim(); - stride = static_cast(dataset.stride()); + stride = static_cast(dataset.data_view().stride()); } else { RAFT_EXPECTS(dim == index->dim(), "Dimension of datasets in indices must be equal."); - RAFT_EXPECTS(stride == static_cast(dataset.stride()), + RAFT_EXPECTS(stride == static_cast(dataset.data_view().stride()), "Row stride of datasets in indices must be equal."); } merged_rows += static_cast(index->size()); @@ -116,12 +116,12 @@ cuvs::neighbors::cagra::index merge_rebuild( "merged_dataset dimension (%u) must equal the input dimension (%u)", unsigned(merged_dataset.dim()), unsigned(dim)); - RAFT_EXPECTS(merged_dataset.stride() == stride, + RAFT_EXPECTS(merged_dataset.data_view().stride() == stride, "merged_dataset stride (%u) must equal the input stride (%ld)", - unsigned(merged_dataset.stride()), + unsigned(merged_dataset.data_view().stride()), long(stride)); - auto output_const_view = merged_dataset.view(); + auto output_const_view = merged_dataset.data_view(); auto output_view = raft::make_device_matrix_view( const_cast(output_const_view.data_handle()), final_rows, stride); @@ -132,7 +132,7 @@ cuvs::neighbors::cagra::index merge_rebuild( std::size_t n_rows = 0; auto const& v = index->dataset(); if constexpr (cuvs::neighbors::is_dense_row_major_dataset_view_v>) { - src_ptr = v.view().data_handle(); + src_ptr = v.data_view().data_handle(); n_rows = static_cast(v.n_rows()); } else { RAFT_FAIL("cagra::merge: unexpected dataset type while copying rows"); @@ -300,14 +300,14 @@ auto preflight_fastener( } if (result.offsets.size() == 1) { result.dim = static_cast(index->dim()); - result.stride = static_cast(dataset.stride()); + result.stride = static_cast(dataset.data_view().stride()); } else { if (result.dim != static_cast(index->dim())) { return reject("all input dimensions must match"); } // The merged dataset has a single row pitch, so mixed input strides cannot be consolidated // without re-padding each input separately. - if (result.stride != static_cast(dataset.stride())) { + if (result.stride != static_cast(dataset.data_view().stride())) { return reject("all input row strides must match"); } } @@ -382,8 +382,8 @@ void copy_input_datasets( auto const& source = indices[i]->dataset(); raft::copy_matrix(destination + offsets[i] * destination_stride, static_cast(destination_stride), - source.view().data_handle(), - static_cast(source.stride()), + source.data_view().data_handle(), + static_cast(source.data_view().stride()), static_cast(dim), static_cast(source.n_rows()), raft::resource::get_cuda_stream(handle)); @@ -399,7 +399,7 @@ auto merge_fastener(raft::resources const& handle, fastener_preflight_result const& preflight) -> cuvs::neighbors::cagra::index { - auto const stride = static_cast(merged_dataset.stride()); + auto const stride = static_cast(merged_dataset.data_view().stride()); RAFT_EXPECTS(merged_dataset.n_rows() == preflight.rows, "merged_dataset rows (%ld) must equal the merged row count (%ld)", long(merged_dataset.n_rows()), @@ -409,7 +409,7 @@ auto merge_fastener(raft::resources const& handle, unsigned(merged_dataset.dim()), long(preflight.dim)); - auto const output_const_view = merged_dataset.view(); + auto const output_const_view = merged_dataset.data_view(); auto* destination = const_cast(output_const_view.data_handle()); { raft::common::nvtx::range scope("cagra::merge/consolidate"); diff --git a/cpp/src/neighbors/detail/cagra/cagra_search.cuh b/cpp/src/neighbors/detail/cagra/cagra_search.cuh index 2d7683816c..81cf407909 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_search.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_search.cuh @@ -101,13 +101,13 @@ void search_main_core( std::unique_ptr> queries_padded_own; if (cuvs::neighbors::matrix_row_width_matches_cagra_required(queries)) { auto v = cuvs::neighbors::make_device_padded_dataset_view(res, queries); - queries_buf = v.view().data_handle(); - query_row_stride = v.stride(); + queries_buf = v.data_view().data_handle(); + query_row_stride = v.data_view().stride(); } else { queries_padded_own = cuvs::neighbors::make_device_padded_dataset(res, queries); auto v = queries_padded_own->as_dataset_view(); - queries_buf = v.view().data_handle(); - query_row_stride = v.stride(); + queries_buf = v.data_view().data_handle(); + query_row_stride = v.data_view().stride(); } const bool can_batch_n_queries = (query_row_stride == query_dim); @@ -245,7 +245,7 @@ void search_main(raft::resources const& res, params.smem_dtype = cuvs::neighbors::cagra::internal_dtype::F16; } auto desc = dataset_descriptor_init_with_cache( - res, params, vv.dset(), index.metric(), nullptr); + res, params, vv, index.metric(), nullptr); search_main_core( res, params, diff --git a/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh b/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh index 461928d4ab..844a20574b 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh @@ -293,8 +293,8 @@ void write_hnswlib_rows_host( "CAGRA dataset rows (%zu) do not match index size (%zu)", static_cast(dataset_view.n_rows()), n_rows); - dataset_data = dataset_view.view().data_handle(); - dataset_stride = dataset_view.stride(); + dataset_data = dataset_view.data_view().data_handle(); + dataset_stride = dataset_view.data_view().stride(); dataset_is_device = is_device_cagra_hnsw_serialize_index_v; } @@ -427,16 +427,16 @@ void write_hnswlib_rows_device(raft::resources const& res, for (size_t first_row = 0; first_row < n_rows; first_row += batch_rows) { auto const rows = std::min(batch_rows, n_rows - first_row); auto const blocks = (rows + warps_per_block - 1) / warps_per_block; - pack_hnswlib_rows - <<(blocks), block_size, 0, stream>>>(output.data_handle(), - row_size, - graph.data_handle(), - dataset.view().data_handle(), - first_row, - rows, - graph_degree, - dim, - dataset.stride()); + pack_hnswlib_rows<<(blocks), block_size, 0, stream>>>( + output.data_handle(), + row_size, + graph.data_handle(), + dataset.data_view().data_handle(), + first_row, + rows, + graph_degree, + dim, + dataset.data_view().stride()); RAFT_CUDA_TRY(cudaPeekAtLastError()); raft::resource::sync_stream(res); diff --git a/cpp/src/neighbors/detail/cagra/compute_distance_standard.hpp b/cpp/src/neighbors/detail/cagra/compute_distance_standard.hpp index 5f7c8efb10..8014d3a070 100644 --- a/cpp/src/neighbors/detail/cagra/compute_distance_standard.hpp +++ b/cpp/src/neighbors/detail/cagra/compute_distance_standard.hpp @@ -28,7 +28,7 @@ struct standard_descriptor_spec : public instance_spec template constexpr static inline bool accepts_dataset() { - return is_padded_dataset_v; + return cuvs::neighbors::is_padded_dataset_view_v; } template @@ -37,11 +37,12 @@ struct standard_descriptor_spec : public instance_spec cuvs::distance::DistanceType metric, const DistanceT* dataset_norms = nullptr) -> host_type { + auto const data_view = dataset.data_view(); return init_(params, - dataset.view().data_handle(), + data_view.data_handle(), IndexT(dataset.n_rows()), dataset.dim(), - dataset.stride(), + data_view.stride(), dataset_norms); } diff --git a/cpp/src/neighbors/detail/cagra/compute_distance_vpq.hpp b/cpp/src/neighbors/detail/cagra/compute_distance_vpq.hpp index 7ae9dc87ba..45bb515f15 100644 --- a/cpp/src/neighbors/detail/cagra/compute_distance_vpq.hpp +++ b/cpp/src/neighbors/detail/cagra/compute_distance_vpq.hpp @@ -41,16 +41,20 @@ struct vpq_descriptor_spec : public instance_spec { using typename base_type::host_type; using typename base_type::index_type; + // `DatasetT` here is the non-owning dataset_view (not the owning dataset) -- callers pass the + // view directly rather than reaching back through a `.dset()`-style owner pointer, so + // classification and the codebook element type are read off the view via its own dictionary + // state (`is_vpq_dataset_view_v`, `DatasetT::value_type`), not an owning-only trait/typedef. template constexpr static inline auto accepts_dataset() - -> std::enable_if_t, bool> + -> std::enable_if_t, bool> { - return std::is_same_v; + return std::is_same_v; } template constexpr static inline auto accepts_dataset() - -> std::enable_if_t, bool> + -> std::enable_if_t, bool> { return false; } @@ -61,11 +65,13 @@ struct vpq_descriptor_spec : public instance_spec { cuvs::distance::DistanceType metric, const DistanceT* dataset_norms = nullptr) -> host_type { + auto const data_view = dataset.data_view(); + auto const dict_view = dataset.dictionary_view(); return init_(params, - dataset.data.data_handle(), - dataset.encoded_row_length(), - dataset.vq_code_book.data_handle(), - dataset.pq_code_book.data_handle(), + data_view.data_handle(), + static_cast(data_view.extent(1)), + dict_view.vq_code_book.data_handle(), + dict_view.pq_code_book.data_handle(), IndexT(dataset.n_rows()), dataset.dim()); } @@ -79,8 +85,9 @@ struct vpq_descriptor_spec : public instance_spec { if (params.team_size != 0 && TeamSize != params.team_size) { return -1.0; } if (cuvs::distance::DistanceType::L2Expanded != metric) { return -1.0; } // Match codebook params - if (dataset.pq_bits() != PqBits) { return -1.0; } - if (dataset.pq_len() != PqLen) { return -1.0; } + auto const dict_view = dataset.dictionary_view(); + if (dict_view.pq_bits() != PqBits) { return -1.0; } + if (dict_view.pq_len() != PqLen) { return -1.0; } if (select_supported_vpq_smem_dtype(params) != SmemDType) { return -1.0; } // Keep auto-selection on the tuned VPQ diagonal while allowing explicit team_size requests to // use the expanded team_size / dataset_block_dim grid. diff --git a/cpp/src/neighbors/detail/cagra/factory.cuh b/cpp/src/neighbors/detail/cagra/factory.cuh index cdcc18867b..3b746ac680 100644 --- a/cpp/src/neighbors/detail/cagra/factory.cuh +++ b/cpp/src/neighbors/detail/cagra/factory.cuh @@ -90,16 +90,19 @@ struct key { uint32_t smem_dtype; }; +// `DatasetT` here is the non-owning dataset_view passed in by the search path, so all state comes +// off the view's own `data_view()`/`dictionary_view()`, not owning-only members. template auto make_key(const cagra::search_params& params, const DatasetT& dataset, cuvs::distance::DistanceType metric) - -> std::enable_if_t, key> + -> std::enable_if_t, key> { - return key{reinterpret_cast(dataset.view().data_handle()), + auto const data_view = dataset.data_view(); + return key{reinterpret_cast(data_view.data_handle()), uint64_t(dataset.n_rows()), dataset.dim(), - dataset.stride(), + data_view.stride(), uint32_t(params.team_size), uint32_t(metric), uint32_t(params.smem_dtype)}; @@ -109,12 +112,14 @@ template auto make_key(const cagra::search_params& params, const DatasetT& dataset, cuvs::distance::DistanceType metric) - -> std::enable_if_t, key> + -> std::enable_if_t, key> { - return key{reinterpret_cast(dataset.data.data_handle()), + auto const data_view = dataset.data_view(); + auto const dict_view = dataset.dictionary_view(); + return key{reinterpret_cast(data_view.data_handle()), uint64_t(dataset.n_rows()), dataset.dim(), - uint32_t(reinterpret_cast(dataset.pq_code_book.data_handle()) >> 6), + uint32_t(reinterpret_cast(dict_view.pq_code_book.data_handle()) >> 6), uint32_t(params.team_size), uint32_t(metric), uint32_t(params.smem_dtype)}; diff --git a/cpp/src/neighbors/detail/dataset_serialize.hpp b/cpp/src/neighbors/detail/dataset_serialize.hpp index e66e60d941..081110eb06 100644 --- a/cpp/src/neighbors/detail/dataset_serialize.hpp +++ b/cpp/src/neighbors/detail/dataset_serialize.hpp @@ -85,13 +85,14 @@ template requires cuvs::neighbors::is_dense_row_major_dataset_view_v void serialize(const raft::resources& res, std::ostream& os, ViewT const& dataset) { - auto n_rows = dataset.n_rows(); - auto dim = dataset.dim(); - auto stride = dataset.stride(); + auto n_rows = dataset.n_rows(); + auto dim = dataset.dim(); + auto data_view = dataset.data_view(); + auto stride = data_view.stride(); raft::serialize_scalar(res, os, n_rows); raft::serialize_scalar(res, os, dim); raft::serialize_scalar(res, os, stride); - auto src = dataset.view(); + auto src = data_view; auto const elements = dense_matrix_elements(n_rows, dim, "serialize_dense_dataset"); raft::numpy_serializer::write_header(os, {raft::numpy_serializer::get_numpy_dtype(), @@ -434,8 +435,9 @@ auto deserialize_vpq(raft::resources const& res, std::istream& is) raft::deserialize_mdspan(res, is, pq_code_book.view()); raft::deserialize_mdspan(res, is, data.view()); - return std::make_unique>( - std::move(vq_code_book), std::move(pq_code_book), std::move(data)); + using owning_t = device_vpq_dataset; + typename owning_t::dictionary_type dictionary{std::move(vq_code_book), std::move(pq_code_book)}; + return std::make_unique(std::move(data), std::move(dictionary)); } template diff --git a/cpp/src/neighbors/detail/hnsw.hpp b/cpp/src/neighbors/detail/hnsw.hpp index bac127294e..c9df2de391 100644 --- a/cpp/src/neighbors/detail/hnsw.hpp +++ b/cpp/src/neighbors/detail/hnsw.hpp @@ -360,8 +360,8 @@ from_cagra(raft::resources const& res, host_dataset = raft::make_host_matrix(dataset_view.n_rows(), dataset_view.dim()); raft::copy_matrix(host_dataset.data_handle(), host_dataset.extent(1), - dataset_view.view().data_handle(), - dataset_view.stride(), + dataset_view.data_view().data_handle(), + dataset_view.data_view().stride(), host_dataset.extent(1), dataset_view.n_rows(), raft::resource::get_cuda_stream(res)); @@ -1003,12 +1003,13 @@ void serialize_to_hnswlib_from_inmem( source_stride = dim; } else if constexpr (is_host_cagra_hnsw_export_index_v) { RAFT_FAIL("serialize_to_hnswlib_from_inmem requires dataset for host CAGRA index"); - } else if (auto dataset_view = index_.dataset(); dataset_view.view().data_handle() != nullptr) { + } else if (auto dataset_view = index_.dataset(); + dataset_view.data_view().data_handle() != nullptr) { n_rows = dataset_view.n_rows(); dim = dataset_view.dim(); device_dataset = true; - source_dataset = dataset_view.view().data_handle(); - source_stride = dataset_view.stride(); + source_dataset = dataset_view.data_view().data_handle(); + source_stride = dataset_view.data_view().stride(); } else { RAFT_FAIL("serialize_to_hnswlib_from_inmem: No dataset provided"); } @@ -1129,12 +1130,12 @@ from_cagra(raft::resources const& res, } else if constexpr (is_host_cagra_hnsw_export_index_v) { RAFT_FAIL("hnsw::from_cagra requires dataset for host CAGRA index"); } else if (auto dataset_view = cagra_index.dataset(); - dataset_view.view().data_handle() != nullptr) { + dataset_view.data_view().data_handle() != nullptr) { n_rows = dataset_view.n_rows(); dim = dataset_view.dim(); device_copy = true; - source_dataset = dataset_view.view().data_handle(); - source_stride = dataset_view.stride(); + source_dataset = dataset_view.data_view().data_handle(); + source_stride = dataset_view.data_view().stride(); } else { RAFT_FAIL("hnsw::from_cagra: No dataset provided"); } diff --git a/cpp/src/neighbors/detail/vamana/vamana_build.cuh b/cpp/src/neighbors/detail/vamana/vamana_build.cuh index fc262b3311..82aa16e4ea 100644 --- a/cpp/src/neighbors/detail/vamana/vamana_build.cuh +++ b/cpp/src/neighbors/detail/vamana/vamana_build.cuh @@ -644,12 +644,14 @@ index build( // process in batches const uint32_t n_rows = dataset.extent(0); + using vpq_owning_t = cuvs::neighbors::device_vpq_dataset; + typename vpq_owning_t::dictionary_type vpq_dictionary{ + raft::make_device_matrix(res, 0, 0), + std::move(pq_codebook)}; auto quantizer = cuvs::preprocessing::quantize::pq::quantizer( pq_params, - cuvs::neighbors::device_vpq_dataset{ - raft::make_device_matrix(res, 0, 0), - std::move(pq_codebook), - raft::make_device_matrix(res, 0, 0)}); + vpq_owning_t(raft::make_device_matrix(res, 0, 0), + std::move(vpq_dictionary))); const int64_t codes_rowlen = cuvs::preprocessing::quantize::pq::get_quantized_dim(pq_params); quantized_vectors = raft::make_device_matrix(res, n_rows, codes_rowlen); diff --git a/cpp/src/neighbors/detail/vamana/vamana_serialize.cuh b/cpp/src/neighbors/detail/vamana/vamana_serialize.cuh index 8a21ed3f1b..c77482c274 100644 --- a/cpp/src/neighbors/detail/vamana/vamana_serialize.cuh +++ b/cpp/src/neighbors/detail/vamana/vamana_serialize.cuh @@ -95,7 +95,7 @@ void serialize_dataset(raft::resources const& res, { if (dataset == nullptr) { return; } try { - serialize_dataset_view(res, dataset->view(), dataset_base_file); + serialize_dataset_view(res, dataset->data_view(), dataset_base_file); } catch (std::bad_alloc& e) { RAFT_LOG_INFO("Failed to serialize dataset"); } catch (raft::logic_error& e) { @@ -172,8 +172,8 @@ void serialize_sector_aligned( auto h_data = raft::make_host_matrix(npts, ndims); raft::copy_matrix(h_data.data_handle(), ndims, - dataset.view().data_handle(), - dataset.stride(), + dataset.data_view().data_handle(), + dataset.data_view().stride(), ndims, npts, raft::resource::get_cuda_stream(res)); diff --git a/cpp/src/neighbors/mg/mg_cagra_inst.cu.in b/cpp/src/neighbors/mg/mg_cagra_inst.cu.in index 077a065e58..0fbe7b926f 100644 --- a/cpp/src/neighbors/mg/mg_cagra_inst.cu.in +++ b/cpp/src/neighbors/mg/mg_cagra_inst.cu.in @@ -14,7 +14,7 @@ void distribute_padded_dataset( cuvs::neighbors::device_padded_dataset_view const& padded_dataset, RankUpdate&& rank_update) { - auto padded_mds = padded_dataset.view(); + auto padded_mds = padded_dataset.data_view(); auto stride = padded_mds.extent(1); const raft::resources& root_res = raft::resource::set_current_device_to_root_rank(res); auto padded_host = @@ -61,7 +61,7 @@ void distribute_padded_dataset( res, \ index, \ static_cast(&index_params), \ - index_dataset.view()); \ + index_dataset.data_view()); \ return index; \ } \ \ @@ -76,7 +76,7 @@ void distribute_padded_dataset( res, \ index, \ static_cast(&index_params), \ - index_dataset.view()); \ + index_dataset.data_view()); \ return index; \ } \ \ @@ -122,7 +122,7 @@ void distribute_padded_dataset( cuvs::neighbors::host_padded_dataset_view new_vectors, \ std::optional> new_indices) \ { \ - cuvs::neighbors::snmg::detail::extend(res, index, new_vectors.view(), new_indices); \ + cuvs::neighbors::snmg::detail::extend(res, index, new_vectors.data_view(), new_indices); \ } \ \ void extend(const raft::resources& res, \ @@ -130,7 +130,7 @@ void distribute_padded_dataset( cuvs::neighbors::host_standard_dataset_view new_vectors, \ std::optional> new_indices) \ { \ - cuvs::neighbors::snmg::detail::extend(res, index, new_vectors.view(), new_indices); \ + cuvs::neighbors::snmg::detail::extend(res, index, new_vectors.data_view(), new_indices); \ } \ \ void search( \ diff --git a/cpp/src/neighbors/scann/detail/scann_build.cuh b/cpp/src/neighbors/scann/detail/scann_build.cuh index c01e50bc83..56611268f7 100644 --- a/cpp/src/neighbors/scann/detail/scann_build.cuh +++ b/cpp/src/neighbors/scann/detail/scann_build.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -290,7 +290,7 @@ index build( // Codebooks from VPQ have the shape [subspace idx, subspace dim, code] // This converts the codebook into matrix format for easy interoperability // with open-source ScaNN search - auto full_codebook_view = pq_quantizer.vpq_codebooks.pq_code_book.view(); + auto full_codebook_view = pq_quantizer.vpq_codebooks.dictionary_view().pq_code_book; raft::linalg::map_offset( res, diff --git a/cpp/src/neighbors/tiered_index.cu b/cpp/src/neighbors/tiered_index.cu index fd7848454f..5d75886930 100644 --- a/cpp/src/neighbors/tiered_index.cu +++ b/cpp/src/neighbors/tiered_index.cu @@ -59,7 +59,7 @@ auto build(raft::resources const& res, -> tiered_index::index> { auto state = detail::build>( - res, params, cagra_build_for_tiered_padded, dataset.view()); + res, params, cagra_build_for_tiered_padded, dataset.data_view()); return cuvs::neighbors::tiered_index::index>(state); } @@ -92,7 +92,7 @@ auto convert_standard_to_padded_index( next_state->ann_index.reset(); if (idx.state->ann_index) { - auto padded_mds = padded_dataset.view(); + auto padded_mds = padded_dataset.data_view(); auto ann_rows = static_cast(idx.state->ann_rows()); auto ann_mds = raft::make_device_matrix_view( padded_mds.data_handle(), ann_rows, static_cast(padded_mds.extent(1))); diff --git a/cpp/src/preprocessing/quantize/detail/pq.cuh b/cpp/src/preprocessing/quantize/detail/pq.cuh index 7fea89461a..f0d749b9a6 100644 --- a/cpp/src/preprocessing/quantize/detail/pq.cuh +++ b/cpp/src/preprocessing/quantize/detail/pq.cuh @@ -192,9 +192,9 @@ quantizer build( pq_code_book = cuvs::neighbors::detail::train_pq( res, filled_params, dataset, raft::make_const_mdspan(vq_code_book.view())); } - return {filled_params, - cuvs::neighbors::device_vpq_dataset{ - std::move(vq_code_book), std::move(pq_code_book), std::move(empty_codes)}}; + using owning_t = cuvs::neighbors::device_vpq_dataset; + typename owning_t::dictionary_type dictionary{std::move(vq_code_book), std::move(pq_code_book)}; + return {filled_params, owning_t(std::move(empty_codes), std::move(dictionary))}; } template @@ -217,7 +217,7 @@ void transform( RAFT_EXPECTS(quantizer.params_quantizer.pq_bits >= 4 && quantizer.params_quantizer.pq_bits <= 16, "PQ bits must be within [4, 16]"); // Encode dataset - auto vq_centers = raft::make_const_mdspan(quantizer.vpq_codebooks.vq_code_book.view()); + auto vq_centers = quantizer.vpq_codebooks.dictionary_view().vq_code_book; auto vq_labels_view = raft::make_device_vector_view(nullptr, 0); if (vq_labels.has_value()) { vq_labels_view = vq_labels.value(); } @@ -226,7 +226,7 @@ void transform( res, to_vpq_params(quantizer.params_quantizer), dataset, - raft::make_const_mdspan(quantizer.vpq_codebooks.pq_code_book.view()), + quantizer.vpq_codebooks.dictionary_view().pq_code_book, vq_centers, vq_labels_view, pq_codes_out); @@ -235,7 +235,7 @@ void transform( res, to_vpq_params(quantizer.params_quantizer), dataset, - raft::make_const_mdspan(quantizer.vpq_codebooks.pq_code_book.view()), + quantizer.vpq_codebooks.dictionary_view().pq_code_book, vq_centers, vq_labels_view, pq_codes_out); @@ -356,30 +356,32 @@ void inverse_transform( "Codes matrix doesn't have the correct number of columns"); RAFT_EXPECTS(quant.params_quantizer.pq_bits >= 4 && quant.params_quantizer.pq_bits <= 16, "PQ bits must be within [4, 16]"); - reconstruct_vectors( - res, - quant.params_quantizer, - codes, - raft::make_const_mdspan(quant.vpq_codebooks.pq_code_book.view()), - raft::make_const_mdspan(quant.vpq_codebooks.vq_code_book.view()), - vq_labels, - out, - quant.params_quantizer.use_subspaces); + auto const quant_dict = quant.vpq_codebooks.dictionary_view(); + reconstruct_vectors(res, + quant.params_quantizer, + codes, + quant_dict.pq_code_book, + quant_dict.vq_code_book, + vq_labels, + out, + quant.params_quantizer.use_subspaces); } template -void vpq_convert_math_type(const raft::resources& res, - const cuvs::neighbors::device_vpq_dataset& src, - cuvs::neighbors::device_vpq_dataset& dst) +void vpq_convert_math_type( + const raft::resources& res, + const cuvs::neighbors::device_vpq_dataset& src, + typename cuvs::neighbors::device_vpq_dataset::dictionary_type& dst_dict) { + auto const src_dict = src.dictionary_view(); raft::linalg::map(res, - dst.vq_code_book.view(), + dst_dict.vq_code_book.view(), cuvs::spatial::knn::detail::utils::mapping{}, - raft::make_const_mdspan(src.vq_code_book.view())); + src_dict.vq_code_book); raft::linalg::map(res, - dst.pq_code_book.view(), + dst_dict.pq_code_book.view(), cuvs::spatial::knn::detail::utils::mapping{}, - raft::make_const_mdspan(src.pq_code_book.view())); + src_dict.pq_code_book); } inline auto make_pq_params_from_vpq(const cuvs::neighbors::vpq_params& in_params, @@ -437,8 +439,9 @@ auto vpq_build(const raft::resources& res, codes.view(), true); - return cuvs::neighbors::device_vpq_dataset{ - std::move(vq_code_book), std::move(pq_code_book), std::move(codes)}; + using owning_t = cuvs::neighbors::device_vpq_dataset; + typename owning_t::dictionary_type dictionary{std::move(vq_code_book), std::move(pq_code_book)}; + return owning_t(std::move(codes), std::move(dictionary)); } template @@ -446,12 +449,13 @@ auto vpq_build_half(const raft::resources& res, const cuvs::neighbors::vpq_params& params, const DatasetT& dataset) -> cuvs::neighbors::device_vpq_dataset { - auto old_type = vpq_build(res, params, dataset); - auto new_type = cuvs::neighbors::device_vpq_dataset{ - raft::make_device_mdarray(res, old_type.vq_code_book.extents()), - raft::make_device_mdarray(res, old_type.pq_code_book.extents()), - std::move(old_type.data)}; - vpq_convert_math_type(res, old_type, new_type); - return new_type; + auto old_type = vpq_build(res, params, dataset); + using new_owning_t = cuvs::neighbors::device_vpq_dataset; + auto const old_dict = old_type.dictionary_view(); + typename new_owning_t::dictionary_type new_dict{ + raft::make_device_mdarray(res, old_dict.vq_code_book.extents()), + raft::make_device_mdarray(res, old_dict.pq_code_book.extents())}; + vpq_convert_math_type(res, old_type, new_dict); + return new_owning_t(old_type.release_data(), std::move(new_dict)); } } // namespace cuvs::preprocessing::quantize::pq::detail diff --git a/cpp/tests/neighbors/ann_cagra.cuh b/cpp/tests/neighbors/ann_cagra.cuh index ebfbab759e..5a033f4b07 100644 --- a/cpp/tests/neighbors/ann_cagra.cuh +++ b/cpp/tests/neighbors/ann_cagra.cuh @@ -1329,7 +1329,7 @@ class AnnCagraIndexFilteredMergeTest : public ::testing::TestWithParam( handle_, ps.n_rows - static_cast(test_cagra_sample_filter::offset), - static_cast(index0.dataset().stride())); + static_cast(index0.dataset().data_view().stride())); auto merged_dataset = cuvs::neighbors::device_padded_dataset( std::move(merged_matrix), static_cast(ps.dim)); auto merge_idx = cuvs::neighbors::cagra::merge( @@ -1575,7 +1575,7 @@ class AnnCagraIndexMergeTest : public ::testing::TestWithParam { auto const merged_rows = static_cast(index0.size()) + static_cast(index1.size()); auto merged_matrix = raft::make_device_matrix( - handle_, merged_rows, static_cast(index0.dataset().stride())); + handle_, merged_rows, static_cast(index0.dataset().data_view().stride())); auto merged_dataset = cuvs::neighbors::device_padded_dataset( std::move(merged_matrix), static_cast(ps.dim)); auto merged_idx = diff --git a/cpp/tests/neighbors/ann_cagra/test_merge_fastener.cu b/cpp/tests/neighbors/ann_cagra/test_merge_fastener.cu index 9dee454d8c..867de36177 100644 --- a/cpp/tests/neighbors/ann_cagra/test_merge_fastener.cu +++ b/cpp/tests/neighbors/ann_cagra/test_merge_fastener.cu @@ -216,11 +216,13 @@ void expect_dataset_order(raft::resources const& res, // honouring that stride rather than as one contiguous block. auto host = raft::make_host_matrix(res, expected.extent(0), expected.extent(1)); auto stream = raft::resource::get_cuda_stream(res); - int64_t const row_stride = static_cast(view.stride()); + int64_t const row_stride = static_cast(view.data_view().stride()); int64_t const dim = static_cast(view.dim()); for (int64_t row = 0; row < view.n_rows(); ++row) { - raft::copy( - host.data_handle() + row * dim, view.view().data_handle() + row * row_stride, dim, stream); + raft::copy(host.data_handle() + row * dim, + view.data_view().data_handle() + row * row_stride, + dim, + stream); } raft::resource::sync_stream(res); for (int64_t row = 0; row < expected.extent(0); ++row) { @@ -237,10 +239,10 @@ void expect_zero_padding(raft::resources const& res, auto view = merged.dataset(); int64_t const rows = view.n_rows(); int64_t const dim = static_cast(view.dim()); - int64_t const stride = static_cast(view.stride()); + int64_t const stride = static_cast(view.data_view().stride()); std::vector host(static_cast(rows * stride)); auto stream = raft::resource::get_cuda_stream(res); - raft::copy(host.data(), view.view().data_handle(), host.size(), stream); + raft::copy(host.data(), view.data_view().data_handle(), host.size(), stream); raft::resource::sync_stream(res); for (int64_t row = 0; row < rows; ++row) { for (int64_t column = dim; column < stride; ++column) { diff --git a/cpp/tests/neighbors/ann_scann.cuh b/cpp/tests/neighbors/ann_scann.cuh index 81ef21c8e2..dff3b52a86 100644 --- a/cpp/tests/neighbors/ann_scann.cuh +++ b/cpp/tests/neighbors/ann_scann.cuh @@ -184,10 +184,11 @@ class scann_test : public ::testing::TestWithParam { vq_codebook.data_handle(), idx.centers().data_handle(), idx.centers().size(), stream_); auto empty_data = raft::make_device_matrix(handle_, 0, 0); + using vpq_owning_t = cuvs::neighbors::device_vpq_dataset; + typename vpq_owning_t::dictionary_type vpq_dictionary{std::move(vq_codebook), + std::move(pq_codebook_copy)}; cuvs::preprocessing::quantize::pq::quantizer quantizer{ - pq_params, - cuvs::neighbors::device_vpq_dataset{ - std::move(vq_codebook), std::move(pq_codebook_copy), std::move(empty_data)}}; + pq_params, vpq_owning_t(std::move(empty_data), std::move(vpq_dictionary))}; auto quantized_residuals_device = raft::make_device_matrix(handle_, ps.num_db_vecs, num_subspaces); diff --git a/cpp/tests/neighbors/vpq_utils.cuh b/cpp/tests/neighbors/vpq_utils.cuh index 23dba3218c..7b5b38e528 100644 --- a/cpp/tests/neighbors/vpq_utils.cuh +++ b/cpp/tests/neighbors/vpq_utils.cuh @@ -50,7 +50,9 @@ void decode_vpq_dataset(raft::device_matrix_view decoded_datase cudaStream_t cuda_stream) { const auto dataset_size = decoded_dataset.extent(0); - RAFT_EXPECTS(vpq_dataset.data.extent(0) == dataset_size, "Dataset sizes mismatch"); + auto const dict_view = vpq_dataset.dictionary_view(); + auto const data_view = vpq_dataset.data_view(); + RAFT_EXPECTS(data_view.extent(0) == dataset_size, "Dataset sizes mismatch"); RAFT_EXPECTS(vpq_dataset.pq_bits() == 8, "decode_vpq_dataset currently only supports pq_bits == 8 (got %u)", vpq_dataset.pq_bits()); @@ -63,14 +65,14 @@ void decode_vpq_dataset(raft::device_matrix_view decoded_datase decode_vpq_dataset_kernel <<>>(decoded_dataset.data_handle(), decoded_dataset.stride(0), - vpq_dataset.vq_code_book.data_handle(), - vpq_dataset.vq_code_book.stride(0), - vpq_dataset.pq_code_book.data_handle(), + dict_view.vq_code_book.data_handle(), + dict_view.vq_code_book.stride(0), + dict_view.pq_code_book.data_handle(), vpq_dataset.pq_len(), 1u << vpq_dataset.pq_bits(), vpq_dataset.dim(), dataset_size, - vpq_dataset.data.data_handle(), - vpq_dataset.data.stride(0)); + data_view.data_handle(), + data_view.stride(0)); } } // namespace cuvs::neighbors diff --git a/cpp/tests/preprocessing/product_quantization.cu b/cpp/tests/preprocessing/product_quantization.cu index a392e7e1db..42a57598bf 100644 --- a/cpp/tests/preprocessing/product_quantization.cu +++ b/cpp/tests/preprocessing/product_quantization.cu @@ -329,7 +329,7 @@ TEST(ProductQuantizationTestF, MakeVpqDatasetFromHost) EXPECT_EQ(vpq.n_rows(), n_rows); EXPECT_EQ(vpq.dim(), dim); - EXPECT_NE(vpq.data.data_handle(), nullptr); + EXPECT_NE(vpq.data_view().data_handle(), nullptr); } TEST(ProductQuantizationTestF, MakeVpqDatasetFromPaddedView) @@ -362,7 +362,7 @@ TEST(ProductQuantizationTestF, MakeVpqDatasetFromPaddedView) EXPECT_EQ(vpq.n_rows(), n_rows); EXPECT_EQ(vpq.dim(), dim); - EXPECT_NE(vpq.data.data_handle(), nullptr); + EXPECT_NE(vpq.data_view().data_handle(), nullptr); } // Define test cases with different parameters From b70daa59199bbaade001a2ba998767dc8a8c90e3 Mon Sep 17 00:00:00 2001 From: Howard Huang Date: Mon, 28 Sep 2026 20:39:05 -0700 Subject: [PATCH 04/15] Fix Spec-based dataset API fallout from the release/26.10 -> main merge Merging upstream main (26.12) pulled in code written against pieces of the old ContainerType-tagged dataset/dataset_view design that this branch already replaced with dataset. Git merged these files without conflict (mostly new/independently-touched code), but the result didn't compile or was silently wrong against the new API. No functional changes beyond restoring correct usage of the existing Spec-based API. - bbq.hpp: BBQ was written as a partial specialization of the old 4-parameter dataset template, which no longer exists. BBQ's shape (a runtime-sized set of alternate quantized encodings, mutated in place by the C API) doesn't fit the dataset two-slot (data + dictionary) model that padded/standard/vpq share, so it gets standalone bbq_dataset/ bbq_dataset_view types instead of a SpecT -- same public API as before, just no longer named as a dataset<>/dataset_view<> specialization. Re-pointed the surrounding trait specializations (owning_dataset_for_view, is_bbq_dataset, dataset_view_kind_of, dataset_view_is_device_accessible) at the new type names; dropped the now-invalid cagra_view_element_type specialization in favor of a plain value_type member. - factory.cuh: the merge dropped key_hash/operator==(key,key) (needed by raft::cache::lru's descriptor cache) entirely, and reverted two of the make_key() overloads back to is_vpq_dataset_v (owning-only trait, always false for a view) and flat dataset.pq_code_book access. Restored key_hash/operator==, fixed the predicate to is_vpq_dataset_view_v, and fixed the accessor to go through dictionary_view().pq_code_book. - cagra.cuh, cagra_build.cuh: reverted back to the old .dset() owner back-pointer (dataset_view no longer holds one) in three spots, a bare dataset.stride() call (moved to dataset.data_view().stride()), and flat vpq_dset.data/.vq_code_book/.pq_code_book member access in reconstruct_vpq_queries(); rewritten to use data_view()/dictionary_view(). - test_iterative_cagra_q.cu: a new test added by the merge still used the old 3-argument VPQ dataset constructor (vq_code_book, pq_code_book, codes); switched to the current 2-argument form (codes, dictionary_type {vq_code_book, pq_code_book}). Verified: full rebuild clean (including the CUB/Thrust version mismatch from a stale CMakeCache CUB_DIR, fixed by a clean reconfigure -- unrelated to source changes); DATASET_C_TEST, CAGRA_C_TEST, PREPROCESSING_TEST, and NEIGHBORS_TEST all pass. Co-Authored-By: Claude Sonnet 5 --- .../cuvs/preprocessing/quantize/bbq.hpp | 66 ++++++++++--------- cpp/src/neighbors/cagra.cuh | 6 +- .../neighbors/detail/cagra/cagra_build.cuh | 22 ++++--- cpp/src/neighbors/detail/cagra/factory.cuh | 23 ++++++- .../ann_cagra/test_iterative_cagra_q.cu | 3 +- .../cpp-api-preprocessing-quantize-bbq.md | 14 ++++ 6 files changed, 85 insertions(+), 49 deletions(-) diff --git a/cpp/include/cuvs/preprocessing/quantize/bbq.hpp b/cpp/include/cuvs/preprocessing/quantize/bbq.hpp index 21366d47d0..e93860f9f7 100644 --- a/cpp/include/cuvs/preprocessing/quantize/bbq.hpp +++ b/cpp/include/cuvs/preprocessing/quantize/bbq.hpp @@ -243,29 +243,33 @@ void resolve_dequant_factors( } // namespace preprocessing::quantize::bbq namespace neighbors { -struct bbq_dataset_container { - template - using owning_storage = cuvs::preprocessing::quantize::bbq::quantizer; - template - using view_storage = cuvs::preprocessing::quantize::bbq::quantizer_view; -}; + +/** + * BBQ doesn't fit the shared Spec-based `dataset`/`dataset_view` shape (one data slot + one + * optional dictionary slot): a BBQ dataset is a small bag of alternate encodings of the *same* + * rows, one per `bbq_code_layout`, selected at query time. So unlike padded/standard/vpq, BBQ + * gets its own standalone owning/view types rather than a `SpecT` plugged into `dataset`/ + * `dataset_view` -- they only need to satisfy `ann_dataset_view` (`n_rows()`/`dim()`) to work + * everywhere a dataset view is expected, plus the handful of trait specializations below. + */ +template +struct bbq_dataset_view; template -struct dataset { - using owning_storage_type = bbq_dataset_container::owning_storage; +struct bbq_dataset { + using value_type = DataT; + using owning_storage_type = cuvs::preprocessing::quantize::bbq::quantizer; std::vector quantizers; - dataset(owning_storage_type&& quantizer) noexcept { add_quantizer(std::move(quantizer)); } + explicit bbq_dataset(owning_storage_type&& quantizer) noexcept + { + add_quantizer(std::move(quantizer)); + } [[nodiscard]] auto as_dataset_view() const noexcept - -> dataset_view> + -> bbq_dataset_view> { - return dataset_view>{quantizers}; + return bbq_dataset_view>{ + quantizers}; } [[nodiscard]] constexpr auto n_rows() const noexcept -> IdxT { @@ -291,14 +295,15 @@ struct dataset { }; template -struct dataset_view { - using owning_storage_type = bbq_dataset_container::owning_storage; - using view_storage_type = bbq_dataset_container::view_storage; +struct bbq_dataset_view { + using value_type = DataT; + using owning_storage_type = cuvs::preprocessing::quantize::bbq::quantizer; + using view_storage_type = cuvs::preprocessing::quantize::bbq::quantizer_view; std::vector quantizers; - dataset_view() noexcept = default; + bbq_dataset_view() noexcept = default; - dataset_view(const std::vector& quantizers) noexcept + bbq_dataset_view(const std::vector& quantizers) noexcept { for (const auto& quantizer : quantizers) { add_quantizer(quantizer); @@ -340,12 +345,10 @@ struct dataset_view { }; template -using device_bbq_dataset = - dataset>; +using device_bbq_dataset = bbq_dataset>; template -using device_bbq_dataset_view = - dataset_view>; +using device_bbq_dataset_view = bbq_dataset_view>; template struct owning_dataset_for_view> { @@ -356,20 +359,19 @@ template struct is_bbq_dataset : std::false_type {}; template -struct is_bbq_dataset> : std::true_type {}; +struct is_bbq_dataset> : std::true_type {}; template inline constexpr bool is_bbq_dataset_v = is_bbq_dataset::value; template -struct dataset_view_kind_of> { +struct dataset_view_kind_of> { static constexpr dataset_view_kind value = dataset_view_kind::bbq; }; -template -struct cagra_view_element_type> { - using type = DataT; -}; +template +struct dataset_view_is_device_accessible> + : std::bool_constant {}; template inline constexpr bool is_device_bbq_dataset_view_v = diff --git a/cpp/src/neighbors/cagra.cuh b/cpp/src/neighbors/cagra.cuh index c7fa171676..851ed83dab 100644 --- a/cpp/src/neighbors/cagra.cuh +++ b/cpp/src/neighbors/cagra.cuh @@ -320,10 +320,10 @@ auto build(raft::resources const& res, const index_params& params, DatasetViewT RAFT_EXPECTS(effective_params.metric == cuvs::distance::DistanceType::L2Expanded, "cagra::build: a VPQ dataset supports only L2Expanded distance"); RAFT_EXPECTS(dataset.n_rows() > 0, "cagra::build: VPQ dataset must not be empty"); - RAFT_EXPECTS(dataset.dset().pq_bits() == 8, + RAFT_EXPECTS(dataset.pq_bits() == 8, "cagra::build: VPQ dataset requires pq_bits == 8, got %u", - dataset.dset().pq_bits()); - auto const pq_len = dataset.dset().pq_len(); + dataset.pq_bits()); + auto const pq_len = dataset.pq_len(); RAFT_EXPECTS(pq_len == 2 || pq_len == 4 || pq_len == 8, "cagra::build: VPQ dataset requires pq_len in {2, 4, 8}, got %u", pq_len); diff --git a/cpp/src/neighbors/detail/cagra/cagra_build.cuh b/cpp/src/neighbors/detail/cagra/cagra_build.cuh index f0625cf0aa..40a30a0bc1 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_build.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_build.cuh @@ -2312,9 +2312,10 @@ void reconstruct_vpq_queries(raft::resources const& res, uint32_t batch_size, raft::device_matrix_view output) { - auto const& vpq_dset = vpq_view.dset(); - const uint32_t dim = vpq_dset.dim(); - const uint32_t pq_len = vpq_dset.pq_len(); + auto const data_view = vpq_view.data_view(); + auto const dict_view = vpq_view.dictionary_view(); + const uint32_t dim = vpq_view.dim(); + const uint32_t pq_len = vpq_view.pq_len(); const uint32_t output_ld = static_cast(output.extent(1)); const uint32_t threads = std::min(dim, 256u); RAFT_EXPECTS(output_ld >= dim, @@ -2324,10 +2325,10 @@ void reconstruct_vpq_queries(raft::resources const& res, kern_reconstruct_vpq_queries <<>>( - vpq_dset.data.data_handle(), - vpq_dset.encoded_row_length(), - vpq_dset.vq_code_book.data_handle(), - vpq_dset.pq_code_book.data_handle(), + data_view.data_handle(), + vpq_view.encoded_row_length(), + dict_view.vq_code_book.data_handle(), + dict_view.pq_code_book.data_handle(), dim, pq_len, offset, @@ -2480,14 +2481,15 @@ auto iterative_build_graph(raft::resources const& res, vpq_dataset = dataset; } else { auto const required_stride = cuvs::neighbors::cagra_required_row_width(dataset.dim()); - RAFT_EXPECTS(dataset.stride() == required_stride, + auto const data_view = dataset.data_view(); + RAFT_EXPECTS(data_view.stride() == required_stride, "iterative CAGRA build requires a CAGRA-aligned device dataset " "(stride %u, required %u). Pass a device_padded_dataset_view, or a " "device_standard_dataset_view whose row width already matches " "cagra_required_row_width.", - dataset.stride(), + data_view.stride(), required_stride); - dev_dataset = dataset.data_view(); + dev_dataset = data_view; logical_dim = dataset.dim(); final_graph_size = static_cast(dataset.n_rows()); } diff --git a/cpp/src/neighbors/detail/cagra/factory.cuh b/cpp/src/neighbors/detail/cagra/factory.cuh index e7f3a50bbc..3278a1111e 100644 --- a/cpp/src/neighbors/detail/cagra/factory.cuh +++ b/cpp/src/neighbors/detail/cagra/factory.cuh @@ -105,7 +105,7 @@ auto make_key(const cagra::search_params& params, reinterpret_cast(dataset_norms), uint64_t(dataset.n_rows()), dataset.dim(), - dataset.stride(), + data_view.stride(), uint32_t(params.team_size), uint32_t(metric), uint32_t(params.smem_dtype)}; @@ -115,7 +115,8 @@ template auto make_key(const cagra::search_params& params, const DatasetT& dataset, cuvs::distance::DistanceType metric, - const void* dataset_norms) -> std::enable_if_t, key> + const void* dataset_norms) + -> std::enable_if_t, key> { auto const data_view = dataset.data_view(); auto const dict_view = dataset.dictionary_view(); @@ -123,12 +124,28 @@ auto make_key(const cagra::search_params& params, reinterpret_cast(dataset_norms), uint64_t(dataset.n_rows()), dataset.dim(), - uint32_t(reinterpret_cast(dataset.pq_code_book.data_handle()) >> 6), + uint32_t(reinterpret_cast(dict_view.pq_code_book.data_handle()) >> 6), uint32_t(params.team_size), uint32_t(metric), uint32_t(params.smem_dtype)}; } +inline auto operator==(const key& a, const key& b) -> bool +{ + return a.data_ptr == b.data_ptr && a.norms_ptr == b.norms_ptr && a.n_rows == b.n_rows && + a.dim == b.dim && a.extra_val == b.extra_val && a.team_size == b.team_size && + a.metric == b.metric && a.smem_dtype == b.smem_dtype; +} + +struct key_hash { + inline auto operator()(const key& x) const noexcept -> std::size_t + { + return size_t{x.data_ptr} + size_t{x.norms_ptr} + + size_t{x.n_rows} * size_t{x.dim} * size_t{x.extra_val} + + (size_t{x.team_size} ^ size_t{x.metric}) + size_t{x.smem_dtype}; + } +}; + template struct store { /** Number of descriptors to cache. */ diff --git a/cpp/tests/neighbors/ann_cagra/test_iterative_cagra_q.cu b/cpp/tests/neighbors/ann_cagra/test_iterative_cagra_q.cu index abc99b2842..ea0cfcabc5 100644 --- a/cpp/tests/neighbors/ann_cagra/test_iterative_cagra_q.cu +++ b/cpp/tests/neighbors/ann_cagra/test_iterative_cagra_q.cu @@ -293,7 +293,8 @@ TEST_F(CagraQContractTest, RejectsEmptyDataset) auto vq_code_book = raft::make_device_matrix(res_, 1, width); auto pq_code_book = raft::make_device_matrix(res_, 256, 2); auto codes = raft::make_device_matrix(res_, 0, 4 + dim / 2); - vpq_dataset_t empty{std::move(vq_code_book), std::move(pq_code_book), std::move(codes)}; + vpq_dataset_t::dictionary_type dictionary{std::move(vq_code_book), std::move(pq_code_book)}; + vpq_dataset_t empty{std::move(codes), std::move(dictionary)}; ASSERT_EQ(empty.n_rows(), 0); EXPECT_THROW(cagra::build(res_, iterative_params(), empty.as_dataset_view()), raft::exception); diff --git a/fern/pages/cpp_api/cpp-api-preprocessing-quantize-bbq.md b/fern/pages/cpp_api/cpp-api-preprocessing-quantize-bbq.md index 67fa8291c1..6e115bec35 100644 --- a/fern/pages/cpp_api/cpp-api-preprocessing-quantize-bbq.md +++ b/fern/pages/cpp_api/cpp-api-preprocessing-quantize-bbq.md @@ -151,3 +151,17 @@ dequant_sum_delta from that delta and quantized_component_sums. **Returns** `void` + +## Types + + +### neighbors::bbq_dataset_view + +BBQ doesn't fit the shared Spec-based `dataset`/`dataset_view` shape (one data slot + one + +optional dictionary slot): a BBQ dataset is a small bag of alternate encodings of the *same* rows, one per `bbq_code_layout`, selected at query time. So unlike padded/standard/vpq, BBQ gets its own standalone owning/view types rather than a `SpecT` plugged into `dataset`/ `dataset_view` -- they only need to satisfy `ann_dataset_view` (`n_rows()`/`dim()`) to work everywhere a dataset view is expected, plus the handful of trait specializations below. + +```cpp +template +struct bbq_dataset_view; +``` From 8bf275f57e655f7cf744382eeddf763906a1ef08 Mon Sep 17 00:00:00 2001 From: Howard Huang Date: Thu, 1 Oct 2026 10:50:06 -0700 Subject: [PATCH 05/15] Remove self-returning view() from dense dataset view storage (#2395) dense_row_major_dataset_view_storage inherits from the mdspan type, so it already is the view; its view() just returned *this. Drop it and convert at the two places that called view() on a dataset view (cagra.cuh, cagra_search.cuh) by using data_view() directly. Co-Authored-By: Claude Sonnet 5.5 --- cpp/include/cuvs/neighbors/common.hpp | 3 --- cpp/src/neighbors/cagra.cuh | 2 +- cpp/src/neighbors/detail/cagra/cagra_search.cuh | 2 +- 3 files changed, 2 insertions(+), 5 deletions(-) diff --git a/cpp/include/cuvs/neighbors/common.hpp b/cpp/include/cuvs/neighbors/common.hpp index 37351edb87..2bf6fb4336 100644 --- a/cpp/include/cuvs/neighbors/common.hpp +++ b/cpp/include/cuvs/neighbors/common.hpp @@ -318,9 +318,6 @@ struct dense_row_major_dataset_view_storage : public ViewT { { return static_cast(ViewT::stride(0) > 0 ? ViewT::stride(0) : this->extent(1)); } - // ViewT (mdspan) has no view() of its own -- it already *is* the view -- so this shrinks to a - // plain upcast instead of reaching into a wrapped field. - [[nodiscard]] auto view() const noexcept -> ViewT { return *this; } }; /** Spec-side implementation shared by `padded_dataset_spec`/`standard_dataset_spec`; those two diff --git a/cpp/src/neighbors/cagra.cuh b/cpp/src/neighbors/cagra.cuh index 851ed83dab..6885ba631f 100644 --- a/cpp/src/neighbors/cagra.cuh +++ b/cpp/src/neighbors/cagra.cuh @@ -44,7 +44,7 @@ CUVS_EXPORT void index::compute_dataset_norms_(raft::reso if constexpr (nb::is_padded_dataset_view_v || nb::is_standard_dataset_view_v) { - rm_dataset = dataset_.data_view().view(); + rm_dataset = dataset_.data_view(); } else if constexpr (nb::is_vpq_dataset_view_v || nb::is_bbq_dataset_view_v) { skip_norms = true; diff --git a/cpp/src/neighbors/detail/cagra/cagra_search.cuh b/cpp/src/neighbors/detail/cagra/cagra_search.cuh index 7dd0d0989c..e389bf6508 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_search.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_search.cuh @@ -270,7 +270,7 @@ void search_main(raft::resources const& res, sample_filter); } else if constexpr (cuvs::neighbors::is_device_standard_dataset_view_v) { RAFT_EXPECTS( - cuvs::neighbors::matrix_row_width_matches_cagra_required(index.dataset().view()), + cuvs::neighbors::matrix_row_width_matches_cagra_required(index.dataset().data_view()), "CAGRA search requires each dataset row to have the CAGRA-aligned stride. Create a padded " "dataset with make_device_padded_dataset() and attach it with cagra::update_dataset()."); run_strided_like(index.dataset()); From 277f441de9c5052656d57de4bc38077d58c63aa9 Mon Sep 17 00:00:00 2001 From: Howard Huang Date: Thu, 1 Oct 2026 10:50:47 -0700 Subject: [PATCH 06/15] Rename dataset data_view() to as_matrix_view() (#2395) data_view() returned the raft mdspan-derived view of a dataset's rows, while as_dataset_view() returns the dataset_view type; the two names did not convey that distinction. Rename data_view() to as_matrix_view() so "matrix view" consistently means an mdspan derivative (e.g. device_matrix_view) and "dataset view" means dataset_view. as_dataset_view() is unchanged. Mechanical rename across dataset/dataset_view, CAGRA, Vamana, HNSW, MG, tiered index, C API, pq.hpp docs and tests. Verified: full rebuild clean; DATASET_C_TEST (9/9), CAGRA_C_TEST (15/15), PREPROCESSING_TEST (226/226), NEIGHBORS_TEST (371/371) pass. Co-Authored-By: Claude Sonnet 5.5 --- c/src/neighbors/cagra.cpp | 8 +++---- cpp/include/cuvs/neighbors/common.hpp | 6 ++--- .../cuvs/preprocessing/quantize/pq.hpp | 6 ++--- cpp/src/neighbors/cagra.cuh | 2 +- cpp/src/neighbors/detail/cagra/add_nodes.cuh | 2 +- .../neighbors/detail/cagra/cagra_build.cuh | 23 +++++++++--------- .../neighbors/detail/cagra/cagra_merge.cuh | 24 +++++++++---------- .../neighbors/detail/cagra/cagra_search.cuh | 10 ++++---- .../detail/cagra/cagra_serialize.cuh | 8 +++---- .../cagra/compute_distance_standard.hpp | 2 +- .../detail/cagra/compute_distance_vpq.hpp | 2 +- cpp/src/neighbors/detail/cagra/factory.cuh | 6 ++--- .../neighbors/detail/dataset_serialize.hpp | 2 +- cpp/src/neighbors/detail/hnsw.hpp | 16 ++++++------- .../detail/vamana/vamana_serialize.cuh | 6 ++--- cpp/src/neighbors/mg/mg_cagra_inst.cu.in | 10 ++++---- cpp/src/neighbors/tiered_index.cu | 4 ++-- cpp/tests/neighbors/ann_cagra.cuh | 4 ++-- .../ann_cagra/test_merge_fastener.cu | 8 +++---- cpp/tests/neighbors/vpq_utils.cuh | 2 +- .../preprocessing/product_quantization.cu | 4 ++-- 21 files changed, 78 insertions(+), 77 deletions(-) diff --git a/c/src/neighbors/cagra.cpp b/c/src/neighbors/cagra.cpp index fafae48d67..90627e6441 100644 --- a/c/src/neighbors/cagra.cpp +++ b/c/src/neighbors/cagra.cpp @@ -179,7 +179,7 @@ static void merge_indices_for_layout( cuvs::neighbors::cagra::detail::merged_dataset_size( *res_ptr, index_ptrs, row_filter); auto const dim = static_cast(index_ptrs.front()->dim()); - auto const stride = static_cast(index_ptrs.front()->dataset().data_view().stride()); + auto const stride = static_cast(index_ptrs.front()->dataset().as_matrix_view().stride()); try { auto matrix = raft::make_device_matrix(*res_ptr, final_row_count, stride); @@ -227,8 +227,8 @@ 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.data_view().data_handle(), - static_cast(input.data_view().stride()), + input.as_matrix_view().data_handle(), + static_cast(input.as_matrix_view().stride()), static_cast(dim), static_cast(input.n_rows()), stream); @@ -1383,7 +1383,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().data_view(), dataset); + cuvs::core::to_dlpack(idx.dataset().as_matrix_view(), dataset); } }); } diff --git a/cpp/include/cuvs/neighbors/common.hpp b/cpp/include/cuvs/neighbors/common.hpp index 2bf6fb4336..5a5270f5b0 100644 --- a/cpp/include/cuvs/neighbors/common.hpp +++ b/cpp/include/cuvs/neighbors/common.hpp @@ -557,7 +557,7 @@ struct dataset { { return spec_type::get_dim(data_, dictionary_); } - [[nodiscard]] auto data_view() const noexcept { return spec_type::get_data_view(data_); } + [[nodiscard]] auto as_matrix_view() const noexcept { return spec_type::get_data_view(data_); } [[nodiscard]] auto dictionary_view() const noexcept { return spec_type::get_dictionary_view(dictionary_); @@ -565,7 +565,7 @@ struct dataset { [[nodiscard]] auto as_dataset_view() const noexcept -> dataset_view { - return dataset_view(data_view(), dictionary_view()); + return dataset_view(as_matrix_view(), dictionary_view()); } // Move the owning storage out (e.g. to reuse an already-encoded codes matrix while rebuilding @@ -660,7 +660,7 @@ struct dataset_view { { return spec_type::get_dim(data_view_, dictionary_view_); } - [[nodiscard]] auto data_view() const noexcept -> view_type { return data_view_; } + [[nodiscard]] auto as_matrix_view() const noexcept -> view_type { return data_view_; } [[nodiscard]] auto dictionary_view() const noexcept -> dictionary_view_type { return dictionary_view_; diff --git a/cpp/include/cuvs/preprocessing/quantize/pq.hpp b/cpp/include/cuvs/preprocessing/quantize/pq.hpp index 2d213cc087..4cc2895c24 100644 --- a/cpp/include/cuvs/preprocessing/quantize/pq.hpp +++ b/cpp/include/cuvs/preprocessing/quantize/pq.hpp @@ -299,13 +299,13 @@ template SrcT const& src) -> cuvs::neighbors::device_vpq_dataset { - // A cuVS dataset keeps its logical width in `dim()` while `data_view()` spans the full row + // A cuVS dataset keeps its logical width in `dim()` while `as_matrix_view()` spans the full row // pitch. if constexpr (requires { - src.data_view(); + src.as_matrix_view(); src.dim(); }) { - auto const rows = src.data_view(); + auto const rows = src.as_matrix_view(); using value_type = typename decltype(rows)::value_type; using extents_type = raft::matrix_extent; return make_vpq_dataset( diff --git a/cpp/src/neighbors/cagra.cuh b/cpp/src/neighbors/cagra.cuh index 6885ba631f..2532ee55fe 100644 --- a/cpp/src/neighbors/cagra.cuh +++ b/cpp/src/neighbors/cagra.cuh @@ -44,7 +44,7 @@ CUVS_EXPORT void index::compute_dataset_norms_(raft::reso if constexpr (nb::is_padded_dataset_view_v || nb::is_standard_dataset_view_v) { - rm_dataset = dataset_.data_view(); + rm_dataset = dataset_.as_matrix_view(); } else if constexpr (nb::is_vpq_dataset_view_v || nb::is_bbq_dataset_view_v) { skip_norms = true; diff --git a/cpp/src/neighbors/detail/cagra/add_nodes.cuh b/cpp/src/neighbors/detail/cagra/add_nodes.cuh index 158338dc41..a77eae3758 100644 --- a/cpp/src/neighbors/detail/cagra/add_nodes.cuh +++ b/cpp/src/neighbors/detail/cagra/add_nodes.cuh @@ -363,7 +363,7 @@ void extend_core(raft::resources const& handle, "cuvs::neighbors::hnsw::deserialize() before calling extend()."); const std::size_t initial_dataset_size = index.size(); - const auto extended_view = extended_dataset.data_view(); + const auto extended_view = extended_dataset.as_matrix_view(); const std::size_t new_dataset_size = static_cast(extended_view.extent(0)); const std::size_t degree = index.graph_degree(); const std::size_t dim = index.dim(); diff --git a/cpp/src/neighbors/detail/cagra/cagra_build.cuh b/cpp/src/neighbors/detail/cagra/cagra_build.cuh index 40a30a0bc1..545b95fb7e 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_build.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_build.cuh @@ -1347,7 +1347,7 @@ auto build_ace(raft::resources const& res, const index_params& params, DatasetVi params.graph_degree, npartitions); - auto dataset_view = dataset.data_view(); + auto dataset_view = dataset.as_matrix_view(); size_t dataset_size = dataset.n_rows(); size_t dataset_dim = dataset.dim(); @@ -2312,7 +2312,7 @@ void reconstruct_vpq_queries(raft::resources const& res, uint32_t batch_size, raft::device_matrix_view output) { - auto const data_view = vpq_view.data_view(); + auto const data_view = vpq_view.as_matrix_view(); auto const dict_view = vpq_view.dictionary_view(); const uint32_t dim = vpq_view.dim(); const uint32_t pq_len = vpq_view.pq_len(); @@ -2481,7 +2481,7 @@ auto iterative_build_graph(raft::resources const& res, vpq_dataset = dataset; } else { auto const required_stride = cuvs::neighbors::cagra_required_row_width(dataset.dim()); - auto const data_view = dataset.data_view(); + auto const data_view = dataset.as_matrix_view(); RAFT_EXPECTS(data_view.stride() == required_stride, "iterative CAGRA build requires a CAGRA-aligned device dataset " "(stride %u, required %u). Pass a device_padded_dataset_view, or a " @@ -2843,7 +2843,7 @@ auto build_from_host_matrix(raft::resources const& res, static_cast(n_rows), intermediate_degree, graph_degree, - dataset.data_view()); + dataset.as_matrix_view()); }(); RAFT_LOG_TRACE("Graph optimized, creating index"); @@ -2890,13 +2890,14 @@ auto build_from_device_matrix(raft::resources const& res, auto cagra_graph = iterative_build_graph(res, params, device_dataset); idx.update_graph(res, std::move(cagra_graph)); } else { - auto cagra_graph = build_cagra_host_graph_from_knn_params(res, - params, - knn_build_params, - device_dataset.n_rows(), - intermediate_degree, - graph_degree, - device_dataset.data_view()); + auto cagra_graph = + build_cagra_host_graph_from_knn_params(res, + params, + knn_build_params, + device_dataset.n_rows(), + intermediate_degree, + graph_degree, + device_dataset.as_matrix_view()); idx.update_graph(res, raft::make_const_mdspan(cagra_graph.view())); } diff --git a/cpp/src/neighbors/detail/cagra/cagra_merge.cuh b/cpp/src/neighbors/detail/cagra/cagra_merge.cuh index 57d0efb435..4de0fbe79e 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_merge.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_merge.cuh @@ -93,10 +93,10 @@ cuvs::neighbors::cagra::index merge_rebuild( "before merge."); if (dim == 0) { dim = index->dim(); - stride = static_cast(dataset.data_view().stride()); + stride = static_cast(dataset.as_matrix_view().stride()); } else { RAFT_EXPECTS(dim == index->dim(), "Dimension of datasets in indices must be equal."); - RAFT_EXPECTS(stride == static_cast(dataset.data_view().stride()), + RAFT_EXPECTS(stride == static_cast(dataset.as_matrix_view().stride()), "Row stride of datasets in indices must be equal."); } merged_rows += static_cast(index->size()); @@ -118,12 +118,12 @@ cuvs::neighbors::cagra::index merge_rebuild( "merged_dataset dimension (%u) must equal the input dimension (%u)", unsigned(merged_dataset.dim()), unsigned(dim)); - RAFT_EXPECTS(merged_dataset.data_view().stride() == stride, + RAFT_EXPECTS(merged_dataset.as_matrix_view().stride() == stride, "merged_dataset stride (%u) must equal the input stride (%ld)", - unsigned(merged_dataset.data_view().stride()), + unsigned(merged_dataset.as_matrix_view().stride()), long(stride)); - auto output_const_view = merged_dataset.data_view(); + auto output_const_view = merged_dataset.as_matrix_view(); auto output_view = raft::make_device_matrix_view( const_cast(output_const_view.data_handle()), final_rows, stride); @@ -134,7 +134,7 @@ cuvs::neighbors::cagra::index merge_rebuild( std::size_t n_rows = 0; auto const& v = index->dataset(); if constexpr (cuvs::neighbors::is_dense_row_major_dataset_view_v>) { - src_ptr = v.data_view().data_handle(); + src_ptr = v.as_matrix_view().data_handle(); n_rows = static_cast(v.n_rows()); } else { RAFT_FAIL("cagra::merge: unexpected dataset type while copying rows"); @@ -302,14 +302,14 @@ auto preflight_fastener( } if (result.offsets.size() == 1) { result.dim = static_cast(index->dim()); - result.stride = static_cast(dataset.data_view().stride()); + result.stride = static_cast(dataset.as_matrix_view().stride()); } else { if (result.dim != static_cast(index->dim())) { return reject("all input dimensions must match"); } // The merged dataset has a single row pitch, so mixed input strides cannot be consolidated // without re-padding each input separately. - if (result.stride != static_cast(dataset.data_view().stride())) { + if (result.stride != static_cast(dataset.as_matrix_view().stride())) { return reject("all input row strides must match"); } } @@ -384,8 +384,8 @@ void copy_input_datasets( auto const& source = indices[i]->dataset(); raft::copy_matrix(destination + offsets[i] * destination_stride, static_cast(destination_stride), - source.data_view().data_handle(), - static_cast(source.data_view().stride()), + source.as_matrix_view().data_handle(), + static_cast(source.as_matrix_view().stride()), static_cast(dim), static_cast(source.n_rows()), raft::resource::get_cuda_stream(handle)); @@ -401,7 +401,7 @@ auto merge_fastener(raft::resources const& handle, fastener_preflight_result const& preflight) -> cuvs::neighbors::cagra::index { - auto const stride = static_cast(merged_dataset.data_view().stride()); + auto const stride = static_cast(merged_dataset.as_matrix_view().stride()); RAFT_EXPECTS(merged_dataset.n_rows() == preflight.rows, "merged_dataset rows (%ld) must equal the merged row count (%ld)", long(merged_dataset.n_rows()), @@ -411,7 +411,7 @@ auto merge_fastener(raft::resources const& handle, unsigned(merged_dataset.dim()), long(preflight.dim)); - auto const output_const_view = merged_dataset.data_view(); + auto const output_const_view = merged_dataset.as_matrix_view(); auto* destination = const_cast(output_const_view.data_handle()); { raft::common::nvtx::range scope("cagra::merge/consolidate"); diff --git a/cpp/src/neighbors/detail/cagra/cagra_search.cuh b/cpp/src/neighbors/detail/cagra/cagra_search.cuh index e389bf6508..ac516c176c 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_search.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_search.cuh @@ -103,13 +103,13 @@ void search_main_core( std::unique_ptr> queries_padded_own; if (cuvs::neighbors::matrix_row_width_matches_cagra_required(queries)) { auto v = cuvs::neighbors::make_device_padded_dataset_view(res, queries); - queries_buf = v.data_view().data_handle(); - query_row_stride = v.data_view().stride(); + queries_buf = v.as_matrix_view().data_handle(); + query_row_stride = v.as_matrix_view().stride(); } else { queries_padded_own = cuvs::neighbors::make_device_padded_dataset(res, queries); auto v = queries_padded_own->as_dataset_view(); - queries_buf = v.data_view().data_handle(); - query_row_stride = v.data_view().stride(); + queries_buf = v.as_matrix_view().data_handle(); + query_row_stride = v.as_matrix_view().stride(); } const bool can_batch_n_queries = (query_row_stride == query_logical_dim); @@ -270,7 +270,7 @@ void search_main(raft::resources const& res, sample_filter); } else if constexpr (cuvs::neighbors::is_device_standard_dataset_view_v) { RAFT_EXPECTS( - cuvs::neighbors::matrix_row_width_matches_cagra_required(index.dataset().data_view()), + cuvs::neighbors::matrix_row_width_matches_cagra_required(index.dataset().as_matrix_view()), "CAGRA search requires each dataset row to have the CAGRA-aligned stride. Create a padded " "dataset with make_device_padded_dataset() and attach it with cagra::update_dataset()."); run_strided_like(index.dataset()); diff --git a/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh b/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh index 6b00b67999..9f4fe27c1e 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh @@ -312,8 +312,8 @@ void write_hnswlib_rows_host( "CAGRA dataset rows (%zu) do not match index size (%zu)", static_cast(dataset_view.n_rows()), n_rows); - dataset_data = dataset_view.data_view().data_handle(); - dataset_stride = dataset_view.data_view().stride(); + dataset_data = dataset_view.as_matrix_view().data_handle(); + dataset_stride = dataset_view.as_matrix_view().stride(); dataset_is_device = is_device_cagra_hnsw_serialize_index_v; } @@ -450,12 +450,12 @@ void write_hnswlib_rows_device(raft::resources const& res, output.data_handle(), row_size, graph.data_handle(), - dataset.data_view().data_handle(), + dataset.as_matrix_view().data_handle(), first_row, rows, graph_degree, dim, - dataset.data_view().stride()); + dataset.as_matrix_view().stride()); RAFT_CUDA_TRY(cudaPeekAtLastError()); raft::resource::sync_stream(res); diff --git a/cpp/src/neighbors/detail/cagra/compute_distance_standard.hpp b/cpp/src/neighbors/detail/cagra/compute_distance_standard.hpp index 8014d3a070..8f5d975449 100644 --- a/cpp/src/neighbors/detail/cagra/compute_distance_standard.hpp +++ b/cpp/src/neighbors/detail/cagra/compute_distance_standard.hpp @@ -37,7 +37,7 @@ struct standard_descriptor_spec : public instance_spec cuvs::distance::DistanceType metric, const DistanceT* dataset_norms = nullptr) -> host_type { - auto const data_view = dataset.data_view(); + auto const data_view = dataset.as_matrix_view(); return init_(params, data_view.data_handle(), IndexT(dataset.n_rows()), diff --git a/cpp/src/neighbors/detail/cagra/compute_distance_vpq.hpp b/cpp/src/neighbors/detail/cagra/compute_distance_vpq.hpp index 45bb515f15..f220d476bf 100644 --- a/cpp/src/neighbors/detail/cagra/compute_distance_vpq.hpp +++ b/cpp/src/neighbors/detail/cagra/compute_distance_vpq.hpp @@ -65,7 +65,7 @@ struct vpq_descriptor_spec : public instance_spec { cuvs::distance::DistanceType metric, const DistanceT* dataset_norms = nullptr) -> host_type { - auto const data_view = dataset.data_view(); + auto const data_view = dataset.as_matrix_view(); auto const dict_view = dataset.dictionary_view(); return init_(params, data_view.data_handle(), diff --git a/cpp/src/neighbors/detail/cagra/factory.cuh b/cpp/src/neighbors/detail/cagra/factory.cuh index 0fa501f3c3..f82f3d998f 100644 --- a/cpp/src/neighbors/detail/cagra/factory.cuh +++ b/cpp/src/neighbors/detail/cagra/factory.cuh @@ -92,7 +92,7 @@ struct key { }; // `DatasetT` here is the non-owning dataset_view passed in by the search path, so all state comes -// off the view's own `data_view()`/`dictionary_view()`, not owning-only members. +// off the view's own `as_matrix_view()`/`dictionary_view()`, not owning-only members. template auto make_key(const cagra::search_params& params, const DatasetT& dataset, @@ -102,7 +102,7 @@ auto make_key(const cagra::search_params& params, cuvs::neighbors::is_standard_dataset_v, key> { - auto const data_view = dataset.data_view(); + auto const data_view = dataset.as_matrix_view(); return key{reinterpret_cast(data_view.data_handle()), reinterpret_cast(dataset_norms), uint64_t(dataset.n_rows()), @@ -120,7 +120,7 @@ auto make_key(const cagra::search_params& params, const void* dataset_norms) -> std::enable_if_t, key> { - auto const data_view = dataset.data_view(); + auto const data_view = dataset.as_matrix_view(); auto const dict_view = dataset.dictionary_view(); return key{reinterpret_cast(data_view.data_handle()), reinterpret_cast(dataset_norms), diff --git a/cpp/src/neighbors/detail/dataset_serialize.hpp b/cpp/src/neighbors/detail/dataset_serialize.hpp index ef4349152b..1743230cd3 100644 --- a/cpp/src/neighbors/detail/dataset_serialize.hpp +++ b/cpp/src/neighbors/detail/dataset_serialize.hpp @@ -87,7 +87,7 @@ void serialize(const raft::resources& res, std::ostream& os, ViewT const& datase { auto n_rows = dataset.n_rows(); auto dim = dataset.dim(); - auto data_view = dataset.data_view(); + auto data_view = dataset.as_matrix_view(); auto stride = data_view.stride(); raft::serialize_scalar(res, os, n_rows); raft::serialize_scalar(res, os, dim); diff --git a/cpp/src/neighbors/detail/hnsw.hpp b/cpp/src/neighbors/detail/hnsw.hpp index 6d32b5008b..046540539d 100644 --- a/cpp/src/neighbors/detail/hnsw.hpp +++ b/cpp/src/neighbors/detail/hnsw.hpp @@ -394,8 +394,8 @@ from_cagra(raft::resources const& res, host_dataset = raft::make_host_matrix(dataset_view.n_rows(), dataset_view.dim()); raft::copy_matrix(host_dataset.data_handle(), host_dataset.extent(1), - dataset_view.data_view().data_handle(), - dataset_view.data_view().stride(), + dataset_view.as_matrix_view().data_handle(), + dataset_view.as_matrix_view().stride(), host_dataset.extent(1), dataset_view.n_rows(), raft::resource::get_cuda_stream(res)); @@ -1840,12 +1840,12 @@ void serialize_to_hnswlib_from_inmem( } else if constexpr (is_host_cagra_hnsw_export_index_v) { RAFT_FAIL("serialize_to_hnswlib_from_inmem requires dataset for host CAGRA index"); } else if (auto dataset_view = index_.dataset(); - dataset_view.data_view().data_handle() != nullptr) { + dataset_view.as_matrix_view().data_handle() != nullptr) { n_rows = dataset_view.n_rows(); dim = dataset_view.dim(); device_dataset = true; - source_dataset = dataset_view.data_view().data_handle(); - source_stride = dataset_view.data_view().stride(); + source_dataset = dataset_view.as_matrix_view().data_handle(); + source_stride = dataset_view.as_matrix_view().stride(); } else { RAFT_FAIL("serialize_to_hnswlib_from_inmem: No dataset provided"); } @@ -1965,12 +1965,12 @@ from_cagra(raft::resources const& res, } else if constexpr (is_host_cagra_hnsw_export_index_v) { RAFT_FAIL("hnsw::from_cagra requires dataset for host CAGRA index"); } else if (auto dataset_view = cagra_index.dataset(); - dataset_view.data_view().data_handle() != nullptr) { + dataset_view.as_matrix_view().data_handle() != nullptr) { n_rows = dataset_view.n_rows(); dim = dataset_view.dim(); device_copy = true; - source_dataset = dataset_view.data_view().data_handle(); - source_stride = dataset_view.data_view().stride(); + source_dataset = dataset_view.as_matrix_view().data_handle(); + source_stride = dataset_view.as_matrix_view().stride(); } else { RAFT_FAIL("hnsw::from_cagra: No dataset provided"); } diff --git a/cpp/src/neighbors/detail/vamana/vamana_serialize.cuh b/cpp/src/neighbors/detail/vamana/vamana_serialize.cuh index c77482c274..78bac296db 100644 --- a/cpp/src/neighbors/detail/vamana/vamana_serialize.cuh +++ b/cpp/src/neighbors/detail/vamana/vamana_serialize.cuh @@ -95,7 +95,7 @@ void serialize_dataset(raft::resources const& res, { if (dataset == nullptr) { return; } try { - serialize_dataset_view(res, dataset->data_view(), dataset_base_file); + serialize_dataset_view(res, dataset->as_matrix_view(), dataset_base_file); } catch (std::bad_alloc& e) { RAFT_LOG_INFO("Failed to serialize dataset"); } catch (raft::logic_error& e) { @@ -172,8 +172,8 @@ void serialize_sector_aligned( auto h_data = raft::make_host_matrix(npts, ndims); raft::copy_matrix(h_data.data_handle(), ndims, - dataset.data_view().data_handle(), - dataset.data_view().stride(), + dataset.as_matrix_view().data_handle(), + dataset.as_matrix_view().stride(), ndims, npts, raft::resource::get_cuda_stream(res)); diff --git a/cpp/src/neighbors/mg/mg_cagra_inst.cu.in b/cpp/src/neighbors/mg/mg_cagra_inst.cu.in index 0fbe7b926f..6820480ebb 100644 --- a/cpp/src/neighbors/mg/mg_cagra_inst.cu.in +++ b/cpp/src/neighbors/mg/mg_cagra_inst.cu.in @@ -14,7 +14,7 @@ void distribute_padded_dataset( cuvs::neighbors::device_padded_dataset_view const& padded_dataset, RankUpdate&& rank_update) { - auto padded_mds = padded_dataset.data_view(); + auto padded_mds = padded_dataset.as_matrix_view(); auto stride = padded_mds.extent(1); const raft::resources& root_res = raft::resource::set_current_device_to_root_rank(res); auto padded_host = @@ -61,7 +61,7 @@ void distribute_padded_dataset( res, \ index, \ static_cast(&index_params), \ - index_dataset.data_view()); \ + index_dataset.as_matrix_view()); \ return index; \ } \ \ @@ -76,7 +76,7 @@ void distribute_padded_dataset( res, \ index, \ static_cast(&index_params), \ - index_dataset.data_view()); \ + index_dataset.as_matrix_view()); \ return index; \ } \ \ @@ -122,7 +122,7 @@ void distribute_padded_dataset( cuvs::neighbors::host_padded_dataset_view new_vectors, \ std::optional> new_indices) \ { \ - cuvs::neighbors::snmg::detail::extend(res, index, new_vectors.data_view(), new_indices); \ + cuvs::neighbors::snmg::detail::extend(res, index, new_vectors.as_matrix_view(), new_indices); \ } \ \ void extend(const raft::resources& res, \ @@ -130,7 +130,7 @@ void distribute_padded_dataset( cuvs::neighbors::host_standard_dataset_view new_vectors, \ std::optional> new_indices) \ { \ - cuvs::neighbors::snmg::detail::extend(res, index, new_vectors.data_view(), new_indices); \ + cuvs::neighbors::snmg::detail::extend(res, index, new_vectors.as_matrix_view(), new_indices); \ } \ \ void search( \ diff --git a/cpp/src/neighbors/tiered_index.cu b/cpp/src/neighbors/tiered_index.cu index 5d75886930..5c457508a3 100644 --- a/cpp/src/neighbors/tiered_index.cu +++ b/cpp/src/neighbors/tiered_index.cu @@ -59,7 +59,7 @@ auto build(raft::resources const& res, -> tiered_index::index> { auto state = detail::build>( - res, params, cagra_build_for_tiered_padded, dataset.data_view()); + res, params, cagra_build_for_tiered_padded, dataset.as_matrix_view()); return cuvs::neighbors::tiered_index::index>(state); } @@ -92,7 +92,7 @@ auto convert_standard_to_padded_index( next_state->ann_index.reset(); if (idx.state->ann_index) { - auto padded_mds = padded_dataset.data_view(); + auto padded_mds = padded_dataset.as_matrix_view(); auto ann_rows = static_cast(idx.state->ann_rows()); auto ann_mds = raft::make_device_matrix_view( padded_mds.data_handle(), ann_rows, static_cast(padded_mds.extent(1))); diff --git a/cpp/tests/neighbors/ann_cagra.cuh b/cpp/tests/neighbors/ann_cagra.cuh index 736c9d5b04..951e3af13b 100644 --- a/cpp/tests/neighbors/ann_cagra.cuh +++ b/cpp/tests/neighbors/ann_cagra.cuh @@ -1330,7 +1330,7 @@ class AnnCagraIndexFilteredMergeTest : public ::testing::TestWithParam( handle_, ps.n_rows - static_cast(test_cagra_sample_filter::offset), - static_cast(index0.dataset().data_view().stride())); + static_cast(index0.dataset().as_matrix_view().stride())); auto merged_dataset = cuvs::neighbors::device_padded_dataset( std::move(merged_matrix), static_cast(ps.dim)); auto merge_idx = cuvs::neighbors::cagra::merge( @@ -1576,7 +1576,7 @@ class AnnCagraIndexMergeTest : public ::testing::TestWithParam { auto const merged_rows = static_cast(index0.size()) + static_cast(index1.size()); auto merged_matrix = raft::make_device_matrix( - handle_, merged_rows, static_cast(index0.dataset().data_view().stride())); + handle_, merged_rows, static_cast(index0.dataset().as_matrix_view().stride())); auto merged_dataset = cuvs::neighbors::device_padded_dataset( std::move(merged_matrix), static_cast(ps.dim)); auto merged_idx = diff --git a/cpp/tests/neighbors/ann_cagra/test_merge_fastener.cu b/cpp/tests/neighbors/ann_cagra/test_merge_fastener.cu index 13548e98fb..b817d540ab 100644 --- a/cpp/tests/neighbors/ann_cagra/test_merge_fastener.cu +++ b/cpp/tests/neighbors/ann_cagra/test_merge_fastener.cu @@ -216,11 +216,11 @@ void expect_dataset_order(raft::resources const& res, // honouring that stride rather than as one contiguous block. auto host = raft::make_host_matrix(res, expected.extent(0), expected.extent(1)); auto stream = raft::resource::get_cuda_stream(res); - int64_t const row_stride = static_cast(view.data_view().stride()); + int64_t const row_stride = static_cast(view.as_matrix_view().stride()); int64_t const dim = static_cast(view.dim()); for (int64_t row = 0; row < view.n_rows(); ++row) { raft::copy(host.data_handle() + row * dim, - view.data_view().data_handle() + row * row_stride, + view.as_matrix_view().data_handle() + row * row_stride, dim, stream); } @@ -239,10 +239,10 @@ void expect_zero_padding(raft::resources const& res, auto view = merged.dataset(); int64_t const rows = view.n_rows(); int64_t const dim = static_cast(view.dim()); - int64_t const stride = static_cast(view.data_view().stride()); + int64_t const stride = static_cast(view.as_matrix_view().stride()); std::vector host(static_cast(rows * stride)); auto stream = raft::resource::get_cuda_stream(res); - raft::copy(host.data(), view.data_view().data_handle(), host.size(), stream); + raft::copy(host.data(), view.as_matrix_view().data_handle(), host.size(), stream); raft::resource::sync_stream(res); for (int64_t row = 0; row < rows; ++row) { for (int64_t column = dim; column < stride; ++column) { diff --git a/cpp/tests/neighbors/vpq_utils.cuh b/cpp/tests/neighbors/vpq_utils.cuh index 7b5b38e528..5fc5d1bba0 100644 --- a/cpp/tests/neighbors/vpq_utils.cuh +++ b/cpp/tests/neighbors/vpq_utils.cuh @@ -51,7 +51,7 @@ void decode_vpq_dataset(raft::device_matrix_view decoded_datase { const auto dataset_size = decoded_dataset.extent(0); auto const dict_view = vpq_dataset.dictionary_view(); - auto const data_view = vpq_dataset.data_view(); + auto const data_view = vpq_dataset.as_matrix_view(); RAFT_EXPECTS(data_view.extent(0) == dataset_size, "Dataset sizes mismatch"); RAFT_EXPECTS(vpq_dataset.pq_bits() == 8, "decode_vpq_dataset currently only supports pq_bits == 8 (got %u)", diff --git a/cpp/tests/preprocessing/product_quantization.cu b/cpp/tests/preprocessing/product_quantization.cu index 8ef2811fc9..bdf4680a3f 100644 --- a/cpp/tests/preprocessing/product_quantization.cu +++ b/cpp/tests/preprocessing/product_quantization.cu @@ -329,7 +329,7 @@ TEST(ProductQuantizationTestF, MakeVpqDatasetFromHost) EXPECT_EQ(vpq.n_rows(), n_rows); EXPECT_EQ(vpq.dim(), dim); - EXPECT_NE(vpq.data_view().data_handle(), nullptr); + EXPECT_NE(vpq.as_matrix_view().data_handle(), nullptr); } TEST(ProductQuantizationTestF, MakeVpqDatasetFromPaddedView) @@ -362,7 +362,7 @@ TEST(ProductQuantizationTestF, MakeVpqDatasetFromPaddedView) EXPECT_EQ(vpq.n_rows(), n_rows); EXPECT_EQ(vpq.dim(), dim); - EXPECT_NE(vpq.data_view().data_handle(), nullptr); + EXPECT_NE(vpq.as_matrix_view().data_handle(), nullptr); } // Define test cases with different parameters From 54d7fc5522fba71862e58699b35bed7938788669 Mon Sep 17 00:00:00 2001 From: Howard Huang Date: Thu, 1 Oct 2026 11:57:39 -0700 Subject: [PATCH 07/15] Strip dataset/dataset_view down to a minimal, kind-agnostic core (#2395) dataset and dataset_view carried VPQ-specific state and methods (a dictionary slot holding the codebooks, and encoded_row_length/vq_n_centers/pq_n_centers/pq_len/pq_bits/pq_dim gated behind `requires` clauses). Padded, standard and empty datasets have no such state, so the shared structs were leaking one kind's abstraction into every kind, and BBQ could not use them at all (it needed standalone types). dataset/dataset_view now hold exactly one payload (data_type / view_type, chosen by the spec) and expose only what every dataset has: n_rows(), dim(), as_matrix_view(), as_dataset_view() and data(). A spec has three functions: get_data_view(), get_n_rows(), get_dim(). Anything else a kind needs is state and methods of that kind's payload, reached through data(). Removed from dataset/dataset_view: - dictionary_type, dictionary_view_type, dictionary_view(), release_dictionary(), release_data() and the std::monostate placeholders that padded/standard/empty used for "no dictionary" - compressed_dataset_spec and the compressed/uncompressed constructor split - the six VPQ helper methods VPQ: vpq_owning_storage / vpq_view_storage are the payloads. Each is the uint8_t codes mdarray/mdspan (so as_matrix_view().data_handle()/extent() keep working) plus vq_code_book and pq_code_book, with the helper methods as members. Construction is (codes, vq_code_book, pq_code_book). Callers that used dictionary_view() now use as_matrix_view(); callers of the helpers use data().pq_bits() etc. BBQ: bbq_dataset_spec plugs the quantizer payloads (bbq_owning_storage / bbq_view_storage, formerly the standalone bbq_dataset / bbq_dataset_view) into dataset/dataset_view like any other kind. quantizers, add_quantizer, has_layout and get_quantizer are reached through data(). This undoes the standalone-BBQ workaround from b70daa591, which only existed because of the dictionary slot. owning_dataset_for_view and dataset_view_is_device_accessible now work for BBQ through the generic templates; dataset_view_kind_of keeps a BBQ partial specialization. Not moved yet: the VPQ and BBQ payloads/specs/traits still live in cuvs::neighbors (common.hpp and bbq.hpp). Relocating them out of cuvs::neighbors is the next step. Also: - ann_cagra_bbq.cuh, ann_nn_descent_bbq.cuh: rmm::cuda_stream_view -> cuda::stream_ref. These did not compile at HEAD (BBQ #2654 merged alongside the stream_ref migration #2521) and are the tests that exercise the BBQ change. - Regenerated Fern pages for the changed headers (fern-api-reference hook). Verified: full rebuild clean. DATASET_C_TEST (9/9), CAGRA_C_TEST (15/15), PREPROCESSING_TEST (226/226), NEIGHBORS_TEST (371/371), NEIGHBORS_ANN_CAGRA_BBQ_UINT32_TEST (105), NEIGHBORS_ANN_NN_DESCENT_TEST (572), NEIGHBORS_ANN_CAGRA_MERGE_TEST (26) pass. NEIGHBORS_ANN_CAGRA_FLOAT_UINT32_TEST was still running at commit time. Co-Authored-By: Claude Sonnet 5.5 --- c/src/neighbors/cagra.cpp | 14 +- c/src/preprocessing/quantize/pq.cpp | 4 +- cpp/include/cuvs/neighbors/common.hpp | 448 +++++++----------- .../cuvs/preprocessing/quantize/bbq.hpp | 85 ++-- .../preprocessing/bbq_cpu_quantize.hpp | 6 +- cpp/src/neighbors/cagra.cuh | 6 +- .../neighbors/detail/cagra/cagra_build.cuh | 11 +- .../detail/cagra/compute_distance_vpq.hpp | 14 +- cpp/src/neighbors/detail/cagra/factory.cuh | 5 +- .../neighbors/detail/cagra/graph_shared.cu | 24 +- .../neighbors/detail/dataset_serialize.hpp | 4 +- cpp/src/neighbors/detail/nn_descent.cuh | 31 +- .../neighbors/detail/vamana/vamana_build.cuh | 8 +- .../neighbors/scann/detail/scann_build.cuh | 2 +- cpp/src/preprocessing/quantize/detail/pq.cuh | 46 +- .../ann_cagra/test_iterative_cagra_q.cu | 9 +- cpp/tests/neighbors/ann_cagra_bbq.cuh | 6 +- cpp/tests/neighbors/ann_nn_descent_bbq.cuh | 2 +- cpp/tests/neighbors/ann_scann.cuh | 5 +- cpp/tests/neighbors/vpq_utils.cuh | 15 +- .../pages/cpp_api/cpp-api-neighbors-common.md | 12 +- .../cpp-api-preprocessing-quantize-bbq.md | 12 +- 22 files changed, 344 insertions(+), 425 deletions(-) diff --git a/c/src/neighbors/cagra.cpp b/c/src/neighbors/cagra.cpp index 90627e6441..e89555c03f 100644 --- a/c/src/neighbors/cagra.cpp +++ b/c/src/neighbors/cagra.cpp @@ -413,8 +413,8 @@ template void validate_bbq_layouts(cuvs::neighbors::device_bbq_dataset_view const& dataset) { using layout_t = cuvs::preprocessing::quantize::bbq::bbq_code_layout; - if (dataset.quantizers.size() == 1) { - auto layout = dataset.quantizers.front().layout; + if (dataset.data().quantizers.size() == 1) { + auto layout = dataset.data().quantizers.front().layout; RAFT_EXPECTS(layout == layout_t::packed_1b || layout == layout_t::transposed_2b || layout == layout_t::packed_4b || layout == layout_t::packed_7b || layout == layout_t::packed_8b, @@ -422,10 +422,10 @@ void validate_bbq_layouts(cuvs::neighbors::device_bbq_dataset_view c return; } - const bool has_1b = dataset.has_layout(layout_t::packed_1b); - const bool has_2bt = dataset.has_layout(layout_t::transposed_2b); - const bool has_4b = dataset.has_layout(layout_t::packed_4b); - const bool has_4bt = dataset.has_layout(layout_t::transposed_4b); + const bool has_1b = dataset.data().has_layout(layout_t::packed_1b); + const bool has_2bt = dataset.data().has_layout(layout_t::transposed_2b); + const bool has_4b = dataset.data().has_layout(layout_t::packed_4b); + const bool has_4bt = dataset.data().has_layout(layout_t::transposed_4b); RAFT_EXPECTS((has_1b && (has_4b || has_2bt || has_4bt)) || (has_2bt && has_4bt), "cuvsDatasetMakeBbqView: unsupported asymmetric BBQ layout pair"); } @@ -469,7 +469,7 @@ auto make_bbq_dataset_view(cuvsBbqQuantizer_t* quantizers, std::size_t num_quant RAFT_EXPECTS(quantizer.metric == expected_metric, "cuvsDatasetMakeBbqView: all quantizers must use the same metric"); } - dataset->add_quantizer(quantizer); + dataset->data().add_quantizer(quantizer); } return dataset; diff --git a/c/src/preprocessing/quantize/pq.cpp b/c/src/preprocessing/quantize/pq.cpp index 9b6a616e0f..74f656e9a2 100644 --- a/c/src/preprocessing/quantize/pq.cpp +++ b/c/src/preprocessing/quantize/pq.cpp @@ -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.dictionary_view() + ->vpq_codebooks.as_matrix_view() .pq_code_book; cuvs::core::to_dlpack(pq_mdspan, pq_codebook); } else { @@ -267,7 +267,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.dictionary_view() + ->vpq_codebooks.as_matrix_view() .vq_code_book; cuvs::core::to_dlpack(pq_mdspan, vq_codebook); } else { diff --git a/cpp/include/cuvs/neighbors/common.hpp b/cpp/include/cuvs/neighbors/common.hpp index 5a5270f5b0..ecfce392e4 100644 --- a/cpp/include/cuvs/neighbors/common.hpp +++ b/cpp/include/cuvs/neighbors/common.hpp @@ -153,15 +153,16 @@ enum class MergeStrategy { /** * @brief Spec-based `dataset` / `dataset_view`. * - * `dataset` and `dataset_view` are single generic templates with zero - * per-kind dispatch inside them: every member is a one-line forward to `spec_type::get_*(...)`, - * and all kind-specific logic lives in the per-kind Spec structs below (`empty_dataset_spec`, - * `padded_dataset_spec`, `standard_dataset_spec`, `vpq_dataset_spec`), which `dataset`/ - * `dataset_view` never name or branch on. `dataset` and `dataset_view` are deliberately two - * independent, non-inheriting types (no shared_ptr, no "sometimes owning" object): `dataset` holds - * owning storage (mdarray-shaped), `dataset_view` holds the corresponding view storage - * (mdspan-shaped). The same `get_n_rows`/`get_dim` spec functions serve both, since - * `raft::mdarray`/`raft::mdspan` both expose `.extent(r)`. + * `dataset` and `dataset_view` are single generic templates that know + * nothing about any particular kind of dataset. They hold exactly one payload (`data_type` / + * `view_type`, chosen by the spec) and expose only what every dataset has: `n_rows()`, `dim()`, + * `as_matrix_view()`, `as_dataset_view()` and `data()`. Each is a one-line forward to one of the + * three spec functions `get_data_view()`, `get_n_rows()` and `get_dim()`. Anything else a kind + * needs (e.g. VPQ codebooks, BBQ quantizers) is state and methods of that kind's payload type, + * reached through `data()`; `dataset`/`dataset_view` never name or branch on it. `dataset` and + * `dataset_view` are deliberately two independent, non-inheriting types (no shared_ptr, no + * "sometimes owning" object): `dataset` holds the owning payload, `dataset_view` the + * corresponding non-owning payload. */ template @@ -170,18 +171,6 @@ struct dataset; template struct dataset_view; -/** - * A spec defines a dictionary iff it needs a second storage slot to interpret the data (e.g. PQ - * codebooks). Non-compressed specs declare `dictionary_type = std::monostate` -- the same - * vocabulary type for "no dictionary," not just an omitted member -- so `dataset`/`dataset_view` - * never need to branch on whether the slot exists; they just always have one, sometimes empty. - */ -template -concept compressed_dataset_spec = requires { - typename SpecT::dictionary_type; - typename SpecT::dictionary_view_type; -} && !std::is_same_v; - namespace detail { // Default owning/view accessors for public dataset aliases. @@ -328,38 +317,134 @@ template struct dense_dataset_spec_impl { template struct apply { - using value_type = std::remove_cv_t; - using index_type = std::remove_cv_t; - using MatrixT = dense_owning_matrix; - using ViewT = dense_view_matrix; - using data_type = dense_row_major_dataset_owning_storage; - using view_type = dense_row_major_dataset_view_storage; - using dictionary_type = std::monostate; - using dictionary_view_type = std::monostate; + using value_type = std::remove_cv_t; + using index_type = std::remove_cv_t; + using MatrixT = dense_owning_matrix; + using ViewT = dense_view_matrix; + using data_type = dense_row_major_dataset_owning_storage; + using view_type = dense_row_major_dataset_view_storage; [[nodiscard]] static auto get_data_view(data_type const& data) noexcept -> view_type { return view_type(data.view(), data.dim()); } - template - [[nodiscard]] static auto get_n_rows(AnyDatasetOrView const& data) noexcept -> index_type + template + [[nodiscard]] static auto get_n_rows(AnyStorage const& data) noexcept -> index_type { return data.n_rows(); } - template - [[nodiscard]] static auto get_dim(AnyDatasetOrView const& data, dictionary_type const&) noexcept - -> uint32_t + template + [[nodiscard]] static auto get_dim(AnyStorage const& data) noexcept -> uint32_t { return data.dim(); } - [[nodiscard]] static auto get_dictionary_view(dictionary_type const&) noexcept - -> dictionary_view_type - { - return {}; - } }; }; +// ----------------------------------------------------------------------------- +// vpq payloads: everything VPQ-specific lives here, not in dataset/dataset_view. +// ----------------------------------------------------------------------------- + +/** Read-only helpers derived from the codebook shapes; shared by the owning and view payloads. + * `Derived` provides `vq_code_book`, `pq_code_book` and the codes' `extent(r)`. */ +template +struct vpq_codebook_helpers { + /** Logical dimension: it comes from the VQ codebook, not from the encoded rows (row padding + * makes the encoded-row width ambiguous as a dimension). */ + [[nodiscard]] auto dim() const noexcept -> uint32_t + { + return static_cast(self().vq_code_book.extent(1)); + } + [[nodiscard]] auto vq_n_centers() const noexcept -> uint32_t + { + return static_cast(self().vq_code_book.extent(0)); + } + [[nodiscard]] auto pq_n_centers() const noexcept -> uint32_t + { + return static_cast(self().pq_code_book.extent(0)); + } + [[nodiscard]] auto pq_len() const noexcept -> uint32_t + { + return static_cast(self().pq_code_book.extent(1)); + } + [[nodiscard]] 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 bits = 0; + while (pq_width > 1) { + bits++; + pq_width >>= 1; + } + return bits; +#endif + } + [[nodiscard]] auto pq_dim() const noexcept -> uint32_t + { + return raft::div_rounding_up_unsafe(dim(), pq_len()); + } + [[nodiscard]] auto encoded_row_length() const noexcept -> uint32_t + { + return static_cast(self().extent(1)); + } + + private: + [[nodiscard]] auto self() const noexcept -> Derived const& + { + return static_cast(*this); + } +}; + +/** Owning VPQ payload: the encoded rows (it *is* the `uint8_t` codes mdarray) plus the VQ and PQ + * codebooks. `Accessor` drives both codebook and code residency. */ +template +struct vpq_owning_storage : public vpq_data_matrix, + public vpq_codebook_helpers> { + using codes_type = vpq_data_matrix; + using vq_book_type = vpq_vq_book_matrix; + using pq_book_type = vpq_vq_book_matrix; + + vq_book_type vq_code_book; + pq_book_type pq_code_book; + + vpq_owning_storage(codes_type&& codes, vq_book_type&& vq_codes, pq_book_type&& pq_codes) noexcept + : codes_type{std::move(codes)}, + vq_code_book{std::move(vq_codes)}, + pq_code_book{std::move(pq_codes)} + { + } +}; + +/** Non-owning VPQ payload: a view of the encoded rows plus views of the VQ and PQ codebooks. */ +template +struct vpq_view_storage : public raft::mdspan, + raft::row_major, + dataset_view_accessor_for_owning>, + public vpq_codebook_helpers> { + using codes_view_type = raft::mdspan, + raft::row_major, + dataset_view_accessor_for_owning>; + using vq_book_view_type = + typename vpq_owning_storage::vq_book_type::const_view_type; + using pq_book_view_type = + typename vpq_owning_storage::pq_book_type::const_view_type; + + vq_book_view_type vq_code_book; + pq_book_view_type pq_code_book; + + vpq_view_storage() noexcept = default; + vpq_view_storage(codes_view_type codes, + vq_book_view_type vq_codes, + pq_book_view_type pq_codes) noexcept + : codes_view_type(codes), vq_code_book(vq_codes), pq_code_book(pq_codes) + { + } +}; + } // namespace detail // ----------------------------------------------------------------------------- @@ -372,12 +457,10 @@ struct empty_dataset_spec { template struct apply { - using value_type = std::remove_cv_t; - using index_type = std::remove_cv_t; - using data_type = detail::empty_dataset_storage; - using view_type = detail::empty_dataset_storage; - using dictionary_type = std::monostate; - using dictionary_view_type = std::monostate; + using value_type = std::remove_cv_t; + using index_type = std::remove_cv_t; + using data_type = detail::empty_dataset_storage; + using view_type = detail::empty_dataset_storage; [[nodiscard]] static auto get_data_view(data_type const& data) noexcept -> view_type { @@ -387,16 +470,10 @@ struct empty_dataset_spec { { return static_cast(data.n_rows()); } - [[nodiscard]] static auto get_dim(data_type const& data, dictionary_type const&) noexcept - -> uint32_t + [[nodiscard]] static auto get_dim(data_type const& data) noexcept -> uint32_t { return data.dim(); } - [[nodiscard]] static auto get_dictionary_view(dictionary_type const&) noexcept - -> dictionary_view_type - { - return {}; - } }; }; @@ -415,9 +492,8 @@ struct standard_dataset_spec { }; /** `Accessor` drives both codebook and code residency, mirroring today's - * single-`Accessor`-per-VPQ-dataset design (`vpq_vq_book_matrix`/`vpq_data_matrix` are both keyed - * off one `Accessor`). Data = encoded rows (uint8_t codes); dictionary = {vq_code_book, - * pq_code_book}. Inlined directly (unlike padded/standard) since no second tag shares this body. */ + * single-`Accessor`-per-VPQ-dataset design. The payload (`detail::vpq_owning_storage` / + * `detail::vpq_view_storage`) holds the encoded rows and the VQ/PQ codebooks. */ template struct vpq_dataset_spec { using accessor_type = Accessor; @@ -428,94 +504,22 @@ struct vpq_dataset_spec { using index_type = std::remove_cv_t; using math_type = MathT; - using data_type = detail::vpq_data_matrix; - using view_type = raft::mdspan, - raft::row_major, - detail::dataset_view_accessor_for_owning>; - - using vq_book_type = detail::vpq_vq_book_matrix; - using pq_book_type = detail::vpq_vq_book_matrix; - - 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; - - [[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 pq_n_centers() const noexcept -> uint32_t - { - return static_cast(pq_code_book.extent(0)); - } - [[nodiscard]] auto pq_len() const noexcept -> uint32_t - { - return static_cast(pq_code_book.extent(1)); - } - [[nodiscard]] 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 bits = 0; - while (pq_width > 1) { - bits++; - pq_width >>= 1; - } - return bits; -#endif - } - [[nodiscard]] auto pq_dim() const noexcept -> uint32_t - { - return raft::div_rounding_up_unsafe(dim(), pq_len()); - } - }; + using data_type = detail::vpq_owning_storage; + using view_type = detail::vpq_view_storage; [[nodiscard]] static auto get_data_view(data_type const& data) noexcept -> view_type { - return data.view(); + return view_type(data.view(), data.vq_code_book.view(), data.pq_code_book.view()); } - template - [[nodiscard]] static auto get_n_rows(AnyExtentShaped const& data) noexcept -> index_type + template + [[nodiscard]] static auto get_n_rows(AnyStorage const& data) noexcept -> index_type { return static_cast(data.extent(0)); } - /* get_dim differs from a plain dense dataset: the dimension comes from the VQ codebook, not - the encoded rows (row padding makes the encoded-row width ambiguous as a dimension). */ - template - [[nodiscard]] static auto get_dim(AnyData const&, dictionary_type const& dict) noexcept - -> uint32_t - { - return static_cast(dict.vq_code_book.extent(1)); - } - template - [[nodiscard]] static auto get_dim(AnyData const&, dictionary_view_type const& dict) noexcept - -> uint32_t - { - return dict.dim(); - } - [[nodiscard]] static auto get_dictionary_view(dictionary_type const& dict) noexcept - -> dictionary_view_type - { - return {dict.vq_code_book.view(), dict.pq_code_book.view()}; - } - [[nodiscard]] static auto get_encoded_row_length(data_type const& data) noexcept -> uint32_t + template + [[nodiscard]] static auto get_dim(AnyStorage const& data) noexcept -> uint32_t { - return static_cast(data.extent(1)); - } - [[nodiscard]] static auto get_encoded_row_length(view_type const& data) noexcept -> uint32_t - { - return static_cast(data.extent(1)); + return data.dim(); } }; }; @@ -524,131 +528,68 @@ struct vpq_dataset_spec { // dataset / dataset_view // ----------------------------------------------------------------------------- -/** Owning dataset: value-held storage (no shared_ptr -- exclusive ownership). Every member is a - * one-line forward to `spec_type::get_*`; all per-kind logic lives in `SpecT`, never inside this - * struct. */ +/** Owning dataset: value-held payload (no shared_ptr -- exclusive ownership). Every member is a + * one-line forward to `spec_type::get_*` or to the payload; all per-kind state and logic lives in + * the spec's `data_type`, never inside this struct. */ 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 dictionary_type = typename spec_type::dictionary_type; - - // Non-compressed: forward constructor args straight to data_type's own constructor (e.g. - // (MatrixT&&, uint32_t logical_dim) for dense, (uint32_t dim) for empty) -- preserves today's - // construction call sites unchanged. + 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; + + // Forward constructor args straight to data_type's own constructor (e.g. (MatrixT&&, uint32_t + // logical_dim) for dense, (uint32_t dim) for empty, (codes&&, vq&&, pq&&) for VPQ). template explicit dataset(Args&&... args) - requires(!compressed_dataset_spec && std::is_constructible_v) - : data_(std::forward(args)...), dictionary_{} - { - } - - // Compressed: data (codes) and dictionary (codebooks) constructed independently. - dataset(data_type&& data, dictionary_type&& dictionary) - requires(compressed_dataset_spec) - : data_(std::move(data)), dictionary_(std::move(dictionary)) + requires(std::is_constructible_v) + : data_(std::forward(args)...) { } [[nodiscard]] auto n_rows() const noexcept -> index_type { return spec_type::get_n_rows(data_); } - [[nodiscard]] auto dim() const noexcept -> uint32_t - { - return spec_type::get_dim(data_, dictionary_); - } + [[nodiscard]] auto dim() const noexcept -> uint32_t { return spec_type::get_dim(data_); } + /** The spec-defined non-owning view of the payload (for dense and VPQ it is an mdspan + * derivative). */ [[nodiscard]] auto as_matrix_view() const noexcept { return spec_type::get_data_view(data_); } - [[nodiscard]] auto dictionary_view() const noexcept - { - return spec_type::get_dictionary_view(dictionary_); - } - [[nodiscard]] auto as_dataset_view() const noexcept -> dataset_view { - return dataset_view(as_matrix_view(), dictionary_view()); + return dataset_view(as_matrix_view()); } - // Move the owning storage out (e.g. to reuse an already-encoded codes matrix while rebuilding - // only the dictionary at a different math_type, as in VPQ's f32->f16 conversion path). - [[nodiscard]] auto release_data() noexcept -> data_type&& { return std::move(data_); } - [[nodiscard]] auto release_dictionary() noexcept -> dictionary_type&& - { - return std::move(dictionary_); - } - - // Dictionary-derived helpers (VPQ: encoded_row_length/vq_n_centers/pq_bits/pq_dim/pq_len/ - // pq_n_centers) forward through dictionary_view() when the dictionary provides them; SFINAE'd - // away for kinds without a dictionary, matching today's VPQ-only surface without dataset<> - // itself branching on which kind it is. - [[nodiscard]] auto encoded_row_length() const noexcept - requires requires(data_type const& d) { spec_type::get_encoded_row_length(d); } - { - return spec_type::get_encoded_row_length(data_); - } - [[nodiscard]] auto vq_n_centers() const noexcept - requires requires(decltype(dictionary_view()) const& d) { d.vq_n_centers(); } - { - return dictionary_view().vq_n_centers(); - } - [[nodiscard]] auto pq_n_centers() const noexcept - requires requires(decltype(dictionary_view()) const& d) { d.pq_n_centers(); } - { - return dictionary_view().pq_n_centers(); - } - [[nodiscard]] auto pq_len() const noexcept - requires requires(decltype(dictionary_view()) const& d) { d.pq_len(); } - { - return dictionary_view().pq_len(); - } - [[nodiscard]] auto pq_bits() const noexcept - requires requires(decltype(dictionary_view()) const& d) { d.pq_bits(); } - { - return dictionary_view().pq_bits(); - } - [[nodiscard]] auto pq_dim() const noexcept - requires requires(decltype(dictionary_view()) const& d) { d.pq_dim(); } - { - return dictionary_view().pq_dim(); - } + /** The owning payload; kind-specific state and methods are reached through it. */ + [[nodiscard]] auto data() const noexcept -> data_type const& { return data_; } + [[nodiscard]] auto data() noexcept -> data_type& { return data_; } private: data_type data_; - [[no_unique_address]] dictionary_type dictionary_; }; -/** Non-owning dataset view: holds only view-shaped storage (mdspan, not mdarray). Deliberately not - * derived from `dataset` -- a view type holds "all view state" with no inheritance and no shared - * ownership tying it to the owning type. Reuses the same `get_n_rows`/`get_dim` spec functions as - * `dataset`, fed view-shaped arguments instead of owning ones. */ +/** Non-owning dataset view: holds only the view-shaped payload. Deliberately not derived from + * `dataset` -- a view type holds "all view state" with no inheritance and no shared ownership tying + * it to the owning type. Reuses the same `get_n_rows`/`get_dim` spec functions as `dataset`, fed + * the view payload instead of the owning one. */ template struct dataset_view { - using spec_type = typename SpecT::template apply; - using value_type = typename spec_type::value_type; - using index_type = typename spec_type::index_type; - using view_type = typename spec_type::view_type; - using dictionary_view_type = typename spec_type::dictionary_view_type; + using spec_type = typename SpecT::template apply; + using value_type = typename spec_type::value_type; + using index_type = typename spec_type::index_type; + using view_type = typename spec_type::view_type; dataset_view() noexcept = default; - // Already-constructed (view_type, dictionary_view_type) pair -- the shape `as_dataset_view()` - // always constructs with, for every kind (dictionary_view_type is std::monostate and - // defaults away when there's no dictionary). Not a template, so it's preferred over the - // forwarding constructor below whenever both could apply. - dataset_view(view_type data_view, dictionary_view_type dictionary_view = {}) noexcept - : data_view_{data_view}, dictionary_view_{dictionary_view} - { - } + // Already-constructed view payload -- the shape `as_dataset_view()` always constructs with, for + // every kind. Not a template, so it's preferred over the forwarding constructor below whenever + // both could apply. + explicit dataset_view(view_type data_view) noexcept : data_view_{data_view} {} // Forward raw constructor args straight to view_type's own constructor (e.g. (ViewT, uint32_t // logical_dim) for dense, (uint32_t dim) for empty) -- preserves today's direct-construction - // call sites (e.g. `device_padded_dataset_view(raw_mdspan, dim)`) unchanged. `view_type` - // is never itself constructible from `(view_type, dictionary_view_type)` (its own constructors - // only take mdspan-shaped args), so this and the plain constructor above never both match the - // same call -- no ambiguity. + // call sites (e.g. `device_padded_dataset_view(raw_mdspan, dim)`) unchanged. template explicit dataset_view(Args&&... args) requires(std::is_constructible_v) - : data_view_(std::forward(args)...), dictionary_view_{} + : data_view_(std::forward(args)...) { } @@ -656,52 +597,15 @@ struct dataset_view { { return spec_type::get_n_rows(data_view_); } - [[nodiscard]] auto dim() const noexcept -> uint32_t - { - return spec_type::get_dim(data_view_, dictionary_view_); - } + [[nodiscard]] auto dim() const noexcept -> uint32_t { return spec_type::get_dim(data_view_); } [[nodiscard]] auto as_matrix_view() const noexcept -> view_type { return data_view_; } - [[nodiscard]] auto dictionary_view() const noexcept -> dictionary_view_type - { - return dictionary_view_; - } - // See dataset<>'s equivalent block: VPQ-only helpers, SFINAE'd away for kinds without a - // dictionary. - [[nodiscard]] auto encoded_row_length() const noexcept - requires requires(view_type const& d) { spec_type::get_encoded_row_length(d); } - { - return spec_type::get_encoded_row_length(data_view_); - } - [[nodiscard]] auto vq_n_centers() const noexcept - requires requires(dictionary_view_type const& d) { d.vq_n_centers(); } - { - return dictionary_view_.vq_n_centers(); - } - [[nodiscard]] auto pq_n_centers() const noexcept - requires requires(dictionary_view_type const& d) { d.pq_n_centers(); } - { - return dictionary_view_.pq_n_centers(); - } - [[nodiscard]] auto pq_len() const noexcept - requires requires(dictionary_view_type const& d) { d.pq_len(); } - { - return dictionary_view_.pq_len(); - } - [[nodiscard]] auto pq_bits() const noexcept - requires requires(dictionary_view_type const& d) { d.pq_bits(); } - { - return dictionary_view_.pq_bits(); - } - [[nodiscard]] auto pq_dim() const noexcept - requires requires(dictionary_view_type const& d) { d.pq_dim(); } - { - return dictionary_view_.pq_dim(); - } + /** The view payload; kind-specific state and methods are reached through it. */ + [[nodiscard]] auto data() const noexcept -> view_type const& { return data_view_; } + [[nodiscard]] auto data() noexcept -> view_type& { return data_view_; } private: view_type data_view_{}; - [[no_unique_address]] dictionary_view_type dictionary_view_{}; }; /** diff --git a/cpp/include/cuvs/preprocessing/quantize/bbq.hpp b/cpp/include/cuvs/preprocessing/quantize/bbq.hpp index e93860f9f7..332994a5f3 100644 --- a/cpp/include/cuvs/preprocessing/quantize/bbq.hpp +++ b/cpp/include/cuvs/preprocessing/quantize/bbq.hpp @@ -244,33 +244,26 @@ void resolve_dequant_factors( namespace neighbors { +namespace detail { + /** - * BBQ doesn't fit the shared Spec-based `dataset`/`dataset_view` shape (one data slot + one - * optional dictionary slot): a BBQ dataset is a small bag of alternate encodings of the *same* - * rows, one per `bbq_code_layout`, selected at query time. So unlike padded/standard/vpq, BBQ - * gets its own standalone owning/view types rather than a `SpecT` plugged into `dataset`/ - * `dataset_view` -- they only need to satisfy `ann_dataset_view` (`n_rows()`/`dim()`) to work - * everywhere a dataset view is expected, plus the handful of trait specializations below. + * BBQ payloads: a BBQ dataset is a small bag of alternate encodings of the *same* rows, one per + * `bbq_code_layout`, selected at query time. The quantizers and the methods that manage them live + * here, in the BBQ payload, not in `dataset`/`dataset_view`; they are reached through `data()`. */ -template -struct bbq_dataset_view; +template +struct bbq_view_storage; -template -struct bbq_dataset { +template +struct bbq_owning_storage { using value_type = DataT; using owning_storage_type = cuvs::preprocessing::quantize::bbq::quantizer; std::vector quantizers; - explicit bbq_dataset(owning_storage_type&& quantizer) noexcept + explicit bbq_owning_storage(owning_storage_type&& quantizer) noexcept { add_quantizer(std::move(quantizer)); } - [[nodiscard]] auto as_dataset_view() const noexcept - -> bbq_dataset_view> - { - return bbq_dataset_view>{ - quantizers}; - } [[nodiscard]] constexpr auto n_rows() const noexcept -> IdxT { return quantizers.size() > 0 ? quantizers[0].n_rows() : 0; @@ -294,16 +287,16 @@ struct bbq_dataset { } }; -template -struct bbq_dataset_view { +template +struct bbq_view_storage { using value_type = DataT; using owning_storage_type = cuvs::preprocessing::quantize::bbq::quantizer; using view_storage_type = cuvs::preprocessing::quantize::bbq::quantizer_view; std::vector quantizers; - bbq_dataset_view() noexcept = default; + bbq_view_storage() noexcept = default; - bbq_dataset_view(const std::vector& quantizers) noexcept + bbq_view_storage(const std::vector& quantizers) noexcept { for (const auto& quantizer : quantizers) { add_quantizer(quantizer); @@ -344,35 +337,61 @@ struct bbq_dataset_view { } }; -template -using device_bbq_dataset = bbq_dataset>; +} // namespace detail + +/** BBQ is just another dataset type: it plugs its payloads into the shared `dataset`/ + * `dataset_view` through a spec, like padded/standard/vpq do. */ +template +struct bbq_dataset_spec { + using accessor_type = Accessor; + + template + struct apply { + using value_type = std::remove_cv_t; + using index_type = std::remove_cv_t; + + using data_type = detail::bbq_owning_storage; + using view_type = detail::bbq_view_storage; + + [[nodiscard]] static auto get_data_view(data_type const& data) noexcept -> view_type + { + return view_type{data.quantizers}; + } + template + [[nodiscard]] static auto get_n_rows(AnyStorage const& data) noexcept -> index_type + { + return data.n_rows(); + } + template + [[nodiscard]] static auto get_dim(AnyStorage const& data) noexcept -> uint32_t + { + return data.dim(); + } + }; +}; template -using device_bbq_dataset_view = bbq_dataset_view>; +using device_bbq_dataset = + dataset>>; template -struct owning_dataset_for_view> { - using type = device_bbq_dataset; -}; +using device_bbq_dataset_view = + dataset_view>>; template struct is_bbq_dataset : std::false_type {}; template -struct is_bbq_dataset> : std::true_type {}; +struct is_bbq_dataset>> : std::true_type {}; template inline constexpr bool is_bbq_dataset_v = is_bbq_dataset::value; template -struct dataset_view_kind_of> { +struct dataset_view_kind_of>> { static constexpr dataset_view_kind value = dataset_view_kind::bbq; }; -template -struct dataset_view_is_device_accessible> - : std::bool_constant {}; - template inline constexpr bool is_device_bbq_dataset_view_v = dataset_view_kind_v == dataset_view_kind::bbq && dataset_view_is_device_accessible_v; diff --git a/cpp/internal/cuvs_internal/preprocessing/bbq_cpu_quantize.hpp b/cpp/internal/cuvs_internal/preprocessing/bbq_cpu_quantize.hpp index 603517e518..0ae1172147 100644 --- a/cpp/internal/cuvs_internal/preprocessing/bbq_cpu_quantize.hpp +++ b/cpp/internal/cuvs_internal/preprocessing/bbq_cpu_quantize.hpp @@ -331,10 +331,10 @@ inline host_quantizer_storage quantize(const float* data, template auto copy_bbq_owning_storage_host_to_device(raft::resources const& res, host_quantizer_storage const& host_storage) -> - typename cuvs::neighbors::device_bbq_dataset::owning_storage_type + typename cuvs::neighbors::device_bbq_dataset::data_type::owning_storage_type { using device_storage = - typename cuvs::neighbors::device_bbq_dataset::owning_storage_type; + typename cuvs::neighbors::device_bbq_dataset::data_type::owning_storage_type; auto stream = raft::resource::get_cuda_stream(res); device_storage device{res, static_cast(host_storage.codes.extent(0)), @@ -389,7 +389,7 @@ auto make_device_bbq_dataset(raft::resources const& res, cuvs::neighbors::device_bbq_dataset device{ copy_bbq_owning_storage_host_to_device(res, host[0])}; for (std::size_t i = 1; i < host.size(); ++i) { - device.add_quantizer(copy_bbq_owning_storage_host_to_device(res, host[i])); + device.data().add_quantizer(copy_bbq_owning_storage_host_to_device(res, host[i])); } return device; } diff --git a/cpp/src/neighbors/cagra.cuh b/cpp/src/neighbors/cagra.cuh index 2532ee55fe..47d6e6a7ac 100644 --- a/cpp/src/neighbors/cagra.cuh +++ b/cpp/src/neighbors/cagra.cuh @@ -320,10 +320,10 @@ auto build(raft::resources const& res, const index_params& params, DatasetViewT RAFT_EXPECTS(effective_params.metric == cuvs::distance::DistanceType::L2Expanded, "cagra::build: a VPQ dataset supports only L2Expanded distance"); RAFT_EXPECTS(dataset.n_rows() > 0, "cagra::build: VPQ dataset must not be empty"); - RAFT_EXPECTS(dataset.pq_bits() == 8, + RAFT_EXPECTS(dataset.data().pq_bits() == 8, "cagra::build: VPQ dataset requires pq_bits == 8, got %u", - dataset.pq_bits()); - auto const pq_len = dataset.pq_len(); + dataset.data().pq_bits()); + auto const pq_len = dataset.data().pq_len(); RAFT_EXPECTS(pq_len == 2 || pq_len == 4 || pq_len == 8, "cagra::build: VPQ dataset requires pq_len in {2, 4, 8}, got %u", pq_len); diff --git a/cpp/src/neighbors/detail/cagra/cagra_build.cuh b/cpp/src/neighbors/detail/cagra/cagra_build.cuh index 545b95fb7e..d89af69dc5 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_build.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_build.cuh @@ -2313,9 +2313,8 @@ void reconstruct_vpq_queries(raft::resources const& res, raft::device_matrix_view output) { auto const data_view = vpq_view.as_matrix_view(); - auto const dict_view = vpq_view.dictionary_view(); const uint32_t dim = vpq_view.dim(); - const uint32_t pq_len = vpq_view.pq_len(); + const uint32_t pq_len = vpq_view.data().pq_len(); const uint32_t output_ld = static_cast(output.extent(1)); const uint32_t threads = std::min(dim, 256u); RAFT_EXPECTS(output_ld >= dim, @@ -2326,9 +2325,9 @@ void reconstruct_vpq_queries(raft::resources const& res, kern_reconstruct_vpq_queries <<>>( data_view.data_handle(), - vpq_view.encoded_row_length(), - dict_view.vq_code_book.data_handle(), - dict_view.pq_code_book.data_handle(), + vpq_view.data().encoded_row_length(), + data_view.vq_code_book.data_handle(), + data_view.pq_code_book.data_handle(), dim, pq_len, offset, @@ -2950,7 +2949,7 @@ auto build_from_bbq_dataset(raft::resources const& res, DatasetViewT const& dataset) -> cuvs::neighbors::cagra::index { - RAFT_EXPECTS(!dataset.quantizers.empty(), "cagra::build: the BBQ dataset is empty."); + RAFT_EXPECTS(!dataset.data().quantizers.empty(), "cagra::build: the BBQ dataset is empty."); RAFT_EXPECTS(params.metric == cuvs::distance::DistanceType::L2Expanded || params.metric == cuvs::distance::DistanceType::L2SqrtExpanded || params.metric == cuvs::distance::DistanceType::CosineExpanded || diff --git a/cpp/src/neighbors/detail/cagra/compute_distance_vpq.hpp b/cpp/src/neighbors/detail/cagra/compute_distance_vpq.hpp index f220d476bf..281a9f84f3 100644 --- a/cpp/src/neighbors/detail/cagra/compute_distance_vpq.hpp +++ b/cpp/src/neighbors/detail/cagra/compute_distance_vpq.hpp @@ -43,8 +43,8 @@ struct vpq_descriptor_spec : public instance_spec { // `DatasetT` here is the non-owning dataset_view (not the owning dataset) -- callers pass the // view directly rather than reaching back through a `.dset()`-style owner pointer, so - // classification and the codebook element type are read off the view via its own dictionary - // state (`is_vpq_dataset_view_v`, `DatasetT::value_type`), not an owning-only trait/typedef. + // classification and the codebook element type are read off the view via its own VPQ payload + // (`is_vpq_dataset_view_v`, `DatasetT::value_type`), not an owning-only trait/typedef. template constexpr static inline auto accepts_dataset() -> std::enable_if_t, bool> @@ -66,12 +66,11 @@ struct vpq_descriptor_spec : public instance_spec { const DistanceT* dataset_norms = nullptr) -> host_type { auto const data_view = dataset.as_matrix_view(); - auto const dict_view = dataset.dictionary_view(); return init_(params, data_view.data_handle(), static_cast(data_view.extent(1)), - dict_view.vq_code_book.data_handle(), - dict_view.pq_code_book.data_handle(), + data_view.vq_code_book.data_handle(), + data_view.pq_code_book.data_handle(), IndexT(dataset.n_rows()), dataset.dim()); } @@ -85,9 +84,8 @@ struct vpq_descriptor_spec : public instance_spec { if (params.team_size != 0 && TeamSize != params.team_size) { return -1.0; } if (cuvs::distance::DistanceType::L2Expanded != metric) { return -1.0; } // Match codebook params - auto const dict_view = dataset.dictionary_view(); - if (dict_view.pq_bits() != PqBits) { return -1.0; } - if (dict_view.pq_len() != PqLen) { return -1.0; } + if (dataset.data().pq_bits() != PqBits) { return -1.0; } + if (dataset.data().pq_len() != PqLen) { return -1.0; } if (select_supported_vpq_smem_dtype(params) != SmemDType) { return -1.0; } // Keep auto-selection on the tuned VPQ diagonal while allowing explicit team_size requests to // use the expanded team_size / dataset_block_dim grid. diff --git a/cpp/src/neighbors/detail/cagra/factory.cuh b/cpp/src/neighbors/detail/cagra/factory.cuh index f82f3d998f..ec1a6f4313 100644 --- a/cpp/src/neighbors/detail/cagra/factory.cuh +++ b/cpp/src/neighbors/detail/cagra/factory.cuh @@ -92,7 +92,7 @@ struct key { }; // `DatasetT` here is the non-owning dataset_view passed in by the search path, so all state comes -// off the view's own `as_matrix_view()`/`dictionary_view()`, not owning-only members. +// off the view's own `as_matrix_view()`, not owning-only members. template auto make_key(const cagra::search_params& params, const DatasetT& dataset, @@ -121,12 +121,11 @@ auto make_key(const cagra::search_params& params, -> std::enable_if_t, key> { auto const data_view = dataset.as_matrix_view(); - auto const dict_view = dataset.dictionary_view(); return key{reinterpret_cast(data_view.data_handle()), reinterpret_cast(dataset_norms), uint64_t(dataset.n_rows()), dataset.dim(), - uint32_t(reinterpret_cast(dict_view.pq_code_book.data_handle()) >> 6), + uint32_t(reinterpret_cast(data_view.pq_code_book.data_handle()) >> 6), uint32_t(params.team_size), uint32_t(metric), uint32_t(params.smem_dtype)}; diff --git a/cpp/src/neighbors/detail/cagra/graph_shared.cu b/cpp/src/neighbors/detail/cagra/graph_shared.cu index 8dcb4d8b6f..5e362f8f5f 100644 --- a/cpp/src/neighbors/detail/cagra/graph_shared.cu +++ b/cpp/src/neighbors/detail/cagra/graph_shared.cu @@ -257,12 +257,14 @@ auto select_sort_quantizers(cuvs::neighbors::device_bbq_dataset_view @@ -285,7 +287,7 @@ void sort_knn_graph_bbq_impl(raft::resources const& res, { namespace bbq = cuvs::preprocessing::quantize::bbq; - RAFT_EXPECTS(!dataset.quantizers.empty(), "the BBQ dataset holds no quantizer"); + RAFT_EXPECTS(!dataset.data().quantizers.empty(), "the BBQ dataset holds no quantizer"); RAFT_EXPECTS(dataset.n_rows() == knn_graph.extent(0), "dataset size is expected to have the same number of graph index size"); RAFT_EXPECTS(metric == cuvs::distance::DistanceType::L2Expanded || diff --git a/cpp/src/neighbors/detail/dataset_serialize.hpp b/cpp/src/neighbors/detail/dataset_serialize.hpp index 1743230cd3..740a16744f 100644 --- a/cpp/src/neighbors/detail/dataset_serialize.hpp +++ b/cpp/src/neighbors/detail/dataset_serialize.hpp @@ -436,8 +436,8 @@ auto deserialize_vpq(raft::resources const& res, std::istream& is) raft::deserialize_mdspan(res, is, data.view()); using owning_t = device_vpq_dataset; - typename owning_t::dictionary_type dictionary{std::move(vq_code_book), std::move(pq_code_book)}; - return std::make_unique(std::move(data), std::move(dictionary)); + return std::make_unique( + std::move(data), std::move(vq_code_book), std::move(pq_code_book)); } template diff --git a/cpp/src/neighbors/detail/nn_descent.cuh b/cpp/src/neighbors/detail/nn_descent.cuh index b1260784b1..7aa70b4800 100644 --- a/cpp/src/neighbors/detail/nn_descent.cuh +++ b/cpp/src/neighbors/detail/nn_descent.cuh @@ -2485,11 +2485,11 @@ void GNND::local_join( // Both kernels take the same (document, query) pair, so there is no symmetric/asymmetric split // here: a single quantizer just means the same one on both operands, which is exactly what // SelfJoin encodes. Picking the two quantizers is all that differs. - const bool self_join = dataset.quantizers.size() == 1; - const bool has_1b = dataset.has_layout(bbq_code_layout::packed_1b); - const bool has_4b = dataset.has_layout(bbq_code_layout::packed_4b); - const bool has_2bt = dataset.has_layout(bbq_code_layout::transposed_2b); - const bool has_4bt = dataset.has_layout(bbq_code_layout::transposed_4b); + const bool self_join = dataset.data().quantizers.size() == 1; + const bool has_1b = dataset.data().has_layout(bbq_code_layout::packed_1b); + const bool has_4b = dataset.data().has_layout(bbq_code_layout::packed_4b); + const bool has_2bt = dataset.data().has_layout(bbq_code_layout::transposed_2b); + const bool has_4bt = dataset.data().has_layout(bbq_code_layout::transposed_4b); // Asymmetric: a packed_4b query selects the tensor-core path, a transposed query the SIMT one. // Only packed_1b promotes to the tensor-core path; transposed_2b is SIMT-only (it would need @@ -2501,13 +2501,14 @@ void GNND::local_join( "packed_1b x packed_4b (tensor core); packed_1b x transposed_2b, " "packed_1b x transposed_4b, transposed_2b x transposed_4b (SIMT)."); auto quantizer_query = - self_join ? dataset.quantizers[0] - : (tc_pair ? dataset.get_quantizer(bbq_code_layout::packed_4b) - : (has_4bt ? dataset.get_quantizer(bbq_code_layout::transposed_4b) - : dataset.get_quantizer(bbq_code_layout::transposed_2b))); - auto quantizer_document = self_join ? dataset.quantizers[0] - : has_1b ? dataset.get_quantizer(bbq_code_layout::packed_1b) - : dataset.get_quantizer(bbq_code_layout::transposed_2b); + self_join + ? dataset.data().quantizers[0] + : (tc_pair ? dataset.data().get_quantizer(bbq_code_layout::packed_4b) + : (has_4bt ? dataset.data().get_quantizer(bbq_code_layout::transposed_4b) + : dataset.data().get_quantizer(bbq_code_layout::transposed_2b))); + auto quantizer_document = self_join ? dataset.data().quantizers[0] + : has_1b ? dataset.data().get_quantizer(bbq_code_layout::packed_1b) + : dataset.data().get_quantizer(bbq_code_layout::transposed_2b); // stage_tile_simt / stage_promoted_tile cast code buffers to uint32_t*, so every plane stride // must be 4-byte aligned. @@ -2982,8 +2983,8 @@ void build(raft::resources const& res, cuvs::neighbors::device_bbq_dataset_view dataset, index& idx) { - RAFT_EXPECTS(dataset.quantizers.size() > 0, "BBQ dataset must not be empty."); - auto front_quantizer = dataset.quantizers[0]; + RAFT_EXPECTS(dataset.data().quantizers.size() > 0, "BBQ dataset must not be empty."); + auto front_quantizer = dataset.data().quantizers[0]; cuvs::common::nvtx::range fun_scope( "neighbors::nn_descent::detail::build-bbq(%zu, %zu, %zu, %zu, %zu)", size_t(dataset.n_rows()), @@ -3001,7 +3002,7 @@ void build(raft::resources const& res, "BBQ dataset metric does not match the NN-Descent metric."); // packed_4b is the only layout dispatched to local_join_kernel_bbq_wmma, and the int4 MMA that // kernel is built around only exists from sm_75 on. - if (dataset.has_layout(bbq_code_layout::packed_4b)) { + if (dataset.data().has_layout(bbq_code_layout::packed_4b)) { auto kernel = local_join_kernel_bbq_wmma build( const uint32_t n_rows = dataset.extent(0); using vpq_owning_t = cuvs::neighbors::device_vpq_dataset; - typename vpq_owning_t::dictionary_type vpq_dictionary{ - raft::make_device_matrix(res, 0, 0), - std::move(pq_codebook)}; - auto quantizer = cuvs::preprocessing::quantize::pq::quantizer( + auto quantizer = cuvs::preprocessing::quantize::pq::quantizer( pq_params, vpq_owning_t(raft::make_device_matrix(res, 0, 0), - std::move(vpq_dictionary))); + raft::make_device_matrix(res, 0, 0), + std::move(pq_codebook))); const int64_t codes_rowlen = cuvs::preprocessing::quantize::pq::get_quantized_dim(pq_params); quantized_vectors = raft::make_device_matrix(res, n_rows, codes_rowlen); diff --git a/cpp/src/neighbors/scann/detail/scann_build.cuh b/cpp/src/neighbors/scann/detail/scann_build.cuh index 634a8d6567..3bf0fa6b30 100644 --- a/cpp/src/neighbors/scann/detail/scann_build.cuh +++ b/cpp/src/neighbors/scann/detail/scann_build.cuh @@ -437,7 +437,7 @@ index build( // Codebooks from VPQ have the shape [subspace idx, subspace dim, code] // This converts the codebook into matrix format for easy interoperability // with open-source ScaNN search - auto full_codebook_view = pq_quantizer.vpq_codebooks.dictionary_view().pq_code_book; + auto full_codebook_view = pq_quantizer.vpq_codebooks.as_matrix_view().pq_code_book; raft::linalg::map_offset( res, diff --git a/cpp/src/preprocessing/quantize/detail/pq.cuh b/cpp/src/preprocessing/quantize/detail/pq.cuh index c52a8a8790..14d5490aca 100644 --- a/cpp/src/preprocessing/quantize/detail/pq.cuh +++ b/cpp/src/preprocessing/quantize/detail/pq.cuh @@ -193,8 +193,8 @@ quantizer build( res, filled_params, dataset, raft::make_const_mdspan(vq_code_book.view())); } using owning_t = cuvs::neighbors::device_vpq_dataset; - typename owning_t::dictionary_type dictionary{std::move(vq_code_book), std::move(pq_code_book)}; - return {filled_params, owning_t(std::move(empty_codes), std::move(dictionary))}; + return {filled_params, + owning_t(std::move(empty_codes), std::move(vq_code_book), std::move(pq_code_book))}; } template @@ -217,7 +217,7 @@ void transform( RAFT_EXPECTS(quantizer.params_quantizer.pq_bits >= 4 && quantizer.params_quantizer.pq_bits <= 16, "PQ bits must be within [4, 16]"); // Encode dataset - auto vq_centers = quantizer.vpq_codebooks.dictionary_view().vq_code_book; + auto vq_centers = quantizer.vpq_codebooks.as_matrix_view().vq_code_book; auto vq_labels_view = raft::make_device_vector_view(nullptr, 0); if (vq_labels.has_value()) { vq_labels_view = vq_labels.value(); } @@ -226,7 +226,7 @@ void transform( res, to_vpq_params(quantizer.params_quantizer), dataset, - quantizer.vpq_codebooks.dictionary_view().pq_code_book, + quantizer.vpq_codebooks.as_matrix_view().pq_code_book, vq_centers, vq_labels_view, pq_codes_out); @@ -235,7 +235,7 @@ void transform( res, to_vpq_params(quantizer.params_quantizer), dataset, - quantizer.vpq_codebooks.dictionary_view().pq_code_book, + quantizer.vpq_codebooks.as_matrix_view().pq_code_book, vq_centers, vq_labels_view, pq_codes_out); @@ -356,7 +356,7 @@ void inverse_transform( "Codes matrix doesn't have the correct number of columns"); RAFT_EXPECTS(quant.params_quantizer.pq_bits >= 4 && quant.params_quantizer.pq_bits <= 16, "PQ bits must be within [4, 16]"); - auto const quant_dict = quant.vpq_codebooks.dictionary_view(); + auto const quant_dict = quant.vpq_codebooks.as_matrix_view(); reconstruct_vectors(res, quant.params_quantizer, codes, @@ -371,17 +371,18 @@ template void vpq_convert_math_type( const raft::resources& res, const cuvs::neighbors::device_vpq_dataset& src, - typename cuvs::neighbors::device_vpq_dataset::dictionary_type& dst_dict) + raft::device_matrix_view dst_vq_code_book, + raft::device_matrix_view dst_pq_code_book) { - auto const src_dict = src.dictionary_view(); + auto const src_view = src.as_matrix_view(); raft::linalg::map(res, - dst_dict.vq_code_book.view(), + dst_vq_code_book, cuvs::spatial::knn::detail::utils::mapping{}, - src_dict.vq_code_book); + src_view.vq_code_book); raft::linalg::map(res, - dst_dict.pq_code_book.view(), + dst_pq_code_book, cuvs::spatial::knn::detail::utils::mapping{}, - src_dict.pq_code_book); + src_view.pq_code_book); } inline auto make_pq_params_from_vpq(const cuvs::neighbors::vpq_params& in_params, @@ -440,8 +441,7 @@ auto vpq_build(const raft::resources& res, true); using owning_t = cuvs::neighbors::device_vpq_dataset; - typename owning_t::dictionary_type dictionary{std::move(vq_code_book), std::move(pq_code_book)}; - return owning_t(std::move(codes), std::move(dictionary)); + return owning_t(std::move(codes), std::move(vq_code_book), std::move(pq_code_book)); } template @@ -449,13 +449,15 @@ auto vpq_build_half(const raft::resources& res, const cuvs::neighbors::vpq_params& params, const DatasetT& dataset) -> cuvs::neighbors::device_vpq_dataset { - auto old_type = vpq_build(res, params, dataset); - using new_owning_t = cuvs::neighbors::device_vpq_dataset; - auto const old_dict = old_type.dictionary_view(); - typename new_owning_t::dictionary_type new_dict{ - raft::make_device_mdarray(res, old_dict.vq_code_book.extents()), - raft::make_device_mdarray(res, old_dict.pq_code_book.extents())}; - vpq_convert_math_type(res, old_type, new_dict); - return new_owning_t(old_type.release_data(), std::move(new_dict)); + auto old_type = vpq_build(res, params, dataset); + using new_owning_t = cuvs::neighbors::device_vpq_dataset; + auto const old_view = old_type.as_matrix_view(); + auto new_vq_code_book = raft::make_device_mdarray(res, old_view.vq_code_book.extents()); + auto new_pq_code_book = raft::make_device_mdarray(res, old_view.pq_code_book.extents()); + vpq_convert_math_type( + res, old_type, new_vq_code_book.view(), new_pq_code_book.view()); + // Reuse the already-encoded codes (the owning payload's mdarray base) with the new codebooks. + return new_owning_t( + std::move(old_type.data()), std::move(new_vq_code_book), std::move(new_pq_code_book)); } } // namespace cuvs::preprocessing::quantize::pq::detail diff --git a/cpp/tests/neighbors/ann_cagra/test_iterative_cagra_q.cu b/cpp/tests/neighbors/ann_cagra/test_iterative_cagra_q.cu index ea0cfcabc5..e0f250a2b0 100644 --- a/cpp/tests/neighbors/ann_cagra/test_iterative_cagra_q.cu +++ b/cpp/tests/neighbors/ann_cagra/test_iterative_cagra_q.cu @@ -179,7 +179,7 @@ class CagraQBuildTest : public CagraQCompressedTestBase, TEST_P(CagraQBuildTest, BuildsAndSearchesAFreshlyCompressedDataset) { auto compressed = compress(res_, dataset(), params_.pq_dim); - ASSERT_EQ(compressed.pq_len(), static_cast(params_.dim / params_.pq_dim)); + ASSERT_EQ(compressed.data().pq_len(), static_cast(params_.dim / params_.pq_dim)); auto idx = cagra::build(res_, iterative_params(), compressed.as_dataset_view()); ASSERT_EQ(idx.size(), params_.n_rows); @@ -272,7 +272,7 @@ TEST_F(CagraQContractTest, RejectsMetricOtherThanL2Expanded) TEST_F(CagraQContractTest, RejectsPqBitsOtherThan8) { auto compressed = compress(res_, dataset(), 32, /* pq_bits */ 6); - ASSERT_EQ(compressed.pq_bits(), 6u); + ASSERT_EQ(compressed.data().pq_bits(), 6u); EXPECT_THROW(cagra::build(res_, iterative_params(), compressed.as_dataset_view()), raft::exception); } @@ -280,7 +280,7 @@ TEST_F(CagraQContractTest, RejectsPqBitsOtherThan8) TEST_F(CagraQContractTest, RejectsPqLenOutsideSupportedSet) { auto compressed = compress(res_, dataset(), /* pq_dim */ 4); // pq_len = 64 / 4 = 16 - ASSERT_EQ(compressed.pq_len(), 16u); + ASSERT_EQ(compressed.data().pq_len(), 16u); EXPECT_THROW(cagra::build(res_, iterative_params(), compressed.as_dataset_view()), raft::exception); } @@ -293,8 +293,7 @@ TEST_F(CagraQContractTest, RejectsEmptyDataset) auto vq_code_book = raft::make_device_matrix(res_, 1, width); auto pq_code_book = raft::make_device_matrix(res_, 256, 2); auto codes = raft::make_device_matrix(res_, 0, 4 + dim / 2); - vpq_dataset_t::dictionary_type dictionary{std::move(vq_code_book), std::move(pq_code_book)}; - vpq_dataset_t empty{std::move(codes), std::move(dictionary)}; + vpq_dataset_t empty{std::move(codes), std::move(vq_code_book), std::move(pq_code_book)}; ASSERT_EQ(empty.n_rows(), 0); EXPECT_THROW(cagra::build(res_, iterative_params(), empty.as_dataset_view()), raft::exception); diff --git a/cpp/tests/neighbors/ann_cagra_bbq.cuh b/cpp/tests/neighbors/ann_cagra_bbq.cuh index 9ecc4bd1d4..0ac52e2265 100644 --- a/cpp/tests/neighbors/ann_cagra_bbq.cuh +++ b/cpp/tests/neighbors/ann_cagra_bbq.cuh @@ -261,7 +261,7 @@ class AnnCagraBbqTest : public ::testing::TestWithParam { ASSERT_EQ(restored.graph_degree(), graph_index.graph_degree()); EXPECT_EQ(restored.metric(), graph_index.metric()); EXPECT_EQ(restored.dataset().n_rows(), 0); - EXPECT_TRUE(restored.dataset().quantizers.empty()); + EXPECT_TRUE(restored.dataset().data().quantizers.empty()); auto original = cagra::update_dataset(handle_, std::move(graph_index), device_padded.view); auto reattached = cagra::update_dataset(handle_, std::move(restored), device_padded.view); @@ -284,7 +284,7 @@ class AnnCagraBbqTest : public ::testing::TestWithParam { ASSERT_EQ(index.graph_size(), static_cast(ps.n_rows)); ASSERT_EQ(index.graph_degree(), static_cast(ps.graph_degree)); EXPECT_EQ(index.dataset().n_rows(), 0); - EXPECT_TRUE(index.dataset().quantizers.empty()); + EXPECT_TRUE(index.dataset().data().quantizers.empty()); } /** Only NN-descent graph construction and the four BBQ metrics are accepted. */ @@ -331,7 +331,7 @@ class AnnCagraBbqTest : public ::testing::TestWithParam { private: raft::resources handle_; - rmm::cuda_stream_view stream_; + cuda::stream_ref stream_; AnnCagraBbqInputs ps; rmm::device_uvector database; rmm::device_uvector search_queries; diff --git a/cpp/tests/neighbors/ann_nn_descent_bbq.cuh b/cpp/tests/neighbors/ann_nn_descent_bbq.cuh index 295cae935f..675b6399ba 100644 --- a/cpp/tests/neighbors/ann_nn_descent_bbq.cuh +++ b/cpp/tests/neighbors/ann_nn_descent_bbq.cuh @@ -153,7 +153,7 @@ class AnnNNDescentBbqTest : public ::testing::TestWithParam database; }; diff --git a/cpp/tests/neighbors/ann_scann.cuh b/cpp/tests/neighbors/ann_scann.cuh index 8809e203ef..02c145e203 100644 --- a/cpp/tests/neighbors/ann_scann.cuh +++ b/cpp/tests/neighbors/ann_scann.cuh @@ -186,10 +186,9 @@ class scann_test : public ::testing::TestWithParam { auto empty_data = raft::make_device_matrix(handle_, 0, 0); using vpq_owning_t = cuvs::neighbors::device_vpq_dataset; - typename vpq_owning_t::dictionary_type vpq_dictionary{std::move(vq_codebook), - std::move(pq_codebook_copy)}; cuvs::preprocessing::quantize::pq::quantizer quantizer{ - pq_params, vpq_owning_t(std::move(empty_data), std::move(vpq_dictionary))}; + pq_params, + vpq_owning_t(std::move(empty_data), std::move(vq_codebook), std::move(pq_codebook_copy))}; auto quantized_residuals_device = raft::make_device_matrix(handle_, ps.num_db_vecs, num_subspaces); diff --git a/cpp/tests/neighbors/vpq_utils.cuh b/cpp/tests/neighbors/vpq_utils.cuh index 5fc5d1bba0..ed6a8d9264 100644 --- a/cpp/tests/neighbors/vpq_utils.cuh +++ b/cpp/tests/neighbors/vpq_utils.cuh @@ -50,12 +50,11 @@ void decode_vpq_dataset(raft::device_matrix_view decoded_datase cudaStream_t cuda_stream) { const auto dataset_size = decoded_dataset.extent(0); - auto const dict_view = vpq_dataset.dictionary_view(); auto const data_view = vpq_dataset.as_matrix_view(); RAFT_EXPECTS(data_view.extent(0) == dataset_size, "Dataset sizes mismatch"); - RAFT_EXPECTS(vpq_dataset.pq_bits() == 8, + RAFT_EXPECTS(vpq_dataset.data().pq_bits() == 8, "decode_vpq_dataset currently only supports pq_bits == 8 (got %u)", - vpq_dataset.pq_bits()); + vpq_dataset.data().pq_bits()); constexpr uint32_t block_size = 256; constexpr uint32_t warp_size = 32; @@ -65,11 +64,11 @@ void decode_vpq_dataset(raft::device_matrix_view decoded_datase decode_vpq_dataset_kernel <<>>(decoded_dataset.data_handle(), decoded_dataset.stride(0), - dict_view.vq_code_book.data_handle(), - dict_view.vq_code_book.stride(0), - dict_view.pq_code_book.data_handle(), - vpq_dataset.pq_len(), - 1u << vpq_dataset.pq_bits(), + data_view.vq_code_book.data_handle(), + data_view.vq_code_book.stride(0), + data_view.pq_code_book.data_handle(), + vpq_dataset.data().pq_len(), + 1u << vpq_dataset.data().pq_bits(), vpq_dataset.dim(), dataset_size, data_view.data_handle(), diff --git a/fern/pages/cpp_api/cpp-api-neighbors-common.md b/fern/pages/cpp_api/cpp-api-neighbors-common.md index 2f8c347e15..457c97c53d 100644 --- a/fern/pages/cpp_api/cpp-api-neighbors-common.md +++ b/fern/pages/cpp_api/cpp-api-neighbors-common.md @@ -64,7 +64,7 @@ enum class MergeStrategy { Spec-based `dataset` / `dataset_view`. -`dataset<T,IdxT,SpecT>` and `dataset_view<T,IdxT,SpecT>` are single generic templates with zero per-kind dispatch inside them: every member is a one-line forward to `spec_type::get_*(...)`, and all kind-specific logic lives in the per-kind Spec structs below (`empty_dataset_spec`, `padded_dataset_spec`, `standard_dataset_spec`, `vpq_dataset_spec`), which `dataset`/ `dataset_view` never name or branch on. `dataset` and `dataset_view` are deliberately two independent, non-inheriting types (no shared_ptr, no "sometimes owning" object): `dataset` holds owning storage (mdarray-shaped), `dataset_view` holds the corresponding view storage (mdspan-shaped). The same `get_n_rows`/`get_dim` spec functions serve both, since `raft::mdarray`/`raft::mdspan` both expose `.extent(r)`. +`dataset<T,IdxT,SpecT>` and `dataset_view<T,IdxT,SpecT>` are single generic templates that know nothing about any particular kind of dataset. They hold exactly one payload (`data_type` / `view_type`, chosen by the spec) and expose only what every dataset has: `n_rows()`, `dim()`, `as_matrix_view()`, `as_dataset_view()` and `data()`. Each is a one-line forward to one of the three spec functions `get_data_view()`, `get_n_rows()` and `get_dim()`. Anything else a kind needs (e.g. VPQ codebooks, BBQ quantizers) is state and methods of that kind's payload type, reached through `data()`; `dataset`/`dataset_view` never name or branch on it. `dataset` and `dataset_view` are deliberately two independent, non-inheriting types (no shared_ptr, no "sometimes owning" object): `dataset` holds the owning payload, `dataset_view` the corresponding non-owning payload. ```cpp template @@ -76,7 +76,7 @@ struct dataset; `Accessor` drives both codebook and code residency, mirroring today's -single-`Accessor`-per-VPQ-dataset design (`vpq_vq_book_matrix`/`vpq_data_matrix` are both keyed off one `Accessor`). Data = encoded rows (uint8_t codes); dictionary = \{vq_code_book, pq_code_book\}. Inlined directly (unlike padded/standard) since no second tag shares this body. +single-`Accessor`-per-VPQ-dataset design. The payload (`detail::vpq_owning_storage` / `detail::vpq_view_storage`) holds the encoded rows and the VQ/PQ codebooks. ```cpp template @@ -86,9 +86,9 @@ struct vpq_dataset_spec; ### neighbors::dataset -Owning dataset: value-held storage (no shared_ptr -- exclusive ownership). Every member is a +Owning dataset: value-held payload (no shared_ptr -- exclusive ownership). Every member is a -one-line forward to `spec_type::get_*`; all per-kind logic lives in `SpecT`, never inside this struct. +one-line forward to `spec_type::get_*` or to the payload; all per-kind state and logic lives in the spec's `data_type`, never inside this struct. ```cpp template @@ -98,9 +98,9 @@ struct dataset; ### neighbors::dataset_view -Non-owning dataset view: holds only view-shaped storage (mdspan, not mdarray). Deliberately not +Non-owning dataset view: holds only the view-shaped payload. Deliberately not derived from -derived from `dataset` -- a view type holds "all view state" with no inheritance and no shared ownership tying it to the owning type. Reuses the same `get_n_rows`/`get_dim` spec functions as `dataset`, fed view-shaped arguments instead of owning ones. +`dataset` -- a view type holds "all view state" with no inheritance and no shared ownership tying it to the owning type. Reuses the same `get_n_rows`/`get_dim` spec functions as `dataset`, fed the view payload instead of the owning one. ```cpp template diff --git a/fern/pages/cpp_api/cpp-api-preprocessing-quantize-bbq.md b/fern/pages/cpp_api/cpp-api-preprocessing-quantize-bbq.md index 6e115bec35..a5656e234e 100644 --- a/fern/pages/cpp_api/cpp-api-preprocessing-quantize-bbq.md +++ b/fern/pages/cpp_api/cpp-api-preprocessing-quantize-bbq.md @@ -154,14 +154,14 @@ dequant_sum_delta from that delta and quantized_component_sums. ## Types - -### neighbors::bbq_dataset_view + +### neighbors::bbq_dataset_spec -BBQ doesn't fit the shared Spec-based `dataset`/`dataset_view` shape (one data slot + one +BBQ is just another dataset type: it plugs its payloads into the shared `dataset`/ -optional dictionary slot): a BBQ dataset is a small bag of alternate encodings of the *same* rows, one per `bbq_code_layout`, selected at query time. So unlike padded/standard/vpq, BBQ gets its own standalone owning/view types rather than a `SpecT` plugged into `dataset`/ `dataset_view` -- they only need to satisfy `ann_dataset_view` (`n_rows()`/`dim()`) to work everywhere a dataset view is expected, plus the handful of trait specializations below. +`dataset_view` through a spec, like padded/standard/vpq do. ```cpp -template -struct bbq_dataset_view; +template +struct bbq_dataset_spec; ``` From 9832a4f26150d1a69e6de2240aae8760d319dd47 Mon Sep 17 00:00:00 2001 From: Howard Huang Date: Thu, 1 Oct 2026 19:43:06 -0700 Subject: [PATCH 08/15] Move dataset/dataset_view out of neighbors/common.hpp into core/dataset.hpp (#2395) Datasets are a core type, not a neighbors one: quantizers (pq.hpp, bbq.hpp) and clustering also need them and should not depend on anything in neighbors. Create cuvs/core/dataset.hpp as the root of the dataset header tree and move the whole dataset section of neighbors/common.hpp into it: dataset and dataset_view, the accessor aliases and dense/empty storage, the empty/padded/standard/vpq specs and their aliases, the kind traits, with_accessor/device_counterpart, the CAGRA row-width helpers, and the make_*_padded_dataset / make_*_standard_dataset factories. neighbors/common.hpp includes the new header, so none of its 67 includers change. This is a pure cut and paste: the 1187 moved lines are byte-identical to the removed block, the namespace is still cuvs::neighbors, and there are no call-site changes. The Fern API reference is regenerated for the new header (fern-api-reference hook). Not done yet (follow-ups): the VPQ and BBQ pieces and the dataset_view_kind enum still live in the moved block and will move into quantize/pq.hpp and quantize/bbq.hpp as children of dataset; the dataset types then move from cuvs::neighbors to cuvs::core. Verified: full rebuild clean; DATASET_C_TEST (9), CAGRA_C_TEST (15), PREPROCESSING_TEST (226), NEIGHBORS_TEST (371), NEIGHBORS_ANN_CAGRA_BBQ_UINT32_TEST (105), NEIGHBORS_ANN_NN_DESCENT_TEST (572), NEIGHBORS_ANN_CAGRA_MERGE_TEST (26), NEIGHBORS_ANN_CAGRA_FLOAT_UINT32_TEST (3594) pass. Co-Authored-By: Claude Sonnet 5.5 --- cpp/include/cuvs/core/dataset.hpp | 1229 +++++++++++++++++ cpp/include/cuvs/neighbors/common.hpp | 1189 +--------------- fern/docs.yml | 2 + fern/pages/cpp_api/cpp-api-core-dataset.md | 129 ++ .../pages/cpp_api/cpp-api-neighbors-common.md | 204 +-- fern/pages/cpp_api/index.md | 1 + 6 files changed, 1404 insertions(+), 1350 deletions(-) create mode 100644 cpp/include/cuvs/core/dataset.hpp 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 new file mode 100644 index 0000000000..aac994b01b --- /dev/null +++ b/cpp/include/cuvs/core/dataset.hpp @@ -0,0 +1,1229 @@ +/* + * 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 // get_device_for_address, copy_matrix +#include // rounding up + +#include +#include + +#include + +#include +#include +#include +#include +#include +#include +#ifdef __cpp_lib_bitops +#include +#endif + +namespace CUVS_EXPORT cuvs { +namespace neighbors { + +/** + * @brief Spec-based `dataset` / `dataset_view`. + * + * `dataset` and `dataset_view` are single generic templates that know + * nothing about any particular kind of dataset. They hold exactly one payload (`data_type` / + * `view_type`, chosen by the spec) and expose only what every dataset has: `n_rows()`, `dim()`, + * `as_matrix_view()`, `as_dataset_view()` and `data()`. Each is a one-line forward to one of the + * three spec functions `get_data_view()`, `get_n_rows()` and `get_dim()`. Anything else a kind + * needs (e.g. VPQ codebooks, BBQ quantizers) is state and methods of that kind's payload type, + * reached through `data()`; `dataset`/`dataset_view` never name or branch on it. `dataset` and + * `dataset_view` are deliberately two independent, non-inheriting types (no shared_ptr, no + * "sometimes owning" object): `dataset` holds the owning payload, `dataset_view` the + * corresponding non-owning payload. + */ + +template +struct dataset; + +template +struct dataset_view; + +namespace detail { + +// Default owning/view accessors for public dataset aliases. +template +using device_owning_accessor = raft::device_accessor>; + +template +using host_owning_accessor = raft::host_accessor>; + +template +using device_view_accessor = raft::device_accessor>; + +template +using host_view_accessor = raft::host_accessor>; + +/** View accessor paired with an owning dataset accessor (same residency). */ +template +using dataset_view_accessor_for_owning = std::conditional_t, + host_view_accessor>; + +/** Owning accessor paired with a view accessor (same residency). */ +template +using dataset_owning_accessor_for_view = std::conditional_t, + host_owning_accessor>; + +// Accessor here is already device_owning_accessor / host_owning_accessor at every +// call site -- exactly the container policy raft::device_mdarray/host_mdarray default to for +// element type DataT -- so pass it straight through instead of re-deriving a +// raft::device_matrix/host_matrix from scratch. +template +using dense_owning_matrix = + raft::mdarray, raft::row_major, Accessor>; + +template +using dense_view_matrix = raft::mdspan, + raft::row_major, + dataset_view_accessor_for_owning>; + +template +using vpq_vq_book_matrix = + raft::mdarray, raft::row_major, Accessor>; + +// VPQ codes are always uint8_t regardless of MathT, so retarget the owning accessor's element +// type instead of re-deriving a device/host matrix; residency is still driven by Accessor. +template +using owning_accessor_with_value_type = std::conditional_t, + host_owning_accessor>; + +template +using vpq_data_matrix = raft::mdarray, + raft::row_major, + owning_accessor_with_value_type>; + +// ----------------------------------------------------------------------------- +// empty +// ----------------------------------------------------------------------------- + +template +struct empty_dataset_storage { + uint32_t suggested_dim{}; + empty_dataset_storage() noexcept = default; + explicit empty_dataset_storage(uint32_t dim) noexcept : suggested_dim(dim) {} + [[nodiscard]] auto n_rows() const noexcept -> IdxT { return 0; } + [[nodiscard]] auto dim() const noexcept -> uint32_t { return suggested_dim; } +}; + +// ----------------------------------------------------------------------------- +// 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 specs. Publicly inherits + * from MatrixT (a `raft::mdarray`) so `view()`/`data_handle()`/`extent()` etc. are reused as-is + * rather than hand-forwarded; `logical_dim_` is the only state this struct adds. + * + * 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 : public MatrixT { + uint32_t logical_dim_; + + // MatrixT (mdarray) also has its own stride(size_t); pull it back into scope since declaring + // our own no-arg stride() below would otherwise hide it entirely (C++ name hiding). + using MatrixT::stride; + + dense_row_major_dataset_owning_storage(MatrixT&& data, uint32_t logical_dim) noexcept + : MatrixT{std::move(data)}, logical_dim_{logical_dim} + { + } + + [[nodiscard]] auto n_rows() const noexcept -> IdxT { return this->extent(0); } + [[nodiscard]] auto dim() const noexcept -> uint32_t { return logical_dim_; } + [[nodiscard]] auto stride() const noexcept -> uint32_t + { + return static_cast(this->extent(1)); + } + // view() and data_handle() are inherited directly from MatrixT (raft::mdarray); no hand-written + // forwarding needed since MatrixT::view() const already returns exactly ViewT. +}; + +template +struct dense_row_major_dataset_view_storage : public ViewT { + uint32_t logical_dim_; + + // ViewT (mdspan) also has its own stride(size_t); pull it back into scope since declaring our + // own no-arg stride() below would otherwise hide it entirely (C++ name hiding), and the body of + // that stride() itself needs to call the inherited one. + using ViewT::stride; + + dense_row_major_dataset_view_storage() noexcept = default; + + explicit dense_row_major_dataset_view_storage(ViewT v) noexcept + : ViewT(v), logical_dim_(static_cast(v.extent(1))) + { + } + + dense_row_major_dataset_view_storage(ViewT v, uint32_t logical_dim) noexcept + : ViewT(v), logical_dim_(logical_dim) + { + } + + [[nodiscard]] auto n_rows() const noexcept -> IdxT { return this->extent(0); } + [[nodiscard]] auto dim() const noexcept -> uint32_t { return logical_dim_; } + [[nodiscard]] auto stride() const noexcept -> uint32_t + { + return static_cast(ViewT::stride(0) > 0 ? ViewT::stride(0) : this->extent(1)); + } +}; + +/** Spec-side implementation shared by `padded_dataset_spec`/`standard_dataset_spec`; those two + * stay distinct top-level types (identical bodies) purely so classification traits can tell them + * apart -- exactly mirroring today's `padded_dataset_container`/`standard_dataset_container`, + * which are likewise two differently-named tags over one shared storage implementation. */ +template +struct dense_dataset_spec_impl { + template + struct apply { + using value_type = std::remove_cv_t; + using index_type = std::remove_cv_t; + using MatrixT = dense_owning_matrix; + using ViewT = dense_view_matrix; + using data_type = dense_row_major_dataset_owning_storage; + using view_type = dense_row_major_dataset_view_storage; + + [[nodiscard]] static auto get_data_view(data_type const& data) noexcept -> view_type + { + return view_type(data.view(), data.dim()); + } + template + [[nodiscard]] static auto get_n_rows(AnyStorage const& data) noexcept -> index_type + { + return data.n_rows(); + } + template + [[nodiscard]] static auto get_dim(AnyStorage const& data) noexcept -> uint32_t + { + return data.dim(); + } + }; +}; + +// ----------------------------------------------------------------------------- +// vpq payloads: everything VPQ-specific lives here, not in dataset/dataset_view. +// ----------------------------------------------------------------------------- + +/** Read-only helpers derived from the codebook shapes; shared by the owning and view payloads. + * `Derived` provides `vq_code_book`, `pq_code_book` and the codes' `extent(r)`. */ +template +struct vpq_codebook_helpers { + /** Logical dimension: it comes from the VQ codebook, not from the encoded rows (row padding + * makes the encoded-row width ambiguous as a dimension). */ + [[nodiscard]] auto dim() const noexcept -> uint32_t + { + return static_cast(self().vq_code_book.extent(1)); + } + [[nodiscard]] auto vq_n_centers() const noexcept -> uint32_t + { + return static_cast(self().vq_code_book.extent(0)); + } + [[nodiscard]] auto pq_n_centers() const noexcept -> uint32_t + { + return static_cast(self().pq_code_book.extent(0)); + } + [[nodiscard]] auto pq_len() const noexcept -> uint32_t + { + return static_cast(self().pq_code_book.extent(1)); + } + [[nodiscard]] 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 bits = 0; + while (pq_width > 1) { + bits++; + pq_width >>= 1; + } + return bits; +#endif + } + [[nodiscard]] auto pq_dim() const noexcept -> uint32_t + { + return raft::div_rounding_up_unsafe(dim(), pq_len()); + } + [[nodiscard]] auto encoded_row_length() const noexcept -> uint32_t + { + return static_cast(self().extent(1)); + } + + private: + [[nodiscard]] auto self() const noexcept -> Derived const& + { + return static_cast(*this); + } +}; + +/** Owning VPQ payload: the encoded rows (it *is* the `uint8_t` codes mdarray) plus the VQ and PQ + * codebooks. `Accessor` drives both codebook and code residency. */ +template +struct vpq_owning_storage : public vpq_data_matrix, + public vpq_codebook_helpers> { + using codes_type = vpq_data_matrix; + using vq_book_type = vpq_vq_book_matrix; + using pq_book_type = vpq_vq_book_matrix; + + vq_book_type vq_code_book; + pq_book_type pq_code_book; + + vpq_owning_storage(codes_type&& codes, vq_book_type&& vq_codes, pq_book_type&& pq_codes) noexcept + : codes_type{std::move(codes)}, + vq_code_book{std::move(vq_codes)}, + pq_code_book{std::move(pq_codes)} + { + } +}; + +/** Non-owning VPQ payload: a view of the encoded rows plus views of the VQ and PQ codebooks. */ +template +struct vpq_view_storage : public raft::mdspan, + raft::row_major, + dataset_view_accessor_for_owning>, + public vpq_codebook_helpers> { + using codes_view_type = raft::mdspan, + raft::row_major, + dataset_view_accessor_for_owning>; + using vq_book_view_type = + typename vpq_owning_storage::vq_book_type::const_view_type; + using pq_book_view_type = + typename vpq_owning_storage::pq_book_type::const_view_type; + + vq_book_view_type vq_code_book; + pq_book_view_type pq_code_book; + + vpq_view_storage() noexcept = default; + vpq_view_storage(codes_view_type codes, + vq_book_view_type vq_codes, + pq_book_view_type pq_codes) noexcept + : codes_view_type(codes), vq_code_book(vq_codes), pq_code_book(pq_codes) + { + } +}; + +} // namespace detail + +// ----------------------------------------------------------------------------- +// Public specs -- the only place per-kind logic lives. +// ----------------------------------------------------------------------------- + +template +struct empty_dataset_spec { + using accessor_type = Accessor; + + template + struct apply { + using value_type = std::remove_cv_t; + using index_type = std::remove_cv_t; + using data_type = detail::empty_dataset_storage; + using view_type = detail::empty_dataset_storage; + + [[nodiscard]] static auto get_data_view(data_type const& data) noexcept -> view_type + { + return data; + } + [[nodiscard]] static auto get_n_rows(data_type const& data) noexcept -> index_type + { + return static_cast(data.n_rows()); + } + [[nodiscard]] static auto get_dim(data_type const& data) noexcept -> uint32_t + { + return data.dim(); + } + }; +}; + +template +struct padded_dataset_spec { + using accessor_type = ContainerPolicy; + template + struct apply : detail::dense_dataset_spec_impl::template apply {}; +}; + +template +struct standard_dataset_spec { + using accessor_type = ContainerPolicy; + template + struct apply : detail::dense_dataset_spec_impl::template apply {}; +}; + +/** `Accessor` drives both codebook and code residency, mirroring today's + * single-`Accessor`-per-VPQ-dataset design. The payload (`detail::vpq_owning_storage` / + * `detail::vpq_view_storage`) holds the encoded rows and the VQ/PQ codebooks. */ +template +struct vpq_dataset_spec { + using accessor_type = Accessor; + + template + struct apply { + using value_type = std::remove_cv_t; + using index_type = std::remove_cv_t; + using math_type = MathT; + + using data_type = detail::vpq_owning_storage; + using view_type = detail::vpq_view_storage; + + [[nodiscard]] static auto get_data_view(data_type const& data) noexcept -> view_type + { + return view_type(data.view(), data.vq_code_book.view(), data.pq_code_book.view()); + } + template + [[nodiscard]] static auto get_n_rows(AnyStorage const& data) noexcept -> index_type + { + return static_cast(data.extent(0)); + } + template + [[nodiscard]] static auto get_dim(AnyStorage const& data) noexcept -> uint32_t + { + return data.dim(); + } + }; +}; + +// ----------------------------------------------------------------------------- +// dataset / dataset_view +// ----------------------------------------------------------------------------- + +/** Owning dataset: value-held payload (no shared_ptr -- exclusive ownership). Every member is a + * one-line forward to `spec_type::get_*` or to the payload; all per-kind state and logic lives in + * the spec's `data_type`, never inside this struct. */ +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; + + // Forward constructor args straight to data_type's own constructor (e.g. (MatrixT&&, uint32_t + // logical_dim) for dense, (uint32_t dim) for empty, (codes&&, vq&&, pq&&) for VPQ). + template + explicit dataset(Args&&... args) + requires(std::is_constructible_v) + : data_(std::forward(args)...) + { + } + + [[nodiscard]] auto n_rows() const noexcept -> index_type { return spec_type::get_n_rows(data_); } + [[nodiscard]] auto dim() const noexcept -> uint32_t { return spec_type::get_dim(data_); } + /** The spec-defined non-owning view of the payload (for dense and VPQ it is an mdspan + * derivative). */ + [[nodiscard]] auto as_matrix_view() const noexcept { return spec_type::get_data_view(data_); } + [[nodiscard]] auto as_dataset_view() const noexcept -> dataset_view + { + return dataset_view(as_matrix_view()); + } + + /** The owning payload; kind-specific state and methods are reached through it. */ + [[nodiscard]] auto data() const noexcept -> data_type const& { return data_; } + [[nodiscard]] auto data() noexcept -> data_type& { return data_; } + + private: + data_type data_; +}; + +/** Non-owning dataset view: holds only the view-shaped payload. Deliberately not derived from + * `dataset` -- a view type holds "all view state" with no inheritance and no shared ownership tying + * it to the owning type. Reuses the same `get_n_rows`/`get_dim` spec functions as `dataset`, fed + * the view payload instead of the owning one. */ +template +struct dataset_view { + using spec_type = typename SpecT::template apply; + using value_type = typename spec_type::value_type; + using index_type = typename spec_type::index_type; + using view_type = typename spec_type::view_type; + + dataset_view() noexcept = default; + + // Already-constructed view payload -- the shape `as_dataset_view()` always constructs with, for + // every kind. Not a template, so it's preferred over the forwarding constructor below whenever + // both could apply. + explicit dataset_view(view_type data_view) noexcept : data_view_{data_view} {} + + // Forward raw constructor args straight to view_type's own constructor (e.g. (ViewT, uint32_t + // logical_dim) for dense, (uint32_t dim) for empty) -- preserves today's direct-construction + // call sites (e.g. `device_padded_dataset_view(raw_mdspan, dim)`) unchanged. + template + explicit dataset_view(Args&&... args) + requires(std::is_constructible_v) + : data_view_(std::forward(args)...) + { + } + + [[nodiscard]] auto n_rows() const noexcept -> index_type + { + return spec_type::get_n_rows(data_view_); + } + [[nodiscard]] auto dim() const noexcept -> uint32_t { return spec_type::get_dim(data_view_); } + [[nodiscard]] auto as_matrix_view() const noexcept -> view_type { return data_view_; } + + /** The view payload; kind-specific state and methods are reached through it. */ + [[nodiscard]] auto data() const noexcept -> view_type const& { return data_view_; } + [[nodiscard]] auto data() noexcept -> view_type& { return data_view_; } + + private: + view_type data_view_{}; +}; + +/** + * @brief Aliases for concrete `dataset` / `dataset_view` layouts. + */ +template +using device_empty_dataset = + dataset>>; + +template +using device_empty_dataset_view = + dataset_view>>; + +template +using host_empty_dataset = + dataset>>; + +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>>; + +template +using host_vpq_dataset_view = + dataset_view>>; + +// Maps a dataset view type to its owning (allocating) dataset counterpart. Trivial and total under +// the Spec design: the owning type for `dataset_view` is always +// `dataset` +// -- no per-kind specialization table needed (unlike the old Container-tagged design). +template +struct owning_dataset_for_view; + +template +struct owning_dataset_for_view> { + using type = dataset; +}; + +template +using owning_dataset_for_view_t = typename owning_dataset_for_view::type; + +// ----------------------------------------------------------------------------- +// Spec-kind classification (all derived from SpecT; dataset/dataset_view never branch on kind). +// ----------------------------------------------------------------------------- + +template +struct is_empty_spec : std::false_type {}; +template +struct is_empty_spec> : std::true_type {}; +template +inline constexpr bool is_empty_spec_v = is_empty_spec::value; + +template +struct is_padded_spec : std::false_type {}; +template +struct is_padded_spec> : std::true_type {}; +template +inline constexpr bool is_padded_spec_v = is_padded_spec::value; + +template +struct is_standard_spec : std::false_type {}; +template +struct is_standard_spec> : std::true_type {}; +template +inline constexpr bool is_standard_spec_v = is_standard_spec::value; + +template +struct is_vpq_spec : std::false_type {}; +template +struct is_vpq_spec> : std::true_type {}; +template +inline constexpr bool is_vpq_spec_v = is_vpq_spec::value; + +template +struct vpq_spec_math_type {}; +template +struct vpq_spec_math_type> { + using type = MathT; +}; +template +using vpq_spec_math_type_t = typename vpq_spec_math_type::type; + +/** Owning-side kind traits (mirror today's `is_padded_dataset_v`/`is_standard_dataset_v`/ + * `is_vpq_dataset_v`, used for SFINAE overload selection in factory.cuh/compute_distance_vpq.hpp). + */ +template +struct is_padded_dataset : std::false_type {}; +template +struct is_padded_dataset> : std::bool_constant> {}; +template +struct is_padded_dataset> + : std::bool_constant> {}; +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::bool_constant> {}; +template +struct is_standard_dataset> + : std::bool_constant> {}; +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::bool_constant> {}; +template +inline constexpr bool is_vpq_dataset_v = is_vpq_dataset::value; + +// ----------------------------------------------------------------------------- +// Dataset view compile-time classification (replaces runtime std::variant dispatch). +// ----------------------------------------------------------------------------- + +/** Any non-owning dataset view exposing row count and logical dimension. */ +template +concept ann_dataset_view = requires(V const& v) { + { v.n_rows() } -> std::convertible_to; + { v.dim() } -> std::convertible_to; +}; + +enum class dataset_view_kind { + // TODO(removal): Remove `unknown` once all deprecated host_matrix_view / device_matrix_view / + // mdspan overloads are deleted. It exists solely so that overload resolution on the deprecated + // build(host_matrix_view) / build(device_matrix_view) shims does not cause a hard error when + // the compiler evaluates is_host/device_dataset_view_v for a plain mdspan type. + unknown, + empty, + padded, + standard, + vpq_f16, + vpq_f32, + bbq, +}; + +template +using dataset_view_type_t = std::remove_cvref_t; + +/** Primary template returns `unknown` so traits safely return `false` for non-dataset-view types. + */ +template +struct dataset_view_kind_of { + static constexpr dataset_view_kind value = dataset_view_kind::unknown; +}; + +template +struct dataset_view_kind_of> { + static constexpr dataset_view_kind value = []() constexpr { + if constexpr (is_empty_spec_v) { + return dataset_view_kind::empty; + } else if constexpr (is_padded_spec_v) { + return dataset_view_kind::padded; + } else if constexpr (is_standard_spec_v) { + return dataset_view_kind::standard; + } else if constexpr (is_vpq_spec_v) { + static_assert(std::is_same_v, half> || + std::is_same_v, float>, + "VPQ dataset_view_kind_of expects MathT to be half or float"); + return std::is_same_v, half> ? dataset_view_kind::vpq_f16 + : dataset_view_kind::vpq_f32; + } else { + return dataset_view_kind::unknown; + } + }(); +}; + +/** True when the dataset view accessor is device-accessible. */ +template +struct dataset_view_is_device_accessible : std::false_type {}; + +template +struct dataset_view_is_device_accessible> + : std::bool_constant {}; + +template +inline constexpr bool dataset_view_is_device_accessible_v = + dataset_view_is_device_accessible>::value; + +template +inline constexpr dataset_view_kind dataset_view_kind_v = + dataset_view_kind_of>::value; + +template +inline constexpr bool is_device_empty_dataset_view_v = + dataset_view_kind_v == dataset_view_kind::empty && dataset_view_is_device_accessible_v; + +template +inline constexpr bool is_host_empty_dataset_view_v = + dataset_view_kind_v == dataset_view_kind::empty && !dataset_view_is_device_accessible_v; + +/** True for any empty dataset view (device or host). */ +template +inline constexpr bool is_empty_dataset_view_v = + is_device_empty_dataset_view_v || is_host_empty_dataset_view_v; + +template +inline constexpr bool is_device_padded_dataset_view_v = + dataset_view_kind_v == dataset_view_kind::padded && dataset_view_is_device_accessible_v; + +template +inline constexpr bool is_host_padded_dataset_view_v = + dataset_view_kind_v == dataset_view_kind::padded && !dataset_view_is_device_accessible_v; + +/** True for either `device_padded_dataset_view` or `host_padded_dataset_view`. */ +template +inline constexpr bool is_padded_dataset_view_v = + is_device_padded_dataset_view_v || is_host_padded_dataset_view_v; + +template +inline constexpr bool is_device_standard_dataset_view_v = + dataset_view_kind_v == dataset_view_kind::standard && dataset_view_is_device_accessible_v; + +template +inline constexpr bool is_host_standard_dataset_view_v = + dataset_view_kind_v == dataset_view_kind::standard && !dataset_view_is_device_accessible_v; + +/** True for either `device_standard_dataset_view` or `host_standard_dataset_view`. */ +template +inline constexpr bool is_standard_dataset_view_v = + is_device_standard_dataset_view_v || is_host_standard_dataset_view_v; + +template +inline constexpr bool is_device_vpq_f16_dataset_view_v = + dataset_view_kind_v == dataset_view_kind::vpq_f16 && dataset_view_is_device_accessible_v; + +template +inline constexpr bool is_host_vpq_f16_dataset_view_v = + dataset_view_kind_v == dataset_view_kind::vpq_f16 && !dataset_view_is_device_accessible_v; + +template +inline constexpr bool is_vpq_f16_dataset_view_v = + is_device_vpq_f16_dataset_view_v || is_host_vpq_f16_dataset_view_v; + +template +inline constexpr bool is_device_vpq_f32_dataset_view_v = + dataset_view_kind_v == dataset_view_kind::vpq_f32 && dataset_view_is_device_accessible_v; + +template +inline constexpr bool is_host_vpq_f32_dataset_view_v = + dataset_view_kind_v == dataset_view_kind::vpq_f32 && !dataset_view_is_device_accessible_v; + +template +inline constexpr bool is_vpq_f32_dataset_view_v = + is_device_vpq_f32_dataset_view_v || is_host_vpq_f32_dataset_view_v; + +template +inline constexpr bool is_device_vpq_dataset_view_v = + is_device_vpq_f16_dataset_view_v || is_device_vpq_f32_dataset_view_v; + +template +inline constexpr bool is_host_vpq_dataset_view_v = + is_host_vpq_f16_dataset_view_v || is_host_vpq_f32_dataset_view_v; + +template +inline constexpr bool is_vpq_dataset_view_v = + is_device_vpq_dataset_view_v || is_host_vpq_dataset_view_v; + +/** True for any device-resident dataset view. */ +template +inline constexpr bool is_device_dataset_view_v = + dataset_view_kind_v != dataset_view_kind::unknown && dataset_view_is_device_accessible_v; + +/** True for any host-resident dataset view. */ +template +inline constexpr bool is_host_dataset_view_v = + dataset_view_kind_v != dataset_view_kind::unknown && !dataset_view_is_device_accessible_v; + +/** + * True when a host view `H` and device view `D` represent the same storage kind and differ + * only in residency (host vs. device). Used by host/device conversion helpers. + */ +template +inline constexpr bool compatible_host_device_dataset_views_v = + is_host_dataset_view_v && is_device_dataset_view_v && + (dataset_view_kind_v == dataset_view_kind_v); + +/** + * Generic accessor retargeting while preserving the dataset tag/layout and value/index types: + * `dataset> -> dataset>` + * `dataset_view> -> dataset_view>` + */ +template +struct with_accessor; + +template +struct with_accessor>, NewAccessor> { + using type = dataset>; +}; + +template +struct with_accessor>, NewAccessor> { + using type = dataset>; +}; + +template +struct with_accessor>, NewAccessor> { + using type = dataset>; +}; + +template +struct with_accessor>, NewAccessor> { + using type = dataset>; +}; + +template +struct with_accessor>, NewAccessor> { + using type = dataset_view>; +}; + +template +struct with_accessor>, NewAccessor> { + using type = dataset_view>; +}; + +template +struct with_accessor>, NewAccessor> { + using type = dataset_view>; +}; + +template +struct with_accessor>, NewAccessor> { + using type = dataset_view>; +}; + +template +using with_accessor_t = + typename with_accessor, NewAccessor>::type; + +/** Map any host accessor to its device counterpart (same payload policy). */ +template +struct to_device_accessor { + using type = Accessor; +}; + +template +struct to_device_accessor> { + using type = detail::device_view_accessor; +}; + +template +struct to_device_accessor> { + using type = detail::device_owning_accessor; +}; + +template +using to_device_accessor_t = typename to_device_accessor::type; + +/** Maps a host dataset view type to its device-resident counterpart. */ +template +struct device_counterpart; + +template +struct device_counterpart> { + using type = with_accessor_t, + to_device_accessor_t>; +}; + +template +using device_counterpart_t = typename device_counterpart>::type; + +/** True for device padded or standard views accepted by dense graph build (VPQ excluded). */ +template +inline constexpr bool is_dense_row_major_device_dataset_view_v = + is_device_padded_dataset_view_v || is_device_standard_dataset_view_v; + +/** True for host or device padded/standard views (dense graph build; VPQ excluded). */ +template +inline constexpr bool is_dense_row_major_dataset_view_v = + is_padded_dataset_view_v || is_standard_dataset_view_v; + +/** Element type `T` for `cagra::build(res, params, dataset_view)` (deduced, not a template arg). + * Trivial under the Spec design: every `dataset_view` already carries `T` directly. + */ +template +using cagra_view_element_type_t = typename dataset_view_type_t::value_type; + +// ----------------------------------------------------------------------------- +// CAGRA row width in elements (same for make_device_padded_dataset* and index layout checks). +// ----------------------------------------------------------------------------- + +/** + * @brief Required row width in elements for CAGRA: minimum leading dimension (LDA) per row for the + * default per-row byte alignment (16 bytes, combined with `sizeof` element type), given + * `logical_columns` feature columns. + */ +[[nodiscard]] inline uint32_t cagra_required_row_width(uint32_t logical_columns, + std::size_t sizeof_value, + uint32_t align_bytes = 16) +{ + return static_cast( + raft::round_up_safe(static_cast(logical_columns) * sizeof_value, + std::lcm(align_bytes, static_cast(sizeof_value))) / + sizeof_value); +} + +template +[[nodiscard]] inline uint32_t cagra_required_row_width(uint32_t logical_columns, + uint32_t align_bytes = 16) +{ + return cagra_required_row_width(logical_columns, sizeof(ValueT), align_bytes); +} + +/** Actual row width in elements (leading dimension) of a 2D row-major matrix view. */ +template +[[nodiscard]] inline uint32_t matrix_actual_row_width(raft::device_matrix_view m) +{ + return m.stride(0) > 0 ? static_cast(m.stride(0)) : static_cast(m.extent(1)); +} + +template +[[nodiscard]] inline uint32_t matrix_actual_row_width(raft::host_matrix_view m) +{ + return m.stride(0) > 0 ? static_cast(m.stride(0)) : static_cast(m.extent(1)); +} + +/** + * @brief True if the matrix's row width in elements matches `cagra_required_row_width` for + * `m.extent(1)` and element type `T` (CAGRA row layout is satisfied for this view). + */ +template +[[nodiscard]] inline bool matrix_row_width_matches_cagra_required( + raft::device_matrix_view m, uint32_t align_bytes = 16) +{ + using value_type = std::remove_const_t; + const uint32_t need = + cagra_required_row_width(static_cast(m.extent(1)), align_bytes); + return matrix_actual_row_width(m) == need; +} + +template +[[nodiscard]] inline bool matrix_row_width_matches_cagra_required(raft::host_matrix_view m, + uint32_t align_bytes = 16) +{ + using value_type = std::remove_const_t; + const uint32_t need = + cagra_required_row_width(static_cast(m.extent(1)), align_bytes); + return matrix_actual_row_width(m) == need; +} + +namespace detail { + +template +[[nodiscard]] inline uint32_t mdspan_row_stride_elements(SrcT const& src) +{ + return src.stride(0) > 0 ? static_cast(src.stride(0)) + : static_cast(src.extent(1)); +} + +template +[[nodiscard]] inline ValueT* expect_device_accessible_data_handle(SrcT const& src, + char const* error_msg) +{ + cudaPointerAttributes ptr_attrs; + RAFT_CUDA_TRY(cudaPointerGetAttributes(&ptr_attrs, src.data_handle())); + // `devicePointer` is relative to the *current* device: it is null for an allocation owned by + // another device without peer access, even though that allocation is perfectly usable once the + // caller switches to the owning device (as the multi-GPU paths do). Accept device and managed + // allocations on their own merit and only consult `devicePointer` for host memory, which needs a + // mapping to be reachable at all. + if (ptr_attrs.type == cudaMemoryTypeDevice || ptr_attrs.type == cudaMemoryTypeManaged) { + return const_cast(src.data_handle()); + } + auto* device_ptr = reinterpret_cast(ptr_attrs.devicePointer); + RAFT_EXPECTS(device_ptr != nullptr, "%s", error_msg); + return device_ptr; +} + +template +[[nodiscard]] inline ViewT make_device_dense_row_major_view_from_src(SrcT const& src, + uint32_t logical_dim) +{ + auto* device_ptr = expect_device_accessible_data_handle( + src, "make_device_*_dataset_view: source must be device-accessible."); + auto v = raft::make_device_matrix_view( + device_ptr, src.extent(0), static_cast(mdspan_row_stride_elements(src))); + return ViewT(v, logical_dim); +} + +template +[[nodiscard]] inline ViewT make_host_dense_row_major_view_from_src(SrcT const& src, + uint32_t logical_dim) +{ + RAFT_EXPECTS(raft::get_device_for_address(src.data_handle()) == -1, + "make_host_*_dataset_view: source must be host-accessible."); + auto v = raft::make_host_matrix_view(const_cast(src.data_handle()), + src.extent(0), + static_cast(mdspan_row_stride_elements(src))); + return ViewT(v, logical_dim); +} + +template +auto make_device_dense_row_major_dataset_from_src(raft::resources const& res, + SrcT const& src, + uint32_t logical_dim, + uint32_t target_stride, + char const* view_factory_name) + -> std::unique_ptr +{ + uint32_t const src_stride = mdspan_row_stride_elements(src); + RAFT_EXPECTS(logical_dim <= target_stride, + "logical dim (%u) must not exceed row stride (%u).", + static_cast(logical_dim), + static_cast(target_stride)); + RAFT_EXPECTS(static_cast(src.extent(1)) <= target_stride, + "Source row length must not exceed required stride."); + cudaPointerAttributes ptr_attrs; + RAFT_CUDA_TRY(cudaPointerGetAttributes(&ptr_attrs, src.data_handle())); + bool const device_src = + (ptr_attrs.type == cudaMemoryTypeDevice) || (ptr_attrs.type == cudaMemoryTypeManaged); + if (device_src && src_stride == target_stride) { + RAFT_EXPECTS(false, + "source is device and stride is already correct. " + "Use %s() to get a view instead.", + view_factory_name); + } + auto out_array = raft::make_device_matrix(res, src.extent(0), target_stride); + RAFT_CUDA_TRY(cudaMemsetAsync(out_array.data_handle(), + 0, + out_array.size() * sizeof(ValueT), + raft::resource::get_cuda_stream(res).get())); + raft::copy_matrix(out_array.data_handle(), + target_stride, + src.data_handle(), + src_stride, + logical_dim, + src.extent(0), + raft::resource::get_cuda_stream(res)); + return std::make_unique(std::move(out_array), logical_dim); +} + +template +auto make_host_dense_row_major_dataset_from_src(raft::resources const& res, + SrcT const& src, + uint32_t logical_dim, + uint32_t target_stride, + char const* view_factory_name) + -> std::unique_ptr +{ + uint32_t const src_stride = mdspan_row_stride_elements(src); + constexpr bool device_src = SrcT::accessor_type::is_device_accessible; + RAFT_EXPECTS(logical_dim <= target_stride, + "logical dim (%u) must not exceed row stride (%u).", + static_cast(logical_dim), + static_cast(target_stride)); + if (!device_src && src_stride == target_stride) { + RAFT_EXPECTS(false, + "source stride is already correct. Use %s() to get a view instead.", + view_factory_name); + } + RAFT_EXPECTS(static_cast(src.extent(1)) <= target_stride, + "Source row length must not exceed required stride."); + auto out_array = raft::make_host_matrix(src.extent(0), target_stride); + std::memset(out_array.data_handle(), 0, out_array.size() * sizeof(ValueT)); + raft::copy_matrix(out_array.data_handle(), + target_stride, + src.data_handle(), + src_stride, + logical_dim, + src.extent(0), + raft::resource::get_cuda_stream(res)); + if (device_src) { raft::resource::sync_stream(res); } + return std::make_unique(std::move(out_array), logical_dim); +} + +} // namespace detail + +template +auto make_device_padded_dataset_view(const raft::resources& res, + SrcT const& src, + uint32_t align_bytes = 16) + -> device_padded_dataset_view +{ + using value_type = typename SrcT::value_type; + using index_type = typename SrcT::index_type; + uint32_t required_stride = + cagra_required_row_width(static_cast(src.extent(1)), align_bytes); + RAFT_EXPECTS( + detail::mdspan_row_stride_elements(src) == required_stride, + "make_device_padded_dataset_view: stride is incorrect (required stride for alignment). " + "Use make_device_padded_dataset() to get an owning padded copy."); + return detail::make_device_dense_row_major_view_from_src< + value_type, + index_type, + device_padded_dataset_view>(src, static_cast(src.extent(1))); +} + +template +auto make_device_padded_dataset(const raft::resources& res, + SrcT const& src, + uint32_t align_bytes = 16) + -> std::unique_ptr> +{ + using value_type = typename SrcT::value_type; + using index_type = typename SrcT::index_type; + uint32_t const logical_dim = static_cast(src.extent(1)); + uint32_t const required_stride = cagra_required_row_width(logical_dim, align_bytes); + return detail::make_device_dense_row_major_dataset_from_src< + device_padded_dataset, + value_type, + index_type>(res, src, logical_dim, required_stride, "make_device_padded_dataset_view"); +} + +template +auto make_host_padded_dataset_view(SrcT const& src, uint32_t align_bytes = 16) + -> host_padded_dataset_view +{ + using value_type = typename SrcT::value_type; + using index_type = typename SrcT::index_type; + uint32_t required_stride = + cagra_required_row_width(static_cast(src.extent(1)), align_bytes); + RAFT_EXPECTS( + detail::mdspan_row_stride_elements(src) == required_stride, + "make_host_padded_dataset_view: stride is incorrect (required stride for alignment). " + "Use make_host_padded_dataset() to get an owning padded copy."); + return detail::make_host_dense_row_major_view_from_src< + value_type, + index_type, + host_padded_dataset_view>(src, static_cast(src.extent(1))); +} + +template +auto make_host_padded_dataset(const raft::resources& res, + SrcT const& src, + uint32_t align_bytes = 16) + -> std::unique_ptr> +{ + using value_type = typename SrcT::value_type; + using index_type = typename SrcT::index_type; + uint32_t const logical_dim = static_cast(src.extent(1)); + uint32_t const required_stride = cagra_required_row_width(logical_dim, align_bytes); + return detail::make_host_dense_row_major_dataset_from_src< + host_padded_dataset, + value_type, + index_type>(res, src, logical_dim, required_stride, "make_host_padded_dataset_view"); +} + +template +auto make_device_standard_dataset_view(SrcT const& src) + -> device_standard_dataset_view +{ + using value_type = typename SrcT::value_type; + using index_type = typename SrcT::index_type; + return detail::make_device_dense_row_major_view_from_src< + value_type, + index_type, + device_standard_dataset_view>(src, + static_cast(src.extent(1))); +} + +/** + * @brief Create an owning device standard dataset with explicit row layout. + * + * Internal use only: the sole call site today is + * `cuvs::neighbors::detail::deserialize_standard()` in `dataset_serialize.hpp`, which must pass + * wire-format `(logical_dim, stride)` because the deserialized host buffer is tight `[n_rows x + * dim]` while the on-disk stride may be larger. Do not call from user code; prefer + * `make_device_standard_dataset_view()` when wrapping existing correctly-strided storage. + */ +template +auto make_device_standard_dataset(const raft::resources& res, + SrcT const& src, + uint32_t logical_dim, + uint32_t target_stride) + -> std::unique_ptr> +{ + using value_type = typename SrcT::value_type; + using index_type = typename SrcT::index_type; + return detail::make_device_dense_row_major_dataset_from_src< + device_standard_dataset, + value_type, + index_type>(res, src, logical_dim, target_stride, "make_device_standard_dataset_view"); +} + +template +auto make_host_standard_dataset_view(SrcT const& src) + -> host_standard_dataset_view +{ + using value_type = typename SrcT::value_type; + using index_type = typename SrcT::index_type; + return detail::make_host_dense_row_major_view_from_src< + value_type, + index_type, + host_standard_dataset_view>(src, static_cast(src.extent(1))); +} + +} // namespace neighbors +} // namespace CUVS_EXPORT cuvs diff --git a/cpp/include/cuvs/neighbors/common.hpp b/cpp/include/cuvs/neighbors/common.hpp index ecfce392e4..f23bc8ea95 100644 --- a/cpp/include/cuvs/neighbors/common.hpp +++ b/cpp/include/cuvs/neighbors/common.hpp @@ -24,6 +24,7 @@ #include #include +#include #include #include @@ -150,1194 +151,6 @@ enum class MergeStrategy { /** @} */ // end group neighbors_index -/** - * @brief Spec-based `dataset` / `dataset_view`. - * - * `dataset` and `dataset_view` are single generic templates that know - * nothing about any particular kind of dataset. They hold exactly one payload (`data_type` / - * `view_type`, chosen by the spec) and expose only what every dataset has: `n_rows()`, `dim()`, - * `as_matrix_view()`, `as_dataset_view()` and `data()`. Each is a one-line forward to one of the - * three spec functions `get_data_view()`, `get_n_rows()` and `get_dim()`. Anything else a kind - * needs (e.g. VPQ codebooks, BBQ quantizers) is state and methods of that kind's payload type, - * reached through `data()`; `dataset`/`dataset_view` never name or branch on it. `dataset` and - * `dataset_view` are deliberately two independent, non-inheriting types (no shared_ptr, no - * "sometimes owning" object): `dataset` holds the owning payload, `dataset_view` the - * corresponding non-owning payload. - */ - -template -struct dataset; - -template -struct dataset_view; - -namespace detail { - -// Default owning/view accessors for public dataset aliases. -template -using device_owning_accessor = raft::device_accessor>; - -template -using host_owning_accessor = raft::host_accessor>; - -template -using device_view_accessor = raft::device_accessor>; - -template -using host_view_accessor = raft::host_accessor>; - -/** View accessor paired with an owning dataset accessor (same residency). */ -template -using dataset_view_accessor_for_owning = std::conditional_t, - host_view_accessor>; - -/** Owning accessor paired with a view accessor (same residency). */ -template -using dataset_owning_accessor_for_view = std::conditional_t, - host_owning_accessor>; - -// Accessor here is already device_owning_accessor / host_owning_accessor at every -// call site -- exactly the container policy raft::device_mdarray/host_mdarray default to for -// element type DataT -- so pass it straight through instead of re-deriving a -// raft::device_matrix/host_matrix from scratch. -template -using dense_owning_matrix = - raft::mdarray, raft::row_major, Accessor>; - -template -using dense_view_matrix = raft::mdspan, - raft::row_major, - dataset_view_accessor_for_owning>; - -template -using vpq_vq_book_matrix = - raft::mdarray, raft::row_major, Accessor>; - -// VPQ codes are always uint8_t regardless of MathT, so retarget the owning accessor's element -// type instead of re-deriving a device/host matrix; residency is still driven by Accessor. -template -using owning_accessor_with_value_type = std::conditional_t, - host_owning_accessor>; - -template -using vpq_data_matrix = raft::mdarray, - raft::row_major, - owning_accessor_with_value_type>; - -// ----------------------------------------------------------------------------- -// empty -// ----------------------------------------------------------------------------- - -template -struct empty_dataset_storage { - uint32_t suggested_dim{}; - empty_dataset_storage() noexcept = default; - explicit empty_dataset_storage(uint32_t dim) noexcept : suggested_dim(dim) {} - [[nodiscard]] auto n_rows() const noexcept -> IdxT { return 0; } - [[nodiscard]] auto dim() const noexcept -> uint32_t { return suggested_dim; } -}; - -// ----------------------------------------------------------------------------- -// 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 specs. Publicly inherits - * from MatrixT (a `raft::mdarray`) so `view()`/`data_handle()`/`extent()` etc. are reused as-is - * rather than hand-forwarded; `logical_dim_` is the only state this struct adds. - * - * 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 : public MatrixT { - uint32_t logical_dim_; - - // MatrixT (mdarray) also has its own stride(size_t); pull it back into scope since declaring - // our own no-arg stride() below would otherwise hide it entirely (C++ name hiding). - using MatrixT::stride; - - dense_row_major_dataset_owning_storage(MatrixT&& data, uint32_t logical_dim) noexcept - : MatrixT{std::move(data)}, logical_dim_{logical_dim} - { - } - - [[nodiscard]] auto n_rows() const noexcept -> IdxT { return this->extent(0); } - [[nodiscard]] auto dim() const noexcept -> uint32_t { return logical_dim_; } - [[nodiscard]] auto stride() const noexcept -> uint32_t - { - return static_cast(this->extent(1)); - } - // view() and data_handle() are inherited directly from MatrixT (raft::mdarray); no hand-written - // forwarding needed since MatrixT::view() const already returns exactly ViewT. -}; - -template -struct dense_row_major_dataset_view_storage : public ViewT { - uint32_t logical_dim_; - - // ViewT (mdspan) also has its own stride(size_t); pull it back into scope since declaring our - // own no-arg stride() below would otherwise hide it entirely (C++ name hiding), and the body of - // that stride() itself needs to call the inherited one. - using ViewT::stride; - - dense_row_major_dataset_view_storage() noexcept = default; - - explicit dense_row_major_dataset_view_storage(ViewT v) noexcept - : ViewT(v), logical_dim_(static_cast(v.extent(1))) - { - } - - dense_row_major_dataset_view_storage(ViewT v, uint32_t logical_dim) noexcept - : ViewT(v), logical_dim_(logical_dim) - { - } - - [[nodiscard]] auto n_rows() const noexcept -> IdxT { return this->extent(0); } - [[nodiscard]] auto dim() const noexcept -> uint32_t { return logical_dim_; } - [[nodiscard]] auto stride() const noexcept -> uint32_t - { - return static_cast(ViewT::stride(0) > 0 ? ViewT::stride(0) : this->extent(1)); - } -}; - -/** Spec-side implementation shared by `padded_dataset_spec`/`standard_dataset_spec`; those two - * stay distinct top-level types (identical bodies) purely so classification traits can tell them - * apart -- exactly mirroring today's `padded_dataset_container`/`standard_dataset_container`, - * which are likewise two differently-named tags over one shared storage implementation. */ -template -struct dense_dataset_spec_impl { - template - struct apply { - using value_type = std::remove_cv_t; - using index_type = std::remove_cv_t; - using MatrixT = dense_owning_matrix; - using ViewT = dense_view_matrix; - using data_type = dense_row_major_dataset_owning_storage; - using view_type = dense_row_major_dataset_view_storage; - - [[nodiscard]] static auto get_data_view(data_type const& data) noexcept -> view_type - { - return view_type(data.view(), data.dim()); - } - template - [[nodiscard]] static auto get_n_rows(AnyStorage const& data) noexcept -> index_type - { - return data.n_rows(); - } - template - [[nodiscard]] static auto get_dim(AnyStorage const& data) noexcept -> uint32_t - { - return data.dim(); - } - }; -}; - -// ----------------------------------------------------------------------------- -// vpq payloads: everything VPQ-specific lives here, not in dataset/dataset_view. -// ----------------------------------------------------------------------------- - -/** Read-only helpers derived from the codebook shapes; shared by the owning and view payloads. - * `Derived` provides `vq_code_book`, `pq_code_book` and the codes' `extent(r)`. */ -template -struct vpq_codebook_helpers { - /** Logical dimension: it comes from the VQ codebook, not from the encoded rows (row padding - * makes the encoded-row width ambiguous as a dimension). */ - [[nodiscard]] auto dim() const noexcept -> uint32_t - { - return static_cast(self().vq_code_book.extent(1)); - } - [[nodiscard]] auto vq_n_centers() const noexcept -> uint32_t - { - return static_cast(self().vq_code_book.extent(0)); - } - [[nodiscard]] auto pq_n_centers() const noexcept -> uint32_t - { - return static_cast(self().pq_code_book.extent(0)); - } - [[nodiscard]] auto pq_len() const noexcept -> uint32_t - { - return static_cast(self().pq_code_book.extent(1)); - } - [[nodiscard]] 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 bits = 0; - while (pq_width > 1) { - bits++; - pq_width >>= 1; - } - return bits; -#endif - } - [[nodiscard]] auto pq_dim() const noexcept -> uint32_t - { - return raft::div_rounding_up_unsafe(dim(), pq_len()); - } - [[nodiscard]] auto encoded_row_length() const noexcept -> uint32_t - { - return static_cast(self().extent(1)); - } - - private: - [[nodiscard]] auto self() const noexcept -> Derived const& - { - return static_cast(*this); - } -}; - -/** Owning VPQ payload: the encoded rows (it *is* the `uint8_t` codes mdarray) plus the VQ and PQ - * codebooks. `Accessor` drives both codebook and code residency. */ -template -struct vpq_owning_storage : public vpq_data_matrix, - public vpq_codebook_helpers> { - using codes_type = vpq_data_matrix; - using vq_book_type = vpq_vq_book_matrix; - using pq_book_type = vpq_vq_book_matrix; - - vq_book_type vq_code_book; - pq_book_type pq_code_book; - - vpq_owning_storage(codes_type&& codes, vq_book_type&& vq_codes, pq_book_type&& pq_codes) noexcept - : codes_type{std::move(codes)}, - vq_code_book{std::move(vq_codes)}, - pq_code_book{std::move(pq_codes)} - { - } -}; - -/** Non-owning VPQ payload: a view of the encoded rows plus views of the VQ and PQ codebooks. */ -template -struct vpq_view_storage : public raft::mdspan, - raft::row_major, - dataset_view_accessor_for_owning>, - public vpq_codebook_helpers> { - using codes_view_type = raft::mdspan, - raft::row_major, - dataset_view_accessor_for_owning>; - using vq_book_view_type = - typename vpq_owning_storage::vq_book_type::const_view_type; - using pq_book_view_type = - typename vpq_owning_storage::pq_book_type::const_view_type; - - vq_book_view_type vq_code_book; - pq_book_view_type pq_code_book; - - vpq_view_storage() noexcept = default; - vpq_view_storage(codes_view_type codes, - vq_book_view_type vq_codes, - pq_book_view_type pq_codes) noexcept - : codes_view_type(codes), vq_code_book(vq_codes), pq_code_book(pq_codes) - { - } -}; - -} // namespace detail - -// ----------------------------------------------------------------------------- -// Public specs -- the only place per-kind logic lives. -// ----------------------------------------------------------------------------- - -template -struct empty_dataset_spec { - using accessor_type = Accessor; - - template - struct apply { - using value_type = std::remove_cv_t; - using index_type = std::remove_cv_t; - using data_type = detail::empty_dataset_storage; - using view_type = detail::empty_dataset_storage; - - [[nodiscard]] static auto get_data_view(data_type const& data) noexcept -> view_type - { - return data; - } - [[nodiscard]] static auto get_n_rows(data_type const& data) noexcept -> index_type - { - return static_cast(data.n_rows()); - } - [[nodiscard]] static auto get_dim(data_type const& data) noexcept -> uint32_t - { - return data.dim(); - } - }; -}; - -template -struct padded_dataset_spec { - using accessor_type = ContainerPolicy; - template - struct apply : detail::dense_dataset_spec_impl::template apply {}; -}; - -template -struct standard_dataset_spec { - using accessor_type = ContainerPolicy; - template - struct apply : detail::dense_dataset_spec_impl::template apply {}; -}; - -/** `Accessor` drives both codebook and code residency, mirroring today's - * single-`Accessor`-per-VPQ-dataset design. The payload (`detail::vpq_owning_storage` / - * `detail::vpq_view_storage`) holds the encoded rows and the VQ/PQ codebooks. */ -template -struct vpq_dataset_spec { - using accessor_type = Accessor; - - template - struct apply { - using value_type = std::remove_cv_t; - using index_type = std::remove_cv_t; - using math_type = MathT; - - using data_type = detail::vpq_owning_storage; - using view_type = detail::vpq_view_storage; - - [[nodiscard]] static auto get_data_view(data_type const& data) noexcept -> view_type - { - return view_type(data.view(), data.vq_code_book.view(), data.pq_code_book.view()); - } - template - [[nodiscard]] static auto get_n_rows(AnyStorage const& data) noexcept -> index_type - { - return static_cast(data.extent(0)); - } - template - [[nodiscard]] static auto get_dim(AnyStorage const& data) noexcept -> uint32_t - { - return data.dim(); - } - }; -}; - -// ----------------------------------------------------------------------------- -// dataset / dataset_view -// ----------------------------------------------------------------------------- - -/** Owning dataset: value-held payload (no shared_ptr -- exclusive ownership). Every member is a - * one-line forward to `spec_type::get_*` or to the payload; all per-kind state and logic lives in - * the spec's `data_type`, never inside this struct. */ -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; - - // Forward constructor args straight to data_type's own constructor (e.g. (MatrixT&&, uint32_t - // logical_dim) for dense, (uint32_t dim) for empty, (codes&&, vq&&, pq&&) for VPQ). - template - explicit dataset(Args&&... args) - requires(std::is_constructible_v) - : data_(std::forward(args)...) - { - } - - [[nodiscard]] auto n_rows() const noexcept -> index_type { return spec_type::get_n_rows(data_); } - [[nodiscard]] auto dim() const noexcept -> uint32_t { return spec_type::get_dim(data_); } - /** The spec-defined non-owning view of the payload (for dense and VPQ it is an mdspan - * derivative). */ - [[nodiscard]] auto as_matrix_view() const noexcept { return spec_type::get_data_view(data_); } - [[nodiscard]] auto as_dataset_view() const noexcept -> dataset_view - { - return dataset_view(as_matrix_view()); - } - - /** The owning payload; kind-specific state and methods are reached through it. */ - [[nodiscard]] auto data() const noexcept -> data_type const& { return data_; } - [[nodiscard]] auto data() noexcept -> data_type& { return data_; } - - private: - data_type data_; -}; - -/** Non-owning dataset view: holds only the view-shaped payload. Deliberately not derived from - * `dataset` -- a view type holds "all view state" with no inheritance and no shared ownership tying - * it to the owning type. Reuses the same `get_n_rows`/`get_dim` spec functions as `dataset`, fed - * the view payload instead of the owning one. */ -template -struct dataset_view { - using spec_type = typename SpecT::template apply; - using value_type = typename spec_type::value_type; - using index_type = typename spec_type::index_type; - using view_type = typename spec_type::view_type; - - dataset_view() noexcept = default; - - // Already-constructed view payload -- the shape `as_dataset_view()` always constructs with, for - // every kind. Not a template, so it's preferred over the forwarding constructor below whenever - // both could apply. - explicit dataset_view(view_type data_view) noexcept : data_view_{data_view} {} - - // Forward raw constructor args straight to view_type's own constructor (e.g. (ViewT, uint32_t - // logical_dim) for dense, (uint32_t dim) for empty) -- preserves today's direct-construction - // call sites (e.g. `device_padded_dataset_view(raw_mdspan, dim)`) unchanged. - template - explicit dataset_view(Args&&... args) - requires(std::is_constructible_v) - : data_view_(std::forward(args)...) - { - } - - [[nodiscard]] auto n_rows() const noexcept -> index_type - { - return spec_type::get_n_rows(data_view_); - } - [[nodiscard]] auto dim() const noexcept -> uint32_t { return spec_type::get_dim(data_view_); } - [[nodiscard]] auto as_matrix_view() const noexcept -> view_type { return data_view_; } - - /** The view payload; kind-specific state and methods are reached through it. */ - [[nodiscard]] auto data() const noexcept -> view_type const& { return data_view_; } - [[nodiscard]] auto data() noexcept -> view_type& { return data_view_; } - - private: - view_type data_view_{}; -}; - -/** - * @brief Aliases for concrete `dataset` / `dataset_view` layouts. - */ -template -using device_empty_dataset = - dataset>>; - -template -using device_empty_dataset_view = - dataset_view>>; - -template -using host_empty_dataset = - dataset>>; - -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>>; - -template -using host_vpq_dataset_view = - dataset_view>>; - -// Maps a dataset view type to its owning (allocating) dataset counterpart. Trivial and total under -// the Spec design: the owning type for `dataset_view` is always -// `dataset` -// -- no per-kind specialization table needed (unlike the old Container-tagged design). -template -struct owning_dataset_for_view; - -template -struct owning_dataset_for_view> { - using type = dataset; -}; - -template -using owning_dataset_for_view_t = typename owning_dataset_for_view::type; - -// ----------------------------------------------------------------------------- -// Spec-kind classification (all derived from SpecT; dataset/dataset_view never branch on kind). -// ----------------------------------------------------------------------------- - -template -struct is_empty_spec : std::false_type {}; -template -struct is_empty_spec> : std::true_type {}; -template -inline constexpr bool is_empty_spec_v = is_empty_spec::value; - -template -struct is_padded_spec : std::false_type {}; -template -struct is_padded_spec> : std::true_type {}; -template -inline constexpr bool is_padded_spec_v = is_padded_spec::value; - -template -struct is_standard_spec : std::false_type {}; -template -struct is_standard_spec> : std::true_type {}; -template -inline constexpr bool is_standard_spec_v = is_standard_spec::value; - -template -struct is_vpq_spec : std::false_type {}; -template -struct is_vpq_spec> : std::true_type {}; -template -inline constexpr bool is_vpq_spec_v = is_vpq_spec::value; - -template -struct vpq_spec_math_type {}; -template -struct vpq_spec_math_type> { - using type = MathT; -}; -template -using vpq_spec_math_type_t = typename vpq_spec_math_type::type; - -/** Owning-side kind traits (mirror today's `is_padded_dataset_v`/`is_standard_dataset_v`/ - * `is_vpq_dataset_v`, used for SFINAE overload selection in factory.cuh/compute_distance_vpq.hpp). - */ -template -struct is_padded_dataset : std::false_type {}; -template -struct is_padded_dataset> : std::bool_constant> {}; -template -struct is_padded_dataset> - : std::bool_constant> {}; -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::bool_constant> {}; -template -struct is_standard_dataset> - : std::bool_constant> {}; -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::bool_constant> {}; -template -inline constexpr bool is_vpq_dataset_v = is_vpq_dataset::value; - -// ----------------------------------------------------------------------------- -// Dataset view compile-time classification (replaces runtime std::variant dispatch). -// ----------------------------------------------------------------------------- - -/** Any non-owning dataset view exposing row count and logical dimension. */ -template -concept ann_dataset_view = requires(V const& v) { - { v.n_rows() } -> std::convertible_to; - { v.dim() } -> std::convertible_to; -}; - -enum class dataset_view_kind { - // TODO(removal): Remove `unknown` once all deprecated host_matrix_view / device_matrix_view / - // mdspan overloads are deleted. It exists solely so that overload resolution on the deprecated - // build(host_matrix_view) / build(device_matrix_view) shims does not cause a hard error when - // the compiler evaluates is_host/device_dataset_view_v for a plain mdspan type. - unknown, - empty, - padded, - standard, - vpq_f16, - vpq_f32, - bbq, -}; - -template -using dataset_view_type_t = std::remove_cvref_t; - -/** Primary template returns `unknown` so traits safely return `false` for non-dataset-view types. - */ -template -struct dataset_view_kind_of { - static constexpr dataset_view_kind value = dataset_view_kind::unknown; -}; - -template -struct dataset_view_kind_of> { - static constexpr dataset_view_kind value = []() constexpr { - if constexpr (is_empty_spec_v) { - return dataset_view_kind::empty; - } else if constexpr (is_padded_spec_v) { - return dataset_view_kind::padded; - } else if constexpr (is_standard_spec_v) { - return dataset_view_kind::standard; - } else if constexpr (is_vpq_spec_v) { - static_assert(std::is_same_v, half> || - std::is_same_v, float>, - "VPQ dataset_view_kind_of expects MathT to be half or float"); - return std::is_same_v, half> ? dataset_view_kind::vpq_f16 - : dataset_view_kind::vpq_f32; - } else { - return dataset_view_kind::unknown; - } - }(); -}; - -/** True when the dataset view accessor is device-accessible. */ -template -struct dataset_view_is_device_accessible : std::false_type {}; - -template -struct dataset_view_is_device_accessible> - : std::bool_constant {}; - -template -inline constexpr bool dataset_view_is_device_accessible_v = - dataset_view_is_device_accessible>::value; - -template -inline constexpr dataset_view_kind dataset_view_kind_v = - dataset_view_kind_of>::value; - -template -inline constexpr bool is_device_empty_dataset_view_v = - dataset_view_kind_v == dataset_view_kind::empty && dataset_view_is_device_accessible_v; - -template -inline constexpr bool is_host_empty_dataset_view_v = - dataset_view_kind_v == dataset_view_kind::empty && !dataset_view_is_device_accessible_v; - -/** True for any empty dataset view (device or host). */ -template -inline constexpr bool is_empty_dataset_view_v = - is_device_empty_dataset_view_v || is_host_empty_dataset_view_v; - -template -inline constexpr bool is_device_padded_dataset_view_v = - dataset_view_kind_v == dataset_view_kind::padded && dataset_view_is_device_accessible_v; - -template -inline constexpr bool is_host_padded_dataset_view_v = - dataset_view_kind_v == dataset_view_kind::padded && !dataset_view_is_device_accessible_v; - -/** True for either `device_padded_dataset_view` or `host_padded_dataset_view`. */ -template -inline constexpr bool is_padded_dataset_view_v = - is_device_padded_dataset_view_v || is_host_padded_dataset_view_v; - -template -inline constexpr bool is_device_standard_dataset_view_v = - dataset_view_kind_v == dataset_view_kind::standard && dataset_view_is_device_accessible_v; - -template -inline constexpr bool is_host_standard_dataset_view_v = - dataset_view_kind_v == dataset_view_kind::standard && !dataset_view_is_device_accessible_v; - -/** True for either `device_standard_dataset_view` or `host_standard_dataset_view`. */ -template -inline constexpr bool is_standard_dataset_view_v = - is_device_standard_dataset_view_v || is_host_standard_dataset_view_v; - -template -inline constexpr bool is_device_vpq_f16_dataset_view_v = - dataset_view_kind_v == dataset_view_kind::vpq_f16 && dataset_view_is_device_accessible_v; - -template -inline constexpr bool is_host_vpq_f16_dataset_view_v = - dataset_view_kind_v == dataset_view_kind::vpq_f16 && !dataset_view_is_device_accessible_v; - -template -inline constexpr bool is_vpq_f16_dataset_view_v = - is_device_vpq_f16_dataset_view_v || is_host_vpq_f16_dataset_view_v; - -template -inline constexpr bool is_device_vpq_f32_dataset_view_v = - dataset_view_kind_v == dataset_view_kind::vpq_f32 && dataset_view_is_device_accessible_v; - -template -inline constexpr bool is_host_vpq_f32_dataset_view_v = - dataset_view_kind_v == dataset_view_kind::vpq_f32 && !dataset_view_is_device_accessible_v; - -template -inline constexpr bool is_vpq_f32_dataset_view_v = - is_device_vpq_f32_dataset_view_v || is_host_vpq_f32_dataset_view_v; - -template -inline constexpr bool is_device_vpq_dataset_view_v = - is_device_vpq_f16_dataset_view_v || is_device_vpq_f32_dataset_view_v; - -template -inline constexpr bool is_host_vpq_dataset_view_v = - is_host_vpq_f16_dataset_view_v || is_host_vpq_f32_dataset_view_v; - -template -inline constexpr bool is_vpq_dataset_view_v = - is_device_vpq_dataset_view_v || is_host_vpq_dataset_view_v; - -/** True for any device-resident dataset view. */ -template -inline constexpr bool is_device_dataset_view_v = - dataset_view_kind_v != dataset_view_kind::unknown && dataset_view_is_device_accessible_v; - -/** True for any host-resident dataset view. */ -template -inline constexpr bool is_host_dataset_view_v = - dataset_view_kind_v != dataset_view_kind::unknown && !dataset_view_is_device_accessible_v; - -/** - * True when a host view `H` and device view `D` represent the same storage kind and differ - * only in residency (host vs. device). Used by host/device conversion helpers. - */ -template -inline constexpr bool compatible_host_device_dataset_views_v = - is_host_dataset_view_v && is_device_dataset_view_v && - (dataset_view_kind_v == dataset_view_kind_v); - -/** - * Generic accessor retargeting while preserving the dataset tag/layout and value/index types: - * `dataset> -> dataset>` - * `dataset_view> -> dataset_view>` - */ -template -struct with_accessor; - -template -struct with_accessor>, NewAccessor> { - using type = dataset>; -}; - -template -struct with_accessor>, NewAccessor> { - using type = dataset>; -}; - -template -struct with_accessor>, NewAccessor> { - using type = dataset>; -}; - -template -struct with_accessor>, NewAccessor> { - using type = dataset>; -}; - -template -struct with_accessor>, NewAccessor> { - using type = dataset_view>; -}; - -template -struct with_accessor>, NewAccessor> { - using type = dataset_view>; -}; - -template -struct with_accessor>, NewAccessor> { - using type = dataset_view>; -}; - -template -struct with_accessor>, NewAccessor> { - using type = dataset_view>; -}; - -template -using with_accessor_t = - typename with_accessor, NewAccessor>::type; - -/** Map any host accessor to its device counterpart (same payload policy). */ -template -struct to_device_accessor { - using type = Accessor; -}; - -template -struct to_device_accessor> { - using type = detail::device_view_accessor; -}; - -template -struct to_device_accessor> { - using type = detail::device_owning_accessor; -}; - -template -using to_device_accessor_t = typename to_device_accessor::type; - -/** Maps a host dataset view type to its device-resident counterpart. */ -template -struct device_counterpart; - -template -struct device_counterpart> { - using type = with_accessor_t, - to_device_accessor_t>; -}; - -template -using device_counterpart_t = typename device_counterpart>::type; - -/** True for device padded or standard views accepted by dense graph build (VPQ excluded). */ -template -inline constexpr bool is_dense_row_major_device_dataset_view_v = - is_device_padded_dataset_view_v || is_device_standard_dataset_view_v; - -/** True for host or device padded/standard views (dense graph build; VPQ excluded). */ -template -inline constexpr bool is_dense_row_major_dataset_view_v = - is_padded_dataset_view_v || is_standard_dataset_view_v; - -/** Element type `T` for `cagra::build(res, params, dataset_view)` (deduced, not a template arg). - * Trivial under the Spec design: every `dataset_view` already carries `T` directly. - */ -template -using cagra_view_element_type_t = typename dataset_view_type_t::value_type; - -// ----------------------------------------------------------------------------- -// CAGRA row width in elements (same for make_device_padded_dataset* and index layout checks). -// ----------------------------------------------------------------------------- - -/** - * @brief Required row width in elements for CAGRA: minimum leading dimension (LDA) per row for the - * default per-row byte alignment (16 bytes, combined with `sizeof` element type), given - * `logical_columns` feature columns. - */ -[[nodiscard]] inline uint32_t cagra_required_row_width(uint32_t logical_columns, - std::size_t sizeof_value, - uint32_t align_bytes = 16) -{ - return static_cast( - raft::round_up_safe(static_cast(logical_columns) * sizeof_value, - std::lcm(align_bytes, static_cast(sizeof_value))) / - sizeof_value); -} - -template -[[nodiscard]] inline uint32_t cagra_required_row_width(uint32_t logical_columns, - uint32_t align_bytes = 16) -{ - return cagra_required_row_width(logical_columns, sizeof(ValueT), align_bytes); -} - -/** Actual row width in elements (leading dimension) of a 2D row-major matrix view. */ -template -[[nodiscard]] inline uint32_t matrix_actual_row_width(raft::device_matrix_view m) -{ - return m.stride(0) > 0 ? static_cast(m.stride(0)) : static_cast(m.extent(1)); -} - -template -[[nodiscard]] inline uint32_t matrix_actual_row_width(raft::host_matrix_view m) -{ - return m.stride(0) > 0 ? static_cast(m.stride(0)) : static_cast(m.extent(1)); -} - -/** - * @brief True if the matrix's row width in elements matches `cagra_required_row_width` for - * `m.extent(1)` and element type `T` (CAGRA row layout is satisfied for this view). - */ -template -[[nodiscard]] inline bool matrix_row_width_matches_cagra_required( - raft::device_matrix_view m, uint32_t align_bytes = 16) -{ - using value_type = std::remove_const_t; - const uint32_t need = - cagra_required_row_width(static_cast(m.extent(1)), align_bytes); - return matrix_actual_row_width(m) == need; -} - -template -[[nodiscard]] inline bool matrix_row_width_matches_cagra_required(raft::host_matrix_view m, - uint32_t align_bytes = 16) -{ - using value_type = std::remove_const_t; - const uint32_t need = - cagra_required_row_width(static_cast(m.extent(1)), align_bytes); - return matrix_actual_row_width(m) == need; -} - -namespace detail { - -template -[[nodiscard]] inline uint32_t mdspan_row_stride_elements(SrcT const& src) -{ - return src.stride(0) > 0 ? static_cast(src.stride(0)) - : static_cast(src.extent(1)); -} - -template -[[nodiscard]] inline ValueT* expect_device_accessible_data_handle(SrcT const& src, - char const* error_msg) -{ - cudaPointerAttributes ptr_attrs; - RAFT_CUDA_TRY(cudaPointerGetAttributes(&ptr_attrs, src.data_handle())); - // `devicePointer` is relative to the *current* device: it is null for an allocation owned by - // another device without peer access, even though that allocation is perfectly usable once the - // caller switches to the owning device (as the multi-GPU paths do). Accept device and managed - // allocations on their own merit and only consult `devicePointer` for host memory, which needs a - // mapping to be reachable at all. - if (ptr_attrs.type == cudaMemoryTypeDevice || ptr_attrs.type == cudaMemoryTypeManaged) { - return const_cast(src.data_handle()); - } - auto* device_ptr = reinterpret_cast(ptr_attrs.devicePointer); - RAFT_EXPECTS(device_ptr != nullptr, "%s", error_msg); - return device_ptr; -} - -template -[[nodiscard]] inline ViewT make_device_dense_row_major_view_from_src(SrcT const& src, - uint32_t logical_dim) -{ - auto* device_ptr = expect_device_accessible_data_handle( - src, "make_device_*_dataset_view: source must be device-accessible."); - auto v = raft::make_device_matrix_view( - device_ptr, src.extent(0), static_cast(mdspan_row_stride_elements(src))); - return ViewT(v, logical_dim); -} - -template -[[nodiscard]] inline ViewT make_host_dense_row_major_view_from_src(SrcT const& src, - uint32_t logical_dim) -{ - RAFT_EXPECTS(raft::get_device_for_address(src.data_handle()) == -1, - "make_host_*_dataset_view: source must be host-accessible."); - auto v = raft::make_host_matrix_view(const_cast(src.data_handle()), - src.extent(0), - static_cast(mdspan_row_stride_elements(src))); - return ViewT(v, logical_dim); -} - -template -auto make_device_dense_row_major_dataset_from_src(raft::resources const& res, - SrcT const& src, - uint32_t logical_dim, - uint32_t target_stride, - char const* view_factory_name) - -> std::unique_ptr -{ - uint32_t const src_stride = mdspan_row_stride_elements(src); - RAFT_EXPECTS(logical_dim <= target_stride, - "logical dim (%u) must not exceed row stride (%u).", - static_cast(logical_dim), - static_cast(target_stride)); - RAFT_EXPECTS(static_cast(src.extent(1)) <= target_stride, - "Source row length must not exceed required stride."); - cudaPointerAttributes ptr_attrs; - RAFT_CUDA_TRY(cudaPointerGetAttributes(&ptr_attrs, src.data_handle())); - bool const device_src = - (ptr_attrs.type == cudaMemoryTypeDevice) || (ptr_attrs.type == cudaMemoryTypeManaged); - if (device_src && src_stride == target_stride) { - RAFT_EXPECTS(false, - "source is device and stride is already correct. " - "Use %s() to get a view instead.", - view_factory_name); - } - auto out_array = raft::make_device_matrix(res, src.extent(0), target_stride); - RAFT_CUDA_TRY(cudaMemsetAsync(out_array.data_handle(), - 0, - out_array.size() * sizeof(ValueT), - raft::resource::get_cuda_stream(res).get())); - raft::copy_matrix(out_array.data_handle(), - target_stride, - src.data_handle(), - src_stride, - logical_dim, - src.extent(0), - raft::resource::get_cuda_stream(res)); - return std::make_unique(std::move(out_array), logical_dim); -} - -template -auto make_host_dense_row_major_dataset_from_src(raft::resources const& res, - SrcT const& src, - uint32_t logical_dim, - uint32_t target_stride, - char const* view_factory_name) - -> std::unique_ptr -{ - uint32_t const src_stride = mdspan_row_stride_elements(src); - constexpr bool device_src = SrcT::accessor_type::is_device_accessible; - RAFT_EXPECTS(logical_dim <= target_stride, - "logical dim (%u) must not exceed row stride (%u).", - static_cast(logical_dim), - static_cast(target_stride)); - if (!device_src && src_stride == target_stride) { - RAFT_EXPECTS(false, - "source stride is already correct. Use %s() to get a view instead.", - view_factory_name); - } - RAFT_EXPECTS(static_cast(src.extent(1)) <= target_stride, - "Source row length must not exceed required stride."); - auto out_array = raft::make_host_matrix(src.extent(0), target_stride); - std::memset(out_array.data_handle(), 0, out_array.size() * sizeof(ValueT)); - raft::copy_matrix(out_array.data_handle(), - target_stride, - src.data_handle(), - src_stride, - logical_dim, - src.extent(0), - raft::resource::get_cuda_stream(res)); - if (device_src) { raft::resource::sync_stream(res); } - return std::make_unique(std::move(out_array), logical_dim); -} - -} // namespace detail - -template -auto make_device_padded_dataset_view(const raft::resources& res, - SrcT const& src, - uint32_t align_bytes = 16) - -> device_padded_dataset_view -{ - using value_type = typename SrcT::value_type; - using index_type = typename SrcT::index_type; - uint32_t required_stride = - cagra_required_row_width(static_cast(src.extent(1)), align_bytes); - RAFT_EXPECTS( - detail::mdspan_row_stride_elements(src) == required_stride, - "make_device_padded_dataset_view: stride is incorrect (required stride for alignment). " - "Use make_device_padded_dataset() to get an owning padded copy."); - return detail::make_device_dense_row_major_view_from_src< - value_type, - index_type, - device_padded_dataset_view>(src, static_cast(src.extent(1))); -} - -template -auto make_device_padded_dataset(const raft::resources& res, - SrcT const& src, - uint32_t align_bytes = 16) - -> std::unique_ptr> -{ - using value_type = typename SrcT::value_type; - using index_type = typename SrcT::index_type; - uint32_t const logical_dim = static_cast(src.extent(1)); - uint32_t const required_stride = cagra_required_row_width(logical_dim, align_bytes); - return detail::make_device_dense_row_major_dataset_from_src< - device_padded_dataset, - value_type, - index_type>(res, src, logical_dim, required_stride, "make_device_padded_dataset_view"); -} - -template -auto make_host_padded_dataset_view(SrcT const& src, uint32_t align_bytes = 16) - -> host_padded_dataset_view -{ - using value_type = typename SrcT::value_type; - using index_type = typename SrcT::index_type; - uint32_t required_stride = - cagra_required_row_width(static_cast(src.extent(1)), align_bytes); - RAFT_EXPECTS( - detail::mdspan_row_stride_elements(src) == required_stride, - "make_host_padded_dataset_view: stride is incorrect (required stride for alignment). " - "Use make_host_padded_dataset() to get an owning padded copy."); - return detail::make_host_dense_row_major_view_from_src< - value_type, - index_type, - host_padded_dataset_view>(src, static_cast(src.extent(1))); -} - -template -auto make_host_padded_dataset(const raft::resources& res, - SrcT const& src, - uint32_t align_bytes = 16) - -> std::unique_ptr> -{ - using value_type = typename SrcT::value_type; - using index_type = typename SrcT::index_type; - uint32_t const logical_dim = static_cast(src.extent(1)); - uint32_t const required_stride = cagra_required_row_width(logical_dim, align_bytes); - return detail::make_host_dense_row_major_dataset_from_src< - host_padded_dataset, - value_type, - index_type>(res, src, logical_dim, required_stride, "make_host_padded_dataset_view"); -} - -template -auto make_device_standard_dataset_view(SrcT const& src) - -> device_standard_dataset_view -{ - using value_type = typename SrcT::value_type; - using index_type = typename SrcT::index_type; - return detail::make_device_dense_row_major_view_from_src< - value_type, - index_type, - device_standard_dataset_view>(src, - static_cast(src.extent(1))); -} - -/** - * @brief Create an owning device standard dataset with explicit row layout. - * - * Internal use only: the sole call site today is - * `cuvs::neighbors::detail::deserialize_standard()` in `dataset_serialize.hpp`, which must pass - * wire-format `(logical_dim, stride)` because the deserialized host buffer is tight `[n_rows x - * dim]` while the on-disk stride may be larger. Do not call from user code; prefer - * `make_device_standard_dataset_view()` when wrapping existing correctly-strided storage. - */ -template -auto make_device_standard_dataset(const raft::resources& res, - SrcT const& src, - uint32_t logical_dim, - uint32_t target_stride) - -> std::unique_ptr> -{ - using value_type = typename SrcT::value_type; - using index_type = typename SrcT::index_type; - return detail::make_device_dense_row_major_dataset_from_src< - device_standard_dataset, - value_type, - index_type>(res, src, logical_dim, target_stride, "make_device_standard_dataset_view"); -} - -template -auto make_host_standard_dataset_view(SrcT const& src) - -> host_standard_dataset_view -{ - using value_type = typename SrcT::value_type; - using index_type = typename SrcT::index_type; - return detail::make_host_dense_row_major_view_from_src< - value_type, - index_type, - host_standard_dataset_view>(src, static_cast(src.extent(1))); -} - namespace filtering { /** 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..86a2a17e80 --- /dev/null +++ b/fern/pages/cpp_api/cpp-api-core-dataset.md @@ -0,0 +1,129 @@ +--- +slug: api-reference/cpp-api-core-dataset +--- + +# Dataset + +_Source header: `cuvs/core/dataset.hpp`_ + +## Types + + +### neighbors::dataset + +Spec-based `dataset` / `dataset_view`. + +`dataset<T,IdxT,SpecT>` and `dataset_view<T,IdxT,SpecT>` are single generic templates that know nothing about any particular kind of dataset. They hold exactly one payload (`data_type` / `view_type`, chosen by the spec) and expose only what every dataset has: `n_rows()`, `dim()`, `as_matrix_view()`, `as_dataset_view()` and `data()`. Each is a one-line forward to one of the three spec functions `get_data_view()`, `get_n_rows()` and `get_dim()`. Anything else a kind needs (e.g. VPQ codebooks, BBQ quantizers) is state and methods of that kind's payload type, reached through `data()`; `dataset`/`dataset_view` never name or branch on it. `dataset` and `dataset_view` are deliberately two independent, non-inheriting types (no shared_ptr, no "sometimes owning" object): `dataset` holds the owning payload, `dataset_view` the corresponding non-owning payload. + +```cpp +template +struct dataset; +``` + + +### neighbors::vpq_dataset_spec + +`Accessor` drives both codebook and code residency, mirroring today's + +single-`Accessor`-per-VPQ-dataset design. The payload (`detail::vpq_owning_storage` / `detail::vpq_view_storage`) holds the encoded rows and the VQ/PQ codebooks. + +```cpp +template +struct vpq_dataset_spec; +``` + + +### neighbors::dataset + +Owning dataset: value-held payload (no shared_ptr -- exclusive ownership). Every member is a + +one-line forward to `spec_type::get_*` or to the payload; all per-kind state and logic lives in the spec's `data_type`, never inside this struct. + +```cpp +template +struct dataset; +``` + + +### neighbors::dataset_view + +Non-owning dataset view: holds only the view-shaped payload. Deliberately not derived from + +`dataset` -- a view type holds "all view state" with no inheritance and no shared ownership tying it to the owning type. Reuses the same `get_n_rows`/`get_dim` spec functions as `dataset`, fed the view payload instead of the owning one. + +```cpp +template +struct dataset_view; +``` + + +### neighbors::is_padded_dataset + +Owning-side kind traits (mirror today's `is_padded_dataset_v`/`is_standard_dataset_v`/ + +`is_vpq_dataset_v`, used for SFINAE overload selection in factory.cuh/compute_distance_vpq.hpp). + +```cpp +template +struct is_padded_dataset; +``` + + +### neighbors::dataset_view_kind_of + +Primary template returns `unknown` so traits safely return `false` for non-dataset-view types. + +```cpp +template +struct dataset_view_kind_of { + static constexpr dataset_view_kind value; +}; +``` + +**Fields** + +| Name | Type | Description | +| --- | --- | --- | +| `value` | `static constexpr dataset_view_kind` | | + + +### neighbors::dataset_view_is_device_accessible + +True when the dataset view accessor is device-accessible. + +```cpp +template +struct dataset_view_is_device_accessible; +``` + + +### neighbors::with_accessor + +Generic accessor retargeting while preserving the dataset tag/layout and value/index types: + +`dataset<T, IdxT, SpecT<..., OldAccessor>> -> dataset<T, IdxT, SpecT<..., NewAccessor>>` `dataset_view<T, IdxT, SpecT<..., OldAccessor>> -> dataset_view<T, IdxT, SpecT<..., NewAccessor>>` + +```cpp +template +struct with_accessor; +``` + + +### neighbors::to_device_accessor + +Map any host accessor to its device counterpart (same payload policy). + +```cpp +template +struct to_device_accessor; +``` + + +### neighbors::device_counterpart + +Maps a host dataset view type to its device-resident counterpart. + +```cpp +template +struct device_counterpart; +``` diff --git a/fern/pages/cpp_api/cpp-api-neighbors-common.md b/fern/pages/cpp_api/cpp-api-neighbors-common.md index 457c97c53d..da72bab316 100644 --- a/fern/pages/cpp_api/cpp-api-neighbors-common.md +++ b/fern/pages/cpp_api/cpp-api-neighbors-common.md @@ -57,168 +57,6 @@ enum class MergeStrategy { | `MERGE_STRATEGY_PHYSICAL` | `0` | | `MERGE_STRATEGY_LOGICAL` | `1` | -## Types - - -### neighbors::dataset - -Spec-based `dataset` / `dataset_view`. - -`dataset<T,IdxT,SpecT>` and `dataset_view<T,IdxT,SpecT>` are single generic templates that know nothing about any particular kind of dataset. They hold exactly one payload (`data_type` / `view_type`, chosen by the spec) and expose only what every dataset has: `n_rows()`, `dim()`, `as_matrix_view()`, `as_dataset_view()` and `data()`. Each is a one-line forward to one of the three spec functions `get_data_view()`, `get_n_rows()` and `get_dim()`. Anything else a kind needs (e.g. VPQ codebooks, BBQ quantizers) is state and methods of that kind's payload type, reached through `data()`; `dataset`/`dataset_view` never name or branch on it. `dataset` and `dataset_view` are deliberately two independent, non-inheriting types (no shared_ptr, no "sometimes owning" object): `dataset` holds the owning payload, `dataset_view` the corresponding non-owning payload. - -```cpp -template -struct dataset; -``` - - -### neighbors::vpq_dataset_spec - -`Accessor` drives both codebook and code residency, mirroring today's - -single-`Accessor`-per-VPQ-dataset design. The payload (`detail::vpq_owning_storage` / `detail::vpq_view_storage`) holds the encoded rows and the VQ/PQ codebooks. - -```cpp -template -struct vpq_dataset_spec; -``` - - -### neighbors::dataset - -Owning dataset: value-held payload (no shared_ptr -- exclusive ownership). Every member is a - -one-line forward to `spec_type::get_*` or to the payload; all per-kind state and logic lives in the spec's `data_type`, never inside this struct. - -```cpp -template -struct dataset; -``` - - -### neighbors::dataset_view - -Non-owning dataset view: holds only the view-shaped payload. Deliberately not derived from - -`dataset` -- a view type holds "all view state" with no inheritance and no shared ownership tying it to the owning type. Reuses the same `get_n_rows`/`get_dim` spec functions as `dataset`, fed the view payload instead of the owning one. - -```cpp -template -struct dataset_view; -``` - - -### neighbors::is_padded_dataset - -Owning-side kind traits (mirror today's `is_padded_dataset_v`/`is_standard_dataset_v`/ - -`is_vpq_dataset_v`, used for SFINAE overload selection in factory.cuh/compute_distance_vpq.hpp). - -```cpp -template -struct is_padded_dataset; -``` - - -### neighbors::dataset_view_kind_of - -Primary template returns `unknown` so traits safely return `false` for non-dataset-view types. - -```cpp -template -struct dataset_view_kind_of { - static constexpr dataset_view_kind value; -}; -``` - -**Fields** - -| Name | Type | Description | -| --- | --- | --- | -| `value` | `static constexpr dataset_view_kind` | | - - -### neighbors::dataset_view_is_device_accessible - -True when the dataset view accessor is device-accessible. - -```cpp -template -struct dataset_view_is_device_accessible; -``` - - -### neighbors::with_accessor - -Generic accessor retargeting while preserving the dataset tag/layout and value/index types: - -`dataset<T, IdxT, SpecT<..., OldAccessor>> -> dataset<T, IdxT, SpecT<..., NewAccessor>>` `dataset_view<T, IdxT, SpecT<..., OldAccessor>> -> dataset_view<T, IdxT, SpecT<..., NewAccessor>>` - -```cpp -template -struct with_accessor; -``` - - -### neighbors::to_device_accessor - -Map any host accessor to its device counterpart (same payload policy). - -```cpp -template -struct to_device_accessor; -``` - - -### neighbors::device_counterpart - -Maps a host dataset view type to its device-resident counterpart. - -```cpp -template -struct device_counterpart; -``` - - -### neighbors::ivf::list_base - -Abstract base class for IVF list data. - -This allows polymorphic access to list data regardless of the underlying layout. - -TODO: Make this struct internal (tracking issue: https://github.com/nvidia/cuvs/issues/1726) - -```cpp -template -struct list_base; -``` - - -### neighbors::ivf::list - -The data for a single IVF list. - -```cpp -template