update_dataset CAGRA (#2427)

This PR introduces the `update_dataset` function in the cpp api.
`update_dataset` should 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. 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_dataset` and other helper functions such as `convert_standard_to_padded_index` and `convert_host_to_device_index`. The callers that use the removed functions should be updated to use `update_dataset` instead.

Resolves #2404

### I've annotated this PR below to make it easier to review.

Authors:
  - Anupam (https://github.com/aamijar)

Approvers:
  - Tarang Jain (https://github.com/tarang-jain)
  - Divye Gala (https://github.com/divyegala)

URL: #2427

Author: aamijar

c/src/neighbors/cagra.cpp

@@ -529,7 +529,7 @@ static void make_host_standard_dataset_view(raft::resources*,
 }
 
 template <typename T>
-static void attach_dataset(raft::resources* res_ptr,
+static void update_dataset(raft::resources* res_ptr,
                            cuvsDataset_t device_padded_dataset,
                            cuvsCagraIndex_t index)
 {
@@ -542,7 +542,7 @@ static void attach_dataset(raft::resources* res_ptr,
   auto* box = reinterpret_cast<sg_cagra_c_api_index_box*>(index->addr);
   RAFT_EXPECTS(device_padded_dataset->mem_type == CUVS_DATASET_MEM_TYPE_DEVICE &&
                  device_padded_dataset->layout == CUVS_DATASET_LAYOUT_PADDED,
-               "cuvsCagraAttachDataset: dataset must be device padded");
+               "cuvsCagraUpdateDataset: dataset must be device padded");
 
   using owner_t = cuvs::neighbors::device_padded_dataset<T, int64_t>;
   using view_t  = cuvs::neighbors::device_padded_dataset_view<T, int64_t>;
@@ -552,7 +552,8 @@ static void attach_dataset(raft::resources* res_ptr,
       "cuvsCagraUpdateDataset: null index handle",
       "cuvsCagraUpdateDataset: host index layout is allowed for this operation",
       [&](auto& idx) {
-        auto padded_idx = cuvs::neighbors::cagra::attach_dataset(*res_ptr, idx, padded_view);
+        auto padded_idx =
+          cuvs::neighbors::cagra::update_dataset(*res_ptr, std::move(idx), padded_view);
         auto* holder =
           new cuvs_cagra_c_api_index_lifetime_holder<T, view_t>{std::move(padded_idx)};
         destroy_sg_cagra_c_api_box(index->addr);
@@ -562,53 +563,6 @@ static void attach_dataset(raft::resources* res_ptr,
   });
 }
 
-template <typename T>
-static void update_device_dataset_same_layout(raft::resources* res_ptr,
-                                              cuvsDataset_t device_dataset,
-                                              cuvsCagraIndex_t index)
-{
-  RAFT_EXPECTS(device_dataset != nullptr, "cuvsCagraUpdateDataset: null dataset");
-  RAFT_EXPECTS(index != nullptr, "cuvsCagraUpdateDataset: null index handle");
-  RAFT_EXPECTS(index->addr != 0, "cuvsCagraUpdateDataset: null index storage");
-  RAFT_EXPECTS(device_dataset->addr != 0, "cuvsCagraUpdateDataset: null dataset storage");
-
-  auto* box = reinterpret_cast<sg_cagra_c_api_index_box*>(index->addr);
-  if (box->layout == sg_cagra_c_api_index_box::dataset_layout::device_padded) {
-    RAFT_EXPECTS(device_dataset->mem_type == CUVS_DATASET_MEM_TYPE_DEVICE &&
-                   device_dataset->layout == CUVS_DATASET_LAYOUT_PADDED,
-                 "cuvsCagraUpdateDeviceDatasetSameLayout: device-padded index "
-                 "requires a "
-                 "device-padded dataset");
-    using owner_t = cuvs::neighbors::device_padded_dataset<T, int64_t>;
-    using view_t  = cuvs::neighbors::device_padded_dataset_view<T, int64_t>;
-    with_dataset_view<owner_t, view_t>(device_dataset, [&](auto const& dataset_view) {
-      auto* idx =
-        reinterpret_cast<cuvs::neighbors::cagra::device_padded_index<T, uint32_t>*>(box->index_ptr);
-      RAFT_EXPECTS(idx != nullptr, "cuvsCagraUpdateDataset: null index handle");
-      idx->update_device_dataset_same_layout(*res_ptr, dataset_view);
-    });
-  } else if (box->layout == sg_cagra_c_api_index_box::dataset_layout::device_standard) {
-    RAFT_EXPECTS(device_dataset->mem_type == CUVS_DATASET_MEM_TYPE_DEVICE &&
-                   device_dataset->layout == CUVS_DATASET_LAYOUT_STANDARD,
-                 "cuvsCagraUpdateDeviceDatasetSameLayout: device-standard "
-                 "index requires a "
-                 "device-standard dataset");
-    using owner_t = cuvs::neighbors::device_standard_dataset<T, int64_t>;
-    using view_t  = cuvs::neighbors::device_standard_dataset_view<T, int64_t>;
-    with_dataset_view<owner_t, view_t>(device_dataset, [&](auto const& dataset_view) {
-      auto* idx =
-        reinterpret_cast<cuvs::neighbors::cagra::device_standard_index<T, uint32_t>*>(box->index_ptr);
-      RAFT_EXPECTS(idx != nullptr, "cuvsCagraUpdateDataset: null index handle");
-      idx->update_device_dataset_same_layout(*res_ptr, dataset_view);
-    });
-  } else {
-    RAFT_FAIL(
-      "cuvsCagraUpdateDataset: C++ "
-      "update_device_dataset_same_layout "
-      "requires a device index and dataset");
-  }
-}
-
 static void _set_graph_build_params(
   std::variant<std::monostate,
                cuvs::neighbors::cagra::graph_build_params::ivf_pq_params,
@@ -716,7 +670,8 @@ void _from_args(cuvsResources_t res,
       auto dataset_view = cuvs::neighbors::make_device_padded_dataset_view(*res_ptr, mds);
       auto* raw         = new cuvs::neighbors::cagra::device_padded_index<T, uint32_t>(
         *res_ptr, metric);
-      raw->update_device_dataset_same_layout(*res_ptr, dataset_view);
+      *raw =
+        cuvs::neighbors::cagra::update_dataset(*res_ptr, std::move(*raw), dataset_view);
       update_graph_from_dlpack(raw);
       wrap_CPP_index_in_lifetime_holder_and_bind_to_C_index<
         T,
@@ -726,7 +681,8 @@ void _from_args(cuvsResources_t res,
       auto dataset_view = cuvs::neighbors::make_device_standard_dataset_view(mds);
       auto* raw         = new cuvs::neighbors::cagra::device_standard_index<T, uint32_t>(
         *res_ptr, metric);
-      raw->update_device_dataset_same_layout(*res_ptr, dataset_view);
+      *raw =
+        cuvs::neighbors::cagra::update_dataset(*res_ptr, std::move(*raw), dataset_view);
       update_graph_from_dlpack(raw);
       wrap_CPP_index_in_lifetime_holder_and_bind_to_C_index<
         T,
@@ -1584,7 +1540,7 @@ extern "C" cuvsError_t cuvsDatasetMakeStandardView(cuvsResources_t res,
   });
 }
 
-static cuvsError_t dispatch_attach_dataset(cuvsResources_t res,
+static cuvsError_t dispatch_update_dataset(cuvsResources_t res,
                                            cuvsDataset_t device_padded_dataset,
                                            cuvsCagraIndex_t index)
 {
@@ -1598,40 +1554,13 @@ static cuvsError_t dispatch_attach_dataset(cuvsResources_t res,
                    index->dtype.bits == device_padded_dataset->dtype.bits,
                  "cuvsCagraUpdateDataset: dtype mismatch between index and dataset");
     if (index->dtype.code == kDLFloat && index->dtype.bits == 32) {
-      attach_dataset<float>(res_ptr, device_padded_dataset, index);
-    } else if (index->dtype.code == kDLFloat && index->dtype.bits == 16) {
-      attach_dataset<half>(res_ptr, device_padded_dataset, index);
-    } else if (index->dtype.code == kDLInt && index->dtype.bits == 8) {
-      attach_dataset<int8_t>(res_ptr, device_padded_dataset, index);
-    } else if (index->dtype.code == kDLUInt && index->dtype.bits == 8) {
-      attach_dataset<uint8_t>(res_ptr, device_padded_dataset, index);
-    } else {
-      RAFT_FAIL("Unsupported index dtype: %d and bits: %d", index->dtype.code, index->dtype.bits);
-    }
-  });
-}
-
-static cuvsError_t dispatch_update_device_dataset_same_layout(cuvsResources_t res,
-                                                              cuvsDataset_t device_dataset,
-                                                              cuvsCagraIndex_t index)
-{
-  return cuvs::core::translate_exceptions([=] {
-    auto* res_ptr = reinterpret_cast<raft::resources*>(res);
-    RAFT_EXPECTS(index != nullptr, "cuvsCagraUpdateDataset: null index handle");
-    RAFT_EXPECTS(device_dataset != nullptr,
-                 "cuvsCagraUpdateDataset: null dataset view");
-    RAFT_EXPECTS(index->dtype.code == device_dataset->dtype.code &&
-                   index->dtype.bits == device_dataset->dtype.bits,
-                 "cuvsCagraUpdateDataset: dtype mismatch "
-                 "between index and dataset");
-    if (index->dtype.code == kDLFloat && index->dtype.bits == 32) {
-      update_device_dataset_same_layout<float>(res_ptr, device_dataset, index);
+      update_dataset<float>(res_ptr, device_padded_dataset, index);
     } else if (index->dtype.code == kDLFloat && index->dtype.bits == 16) {
-      update_device_dataset_same_layout<half>(res_ptr, device_dataset, index);
+      update_dataset<half>(res_ptr, device_padded_dataset, index);
     } else if (index->dtype.code == kDLInt && index->dtype.bits == 8) {
-      update_device_dataset_same_layout<int8_t>(res_ptr, device_dataset, index);
+      update_dataset<int8_t>(res_ptr, device_padded_dataset, index);
     } else if (index->dtype.code == kDLUInt && index->dtype.bits == 8) {
-      update_device_dataset_same_layout<uint8_t>(res_ptr, device_dataset, index);
+      update_dataset<uint8_t>(res_ptr, device_padded_dataset, index);
     } else {
       RAFT_FAIL("Unsupported index dtype: %d and bits: %d", index->dtype.code, index->dtype.bits);
     }
@@ -1656,12 +1585,7 @@ extern "C" cuvsError_t cuvsCagraUpdateDataset(cuvsResources_t res,
                  "cuvsCagraUpdateDataset: dtype mismatch between index and dataset");
   });
   if (status != CUVS_SUCCESS) { return status; }
-
-  auto* box = reinterpret_cast<sg_cagra_c_api_index_box*>(index->addr);
-  if (box->layout == sg_cagra_c_api_index_box::dataset_layout::device_padded) {
-    return dispatch_update_device_dataset_same_layout(res, device_padded_dataset, index);
-  }
-  return dispatch_attach_dataset(res, device_padded_dataset, index);
+  return dispatch_update_dataset(res, device_padded_dataset, index);
 }
 
 /**

c/src/neighbors/mg_cagra.cpp

@@ -229,16 +229,16 @@ void _mg_update_dataset(cuvsResources_t res,
       using padded_ann_t   = cuvs::neighbors::cagra::device_padded_index<T, uint32_t>;
       auto* standard_index = reinterpret_cast<mg_cagra_index_t<T, standard_ann_t>*>(box->index_ptr);
       auto* padded_index   = new mg_cagra_index_t<T, padded_ann_t>(
-        cuvs::neighbors::cagra::attach_dataset(*res_ptr, *standard_index, padded_view));
+        cuvs::neighbors::cagra::update_dataset(
+          *res_ptr, std::move(*standard_index), padded_view));
       auto* padded_box =
         make_mg_cagra_box<T, padded_ann_t>(padded_index, mg_cagra_dataset_layout::device_padded);
       destroy_mg_cagra_c_api_box(index->addr);
       index->addr = reinterpret_cast<uintptr_t>(padded_box);
     } else if (box->layout == mg_cagra_dataset_layout::device_padded) {
       using padded_ann_t = cuvs::neighbors::cagra::device_padded_index<T, uint32_t>;
       auto* padded_index = reinterpret_cast<mg_cagra_index_t<T, padded_ann_t>*>(box->index_ptr);
-      cuvs::neighbors::cagra::update_device_dataset_same_layout(
-        *res_ptr, *padded_index, padded_view);
+      cuvs::neighbors::cagra::update_dataset(*res_ptr, *padded_index, padded_view);
     } else {
       RAFT_FAIL("cuvsMultiGpuCagraUpdateDataset: unsupported index dataset layout");
     }

cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h

@@ -418,8 +418,8 @@ void cuvs_cagra<T, IdxT>::compress_dataset(const T* dataset, size_t nrow)
   // Search runs on the compressed rows and the graph, so release the dense copy of the dataset.
   cuvs::neighbors::device_padded_dataset_view<T, int64_t> empty_dv(
     raft::make_device_matrix_view(static_cast<T const*>(nullptr), 0, this->dim_), this->dim_);
-  index_->update_device_dataset_same_layout(handle_, empty_dv);
-  *dataset_            = raft::make_device_matrix<T, int64_t>(handle_, 0, 0);
+  *index_   = cuvs::neighbors::cagra::update_dataset(handle_, std::move(*index_), empty_dv);
+  *dataset_ = raft::make_device_matrix<T, int64_t>(handle_, 0, 0);
   need_dataset_update_ = false;
 }
 
@@ -485,7 +485,7 @@ void cuvs_cagra<T, IdxT>::set_search_param(const search_param_base& param,
     *dataset_ = raft::make_device_matrix<T, int64_t>(handle_, 0, 0);
     cuvs::neighbors::device_padded_dataset_view<T, int64_t> empty_dv(
       raft::make_device_matrix_view(static_cast<T const*>(nullptr), 0, this->dim_), this->dim_);
-    index_->update_device_dataset_same_layout(handle_, empty_dv);
+    *index_ = cuvs::neighbors::cagra::update_dataset(handle_, std::move(*index_), empty_dv);
 
     // Allocate space using the correct memory resource.
     RAFT_LOG_DEBUG("moving dataset to new memory space: %s",
@@ -498,7 +498,7 @@ void cuvs_cagra<T, IdxT>::set_search_param(const search_param_base& param,
       raft::make_device_matrix_view(
         dataset_->data_handle(), dataset_->extent(0), dataset_->extent(1)),
       this->dim_);
-    index_->update_device_dataset_same_layout(handle_, dv);
+    *index_ = cuvs::neighbors::cagra::update_dataset(handle_, std::move(*index_), dv);
 
     need_dataset_update_         = false;
     needs_dynamic_batcher_update = true;
@@ -562,11 +562,15 @@ void cuvs_cagra<T, IdxT>::set_search_dataset(const T* dataset, size_t nrow)
       auto& sub_dataset_buffer = (*sub_dataset_buffers_)[i];
       sub_dataset_buffer       = raft::make_device_matrix<T, int64_t>(handle_, 0, 0);
       if (dataset_is_on_host) {
-        sub_index->update_device_dataset_same_layout(
-          handle_, detail::make_padded_view<T>(handle_, sub_host, sub_dataset_buffer));
+        *sub_index = cuvs::neighbors::cagra::update_dataset(
+          handle_,
+          std::move(*sub_index),
+          detail::make_padded_view<T>(handle_, sub_host, sub_dataset_buffer));
       } else {
-        sub_index->update_device_dataset_same_layout(
-          handle_, detail::make_padded_view<T>(handle_, sub_dev, sub_dataset_buffer));
+        *sub_index = cuvs::neighbors::cagra::update_dataset(
+          handle_,
+          std::move(*sub_index),
+          detail::make_padded_view<T>(handle_, sub_dev, sub_dataset_buffer));
       }
     }
     need_dataset_update_ = false;