[REVIEW] New Dataset API Clarifying Ownership

[HowardHuang1]

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:

1. Storage

2. Container tags

3. Public dataset/view aliases

Ownership

The index and cagra::build / cagra::index do not own raw vector storage, they only take views.

• callers (or the C merged holder) must keep backing memory alive for as long as the index is used.

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

• unique_ptr<vpq_dataset> vpq_owner
• unique_ptr padded_dataset_owner
• raft::device_matrix dataset
• cagra::index idx
  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:

• make_device_padded_dataset_view
• make_host_padded_dataset_view
• make_device_padded_dataset
• make_host_padded_dataset
• make_vpq_dataset
  • in pq.hpp and pq.cu
• make_merged_dataset
  • in cagra.hpp

Places where make_padded_dataset/view are called internally (not by user):

Host non-ACE path

• cpp/src/neighbors/cagra_build_inst.cu.in
• cagra_from_host_padded in cpp/src/neighbors/iface/iface.hpp
• c/src/neighbors/cagra.cpp

Tiered CAGRA

• update_cagra_ann_dataset_for_stride
• build_upstream_ann

Ownership in Downstream Functions:

• build() takes dataset_view only.
• Downstream functions search / serialize / deserialize / merge only take views.

Improvements:

• build() function previously only supported device dataset inputs. It now supports host dataset inputs.

Breaking Changes for Dataset API:

The following functions are removed since index no longer owns the dataset, index only takes views:

• Removed all owning dataset based builds. Build only takes views.
• Removed all update_dataset() overloads that take owning dataset. Update_dataset() 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:

• Non-ACE / typical padded attach: rows live under index_owning_dataset_storage_ (type-erased owning wrapper, commonly device_padded_dataset).
• ACE in-memory device_matrix attach: rows live under index_owning_dataset_storage_ (optional raw device_matrix).

Compression Param:

• implicit vpq dataset creation within build() when compression==True
• this forces index to own new vpq dataset which violates our new contract that we want index to never own dataset.

Merge:

• MERGE path: merge() internally creates merged_dataset on a deprecated internal merged dataset creation path. Here, index takes ownership of merged_dataset by storing it in index_owning_dataset_storage_ .

These paths have since been removed.

Attach Dataset

• Previously, in the old code, ACE attach_dataset_on_build called make_device_padded_dataset on host dataset which did a H2D copy in order to attach dataset to final index.
• cpp/src/neighbors/detail/cagra/cagra_build.cuh
• This has since been removed. attach_dataset_on_build is disabled for host build paths. This avoids a H2D copy.

Compressed Dataset

• Removed old code that did compressed dataset creation within build. This should only happen in the factory.

Merged Dataset

• Removed implicit memory allocation within merge(), memory allocation now delegated to make_merged_dataset() factory. Removed index ownership within merge.

Deserialize

• Removed index ownership of dataset during deserialization. Now users are expected to create/declare the dataset type to be deserialized and then pass it as a reference to the deserialize() function which will then populate this dataset and return it to the caller.

Helpers

• cagra_required_row_width
• matrix_actual_row_width
• matrix_row_width_matches_cagra_required
• convert_dataset_view_to_padded_for_graph_build
• convert_host_to_device_index
• attach_device_dataset_on_host_index

How to attach a compressed dataset onto an uncompressed index?

1. construct a new compressed index
2. copy over the graph and other params from the uncompressed index
3. delete the old uncompressed index
4. attach the vpq dataset onto the compressed index

How to attach a searchable device dataset onto an index built with host build?

1. Convert host index to device index with helper function convert_host_to_device_index
   a. Utilizes map of host dataset type to device dataset type counterpart

TODOs:

• Bring back Host functions [DONE]
• Mark any old functions that are no longer used as [DONE]
• Use templates wherever possible. Shift towards composition rather than inheritance [DONE]

Recent Updates:

• build_ace() and build() functions merged on Public API surface
• removed build_result(), ace_build_result(), and merge_result()
• deprecated internal vpq dataset creation inside build() when index_params::compression == true --> moved to make_vpq_dataset() factory
• deprecated internal merged dataset creation inside merge() --> moved to make_merged_dataset() factory. For backwards compatibility, have index take ownership of deprecated internal merged dataset creation path ONLY.
• build() and build_ace() both had a attach_dataset_on_build which requires ownership of dataset. Ownership is given to index temporarily. This will later be deprecated. Users will be expected to pass padded dataset on device and call search() directly. Attach_dataset_on_build will no longer be supported for host builds.
• added host versions of dataset API
• templated build() and downstream functions to work on host datasets
• added index template type conversion helpers

Future PRs:

PR#2: Add Support for Compressed Datasets

• pq_dataset
• bbq_dataset
• rabitq_dataset
• sq_dataset

PR#3: Migrate Rest of Algorithms to use Dataset API

• HNSW
• IVF
• Vamana
• Scann
• Brute Force

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[aamijar]

/ok to test 5447a4c

[aamijar]

/ok to test 17ab09d

[achirkin]

NB: I updated the label to breaking, since the description implies removal of a publicly visible class strided_dataset

[cjnolet]

Does the dataset(_view) type bring anything on top of mdarray/mdspan in that case?

@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:

1. The user should not have to know that they need to pad a dataset in order to use cagra without the additional copy. They should not need to know how any of these algorithms work internally. They should, however, need to know that CAGRA expects a padded dataset, and they should have an API to construct one so that they can own the dataset class and not have cagra creating one under the hood.
2. APIs, especially the graph-based APIs, should be able to accept as inputs data which has been quantized using a metod like PQ, which carries with additional information. In the case of PQ, the codebooks are needed to compute the distances. This again decouples the quantization from the algorithm (CAGRA-Q does not need to do its own quantization. It should just accept the quantized vectors). We're being asked for the same behavior with Vamana.

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).

[HowardHuang1]

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:

• cagra_build_view_only.cu
• cagra_padded_dataset.cu
• cagra_vpq_build_result.cu
• dataset_compression.cu
• dataset_types.cu

[cjnolet]

/ok to test 70b6b58

[achirkin]

Sorry for being so late for the second round of reviews! I have a few suggestions for the dataset type hierarchy.

[aamijar]

/ok to test 2b753ac

[HowardHuang1]

/ok to test b4f7d3e

[HowardHuang1]

/ok to test 2ffe67b

[HowardHuang1]

/ok to test 80905ea

My change requests are outdated or postponed to follow-up PRs.

[HowardHuang1]

/ok to test a6dd4eb