update_dataset CAGRA - #2427
Conversation
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
| #include <numeric> | ||
| #include <optional> | ||
| #include <string> | ||
| #include <type_traits> |
There was a problem hiding this comment.
Cleaning up some unused headers and reordering.
| typename IdxT, | ||
| cuvs::neighbors::ann_dataset_view DatasetViewT = | ||
| cuvs::neighbors::device_padded_dataset_view<T, int64_t>> | ||
| struct index; |
There was a problem hiding this comment.
Removing this unnecessary forward declaration.
| index(raft::resources const& res, | ||
| cuvs::distance::DistanceType metric = cuvs::distance::DistanceType::L2Expanded) | ||
| explicit index(raft::resources const& res, | ||
| cuvs::distance::DistanceType metric = cuvs::distance::DistanceType::L2Expanded) |
There was a problem hiding this comment.
Adding explicit keyword to avoid unintentional implicit conversions.
| graph_degree_(other.graph_degree_) | ||
| { | ||
| update_dataset(res, dataset); | ||
| } |
There was a problem hiding this comment.
The core implementation of the free function update_dataset is in this constructor. The update_dataset is really a wrapper around this.
This constructor creates a cagra index and takes ownership of the resources of the old one (specifically other.graph_ and other.source_indices_).
We reuse the member function index.update_dataset() to avoid duplicating code.
The member function index.update_dataset() is compatible with the same type only.
The update_dataset free function is for updating datasets of different types (ex. standard to padded).
|
|
||
| private: | ||
| template <typename, typename, ann_dataset_view> | ||
| friend struct index; |
There was a problem hiding this comment.
We need this because one index specialization is cannot access the private variables of another. This is the functionality that we need in our new constructor.
| * This returns a new padded index because standard and padded MG indexes have different C++ types. | ||
| * This moves each rank-local CAGRA graph into the returned padded MG index. | ||
| */ | ||
| auto attach_dataset( |
There was a problem hiding this comment.
attach_dataset has been deleted in favor update_dataset
| /** | ||
| * @brief Update an existing padded MG CAGRA index with a padded dataset of the same layout. | ||
| */ | ||
| void update_device_dataset_same_layout( |
There was a problem hiding this comment.
update_device_dataset_same_layout has been deleted in favor of using the member function index.update_dataset().
| typename IdxT, | ||
| ann_dataset_view SrcDatasetViewT, | ||
| ann_dataset_view DstDatasetViewT> | ||
| auto update_dataset(raft::resources const& res, |
There was a problem hiding this comment.
This is where the update_dataset implementation lives. Again, it is mostly a wrapper around the newly added constructor.
| { \ | ||
| return cuvs::neighbors::cagra::update_dataset<T, IdxT, SrcDatasetViewT, DstDatasetViewT>( \ | ||
| res, std::move(cagra_index), dataset); \ | ||
| } |
There was a problem hiding this comment.
I am not sure whether this is the correct file where the concrete definitions should live. But it seemed repetitive to create another file like cagra_update_dataset_inst.cu.in.
| // call cagra::index::update_device_dataset_same_layout on it to update the ann_index to point | ||
| // to the | ||
| // new dataset | ||
| // call cagra::index::update_dataset on it to update the ann_index to point to the new dataset |
There was a problem hiding this comment.
We need to rework tiered_index IMO. I've added my thoughts here #2434. We will properly get rid of the convert_standard_to_padded_index from cagra.hpp that is only used by tiered index in a follow up.
| auto update_dataset(raft::resources const& res, | ||
| index<half, uint32_t, host_standard_dataset_view<half, int64_t>>&& cagra_index, | ||
| device_standard_dataset_view<half, int64_t> dataset) | ||
| -> index<half, uint32_t, device_standard_dataset_view<half, int64_t>>; |
There was a problem hiding this comment.
The four instantiations for host_standard_dataset_view -> device_standard_dataset_view are to support the MG build implementation.
However, this use case is puzzling to me. Would a user ever need to do this? And why are we doing it in the MG code? We should be doing host_standard_dataset_view -> device_padded_dataset_view if anything.
Another question is why are we doing copying from H2D within the MG code anyway? Why not move that responsibility to the user so that they have ownership of the device dataset?
This PR introduces the
update_datasetfunction in the cpp api.update_datasetshould swap out the dataset from an index. For instance, from a device_standard_dataset_view -> device_padded_dataset_view.The major use case here is when a user builds a cagra index with a certain type of dataset. But during search they want to use a different type of dataset. In the future this may look like device_standard_dataset_view -> device_pq_dataset_view.
The type of dataset is coupled to the index type. So we need to construct a new index object altogether. To avoid copying expensive member variables like the cagra graph or source indices we simply move them to the new object instead. This means the user relinquishes the old index object.
A new constructor in cagra.hpp has been introduced to move and take ownership of the resources from an existing cagra index and assign the new dataset.
We intend to remove the
attach_datasetand other helper functions such asconvert_standard_to_padded_indexandconvert_host_to_device_index. The callers that use the removed functions should be updated to useupdate_datasetinstead.Resolves #2404
I've annotated this PR below to make it easier to review.