Skip to content

Commit

Permalink
More updates
Browse files Browse the repository at this point in the history
  • Loading branch information
@cjnolet
cjnolet committed Feb 5, 2024
1 parent c3c6780 commit 1d337e1
Show file tree
Hide file tree
Showing 5 changed files with 69 additions and 69 deletions.
4 changes: 2 additions & 2 deletions ci/test_wheel_cuvs.sh
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,8 +4,8 @@
set -euo pipefail

mkdir -p ./dist
RAPIDS_PY_CUDA_SUFFIX="$(rapids-wheel-ctk-name-gen ${RAPIDS_CUDA_VERSION})"
RAPIDS_PY_WHEEL_NAME="cuvs_${RAPIDS_PY_CUDA_SUFFIX}" rapids-download-wheels-from-s3 ./dist
#RAPIDS_PY_CUDA_SUFFIX="$(rapids-wheel-ctk-name-gen ${RAPIDS_CUDA_VERSION})"
#RAPIDS_PY_WHEEL_NAME="cuvs_${RAPIDS_PY_CUDA_SUFFIX}" rapids-download-wheels-from-s3 ./dist

## echo to expand wildcard before adding `[extra]` requires for pip
#python -m pip install $(echo ./dist/cuvs*.whl)[test]
Expand Down
68 changes: 34 additions & 34 deletions cpp/include/cuvs/neighbors/cagra_c.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,7 @@ extern "C" {
* @brief Enum to denote which ANN algorithm is used to build CAGRA graph
*
*/
enum cagraGraphBuildAlgo {
enum cuvsCagraGraphBuildAlgo {
/* Use IVF-PQ to build all-neighbors knn graph */
IVF_PQ,
/* Experimental, use NN-Descent to build all-neighbors knn graph */
Expand All @@ -44,18 +44,18 @@ enum cagraGraphBuildAlgo {
* @brief Supplemental parameters to build CAGRA Index
*
*/
struct cagraIndexParams {
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 cagraGraphBuildAlgo build_algo;
enum cuvsCagraGraphBuildAlgo build_algo;
/** Number of Iterations to run if building with NN_DESCENT */
size_t nn_descent_niter;
};

typedef struct cagraIndexParams* cuvsCagraIndexParams_t;
typedef struct cuvsCagraIndexParams* cuvsCagraIndexParams_t;

/**
* @brief Allocate CAGRA Index params, and populate with default values
Expand All @@ -77,7 +77,7 @@ cuvsError_t cuvsCagraIndexParamsDestroy(cuvsCagraIndexParams_t index);
* @brief Enum to denote algorithm used to search CAGRA Index
*
*/
enum cagraSearchAlgo {
enum cuvsCagraSearchAlgo {
/** For large batch sizes. */
SINGLE_CTA,
/** For small batch sizes. */
Expand All @@ -90,13 +90,13 @@ enum cagraSearchAlgo {
* @brief Enum to denote Hash Mode used while searching CAGRA index
*
*/
enum cagraHashMode { HASH, SMALL, AUTO_HASH };
enum cuvsCagraHashMode { HASH, SMALL, AUTO_HASH };

/**
* @brief Supplemental parameters to search CAGRA index
*
*/
struct cagraSearchParams {
struct cuvsCagraSearchParams {
/** Maximum number of queries to search at the same time (batch size). Auto select when 0.*/
size_t max_queries;

Expand All @@ -114,7 +114,7 @@ struct cagraSearchParams {
// Reasonable default values are automatically chosen.

/** Which search implementation to use. */
enum cagraSearchAlgo algo;
enum cuvsCagraSearchAlgo algo;

/** Number of threads used to calculate a single distance. 4, 8, 16, or 32. */
size_t team_size;
Expand All @@ -128,7 +128,7 @@ struct cagraSearchParams {
/** Thread block size. 0, 64, 128, 256, 512, 1024. Auto selection when 0. */
size_t thread_block_size;
/** Hashmap type. Auto selection when AUTO. */
enum cagraHashMode hashmap_mode;
enum cuvsCagraHashMode hashmap_mode;
/** Lower limit of hashmap bit length. More than 8. */
size_t hashmap_min_bitlen;
/** Upper limit of hashmap fill rate. More than 0.1, less than 0.9.*/
Expand All @@ -140,7 +140,7 @@ struct cagraSearchParams {
uint64_t rand_xor_mask;
};

typedef struct cagraSearchParams* cuvsCagraSearchParams_t;
typedef struct cuvsCagraSearchParams* cuvsCagraSearchParams_t;

/**
* @brief Allocate CAGRA search params, and populate with default values
Expand All @@ -166,24 +166,24 @@ typedef struct {
uintptr_t addr;
DLDataType dtype;

} cagraIndex;
} cuvsCagraIndex;

typedef cagraIndex* cagraIndex_t;
typedef cagraIndex* cuvsCagraIndex_t;

/**
* @brief Allocate CAGRA index
*
* @param[in] index cagraIndex_t to allocate
* @param[in] index cuvsCagraIndex_t to allocate
* @return cagraError_t
*/
cuvsError_t cagraIndexCreate(cagraIndex_t* index);
cuvsError_t cuvsCagraIndexCreate(cuvsCagraIndex_t* index);

/**
* @brief De-allocate CAGRA index
*
* @param[in] index cagraIndex_t to de-allocate
* @param[in] index cuvsCagraIndex_t to de-allocate
*/
cuvsError_t cagraIndexDestroy(cagraIndex_t index);
cuvsError_t cuvsCagraIndexDestroy(cuvsCagraIndex_t index);

/**
* @brief Build a CAGRA index with a `DLManagedTensor` which has underlying
Expand All @@ -209,28 +209,28 @@ cuvsError_t cagraIndexDestroy(cagraIndex_t index);
* cuvsError_t params_create_status = cuvsCagraIndexParamsCreate(&params);
*
* // Create CAGRA index
* cagraIndex_t index;
* cuvsError_t index_create_status = cagraIndexCreate(&index);
* cuvsCagraIndex_t index;
* cuvsError_t index_create_status = cuvsCagraIndexCreate(&index);
*
* // Build the CAGRA Index
* cuvsError_t build_status = cagraBuild(res, params, &dataset, index);
* cuvsError_t build_status = cuvsCagraBuild(res, params, &dataset, index);
*
* // de-allocate `params`, `index` and `res`
* cuvsError_t params_destroy_status = cuvsCagraIndexParamsDestroy(params);
* cuvsError_t index_destroy_status = cagraIndexDestroy(index);
* cuvsError_t index_destroy_status = cuvsCagraIndexDestroy(index);
* cuvsError_t res_destroy_status = cuvsResourcesDestroy(res);
* @endcode
*
* @param[in] res cuvsResources_t opaque C handle
* @param[in] params cuvsCagraIndexParams_t used to build CAGRA index
* @param[in] dataset DLManagedTensor* training dataset
* @param[out] index cagraIndex_t Newly built CAGRA index
* @param[out] index cuvsCagraIndex_t Newly built CAGRA index
* @return cuvsError_t
*/
cuvsError_t cagraBuild(cuvsResources_t res,
cuvsCagraIndexParams_t params,
DLManagedTensor* dataset,
cagraIndex_t index);
cuvsError_t cuvsCagraBuild(cuvsResources_t res,
cuvsCagraIndexParams_t params,
DLManagedTensor* dataset,
cuvsCagraIndex_t index);

/**
* @brief Search a CAGRA index with a `DLManagedTensor` which has underlying
Expand Down Expand Up @@ -259,8 +259,8 @@ cuvsError_t cagraBuild(cuvsResources_t res,
* cuvsCagraSearchParams_t params;
* cuvsError_t params_create_status = cuvsCagraSearchParamsCreate(&params);
*
* // Search the `index` built using `cagraBuild`
* cuvsError_t search_status = cagraSearch(res, params, index, queries, neighbors, distances);
* // Search the `index` built using `cuvsCagraBuild`
* cuvsError_t search_status = cuvsCagraSearch(res, params, index, queries, neighbors, distances);
*
* // de-allocate `params` and `res`
* cuvsError_t params_destroy_status = cuvsCagraSearchParamsDestroy(params);
Expand All @@ -269,17 +269,17 @@ cuvsError_t cagraBuild(cuvsResources_t res,
*
* @param[in] res cuvsResources_t opaque C handle
* @param[in] params cuvsCagraSearchParams_t used to search CAGRA index
* @param[in] index cagraIndex which has been returned by `cagraBuild`
* @param[in] index cuvsCagraIndex which has been returned by `cuvsCagraBuild`
* @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,
cuvsCagraSearchParams_t params,
cagraIndex_t index,
DLManagedTensor* queries,
DLManagedTensor* neighbors,
DLManagedTensor* distances);
cuvsError_t cuvsCagraSearch(cuvsResources_t res,
cuvsCagraSearchParams_t params,
cuvsCagraIndex_t index,
DLManagedTensor* queries,
DLManagedTensor* neighbors,
DLManagedTensor* distances);

#ifdef __cplusplus
}
Expand Down
50 changes: 25 additions & 25 deletions cpp/src/neighbors/cagra_c.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@
namespace {

template <typename T>
void* _build(cuvsResources_t res, cagraIndexParams params, DLManagedTensor* dataset_tensor)
void* _build(cuvsResources_t res, cuvsCagraIndexParams params, DLManagedTensor* dataset_tensor)
{
auto dataset = dataset_tensor->dl_tensor;

Expand Down Expand Up @@ -59,8 +59,8 @@ void* _build(cuvsResources_t res, cagraIndexParams params, DLManagedTensor* data

template <typename T>
void _search(cuvsResources_t res,
cagraSearchParams params,
cagraIndex index,
cuvsCagraSearchParams params,
cuvsCagraIndex index,
DLManagedTensor* queries_tensor,
DLManagedTensor* neighbors_tensor,
DLManagedTensor* distances_tensor)
Expand Down Expand Up @@ -95,17 +95,17 @@ void _search(cuvsResources_t res,

} // namespace

extern "C" cuvsError_t cagraIndexCreate(cagraIndex_t* index)
extern "C" cuvsError_t cuvsCagraIndexCreate(cuvsCagraIndex_t* index)
{
try {
*index = new cagraIndex{};
*index = new cuvsCagraIndex{};
return CUVS_SUCCESS;
} catch (...) {
return CUVS_ERROR;
}
}

extern "C" cuvsError_t cagraIndexDestroy(cagraIndex_t index_c_ptr)
extern "C" cuvsError_t cuvsCagraIndexDestroy(cuvsCagraIndex_t index_c_ptr)
{
try {
auto index = *index_c_ptr;
Expand All @@ -130,10 +130,10 @@ extern "C" cuvsError_t cagraIndexDestroy(cagraIndex_t index_c_ptr)
}
}

extern "C" cuvsError_t cagraBuild(cuvsResources_t res,
cuvsCagraIndexParams_t params,
DLManagedTensor* dataset_tensor,
cagraIndex_t index)
extern "C" cuvsError_t cuvsCagraBuild(cuvsResources_t res,
cuvsCagraIndexParams_t params,
DLManagedTensor* dataset_tensor,
cuvsCagraIndex_t index)
{
try {
auto dataset = dataset_tensor->dl_tensor;
Expand All @@ -158,12 +158,12 @@ extern "C" cuvsError_t cagraBuild(cuvsResources_t res,
}
}

extern "C" cuvsError_t cagraSearch(cuvsResources_t res,
cuvsCagraSearchParams_t params,
cagraIndex_t index_c_ptr,
DLManagedTensor* queries_tensor,
DLManagedTensor* neighbors_tensor,
DLManagedTensor* distances_tensor)
extern "C" cuvsError_t cuvsCagraSearch(cuvsResources_t res,
cuvsCagraSearchParams_t params,
cuvsCagraIndex_t index_c_ptr,
DLManagedTensor* queries_tensor,
DLManagedTensor* neighbors_tensor,
DLManagedTensor* distances_tensor)
{
try {
auto queries = queries_tensor->dl_tensor;
Expand Down Expand Up @@ -205,10 +205,10 @@ extern "C" cuvsError_t cagraSearch(cuvsResources_t res,
extern "C" cuvsError_t cuvsCagraIndexParamsCreate(cuvsCagraIndexParams_t* params)
{
try {
*params = new cagraIndexParams{.intermediate_graph_degree = 128,
.graph_degree = 64,
.build_algo = IVF_PQ,
.nn_descent_niter = 20};
*params = new cuvsCagraIndexParams{.intermediate_graph_degree = 128,
.graph_degree = 64,
.build_algo = IVF_PQ,
.nn_descent_niter = 20};
return CUVS_SUCCESS;
} catch (...) {
return CUVS_ERROR;
Expand All @@ -228,11 +228,11 @@ extern "C" cuvsError_t cuvsCagraIndexParamsDestroy(cuvsCagraIndexParams_t params
extern "C" cuvsError_t cuvsCagraSearchParamsCreate(cuvsCagraSearchParams_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};
*params = new cuvsCagraSearchParams{.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;
Expand Down
10 changes: 5 additions & 5 deletions cpp/test/neighbors/ann_cagra_c.cu
View file
Original file line number Diff line number Diff line change
Expand Up @@ -56,13 +56,13 @@ TEST(CagraC, BuildSearch)
dataset_tensor.dl_tensor.strides = nullptr;

// create index
cagraIndex_t index;
cagraIndexCreate(&index);
cuvsCagraIndex_t index;
cuvsCagraIndexCreate(&index);

// build index
cuvsCagraIndexParams_t build_params;
cuvsCagraIndexParamsCreate(&build_params);
cagraBuild(res, build_params, &dataset_tensor, index);
cuvsCagraBuild(res, build_params, &dataset_tensor, index);

// create queries DLTensor
float* queries_d;
Expand Down Expand Up @@ -113,7 +113,7 @@ TEST(CagraC, BuildSearch)
// search index
cuvsCagraSearchParams_t search_params;
cuvsCagraSearchParamsCreate(&search_params);
cagraSearch(res, search_params, index, &queries_tensor, &neighbors_tensor, &distances_tensor);
cuvsCagraSearch(res, search_params, index, &queries_tensor, &neighbors_tensor, &distances_tensor);

// verify output
ASSERT_TRUE(cuvs::devArrMatchHost(neighbors_exp, neighbors_d, 4, cuvs::Compare<uint32_t>()));
Expand All @@ -128,6 +128,6 @@ TEST(CagraC, BuildSearch)
// de-allocate index and res
cuvsCagraSearchParamsDestroy(search_params);
cuvsCagraIndexParamsDestroy(build_params);
cagraIndexDestroy(index);
cuvsCagraIndexDestroy(index);
cuvsResourcesDestroy(res);
}
6 changes: 3 additions & 3 deletions cpp/test/neighbors/c_api.c
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,8 +24,8 @@ int main()
// simple smoke test to make sure that we can compile the cagra_c.h API
// using a c compiler. This isn't aiming to be a full test, just checking
// that the exposed C-API is valid C code and doesn't contain C++ features
cagraIndex_t index;
cagraIndexCreate(&index);
cagraIndexDestroy(index);
cuvsCagraIndex_t index;
cuvsCagraIndexCreate(&index);
cuvsCagraIndexDestroy(index);
return 0;
}

0 comments on commit 1d337e1

Please sign in to comment.