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 1b943afe30..46ef145412 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 @@ -147,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. @@ -189,26 +206,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 @@ -223,18 +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. + * 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). @@ -243,377 +266,442 @@ 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) + : ViewT(v), logical_dim_(logical_dim) { } - dense_row_major_dataset_view_storage(dense_row_major_dataset_view_storage const& other) noexcept - : data_(other.data_), logical_dim_(other.logical_dim_) - { - } - - [[nodiscard]] auto n_rows() const noexcept -> IdxT { return data_.extent(0); } + [[nodiscard]] auto 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 -using padded_dataset_owning_storage = - dense_row_major_dataset_owning_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; + + [[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 padded_dataset_view_storage = dense_row_major_dataset_view_storage; +} // namespace detail -template -using standard_dataset_owning_storage = - dense_row_major_dataset_owning_storage; +// ----------------------------------------------------------------------------- +// Public specs -- the only place per-kind logic lives. +// ----------------------------------------------------------------------------- -template -using standard_dataset_view_storage = dense_row_major_dataset_view_storage; +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; + 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(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 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 (`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; + + 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_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()); + } + }; + + [[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 // ----------------------------------------------------------------------------- -/** - * 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)} +/** 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 { + 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. + template + explicit dataset(Args&&... args) + requires(!compressed_dataset_spec && std::is_constructible_v) + : data_(std::forward(args)...), dictionary_{} { } - [[nodiscard]] auto n_rows() const noexcept -> IdxT { return data.extent(0); } - [[nodiscard]] auto dim() const noexcept -> uint32_t { return vq_code_book.extent(1); } + // 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)) + { + } - [[nodiscard]] constexpr inline auto encoded_row_length() const noexcept -> uint32_t + [[nodiscard]] auto n_rows() const noexcept -> index_type { return spec_type::get_n_rows(data_); } + [[nodiscard]] auto dim() const noexcept -> uint32_t { - return data.extent(1); + return spec_type::get_dim(data_, dictionary_); } - [[nodiscard]] constexpr inline auto vq_n_centers() const noexcept -> uint32_t + [[nodiscard]] auto data_view() const noexcept { return spec_type::get_data_view(data_); } + [[nodiscard]] auto dictionary_view() const noexcept { - return vq_code_book.extent(0); + return spec_type::get_dictionary_view(dictionary_); } - [[nodiscard]] constexpr inline auto pq_bits() const noexcept -> uint32_t + + [[nodiscard]] auto as_dataset_view() const noexcept -> dataset_view { - auto pq_width = pq_n_centers(); -#ifdef __cpp_lib_bitops - return std::countr_zero(pq_width); -#else - uint32_t pq_bits = 0; - while (pq_width > 1) { - pq_bits++; - pq_width >>= 1; - } - return pq_bits; -#endif + return dataset_view(data_view(), dictionary_view()); } - [[nodiscard]] constexpr inline auto pq_dim() const noexcept -> uint32_t + + // 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 raft::div_rounding_up_unsafe(dim(), pq_len()); + return std::move(dictionary_); } - [[nodiscard]] constexpr inline auto pq_len() const noexcept -> uint32_t + + // 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 pq_code_book.extent(1); + return spec_type::get_encoded_row_length(data_); } - [[nodiscard]] constexpr inline auto pq_n_centers() const noexcept -> uint32_t + [[nodiscard]] auto vq_n_centers() const noexcept + requires requires(decltype(dictionary_view()) const& d) { d.vq_n_centers(); } { - return pq_code_book.extent(0); + return dictionary_view().vq_n_centers(); } -}; - -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) + [[nodiscard]] auto pq_n_centers() const noexcept + requires requires(decltype(dictionary_view()) const& d) { d.pq_n_centers(); } { - RAFT_EXPECTS(ptr != nullptr, "vpq_dataset_view: null dataset pointer"); + return dictionary_view().pq_n_centers(); } - - [[nodiscard]] auto n_rows() const noexcept + [[nodiscard]] auto pq_len() const noexcept + requires requires(decltype(dictionary_view()) const& d) { d.pq_len(); } { - using idx_type = decltype(std::declval().n_rows()); - return dataset_ != nullptr ? dataset_->n_rows() : idx_type{0}; + return dictionary_view().pq_len(); } - [[nodiscard]] auto dim() const noexcept -> uint32_t + [[nodiscard]] auto pq_bits() const noexcept + requires requires(decltype(dictionary_view()) const& d) { d.pq_bits(); } { - return dataset_ != nullptr ? dataset_->dim() : uint32_t{0}; + 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(); } - [[nodiscard]] owning_dataset_type const& dset() const noexcept { return *dataset_; } -}; - -} // namespace detail - -// ----------------------------------------------------------------------------- -// empty -// ----------------------------------------------------------------------------- - -struct empty_dataset_container { - template - using owning_storage = detail::empty_dataset_owning_storage; - template - using view_storage = detail::empty_dataset_view_storage; -}; - -// ----------------------------------------------------------------------------- -// padded (row-major with logical dim vs stride) -// ----------------------------------------------------------------------------- - -struct padded_dataset_container { - template - using owning_storage = - detail::padded_dataset_owning_storage, - detail::dense_view_matrix, - DataT, - IdxT>; - template - using view_storage = detail:: - padded_dataset_view_storage, DataT, IdxT>; -}; - -// ----------------------------------------------------------------------------- -// standard (row-major with arbitrary stride; no CAGRA alignment requirement) -// ----------------------------------------------------------------------------- - -struct standard_dataset_container { - template - using owning_storage = - detail::standard_dataset_owning_storage, - detail::dense_view_matrix, - DataT, - IdxT>; - template - using view_storage = detail:: - standard_dataset_view_storage, DataT, IdxT>; -}; - -// ----------------------------------------------------------------------------- -// VPQ compressed -// ----------------------------------------------------------------------------- - -struct vpq_dataset_container { - template - using owning_storage = - detail::vpq_dataset_owning_storage, - detail::vpq_vq_book_matrix, - detail::vpq_data_matrix, - MathT, - IdxT>; - template - using view_storage = - detail::vpq_dataset_view_storage; -}; -template -struct dataset { - static_assert(!std::is_same_v, - "dataset: unsupported ContainerType / type-parameter combination"); + private: + data_type data_; + [[no_unique_address]] dictionary_type dictionary_; }; -template +/** 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 { - static_assert(!std::is_same_v, - "dataset_view: unsupported ContainerType / type-parameter combination"); -}; - -// ----------------------------------------------------------------------------- -// empty -// ----------------------------------------------------------------------------- - -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> + 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; + + 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->dim()}; } -}; -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; -}; - -// ----------------------------------------------------------------------------- -// standard (row-major with arbitrary stride) -// ----------------------------------------------------------------------------- - -template -struct dataset - : standard_dataset_container::template owning_storage { - using container_type = standard_dataset_container; - using owning_storage_type = - typename container_type::template owning_storage; - using owning_storage_type::owning_storage_type; - - [[nodiscard]] auto as_dataset_view() const noexcept - -> dataset_view> + // 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_{} { - return dataset_view>(this->view(), - this->dim()); } -}; - -template -struct dataset_view - : standard_dataset_container::template view_storage { - using container_type = standard_dataset_container; - using view_storage_type = typename container_type::template view_storage; - using view_storage_type::view_storage_type; -}; - -// ----------------------------------------------------------------------------- -// padded (row-major with logical dim vs stride) -// ----------------------------------------------------------------------------- -template -struct dataset - : padded_dataset_container::template owning_storage { - using container_type = padded_dataset_container; - using owning_storage_type = - typename container_type::template owning_storage; - using owning_storage_type::owning_storage_type; - - [[nodiscard]] auto as_dataset_view() const noexcept - -> dataset_view> + [[nodiscard]] auto n_rows() const noexcept -> index_type { - return dataset_view>(this->view(), - this->dim()); + 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_view - : padded_dataset_container::template view_storage { - using container_type = padded_dataset_container; - using view_storage_type = typename container_type::template view_storage; - using view_storage_type::view_storage_type; -}; - -// ----------------------------------------------------------------------------- -// VPQ compressed (view holds non-owning pointer to owning dataset) -// ----------------------------------------------------------------------------- -template -struct dataset - : vpq_dataset_container::template owning_storage { - using container_type = vpq_dataset_container; - using owning_storage_type = - typename container_type::template owning_storage; - using owning_storage_type::owning_storage_type; - - [[nodiscard]] auto as_dataset_view() const - -> dataset_view> + // 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_{}; }; /** @@ -621,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 = device_standard_dataset; +template +struct owning_dataset_for_view> { + using type = dataset; }; -template -struct owning_dataset_for_view> { - using type = host_padded_dataset; -}; +template +using owning_dataset_for_view_t = typename owning_dataset_for_view::type; -template -struct owning_dataset_for_view> { - using type = host_standard_dataset; -}; +// ----------------------------------------------------------------------------- +// Spec-kind classification (all derived from SpecT; dataset/dataset_view never branch on kind). +// ----------------------------------------------------------------------------- -template -struct owning_dataset_for_view> { - using type = device_vpq_dataset; +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; -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; @@ -778,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 @@ -785,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 = @@ -923,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_view; +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 @@ -974,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 @@ -993,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). @@ -1306,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, 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