From 2fcbb94c431196c81649f788143a4b38ea6a01d3 Mon Sep 17 00:00:00 2001 From: achirkin Date: Fri, 22 Mar 2024 15:12:30 +0100 Subject: [PATCH] C API: move the compression params into the index params --- cpp/include/cuvs/neighbors/cagra.h | 89 +++++++++++++----------------- cpp/src/neighbors/cagra_c.cpp | 37 +++---------- 2 files changed, 48 insertions(+), 78 deletions(-) diff --git a/cpp/include/cuvs/neighbors/cagra.h b/cpp/include/cuvs/neighbors/cagra.h index 7ae6784d4..3a1a8c9fe 100644 --- 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 index 5cb0b2e9c..2e96e8db6 100644 --- a/cpp/src/neighbors/cagra_c.cpp +++ b/cpp/src/neighbors/cagra_c.cpp @@ -28,15 +28,10 @@ #include #include -#include - namespace { template -void* _build(cuvsResources_t res, - cuvsCagraIndexParams params, - std::optional 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(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(_build(res, *params, cparams, dataset_tensor)); + index->addr = reinterpret_cast(_build(res, *params, dataset_tensor)); index->dtype.code = kDLFloat; } else if (dataset.dtype.code == kDLInt && dataset.dtype.bits == 8) { - index->addr = - reinterpret_cast(_build(res, *params, cparams, dataset_tensor)); + index->addr = reinterpret_cast(_build(res, *params, dataset_tensor)); index->dtype.code = kDLInt; } else if (dataset.dtype.code == kDLUInt && dataset.dtype.bits == 8) { - index->addr = - reinterpret_cast(_build(res, *params, cparams, dataset_tensor)); + index->addr = reinterpret_cast(_build(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,