Reuse raft::mdarray/mdspan directly in dense dataset storage - #2530
Open
HowardHuang1 wants to merge 4 commits into
Open
Conversation
…2395) dense_owning_matrix/dense_view_matrix (and the VPQ codebook/code matrix aliases) picked between raft::device_matrix and raft::host_matrix via std::conditional_t, even though those are themselves just aliases for raft::mdarray/raft::mdspan with the exact accessor already computed as Accessor. Point the aliases at raft::mdarray/raft::mdspan directly instead. dense_row_major_dataset_owning_storage/_view_storage also wrapped their matrix/view as a field and hand-forwarded view()/data_handle(), which raft::mdarray/raft::mdspan already provide natively. They now inherit from the matrix/view type instead, so those forwards are no longer needed. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
…dataset-functions-one-level-up-and-reuse-mdarray-26_10
as_dataset_view()/n_rows()/dim() can be shared across all dataset
kinds with zero per-kind dispatch inside dataset/dataset_view
themselves. Adds cuvs::neighbors::experimental::{dataset,dataset_view}
as a Spec/policy-based prototype: every member is a one-line forward
to spec_type::get_*(), with all kind-specific logic (empty_spec,
mdarray_spec, vpq_spec) living outside dataset/dataset_view, which
never name or branch on a concrete kind. dataset and dataset_view
stay two independent, non-inheriting types (no shared_ptr, no
"sometimes owning" object).
Purely additive: lives in its own namespace, not referenced by any
existing type, alias, or call site. Verified via a standalone sandbox
(all three specs, both dataset and dataset_view, asserting n_rows/dim/
as_dataset_view) and a full rebuild + DATASET_C_TEST (7/7) +
CAGRA_C_TEST (14/14) + PREPROCESSING_TEST (226/226) + NEIGHBORS_TEST
(339/339), all passing unchanged from before this commit.
Limitations / not yet done:
- Not wired up: the real dataset/dataset_view types (Container-tagged)
and every downstream consumer (CAGRA build/search, serialization,
the C API, compile-time classification traits like
dataset_view_kind_of/is_padded_dataset_view_v) are untouched. This
prototype does not replace them yet.
- mdarray_spec's layout choice for the padded case is unverified
against CAGRA's actual alignment requirements -- needs confirmation
before real use.
- No sparse (CSR/COO) or scalar-quantized specs; only empty/dense/vpq,
matching today's four kinds minus the empty/dense split.
- This design is still an open thread with the team, not finalized;
land as prototype only, pending further review.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
…VIDIA#2395) common.hpp's dataset<ContainerType,DataT,IdxT,Accessor>/dataset_view<...> (four hand-specialized structs tagged by ContainerType) are replaced by dataset<T,IdxT,SpecT>/dataset_view<T,IdxT,SpecT>: two independent, non-inheriting types where every member is a one-line forward to a per-kind Spec (empty_dataset_spec, padded_dataset_spec, standard_dataset_spec, vpq_dataset_spec) -- no dispatch inside dataset/dataset_view themselves. VPQ's codebooks move from three flat, ungrouped matrices (vq_code_book, pq_code_book, data as direct members) into a dictionary_type{vq_code_book, pq_code_book} bundle alongside data, giving every kind the same shape (one data slot + one optional dictionary slot). The old dataset_view -> owner back-pointer (.dset()) is removed; a VPQ view now holds its own dictionary_view() directly, matching padded/standard views holding their own state. Public aliases (device_padded_dataset<T,IdxT> etc.) keep their names and 2-arg construction signatures, so most call sites are unaffected. Call sites that reached into dataset internals directly needed updating: .view()/.stride() on a dataset or dataset_view -> .data_view(); raw .vq_code_book/.pq_code_book/.data member access -> .dictionary_view()/ .data_view(); old 3-arg VPQ construction -> 2-arg (codes, dictionary_type{vq,pq}). Touches CAGRA build/search/serialize/merge, VPQ training (pq.cuh), SCANN, Vamana, HNSW export, multi-GPU CAGRA, and the C API's product-quantizer codebook accessors (pq.cpp) -- the one place the C API reaches into dataset internals rather than going through the opaque cuvsDataset handle. Also: removes detail::vpq_dataset_spec_impl (only empty/padded/standard/ vpq_dataset_spec exist as public specs; padded and standard share dense_dataset_spec_impl since two tags need the same body, but VPQ had no second tag to share with, so its body is now inlined directly into vpq_dataset_spec, consistent with how empty_dataset_spec is already written). Verified: full rebuild clean; DATASET_C_TEST (7/7), CAGRA_C_TEST (14/14), PREPROCESSING_TEST (226/226), NEIGHBORS_TEST (339/339) all pass. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Addresses #2395.
dense_owning_matrix/dense_view_matrix (and the VPQ codebook/code matrix aliases) picked between raft::device_matrix and raft::host_matrix via std::conditional_t, even though those are themselves just aliases for raft::mdarray/raft::mdspan with the exact accessor already computed as Accessor. Point the aliases at raft::mdarray/raft::mdspan directly instead.
dense_row_major_dataset_owning_storage/_view_storage also wrapped their matrix/view as a field and hand-forwarded view()/data_handle(), which raft::mdarray/raft::mdspan already provide natively. They now inherit from the matrix/view type instead, so those forwards are no longer needed.