C API: move the compression params into the index params

rapidsai · Mar 22, 2024 · 2fcbb94 · 2fcbb94
1 parent 597e8e4
commit 2fcbb94
Show file tree

Hide file tree

Showing 2 changed files with 48 additions and 78 deletions.
diff --git a/cpp/include/cuvs/neighbors/cagra.h b/cpp/include/cuvs/neighbors/cagra.h
@@ -40,39 +40,6 @@ enum cuvsCagraGraphBuildAlgo {
   NN_DESCENT
 };
 
-/**
- * @brief Supplemental parameters to build CAGRA Index
- *
- */
-struct cuvsCagraIndexParams {
-  /** Degree of input graph for pruning. */
-  size_t intermediate_graph_degree;
-  /** Degree of output graph. */
-  size_t graph_degree;
-  /** ANN algorithm to build knn graph. */
-  enum cuvsCagraGraphBuildAlgo build_algo;
-  /** Number of Iterations to run if building with NN_DESCENT */
-  size_t nn_descent_niter;
-};
-
-typedef struct cuvsCagraIndexParams* cuvsCagraIndexParams_t;
-
-/**
- * @brief Allocate CAGRA Index params, and populate with default values
- *
- * @param[in] params cuvsCagraIndexParams_t to allocate
- * @return cuvsError_t
- */
-cuvsError_t cuvsCagraIndexParamsCreate(cuvsCagraIndexParams_t* params);
-
-/**
- * @brief De-allocate CAGRA Index params
- *
- * @param[in] params
- * @return cuvsError_t
- */
-cuvsError_t cuvsCagraIndexParamsDestroy(cuvsCagraIndexParams_t params);
-
 /** Parameters for VPQ compression. */
 struct cuvsCagraCompressionParams {
   /**
@@ -112,6 +79,45 @@ struct cuvsCagraCompressionParams {
 
 typedef struct cuvsCagraCompressionParams* cuvsCagraCompressionParams_t;
 
+/**
+ * @brief Supplemental parameters to build CAGRA Index
+ *
+ */
+struct cuvsCagraIndexParams {
+  /** Degree of input graph for pruning. */
+  size_t intermediate_graph_degree;
+  /** Degree of output graph. */
+  size_t graph_degree;
+  /** ANN algorithm to build knn graph. */
+  enum cuvsCagraGraphBuildAlgo build_algo;
+  /** Number of Iterations to run if building with NN_DESCENT */
+  size_t nn_descent_niter;
+  /**
+   * Optional: specify compression parameters if compression is desired.
+   *
+   * NOTE: this is experimental new API, consider it unsafe.
+   */
+  cuvsCagraCompressionParams_t compression;
+};
+
+typedef struct cuvsCagraIndexParams* cuvsCagraIndexParams_t;
+
+/**
+ * @brief Allocate CAGRA Index params, and populate with default values
+ *
+ * @param[in] params cuvsCagraIndexParams_t to allocate
+ * @return cuvsError_t
+ */
+cuvsError_t cuvsCagraIndexParamsCreate(cuvsCagraIndexParams_t* params);
+
+/**
+ * @brief De-allocate CAGRA Index params
+ *
+ * @param[in] params
+ * @return cuvsError_t
+ */
+cuvsError_t cuvsCagraIndexParamsDestroy(cuvsCagraIndexParams_t params);
+
 /**
  * @brief Allocate CAGRA Compression params, and populate with default values
  *
@@ -314,23 +320,6 @@ cuvsError_t cuvsCagraBuild(cuvsResources_t res,
                            DLManagedTensor* dataset,
                            cuvsCagraIndex_t index);
 
-/**
- * @brief Build a CAGRA index with compression
- *
- * @param[in] res cuvsResources_t opaque C handle
- * @param[in] params cuvsCagraIndexParams_t used to build CAGRA index
- * @param[in] compression_params cuvsCagraCompressionParams_t used to compress the input dataset for
- * CAGRA index
- * @param[in] dataset DLManagedTensor* training dataset
- * @param[out] index cuvsCagraIndex_t Newly built CAGRA index
- * @return cuvsError_t
- */
-cuvsError_t cuvsCagraBuildCompressed(cuvsResources_t res,
-                                     cuvsCagraIndexParams_t params,
-                                     cuvsCagraCompressionParams_t compression_params,
-                                     DLManagedTensor* dataset,
-                                     cuvsCagraIndex_t index);
-
 /**
  * @}
  */

diff --git a/cpp/src/neighbors/cagra_c.cpp b/cpp/src/neighbors/cagra_c.cpp
@@ -28,15 +28,10 @@
 #include <cuvs/neighbors/cagra.h>
 #include <cuvs/neighbors/cagra.hpp>
 
-#include <optional>
-
 namespace {
 
 template <typename T>
-void* _build(cuvsResources_t res,
-             cuvsCagraIndexParams params,
-             std::optional<cuvsCagraCompressionParams> cparams,
-             DLManagedTensor* dataset_tensor)
+void* _build(cuvsResources_t res, cuvsCagraIndexParams params, DLManagedTensor* dataset_tensor)
 {
   auto dataset = dataset_tensor->dl_tensor;
 
@@ -50,7 +45,7 @@ void* _build(cuvsResources_t res,
     static_cast<cuvs::neighbors::cagra::graph_build_algo>(params.build_algo);
   build_params.nn_descent_niter = params.nn_descent_niter;
 
-  if (cparams.has_value()) {
+  if (auto* cparams = params.compression; cparams != nullptr) {
     auto compression_params                        = cuvs::neighbors::cagra::vpq_params();
     compression_params.pq_bits                     = cparams->pq_bits;
     compression_params.pq_dim                      = cparams->pq_dim;
@@ -138,27 +133,21 @@ extern "C" cuvsError_t cuvsCagraIndexDestroy(cuvsCagraIndex_t index_c_ptr)
   });
 }
 
-extern "C" cuvsError_t cuvsCagraBuildCompressed(cuvsResources_t res,
-                                                cuvsCagraIndexParams_t params,
-                                                cuvsCagraCompressionParams_t compression_params,
-                                                DLManagedTensor* dataset_tensor,
-                                                cuvsCagraIndex_t index)
+extern "C" cuvsError_t cuvsCagraBuild(cuvsResources_t res,
+                                      cuvsCagraIndexParams_t params,
+                                      DLManagedTensor* dataset_tensor,
+                                      cuvsCagraIndex_t index)
 {
   return cuvs::core::translate_exceptions([=] {
     auto dataset = dataset_tensor->dl_tensor;
-    auto cparams =
-      compression_params == nullptr ? std::nullopt : std::make_optional(*compression_params);
     if (dataset.dtype.code == kDLFloat && dataset.dtype.bits == 32) {
-      index->addr =
-        reinterpret_cast<uintptr_t>(_build<float>(res, *params, cparams, dataset_tensor));
+      index->addr       = reinterpret_cast<uintptr_t>(_build<float>(res, *params, dataset_tensor));
       index->dtype.code = kDLFloat;
     } else if (dataset.dtype.code == kDLInt && dataset.dtype.bits == 8) {
-      index->addr =
-        reinterpret_cast<uintptr_t>(_build<int8_t>(res, *params, cparams, dataset_tensor));
+      index->addr       = reinterpret_cast<uintptr_t>(_build<int8_t>(res, *params, dataset_tensor));
       index->dtype.code = kDLInt;
     } else if (dataset.dtype.code == kDLUInt && dataset.dtype.bits == 8) {
-      index->addr =
-        reinterpret_cast<uintptr_t>(_build<uint8_t>(res, *params, cparams, dataset_tensor));
+      index->addr = reinterpret_cast<uintptr_t>(_build<uint8_t>(res, *params, dataset_tensor));
       index->dtype.code = kDLUInt;
     } else {
       RAFT_FAIL("Unsupported dataset DLtensor dtype: %d and bits: %d",
@@ -168,14 +157,6 @@ extern "C" cuvsError_t cuvsCagraBuildCompressed(cuvsResources_t res,
   });
 }
 
-extern "C" cuvsError_t cuvsCagraBuild(cuvsResources_t res,
-                                      cuvsCagraIndexParams_t params,
-                                      DLManagedTensor* dataset_tensor,
-                                      cuvsCagraIndex_t index)
-{
-  return cuvsCagraBuildCompressed(res, params, nullptr, dataset_tensor, index);
-}
-
 extern "C" cuvsError_t cuvsCagraSearch(cuvsResources_t res,
                                        cuvsCagraSearchParams_t params,
                                        cuvsCagraIndex_t index_c_ptr,