diff --git a/cpp/include/cuvs/neighbors/common.hpp b/cpp/include/cuvs/neighbors/common.hpp index 1b943afe30..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 @@ -189,26 +190,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 +245,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 +256,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) + : 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 @@ -1339,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 { /**