Expose extend() in C API (#276)

Add functionality to add additional vectors after build to C API

Authors:
  - Ajit Mistry (https://github.com/ajit283)
  - Corey J. Nolet (https://github.com/cjnolet)
  - Ben Frederickson (https://github.com/benfred)
  - Micka (https://github.com/lowener)

Approvers:
  - Ben Frederickson (https://github.com/benfred)

URL: #276

cpp/include/cuvs/neighbors/cagra.h

@@ -141,6 +141,45 @@ cuvsError_t cuvsCagraCompressionParamsCreate(cuvsCagraCompressionParams_t* param
  */
 cuvsError_t cuvsCagraCompressionParamsDestroy(cuvsCagraCompressionParams_t params);
 
+/**
+ * @}
+ */
+
+/**
+ * @defgroup cagra_c_extend_params C API for CUDA ANN Graph-based nearest neighbor search
+ * @{
+ */
+/**
+ * @brief Supplemental parameters to extend CAGRA Index
+ *
+ */
+struct cuvsCagraExtendParams {
+  /** The additional dataset is divided into chunks and added to the graph. This is the knob to
+   * adjust the tradeoff between the recall and operation throughput. Large chunk sizes can result
+   * in high throughput, but use more working memory (O(max_chunk_size*degree^2)). This can also
+   * degrade recall because no edges are added between the nodes in the same chunk. Auto select when
+   * 0. */
+  uint32_t max_chunk_size;
+};
+
+typedef struct cuvsCagraExtendParams* cuvsCagraExtendParams_t;
+
+/**
+ * @brief Allocate CAGRA Extend params, and populate with default values
+ *
+ * @param[in] params cuvsCagraExtendParams_t to allocate
+ * @return cuvsError_t
+ */
+cuvsError_t cuvsCagraExtendParamsCreate(cuvsCagraExtendParams_t* params);
+
+/**
+ * @brief De-allocate CAGRA Extend params
+ *
+ * @param[in] params
+ * @return cuvsError_t
+ */
+cuvsError_t cuvsCagraExtendParamsDestroy(cuvsCagraExtendParams_t params);
+
 /**
  * @}
  */
@@ -340,6 +379,36 @@ cuvsError_t cuvsCagraBuild(cuvsResources_t res,
  * @}
  */
 
+/**
+ * @defgroup cagra_c_extend_params C API for CUDA ANN Graph-based nearest neighbor search
+ * @{
+ */
+
+/**
+ * @brief Extend a CAGRA index with a `DLManagedTensor` which has underlying
+ *        `DLDeviceType` equal to `kDLCUDA`, `kDLCUDAHost`, `kDLCUDAManaged`,
+ *        or `kDLCPU`. Also, acceptable underlying types are:
+ *        1. `kDLDataType.code == kDLFloat` and `kDLDataType.bits = 32`
+ *        2. `kDLDataType.code == kDLInt` and `kDLDataType.bits = 8`
+ *        3. `kDLDataType.code == kDLUInt` and `kDLDataType.bits = 8`
+ *
+ * @param[in] res cuvsResources_t opaque C handle
+ * @param[in] params cuvsCagraExtendParams_t used to extend CAGRA index
+ * @param[in] additional_dataset DLManagedTensor* additional dataset
+ * @param[in,out] index cuvsCagraIndex_t CAGRA index
+ * @param[out] return_dataset DLManagedTensor* extended dataset
+ * @return cuvsError_t
+ */
+cuvsError_t cuvsCagraExtend(cuvsResources_t res,
+                            cuvsCagraExtendParams_t params,
+                            DLManagedTensor* additional_dataset,
+                            cuvsCagraIndex_t index,
+                            DLManagedTensor* return_dataset);
+
+/**
+ * @}
+ */
+
 /**
  * @defgroup cagra_c_index_search C API for CUDA ANN Graph-based nearest neighbor search
  * @{

cpp/src/neighbors/cagra_c.cpp

@@ -87,6 +87,43 @@ void* _build(cuvsResources_t res, cuvsCagraIndexParams params, DLManagedTensor*
   return index;
 }
 
+template <typename T>
+void _extend(cuvsResources_t res,
+             cuvsCagraExtendParams params,
+             cuvsCagraIndex index,
+             DLManagedTensor* additional_dataset_tensor,
+             DLManagedTensor* return_tensor)
+{
+  auto dataset          = additional_dataset_tensor->dl_tensor;
+  auto return_dl_tensor = return_tensor->dl_tensor;
+  auto index_ptr        = reinterpret_cast<cuvs::neighbors::cagra::index<T, uint32_t>*>(index.addr);
+  auto res_ptr          = reinterpret_cast<raft::resources*>(res);
+
+  // TODO: use C struct here (see issue #487)
+  auto extend_params           = cuvs::neighbors::cagra::extend_params();
+  extend_params.max_chunk_size = params.max_chunk_size;
+
+  if (cuvs::core::is_dlpack_device_compatible(dataset) &&
+      cuvs::core::is_dlpack_device_compatible(return_dl_tensor)) {
+    using mdspan_type        = raft::device_matrix_view<T const, int64_t, raft::row_major>;
+    using mdspan_return_type = raft::device_matrix_view<T, int64_t, raft::row_major>;
+    auto mds                 = cuvs::core::from_dlpack<mdspan_type>(additional_dataset_tensor);
+    auto return_mds          = cuvs::core::from_dlpack<mdspan_return_type>(return_tensor);
+    cuvs::neighbors::cagra::extend(*res_ptr, extend_params, mds, *index_ptr, return_mds);
+  } else if (cuvs::core::is_dlpack_host_compatible(dataset) &&
+             cuvs::core::is_dlpack_host_compatible(return_dl_tensor)) {
+    using mdspan_type        = raft::host_matrix_view<T const, int64_t, raft::row_major>;
+    using mdspan_return_type = raft::device_matrix_view<T, int64_t, raft::row_major>;
+    auto mds                 = cuvs::core::from_dlpack<mdspan_type>(additional_dataset_tensor);
+    auto return_mds          = cuvs::core::from_dlpack<mdspan_return_type>(return_tensor);
+    cuvs::neighbors::cagra::extend(*res_ptr, extend_params, mds, *index_ptr, return_mds);
+  } else {
+    RAFT_FAIL("Unsupported dataset DLtensor dtype: %d and bits: %d",
+              dataset.dtype.code,
+              dataset.dtype.bits);
+  }
+}
+
 template <typename T>
 void _search(cuvsResources_t res,
              cuvsCagraSearchParams params,
@@ -229,6 +266,30 @@ extern "C" cuvsError_t cuvsCagraBuild(cuvsResources_t res,
   });
 }
 
+extern "C" cuvsError_t cuvsCagraExtend(cuvsResources_t res,
+                                       cuvsCagraExtendParams_t params,
+                                       DLManagedTensor* additional_dataset_tensor,
+                                       cuvsCagraIndex_t index_c_ptr,
+                                       DLManagedTensor* return_dataset_tensor)
+{
+  return cuvs::core::translate_exceptions([=] {
+    auto dataset = additional_dataset_tensor->dl_tensor;
+    auto index   = *index_c_ptr;
+
+    if ((dataset.dtype.code == kDLFloat) && (dataset.dtype.bits == 32)) {
+      _extend<float>(res, *params, index, additional_dataset_tensor, return_dataset_tensor);
+    } else if (dataset.dtype.code == kDLInt && dataset.dtype.bits == 8) {
+      _extend<int8_t>(res, *params, index, additional_dataset_tensor, return_dataset_tensor);
+    } else if (dataset.dtype.code == kDLUInt && dataset.dtype.bits == 8) {
+      _extend<uint8_t>(res, *params, index, additional_dataset_tensor, return_dataset_tensor);
+    } else {
+      RAFT_FAIL("Unsupported dataset DLtensor dtype: %d and bits: %d",
+                dataset.dtype.code,
+                dataset.dtype.bits);
+    }
+  });
+}
+
 extern "C" cuvsError_t cuvsCagraSearch(cuvsResources_t res,
                                        cuvsCagraSearchParams_t params,
                                        cuvsCagraIndex_t index_c_ptr,
@@ -309,6 +370,17 @@ extern "C" cuvsError_t cuvsCagraCompressionParamsDestroy(cuvsCagraCompressionPar
   return cuvs::core::translate_exceptions([=] { delete params; });
 }
 
+extern "C" cuvsError_t cuvsCagraExtendParamsCreate(cuvsCagraExtendParams_t* params)
+{
+  return cuvs::core::translate_exceptions(
+    [=] { *params = new cuvsCagraExtendParams{.max_chunk_size = 0}; });
+}
+
+extern "C" cuvsError_t cuvsCagraExtendParamsDestroy(cuvsCagraExtendParams_t params)
+{
+  return cuvs::core::translate_exceptions([=] { delete params; });
+}
+
 extern "C" cuvsError_t cuvsCagraSearchParamsCreate(cuvsCagraSearchParams_t* params)
 {
   return cuvs::core::translate_exceptions([=] {