[REVIEW] New Dataset API Clarifying Ownership#1846
Conversation
aamijar
added
non-breaking
|
/ok to test 5447a4c |
|
/ok to test 17ab09d |
achirkin
added
breaking
|
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:
|
…taset_view The device path in BinaryCuvsCagra::train() was passing a raw mdspan to cagra::build() instead of a device_padded_dataset_view, causing: error: no instance of overloaded function "cuvs::neighbors::cagra::build" matches the argument list Fix by wrapping the device matrix view with make_device_padded_dataset_view before passing to build(), matching the pattern used in CuvsCagra<data_t>::train(). Also corrects the downstream hunk offsets (+183/+225/+297) in the BinaryCuvsCagra search() and reset() sections that were shifted by the 3-line addition.
update_graph(device_matrix_view) stores only a VIEW with no ownership transfer. When convert_host_to_device_index is called with a temporary host_idx (e.g. inside iface.hpp for the MG host-data build path), the returned device index ends up with a dangling graph_view_ once host_idx is destroyed, causing a segfault on first access. Fix: copy the graph device→host→device inside convert_host_to_device_index so the returned index always owns its graph memory. Remove the now-redundant workaround that was doing the same extra copy in cagra.cpp.
|
/ok to test e795819 |
…ed update_dataset() supporting host_matrix and device_matrix views were removed
|
/ok to test 26b87e5 |
…iew. All call sites updated to use Dataset API instead
…The main library already uses C++20
…ncepts. The main library already uses C++20" This reverts commit 48fc8d5.
|
/ok to test 217e5bd |
… to use the new Dataset API
|
/ok to test 8001bf4 |
|
/ok to test 3803e00 |
|
/ok to test c31ea7c |
divyegala
left a comment
divyegala
left a comment
There was a problem hiding this comment.
First pass review focussing on CAGRA header and dataset related headers
| auto v = raft::make_device_matrix_view<const T, int64_t>( | ||
| static_cast<const T*>(nullptr), int64_t{0}, uint32_t{0}); | ||
| return DatasetViewT(v, uint32_t{0}); | ||
| } else if constexpr (cuvs::neighbors::is_host_padded_dataset_view_v<DatasetViewT>) { |
There was a problem hiding this comment.
Why are we limiting to only host_padded and device_padded builds? You can build just fine without padding, and it is only search that needs device_padded
There was a problem hiding this comment.
I think this is a very valid concern! For this, one thing I can think of is we can create device_non_padded_dataset/view container types so that we can have build() take non_padded datasets to match the behavior in upstream main. However, is there any real use case/benefit of allowing users to build() with a non-padded dataset?
Enabling this flexibility makes the workflow complicated in 2 ways:
- The users that build with non-padded datasets will need to call the make_padded_dataset() factory + update_dataset() to attach the padded dataset after build and before search anyways. Why not just have them call the factory at the start of the pipeline and restrict build() to taking padded datasets as input only to avoid this 2nd code path?
- With templates, this gets messy. Cagra search currently operates on device_padded_index = index<device_padded_dataset_view>. However, if we allow builds with non_padded_datasets we will need to introduce a device_non_padded_index = index<device_non_padded_dataset_view>. Downstream functions like update_dataset() and merge() can't attach a padded view to a non-padded index. They must call the conversion helper which can make the code more convoluted. Functions like search(), extend(), and serialize/deserialize must add further overloads or dispatches on DatasetViewT.
Please let me know what you think!
| if constexpr (cuvs::neighbors::is_device_padded_dataset_view_v<DatasetViewT>) { | ||
| auto v = raft::make_device_matrix_view<const T, int64_t>( | ||
| static_cast<const T*>(nullptr), int64_t{0}, dim_); | ||
| dataset_ = DatasetViewT(v, dim_); | ||
| } else if constexpr (cuvs::neighbors::is_host_padded_dataset_view_v<DatasetViewT>) { | ||
| auto v = raft::make_host_matrix_view<const T, int64_t>( | ||
| static_cast<const T*>(nullptr), int64_t{0}, dim_); | ||
| dataset_ = DatasetViewT(v, dim_); | ||
| } else if constexpr (cuvs::neighbors::is_empty_dataset_view_v<DatasetViewT>) { |
There was a problem hiding this comment.
This doesn't seem like it is building a padded dataset. To clarify, padding adds 0s.
There was a problem hiding this comment.
Yes you're correct. Here, update_dataset() doesn't build a padded_dataset. Instead, it configures disk-backed mode. The null host/device_padded_dataset_view is only a type placeholder here to satisfy the DatasetViewT template parameter on index since index type needs to match the type of the dataset to be attached.
The real state here is: dataset_fd_ + n_rows_/dim_
At all times we have 3 pieces of state stored in the private member variables of index:
- dataset_ (DatasetViewT) — typed, required
- dataset_fd_ — actual disk backing
- n_rows_ / dim_ — metadata when fd is set
In this case, since this is in the file descriptor overload of update_dataset(), only states (2) and (3) are active. State (1) is a dummy.
I get how this can be confusing since it seems like we are declaring it to be a padded dataset when in reality, it isn't padded at all (rather it's empty). However, a fix for this would require making the codebase more convoluted.
One way is creating a new type let's call it host_empty_index. We will need to re-wire ACE build to return this index type instead of the original host_padded_index. To do this, we may need to potentially wrangle with a function returning 2 different return types. However, in C++ we can't have a single build() function return 2 different return types depending on code path.
One fix would be having the ENTIRE host build path return this new host_empty_index (we need to have disk-backed ACE return host_empty_dataset. This forces us to have in-memory ACE updated to using host_empty_dataset as well to avoid 2 return types. Then we hit non-ACE host returning a different return type so it forces us to change that as well). Furthermore, this involves updating all call sites. However, this can be done since host build is not searchable anyways, so effectively it is an empty dataset but in this process we lose the type of dataset the index was originally built with and will cause confusion amongst users when they call downstream functions that now work on index with empty dataset type.
Is there a simpler solution for this? Or can we leave this as is?
| // Concrete non-template overloads for all supported index types. | ||
| // Previously a single template <T, IdxT, DatasetViewT, OutputIdxT> covered all index types; it | ||
| // has been replaced with explicit overloads to maintain a stable non-template ABI. When a new | ||
| // index type is added, add corresponding overloads here. Index types whose search is not yet | ||
| // implemented (e.g. vpq_f32_index) are still declared so the symbols exist when the | ||
| // implementation lands. |
| * @param[in] sample_filter an optional device filter function object that greenlights samples | ||
| * for a given query. (none_sample_filter for no filtering) | ||
| */ | ||
| void search(raft::resources const& res, |
There was a problem hiding this comment.
All the overloads seem to be missing docs
| } | ||
|
|
||
| template <typename T, typename IdxT> | ||
| auto dataset_view_to_strided_device_matrix(device_padded_dataset_view<T, IdxT> const& view) |
There was a problem hiding this comment.
Remind me please where strided_device_matrix is used?
There was a problem hiding this comment.
Should this file be in detail namespace?
| struct device_empty_dataset_container {}; | ||
| struct host_padded_dataset_container {}; | ||
| struct device_padded_dataset_container {}; | ||
| struct host_vpq_dataset_container {}; |
There was a problem hiding this comment.
Let's not add tags for containers that don't currently exist
| * | ||
| * The first template parameter `containertype` on `dataset` / `dataset_view` is one of these types. | ||
| */ | ||
| struct host_empty_dataset_container {}; |
There was a problem hiding this comment.
These tags shouldn't just be an empty type, it can at least store storage_type and view_type
| [[nodiscard]] constexpr inline auto encoded_row_length() const noexcept -> uint32_t | ||
| { | ||
| return data.extent(1); | ||
| } | ||
| [[nodiscard]] constexpr inline auto vq_n_centers() const noexcept -> uint32_t | ||
| { | ||
| return vq_code_book.extent(0); | ||
| } | ||
| [[nodiscard]] constexpr inline auto pq_bits() const noexcept -> uint32_t | ||
| { | ||
| auto pq_width = pq_n_centers(); | ||
| #ifdef __cpp_lib_bitops | ||
| return std::countr_zero(pq_width); | ||
| #else | ||
| uint32_t pq_bits = 0; | ||
| while (pq_width > 1) { | ||
| pq_bits++; | ||
| pq_width >>= 1; | ||
| } | ||
| return pq_bits; | ||
| #endif | ||
| } | ||
| [[nodiscard]] constexpr inline auto pq_dim() const noexcept -> uint32_t | ||
| { | ||
| return raft::div_rounding_up_unsafe(dim(), pq_len()); | ||
| } | ||
| [[nodiscard]] constexpr inline auto pq_len() const noexcept -> uint32_t | ||
| { | ||
| return pq_code_book.extent(1); | ||
| } | ||
| [[nodiscard]] constexpr inline auto pq_n_centers() const noexcept -> uint32_t | ||
| { | ||
| return pq_code_book.extent(0); | ||
| } |
There was a problem hiding this comment.
These member functions are actually a property of the container type and that's where they should be
achirkin
left a comment
achirkin
left a comment
There was a problem hiding this comment.
Thanks for the updates! I believe we still have a room for improvement on the dataset definition side (I haven't looked past the public headers yet).
| device_empty, | ||
| host_empty, | ||
| device_padded, | ||
| host_padded, | ||
| device_vpq_f16, | ||
| host_vpq_f16, | ||
| device_vpq_f32, | ||
| host_vpq_f32, |
There was a problem hiding this comment.
We have the accessor template parameter to distinguish between host/device/both accessible dataset. We shouldn't need to duplicate that information in the dataset kind enum.
| template <typename HostViewT> | ||
| struct device_counterpart; | ||
|
|
||
| template <typename DataT, typename IdxT> | ||
| struct device_counterpart<host_padded_dataset_view<DataT, IdxT>> { | ||
| using type = device_padded_dataset_view<DataT, IdxT>; | ||
| }; | ||
|
|
||
| template <typename MathT, typename IdxT> | ||
| struct device_counterpart<host_vpq_dataset_view<MathT, IdxT>> { | ||
| using type = device_vpq_dataset_view<MathT, IdxT>; | ||
| }; | ||
|
|
||
| template <typename IdxT> | ||
| struct device_counterpart<host_empty_dataset_view<IdxT>> { | ||
| using type = device_empty_dataset_view<IdxT>; | ||
| }; | ||
|
|
||
| template <typename HostViewT> | ||
| using device_counterpart_t = typename device_counterpart<dataset_view_type_t<HostViewT>>::type; |
There was a problem hiding this comment.
I think, implementing this logic via ad-hoc polymorphism is error-prone and not very maintainable.
We should be able to cover the conversion via a single mapping/function like "dataset<Tag, DataT, IdxT, AccessorA> -> dataset<Tag, DataT, IdxT, AccessorB>".
Let's try to avoid mixing the accessibility (device/host) with the data layout tags (empty/plain/vpq/padded etc).
| template <typename containertype, typename DataT, typename IdxT> | ||
| struct dataset { | ||
| static_assert(!std::is_same_v<containertype, containertype>, | ||
| "dataset: unsupported containertype / type-parameter combination"); | ||
| }; | ||
|
|
||
| template <typename containertype, typename DataT, typename IdxT> | ||
| struct dataset_view { | ||
| static_assert(!std::is_same_v<containertype, containertype>, | ||
| "dataset_view: unsupported containertype / type-parameter combination"); | ||
| }; |
There was a problem hiding this comment.
We've been discussing for a while that one of the main goals of this refactor is to untangle the accessibility property (device/host/etc) from the container types (empty/padded/vpq). We're not there yet. The description mentions the host/device accessible template parameters, but they are missing in the definitions here.
template <typename DatasetTag, typename DataT, typename IdxT, AccessorPolicy Accessor>
struct dataset {
...
};Then, as an example, you get this specialization for the padded dataset:
template <typename DataT, typename IdxT, AccessoryPolicy Accessor>
struct dataset<device_padded_dataset_container, DataT, IdxT, Accessor> {
using storage_type = raft::mdarray<DataT, raft::matrix_extent<IdxT>, std::layout_left_padded, Accessor>;
...
}Here accessor can be anything: device, host, managed, etc; you don't need to copy it twice for every type of the accessor.
*Also note: we can now use std::layout_left_padded or std::layout_stride to natively support the padded/strided layout rather than using a workaround with layout_c_contiguous + dim.
| template <typename containertype, typename DataT, typename IdxT> | ||
| struct dataset { | ||
| static_assert(!std::is_same_v<containertype, containertype>, | ||
| "dataset: unsupported containertype / type-parameter combination"); |
There was a problem hiding this comment.
I'm afraid, the current empty dataset definition as-is serves no purpose at all, because it has no members.
Please investigate how to use composition (like we do in raft mdarrays) or (as a last resort) CRTP to define the dataset with useful member definitions and define the per-implementation functionality without repeating all members in all instances manually. You can have a look at my past experiments for an example of how PQ compression can be implemented as as container policy without modifying the parent (in my case it's mdarray, in your case - dataset).
10 participants
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):
Intermediate Class Structure (Classes + Inheritance):
dataset and dataset_view are now 2 separate parent classes.
Intermediate Class Structure (ContainerType Tags + Composition + Variants):
New Class Structure (ContainerType Tags + Composition):
Inheritance is removed entirely and all dataset types are on the same level of the inheritance tree.
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.
cuvs_cagra_c_api_lifetime_holder is a separate heap object from cagra::index. It is heap-allocated in cagra.cpp with new cuvs_cagra_c_api_lifetime_holder<...>. The C API keeps a raw pointer to it in cuvsCagraIndex.cuvs_cagra_c_api_lifetime_holder It is not embedded in the index, which is why the C layer needs that second field to delete the holder on destroy.
Heap-allocated bundle for the C API: owns
cagra::indexand any co-owned device storage (VPQ, padded dataset copy, merge/de-serialize/extend buffers) when the index is not standalone.cuvsCagraIndex.c_api_lifetime_ownerpoints at this. Used for merge, build, deserialize, from_args, extend.The holder moves the owning device_padded_dataset (as unique_ptr<dataset<>> in padded_dataset_owner) to the heap, and cuvsCagraIndex.merged_owner points at the holder. Destroying the C index later destroys the holder first, so the dataset outlives the index’s use of the view, or the ordering is set up so the view is not used after free.
In cuvsCagraIndexFromArgs in cagra.cpp (C API) where callers are things like the Python cagra.from_graph (via Cython) and the Java CagraIndexImpl, and any C code that uses that function:
The flow is: caller → cuvsCagraIndexFromArgs → _from_args, which writes into the cuvsCagraIndex struct the user passed
The holder is not returned as a separate C return value. It is allocated on the heap and its address is stored in output_index->merged_owner, and output_index->addr points at the index inside that holder (or at a freestanding index when merged_owner == 0).
So when _from_args returns, the user’s cuvsCagraIndex already holds the pointers that describe where everything lives.
The unique_ptr to the copy of the dataset from make_padded_dataset is not local to _from_args—it is a member of the holder, which is on the heap and stays alive.
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.
3 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:
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