Skip to content

Commit

Permalink
create/destroy methods for params objects
Browse files Browse the repository at this point in the history
  • Loading branch information
@benfred
benfred committed Jan 31, 2024
1 parent e05040d commit 106fe15
Show file tree
Hide file tree
Showing 3 changed files with 105 additions and 34 deletions.
61 changes: 46 additions & 15 deletions cpp/include/cuvs/neighbors/cagra_c.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -181,31 +204,36 @@ 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(&params);
*
* // Create CAGRA index
* cagraIndex_t index;
* cuvsError_t index_create_status = cagraIndexCreate(&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 ==
Expand All @@ -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(&params);
*
* // 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,
Expand Down
70 changes: 53 additions & 17 deletions cpp/src/neighbors/cagra_c.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -131,21 +131,21 @@ 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)
{
try {
auto dataset = dataset_tensor->dl_tensor;

if (dataset.dtype.code == kDLFloat && dataset.dtype.bits == 32) {
index->addr = reinterpret_cast<uintptr_t>(_build<float>(res, params, 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, 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, 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",
Expand All @@ -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,
Expand All @@ -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<float>(res, params, index, queries_tensor, neighbors_tensor, distances_tensor);
_search<float>(res, *params, index, queries_tensor, neighbors_tensor, distances_tensor);
} else if (queries.dtype.code == kDLInt && queries.dtype.bits == 8) {
_search<int8_t>(res, params, index, queries_tensor, neighbors_tensor, distances_tensor);
_search<int8_t>(res, *params, index, queries_tensor, neighbors_tensor, distances_tensor);
} else if (queries.dtype.code == kDLUInt && queries.dtype.bits == 8) {
_search<uint8_t>(res, params, index, queries_tensor, neighbors_tensor, distances_tensor);
_search<uint8_t>(res, *params, index, queries_tensor, neighbors_tensor, distances_tensor);
} else {
RAFT_FAIL("Unsupported queries DLtensor dtype: %d and bits: %d",
queries.dtype.code,
Expand All @@ -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;
}
}
8 changes: 6 additions & 2 deletions cpp/test/neighbors/ann_cagra_c.cu
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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);
}

0 comments on commit 106fe15

Please sign in to comment.