diff --git a/cpp/include/cuvs/neighbors/cagra_c.h b/cpp/include/cuvs/neighbors/cagra_c.h index 2d6ce74b6..b7374acbb 100644 --- a/cpp/include/cuvs/neighbors/cagra_c.h +++ b/cpp/include/cuvs/neighbors/cagra_c.h @@ -55,12 +55,23 @@ struct cagraIndexParams { size_t nn_descent_niter; }; -typedef struct cagraIndexParams cagraIndexParams; +typedef struct cagraIndexParams* cagraIndexParams_t; /** - * @brief Default cagraIndexParams to use + * @brief Allocate CAGRA Index params, and populate with default values + * + * @param[in] params cagraIndexParams_t to allocate + * @return cuvsError_t + */ +cuvsError_t cagraIndexParamsCreate(cagraIndexParams_t* params); + +/** + * @brief De-allocate CAGRA Index params + * + * @param[in] params + * @return cuvsError_t */ -extern const cagraIndexParams cagraDefaultIndexParams; +cuvsError_t cagraIndexParamsDestroy(cagraIndexParams_t index); /** * @brief Enum to denote algorithm used to search CAGRA Index @@ -128,12 +139,24 @@ struct cagraSearchParams { /** Bit mask used for initial random seed node selection. */ uint64_t rand_xor_mask; }; -typedef struct cagraSearchParams cagraSearchParams; + +typedef struct cagraSearchParams* cagraSearchParams_t; /** - * @brief Default cagraSearchParams to use + * @brief Allocate CAGRA search params, and populate with default values + * + * @param[in] params cagraSearchParams_t to allocate + * @return cuvsError_t */ -extern const cagraSearchParams cagraDefaultSearchParams; +cuvsError_t cagraSearchParamsCreate(cagraSearchParams_t* params); + +/** + * @brief De-allocate CAGRA search params + * + * @param[in] params + * @return cuvsError_t + */ +cuvsError_t cagraSearchParamsDestroy(cagraSearchParams_t params); /** * @brief Struct to hold address of cuvs::neighbors::cagra::index and its active trained dtype @@ -181,6 +204,10 @@ cuvsError_t cagraIndexDestroy(cagraIndex_t index); * // Assume a populated `DLManagedTensor` type here * DLManagedTensor dataset; * + * // Create default index params + * cagraIndexParams_t params; + * cuvsError_t params_create_status = cagraIndexParamsCreate(¶ms); + * * // Create CAGRA index * cagraIndex_t index; * cuvsError_t index_create_status = cagraIndexCreate(&index); @@ -188,24 +215,25 @@ cuvsError_t cagraIndexDestroy(cagraIndex_t index); * // Build the CAGRA Index * cuvsError_t build_status = cagraBuild(res, params, &dataset, index); * - * // de-allocate `index` and `res` + * // de-allocate `params`, `index` and `res` + * cuvsError_t params_destroy_status = cagraIndexParamsDestroy(params); * cuvsError_t index_destroy_status = cagraIndexDestroy(index); * cuvsError_t res_destroy_status = cuvsResourcesDestroy(res); * @endcode * * @param[in] res cuvsResources_t opaque C handle - * @param[in] params cagraIndexParams used to build CAGRA index + * @param[in] params cagraIndexParams_t used to build CAGRA index * @param[in] dataset DLManagedTensor* training dataset * @param[out] index cagraIndex_t Newly built CAGRA index * @return cuvsError_t */ cuvsError_t cagraBuild(cuvsResources_t res, - struct cagraIndexParams params, + cagraIndexParams_t params, DLManagedTensor* dataset, cagraIndex_t index); /** - * @brief Build a CAGRA index with a `DLManagedTensor` which has underlying + * @brief Search a CAGRA index with a `DLManagedTensor` which has underlying * `DLDeviceType` equal to `kDLCUDA`, `kDLCUDAHost`, `kDLCUDAManaged`. * It is also important to note that the CAGRA Index must have been built * with the same type of `queries`, such that `index.dtype.code == @@ -227,24 +255,27 @@ cuvsError_t cagraBuild(cuvsResources_t res, * DLManagedTensor queries; * DLManagedTensor neighbors; * + * // Create default search params + * cagraSearchParams_t params; + * cuvsError_t params_create_status = cagraSearchParamsCreate(¶ms); + * * // Search the `index` built using `cagraBuild` - * cagraSearchParams params; * cuvsError_t search_status = cagraSearch(res, params, index, queries, neighbors, distances); * - * // de-allocate `index` and `res` - * cuvsError_t index_destroy_status = cagraIndexDestroy(index); + * // de-allocate `params` and `res` + * cuvsError_t params_destroy_status = cagraSearchParamsDestroy(params); * cuvsError_t res_destroy_status = cuvsResourcesDestroy(res); * @endcode * * @param[in] res cuvsResources_t opaque C handle - * @param[in] params cagraSearchParams used to search CAGRA index + * @param[in] params cagraSearchParams_t used to search CAGRA index * @param[in] index cagraIndex which has been returned by `cagraBuild` * @param[in] queries DLManagedTensor* queries dataset to search * @param[out] neighbors DLManagedTensor* output `k` neighbors for queries * @param[out] distances DLManagedTensor* output `k` distances for queries */ cuvsError_t cagraSearch(cuvsResources_t res, - cagraSearchParams params, + cagraSearchParams_t params, cagraIndex_t index, DLManagedTensor* queries, DLManagedTensor* neighbors, diff --git a/cpp/src/neighbors/cagra_c.cpp b/cpp/src/neighbors/cagra_c.cpp index eaaf3cc98..e9b9860fb 100644 --- a/cpp/src/neighbors/cagra_c.cpp +++ b/cpp/src/neighbors/cagra_c.cpp @@ -131,7 +131,7 @@ extern "C" cuvsError_t cagraIndexDestroy(cagraIndex_t index_c_ptr) } extern "C" cuvsError_t cagraBuild(cuvsResources_t res, - cagraIndexParams params, + cagraIndexParams_t params, DLManagedTensor* dataset_tensor, cagraIndex_t index) { @@ -139,13 +139,13 @@ extern "C" cuvsError_t cagraBuild(cuvsResources_t res, auto dataset = dataset_tensor->dl_tensor; if (dataset.dtype.code == kDLFloat && dataset.dtype.bits == 32) { - index->addr = reinterpret_cast(_build(res, params, 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, 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, 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", @@ -159,7 +159,7 @@ extern "C" cuvsError_t cagraBuild(cuvsResources_t res, } extern "C" cuvsError_t cagraSearch(cuvsResources_t res, - cagraSearchParams params, + cagraSearchParams_t params, cagraIndex_t index_c_ptr, DLManagedTensor* queries_tensor, DLManagedTensor* neighbors_tensor, @@ -186,11 +186,11 @@ extern "C" cuvsError_t cagraSearch(cuvsResources_t res, RAFT_EXPECTS(queries.dtype.code == index.dtype.code, "type mismatch between index and queries"); if (queries.dtype.code == kDLFloat && queries.dtype.bits == 32) { - _search(res, params, index, queries_tensor, neighbors_tensor, distances_tensor); + _search(res, *params, index, queries_tensor, neighbors_tensor, distances_tensor); } else if (queries.dtype.code == kDLInt && queries.dtype.bits == 8) { - _search(res, params, index, queries_tensor, neighbors_tensor, distances_tensor); + _search(res, *params, index, queries_tensor, neighbors_tensor, distances_tensor); } else if (queries.dtype.code == kDLUInt && queries.dtype.bits == 8) { - _search(res, params, index, queries_tensor, neighbors_tensor, distances_tensor); + _search(res, *params, index, queries_tensor, neighbors_tensor, distances_tensor); } else { RAFT_FAIL("Unsupported queries DLtensor dtype: %d and bits: %d", queries.dtype.code, @@ -202,13 +202,49 @@ extern "C" cuvsError_t cagraSearch(cuvsResources_t res, } } -extern "C" const cagraSearchParams cagraDefaultSearchParams = {.itopk_size = 64, - .search_width = 1, - .hashmap_max_fill_rate = 0.5, - .num_random_samplings = 1, - .rand_xor_mask = 0x128394}; +extern "C" cuvsError_t cagraIndexParamsCreate(cagraIndexParams_t* params) +{ + try { + *params = new cagraIndexParams{.intermediate_graph_degree = 128, + .graph_degree = 64, + .build_algo = IVF_PQ, + .nn_descent_niter = 20}; + return CUVS_SUCCESS; + } catch (...) { + return CUVS_ERROR; + } +} + +extern "C" cuvsError_t cagraIndexParamsDestroy(cagraIndexParams_t params) +{ + try { + delete params; + return CUVS_SUCCESS; + } catch (...) { + return CUVS_ERROR; + } +} -extern "C" const cagraIndexParams cagraDefaultIndexParams = {.intermediate_graph_degree = 128, - .graph_degree = 64, - .build_algo = IVF_PQ, - .nn_descent_niter = 20}; +extern "C" cuvsError_t cagraSearchParamsCreate(cagraSearchParams_t* params) +{ + try { + *params = new cagraSearchParams{.itopk_size = 64, + .search_width = 1, + .hashmap_max_fill_rate = 0.5, + .num_random_samplings = 1, + .rand_xor_mask = 0x128394}; + return CUVS_SUCCESS; + } catch (...) { + return CUVS_ERROR; + } +} + +extern "C" cuvsError_t cagraSearchParamsDestroy(cagraSearchParams_t params) +{ + try { + delete params; + return CUVS_SUCCESS; + } catch (...) { + return CUVS_ERROR; + } +} diff --git a/cpp/test/neighbors/ann_cagra_c.cu b/cpp/test/neighbors/ann_cagra_c.cu index b866c4df1..a41498490 100644 --- a/cpp/test/neighbors/ann_cagra_c.cu +++ b/cpp/test/neighbors/ann_cagra_c.cu @@ -60,7 +60,8 @@ TEST(CagraC, BuildSearch) cagraIndexCreate(&index); // build index - cagraIndexParams build_params = cagraDefaultIndexParams; + cagraIndexParams_t build_params; + cagraIndexParamsCreate(&build_params); cagraBuild(res, build_params, &dataset_tensor, index); // create queries DLTensor @@ -110,7 +111,8 @@ TEST(CagraC, BuildSearch) distances_tensor.dl_tensor.strides = nullptr; // search index - cagraSearchParams search_params = cagraDefaultSearchParams; + cagraSearchParams_t search_params; + cagraSearchParamsCreate(&search_params); cagraSearch(res, search_params, index, &queries_tensor, &neighbors_tensor, &distances_tensor); // verify output @@ -124,6 +126,8 @@ TEST(CagraC, BuildSearch) cudaFree(distances_d); // de-allocate index and res + cagraSearchParamsDestroy(search_params); + cagraIndexParamsDestroy(build_params); cagraIndexDestroy(index); cuvsResourcesDestroy(res); }