[REVIEW] New Dataset API Clarifying Ownership#1846
Conversation
|
/ok to test 5447a4c |
|
/ok to test 17ab09d |
|
NB: I updated the label to |
@achirkin The problem w/ using mdspan/mdarray for this is that it's not carrying along the proper information to either the algorithms nor the user (which is why we created this specialized class for this in the first place!). Two immediate reasons why this API is necessary:
This new API solves both of these problems while leaving the control over the memory ownership entirely in the user's hands. We've discussed this for a long time. We've known this is needed for a long time. it's time to prioritize this and get it done. I agree that an anstract class might make more sense, but ultimately we should not be moving any owneship over to the algorithm (the user should maintain ownership over the class and underlying memory the entire time). |
…tion between make host/device padded dataset in factory
… of dataset + create build_result struct which returns both index and vpq_dataset to prevent automatic out of scope destruction of dataset for vpq case
…rt for cases where we DO need to own the dataset (in order to keep view alive for index). All cases where we build() from dataset already on device --> we don't need to own. Merge + All cases when data is on host --> we DO need to own the device copy we create. This includes within ACE build and C API build from host and from_args with host dataset
|
The doc that outlines some of the API design choices can be found in slack. Let me know if there are any parts of the design that can be altered to better suit our users' needs. The following files are test case files I've added and can be ignored for now. They will be removed before the final merge with upstream repo:
|
|
|
||
| // Default owning/view accessors for public dataset aliases. | ||
| template <typename T> | ||
| using device_owning_accessor = raft::device_accessor<raft::device_container_policy<T>>; |
There was a problem hiding this comment.
Can we split this stuff out into include/core/dataset.hpp please? The dataset APIs are a common vocabulary type and it would be good to have them in core/ instead of neighbors/.
This way the quantize APIs can provide their dataset types and we can avoid any awkward circular references for files outside of neighbors/ to have to include or otherwise reference types within neighbors/.
I see it like this:
core/datasets.hpp -> core dataset types (for e.g. strided, empty, standard), factories, and aliasa
preprocessing/quantize/pq.hpp -> vpq_dataset types, factories, and aliases
preprocessing/quantize/rabitq.hpp -> rabitq_dataset types, factories, and aliases
neighbors/cagra.hpp -> uses datasets from core and preprocessing.
The layout above avoids circular dependencies across namespaces. If we don't do this, we'll ultimately have vpq_dataset in preprocessing/quantize depending upon dataset abstractions in neighbors/common.hpp and that makes the dependency tree awkward (as it's no longer a tree, but a graph).
|
/ok to test 1aeeb85 |
…r chain in index() within cagra.hpp
Co-authored-by: Divye Gala <divyegala@gmail.com>
Co-authored-by: Divye Gala <divyegala@gmail.com>
…ust checks row count and logical dim and applies to all datasets regardless of index type
|
/ok to test ed81c08 |
…dex. Remove cagra_build_dataset_ from common.hpp which is no longer used. Remove ann_build_pad_ from tiered_index.cu which is no longer used
|
/ok to test 5cd5d10 |
|
/ok to test 27bd293 |
| -10.0f, | ||
| 10.0f, | ||
| 42ULL, | ||
| raft::random::GenPC); |
There was a problem hiding this comment.
What is the reason for this change, is it related to the PR?
| out->sub_dataset_buffers_ = sub_dataset_buffers_; | ||
| out->deserialized_dataset_ = deserialized_dataset_; | ||
| out->sub_deserialized_datasets_ = sub_deserialized_datasets_; | ||
| return out; |
There was a problem hiding this comment.
Is there a better way to express this? It seems odd to copy a bunch of stuff line by line.
… field used later for distinguishing between padded and standard datasets. This index box also centralizes dispatch so dispatch can happen in one place rather than requiring a separate dispatch at every call site. We now have 3 places that use index box concept so duplicate code is extracted into a c_api_box.hpp file
…ructor for mg_index so deserialize can have index allocated in memory and passed in by reference to be populated by the deserialize function
| // Default owning/view accessors for public dataset aliases. | ||
| template <typename T> | ||
| using device_owning_accessor = raft::device_accessor<raft::device_container_policy<T>>; | ||
|
|
||
| template <typename T> | ||
| using host_owning_accessor = raft::host_accessor<raft::host_container_policy<T>>; | ||
|
|
||
| template <typename T> | ||
| using device_view_accessor = raft::device_accessor<cuda::std::default_accessor<const T>>; | ||
|
|
||
| template <typename T> | ||
| using host_view_accessor = raft::host_accessor<cuda::std::default_accessor<const T>>; |
There was a problem hiding this comment.
Couple issues with these aliases:
- Why do we need to differentiate the accessors for owning/non-owning?
- Please don't mix the accessor and container policy; use cuda::std::default_accessor
- Why some of the parameters T are const?
- why do we need these aliases in the detail namespace at all? I'd suggest to only define convenience aliases once in public namespaces (otherwise you quickly run out of unique names)
| */ | ||
|
|
||
| template <typename ContainerType, typename DataT, typename IdxT, typename Accessor> | ||
| struct dataset; |
There was a problem hiding this comment.
Coming back to the discussion at #1846 (comment)
Please try to explore the options to define the dataset with its common members once to avoid duplicating the code. In the current state, the dataset template itself still carries no information.
| * - IdxT: index type used for row counts (`n_rows()` return type). | ||
| */ | ||
| template <typename MatrixT, typename ViewT, typename DataT, typename IdxT> | ||
| struct dense_row_major_dataset_owning_storage { |
There was a problem hiding this comment.
We have already established that the dataset functionality mirrors the mdarray functionality for the case of dense array. Then, maybe, just use that fact and copy/reuse the code rather than trying to solve the same problem one more time? This could simplify the implementation of the dataset api and make it easier to convert between mdarrays and datasets.
More precisely, I'm suggesting to reuse layout/container/accessor policies from mdarrays in of the two ways:
a) copy the relevant code from mdarrays, so that transforming between the two is straightforward
b) use mdarray as the sole member of dense dataset struct and forward the relevant members.
…tend(). User must now call factory to allocate memory for the extended dataset and pass that in by reference to extend() function which will populate that buffer. Previously, extend did implicit memory allocation as a fallback path internally if optional buffer was not passed in. We don't want any internal implicit memory allocations so we removed that fallback
Overview
Addressing #1574 and #1571.
Replaced strided_dataset with padded_dataset class. Added support all the way up to CAGRA code.
Old class structure (Classes + Inheritance):
New Class Structure (ContainerType Tags + Composition):
Inheritance is removed entirely and all dataset types are on the same level of the inheritance tree.
3 Levels:
Ownership
The index and cagra::build / cagra::index do not own raw vector storage, they only take views.
The old code had a type-erased std::unique_ptr<dataset_view<...>>, i.e. non-owning view handles. The new code uses templates on the index type which determines the type of dataset_view the index holds.
ACE v.s. non-ACE paths on Host
ACE path copies datasets that can't entirely fit in CPU memory in chunks onto GPU memory by calling make_padded_dataset. This is 1x memory on CPU and 1x memory on GPU.
Return types:
Used mainly to maintain lifetime of dataset.
cuvs_cagra_c_api_lifetime_holder
It is a single C++ struct in cagra.cpp that groups the real cagra::index with any extra heap-owned things the C API had to create so the index’s non-owning views stay valid.
Miscellaneous: Extend Serialize Deserialize
Will fill in later
Factories:
Places where make_padded_dataset/view are called internally (not by user):
Host non-ACE path
Tiered CAGRA
Ownership in Downstream Functions:
Improvements:
Breaking Changes for Dataset API:
The following functions are removed since index no longer owns the dataset, index only takes views:
Removed old functions that took mdspan or derivatives of mdspan.
4 cases where index previously owned dataset [all deprecated paths]:
2 edge case build() paths when attach_dataset_on_build == true and a successful dense attach:
Compression Param:
Merge:
These paths have since been removed.
Attach Dataset
Compressed Dataset
Merged Dataset
Deserialize
Helpers
How to attach a compressed dataset onto an uncompressed index?
How to attach a searchable device dataset onto an index built with host build?
a. Utilizes map of host dataset type to device dataset type counterpart
TODOs:
Recent Updates:
Future PRs:
PR#2: Add Support for Compressed Datasets
PR#3: Migrate Rest of Algorithms to use Dataset API