[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). |
|
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:
|
|
/ok to test 70b6b58 |
achirkin
left a comment
There was a problem hiding this comment.
Sorry for being so late for the second round of reviews! I have a few suggestions for the dataset type hierarchy.
| * A populated instance is carried inside `merged_dataset_storage` together with the owning | ||
| * device matrices allocated by `make_merged_storage`. | ||
| */ | ||
| struct merged_dataset { |
There was a problem hiding this comment.
We're removing this aren't we? A merged dataset is just a new dataset with the rows merged. We don't need a new vocab type for this.
| */ | ||
| template <typename T, typename IdxT> | ||
| struct merged_dataset_storage { | ||
| merged_dataset layout{}; |
There was a problem hiding this comment.
Need to remove this too... it's dangerous and not maintainable to keep introducing new vocab types for things like this (especially when we already have things thast can store this).
…VIDIA#2128) ## Summary Small perf cleanup for the filtered brute-force CSR path in `knn_brute_force.cuh`. - When sparsity ≥ 0.9, filtered brute-force takes the sparse CSR / `masked_matmul` path. As part of setup it builds a per-nonzero `rows` array via `raft::sparse::convert::csr_to_coo`. - That `rows` array is only consumed by `cuvs::neighbors::detail::epilogue_on_csr`, which only runs for `L2Expanded` / `L2SqrtExpanded` / `CosineExpanded` (it combines masked inner products with precomputed norms). - For `InnerProduct`, the epilogue is skipped — so the allocation and the `csr_to_coo` kernel launch were dead work on every IP-metric filtered search hitting the CSR path. This PR hoists the `rmm::device_uvector<IdxT> rows(...)` allocation and the `csr_to_coo` call inside the L2/Cosine branch, next to their only consumer. ## What changes - No behavior change for `L2Expanded`, `L2SqrtExpanded`, `CosineExpanded` — same kernels, same order. - For `InnerProduct` filtered searches that hit the sparse CSR branch: saves one device allocation of `nnz * sizeof(IdxT)` and one kernel launch + write pass per call. ## Test plan - [ ] Existing filtered brute-force tests pass (`BRUTE_FORCE_TEST`, prefiltered variants) for `InnerProduct`, `L2Expanded`, `L2SqrtExpanded`, `CosineExpanded`. - [x] No diff in output for L2 / Cosine vs. base. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Authors: - Max Buckley (https://github.com/maxwbuckley) - Micka (https://github.com/lowener) - Anupam (https://github.com/aamijar) Approvers: - Micka (https://github.com/lowener) URL: NVIDIA#2128
| cuvs::neighbors::filtering::none_sample_filter{}) | ||
| -> cuvs::neighbors::cagra::index<half, uint32_t>; | ||
| template <typename T, typename IdxT, cuvs::neighbors::ann_dataset_view DatasetViewT> | ||
| merged_dataset_storage<T, IdxT> make_merged_storage( |
There was a problem hiding this comment.
Please remove. We don't need special "storage" objects. We discussed this offline.
| auto merge(raft::resources const& res, | ||
| const cuvs::neighbors::cagra::index_params& params, | ||
| std::vector<cuvs::neighbors::cagra::index<int8_t, uint32_t>*>& indices, | ||
| std::vector<cuvs::neighbors::cagra::index<T, IdxT, DatasetViewT>*>& indices, | ||
| merged_dataset_storage<T, IdxT>& storage, | ||
| const cuvs::neighbors::filtering::base_filter& row_filter = |
There was a problem hiding this comment.
@HowardHuang1 this looks like a proper row filter (input) not an output like I mentioned offline. In that case, the user should be able to apply the filter and know exactly how many items are going in the final dataset.
| * | ||
| * @internal | ||
| */ | ||
| struct fd_transfer { |
There was a problem hiding this comment.
why do we have detail things in the public apis? Please move this to cpp. If the user shouldn't be interacting with it, it doesn't belong in hpp.
| typename container_type::template owning_storage<DataT, IdxT, Accessor>; | ||
| using owning_storage_type::owning_storage_type; | ||
|
|
||
| [[nodiscard]] auto as_dataset_view() const noexcept |
There was a problem hiding this comment.
please just call this "view()" to match all the other vocab type.
| /** Compressed dataset. */ | ||
| raft::device_matrix<uint8_t, index_type, raft::row_major> data; | ||
|
|
||
| vpq_dataset(raft::device_matrix<math_type, uint32_t, raft::row_major>&& vq_code_book, |
There was a problem hiding this comment.
This is not cagra specific. Please move this to the respective preprocessing/pq.hpp file. Each different quantizer sohuld provide their respective dataset implementation so that when imported, the dataset comes w/ the quantizer.
| * - 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.
Agreed- for a single dense contiguous owning matrix type, we should be using mdarray... nothing else. If we need some type of composite, we should define that upon use... but i'm not seeing an immediate need for that in this PR.
| cuvsDatasetPaddedView_t padded_dataset, | ||
| cuvsCagraIndex_t index); | ||
|
|
||
| /** |
There was a problem hiding this comment.
What is a host cagra index? CAGRA is a GPU algorithn. This naming could be described with more clarity. Does "host cagra index" just imply the underlying dataset is on host? Because there is a potential configuration for CAGRA where the graph is always on host too.
Also, in C, we tend to put the verb last- "Make" "Create" "Destroy" "Attach" should all be last.
I would
There was a problem hiding this comment.
Yes host cagra index just refers to an index where the attached dataset is on host but the index graph itself is on device. Such a dataset is not immediately searchable.
…ng 26.10 to 26.08 versions for release
… not correct. Use owning factory first and then create view from that.
…taset() function. Attach_dataset() takes in any index type as input (host vs device, standard vs padded) and a device padded device and always returns a device padded index which is searchable. It internally dispatches on every combination of input index and input dataset type allowed and does internal host to device or standard to padded conversion if needed. We introduce 2 helper functions: convert_host_to_device_index() and convert_standard_to_padded_index(). These 2 helpers along with the original update_dataset(), which only takes device_padded_index and device_padded_dataset() as input, are internally called within attach_dataset() within each of the dispatch cases.
| extern "C" cuvsError_t cuvsDatasetMakeDevicePadded(cuvsResources_t res, | ||
| DLManagedTensor* dataset_tensor, | ||
| cuvsDatasetPadded_t* padded_dataset) |
There was a problem hiding this comment.
All those new cuvsDataset... functions should be in a separate compilation unit (src/core/dataset.cpp) (Not blocking this PR though, it can be a follow-up)
| DLDataType dtype; | ||
| cuvsDatasetLayout_t layout; | ||
| } cuvsDatasetStandard; | ||
| typedef cuvsDatasetStandard* cuvsDatasetStandard_t; |
There was a problem hiding this comment.
The two structs are similar, can't there be just one struct definition and two typedefs? They already represent the same thing.
There was a problem hiding this comment.
Same with the view struct just below
| void (*destroy_addr)(void*); | ||
| cuvsDatasetViewKind_t kind; | ||
| DLDataType dtype; | ||
| cuvsDatasetLayout_t layout; |
There was a problem hiding this comment.
layout is redundant if kind is already present in the struct
| cuvsCagraIndexParams_t params, | ||
| DLManagedTensor* dataset, | ||
| cuvsCagraIndex_t index); | ||
| CUVS_EXPORT cuvsError_t cuvsCagraGetDatasetViewKind(DLManagedTensor* dataset, |
There was a problem hiding this comment.
Any reason this needs to be exposed to the public?
| CUVS_EXPORT cuvsError_t cuvsCagraDeserializePadded(cuvsResources_t res, | ||
| const char* filename, | ||
| cuvsCagraIndex_t index, | ||
| cuvsDatasetPadded_t* out_padded_dataset); |
There was a problem hiding this comment.
Question- why isn't the dataset type serialized with the cagra graph (in the even the user wants that serialized with the cagra graph)? This way we can just do "cuvsCagraDeserializeWithDataset" (which throws an error if the serialized file doesn't contain the dataset) and cuvsCagraDeserialize (which doesn't deserialize the dataset, even if it's in the file).
There was a problem hiding this comment.
Pre-allocating the dataset buffer gives us 2 advantages: (1) no hidden memory allocations, (2) deterministic dataset type.
If we are pre-allocating the dataset buffer, I think we can just parse the type from that directly so there wouldn't be a need to modify the serialize algo itself to serialize the dataset type along with the graph.
If the proposal is to remove the passing in of the dataset buffer entirely in favor of serializing the dataset type with the cagra graph (which more closely mirrors the original serialize function), we would lose the beenfit of advantage (1).
| typedef enum { | ||
| CUVS_DATASET_VIEW_KIND_DEVICE_PADDED = 0, | ||
| CUVS_DATASET_VIEW_KIND_HOST_PADDED = 1, | ||
| CUVS_DATASET_VIEW_KIND_DEVICE_STANDARD = 2, |
There was a problem hiding this comment.
Instead of view kind containing both host and device for each dataset type (this is going to quickly become unmaintainable), please just define the accessor type (HOST and DEVICE) here and reuse the dataset layout enum above.
| DLDataType dtype; | ||
| cuvsDatasetLayout_t layout; | ||
| } cuvsDatasetPadded; | ||
| typedef cuvsDatasetPadded* cuvsDatasetPadded_t; |
There was a problem hiding this comment.
We discussed this offline- we should probably just make a single dataset struct and use the enums to parameterize it so that we know what kind to cast it to in the C++ layer.
… NEIGHBORS_ANN_CAGRA_tests
…update_dataset() that only updated standard index with standard dataset or padded index with padded dataset() was renamed update_device_dataset_same_layout(). New centralized cuvsCagraUpdateDataset at C API layer was created which internally dispatches to cuvsCagraUpdateDeviceDatasetSameLayout if index type is device_padded_index. For all other index types, it dispatches internally to cuvsCagraAttachDataset. This resolves the mutability issue we were previously experiencing where cuvsCagraAttachDataset guarantees original input index is immutable and returns a new index of new type while cuvsCagraUpdateDeviceDatasetSameLayout requires mutating the input dataset and makes no such index copy. This is not implemented at the C++ API layer yet, it only exists in the C API layer since C API doesn't make immutability guarantees while C++ API does and tries to keep C++ logic immutable to prevent thread conflicts. Downstream langauge wrappers have been updated to call centralized cuvsCagraUpdateDataset as the single sole entry point for all cases of updating/attaching a dataset
…ut from exports in C API layer and remove from downstream wrappers. They are not meant to be exposed to users/callers. Users/callers should use cuvsCagraUpdateDataset as the only C API function they interface with when updating/attaching a dataset
| _deserialize_padded<float>(res, filename, index, out_padded_dataset); | ||
| index->dtype.code = kDLFloat; | ||
| } else if (dtype.kind == 'e' && dtype.itemsize == 2) { | ||
| _deserialize_padded<half>(res, filename, index, out_padded_dataset); | ||
| index->dtype.code = kDLFloat; | ||
| } else if (dtype.kind == 'i' && dtype.itemsize == 1) { | ||
| _deserialize_padded<int8_t>(res, filename, index, out_padded_dataset); | ||
| index->dtype.code = kDLInt; | ||
| } else if (dtype.kind == 'u' && dtype.itemsize == 1) { | ||
| _deserialize_padded<uint8_t>(res, filename, index, out_padded_dataset); | ||
| index->dtype.code = kDLUInt; |
There was a problem hiding this comment.
I think there's a bug here in the order of these ops. The _deserialize_padded copies the index's dtype into the handle it returns:
out->dtype = output_index->dtype;So the dtype should already be set when it's called. Therefore, we should flip the two lines for each dtype kind branch:
index->dtype.code = ...
_deserialize_padded...
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