forked from rapidsai/cuvs
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
KNN bruteforce and IVF methods C/C++ API (rapidsai#33)
Authors: - Victor Lafargue (https://github.com/viclafargue) - Corey J. Nolet (https://github.com/cjnolet) Approvers: - Ben Frederickson (https://github.com/benfred) - Corey J. Nolet (https://github.com/cjnolet) - Jake Awe (https://github.com/AyodeAwe) URL: rapidsai#33
- Loading branch information
1 parent
3975ca3
commit 0db75d0
Showing
71 changed files
with
9,135 additions
and
15 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,70 @@ | ||
| /* | ||
| * Copyright (c) 2024, NVIDIA CORPORATION. | ||
| * | ||
| * Licensed under the Apache License, Version 2.0 (the "License"); | ||
| * you may not use this file except in compliance with the License. | ||
| * You may obtain a copy of the License at | ||
| * | ||
| * http://www.apache.org/licenses/LICENSE-2.0 | ||
| * | ||
| * Unless required by applicable law or agreed to in writing, software | ||
| * distributed under the License is distributed on an "AS IS" BASIS, | ||
| * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| * See the License for the specific language governing permissions and | ||
| * limitations under the License. | ||
| */ | ||
|
|
||
| #ifdef __cplusplus | ||
| extern "C" { | ||
| #endif | ||
|
|
||
| /** enum to tell how to compute distance */ | ||
| enum DistanceType { | ||
|
|
||
| /** evaluate as dist_ij = sum(x_ik^2) + sum(y_ij)^2 - 2*sum(x_ik * y_jk) */ | ||
| L2Expanded = 0, | ||
| /** same as above, but inside the epilogue, perform square root operation */ | ||
| L2SqrtExpanded = 1, | ||
| /** cosine distance */ | ||
| CosineExpanded = 2, | ||
| /** L1 distance */ | ||
| L1 = 3, | ||
| /** evaluate as dist_ij += (x_ik - y-jk)^2 */ | ||
| L2Unexpanded = 4, | ||
| /** same as above, but inside the epilogue, perform square root operation */ | ||
| L2SqrtUnexpanded = 5, | ||
| /** basic inner product **/ | ||
| InnerProduct = 6, | ||
| /** Chebyshev (Linf) distance **/ | ||
| Linf = 7, | ||
| /** Canberra distance **/ | ||
| Canberra = 8, | ||
| /** Generalized Minkowski distance **/ | ||
| LpUnexpanded = 9, | ||
| /** Correlation distance **/ | ||
| CorrelationExpanded = 10, | ||
| /** Jaccard distance **/ | ||
| JaccardExpanded = 11, | ||
| /** Hellinger distance **/ | ||
| HellingerExpanded = 12, | ||
| /** Haversine distance **/ | ||
| Haversine = 13, | ||
| /** Bray-Curtis distance **/ | ||
| BrayCurtis = 14, | ||
| /** Jensen-Shannon distance**/ | ||
| JensenShannon = 15, | ||
| /** Hamming distance **/ | ||
| HammingUnexpanded = 16, | ||
| /** KLDivergence **/ | ||
| KLDivergence = 17, | ||
| /** RusselRao **/ | ||
| RusselRaoExpanded = 18, | ||
| /** Dice-Sorensen distance **/ | ||
| DiceExpanded = 19, | ||
| /** Precomputed (special value) **/ | ||
| Precomputed = 100 | ||
| }; | ||
|
|
||
| #ifdef __cplusplus | ||
| } | ||
| #endif |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,163 @@ | ||
| /* | ||
| * Copyright (c) 2024, NVIDIA CORPORATION. | ||
| * | ||
| * Licensed under the Apache License, Version 2.0 (the "License"); | ||
| * you may not use this file except in compliance with the License. | ||
| * You may obtain a copy of the License at | ||
| * | ||
| * http://www.apache.org/licenses/LICENSE-2.0 | ||
| * | ||
| * Unless required by applicable law or agreed to in writing, software | ||
| * distributed under the License is distributed on an "AS IS" BASIS, | ||
| * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| * See the License for the specific language governing permissions and | ||
| * limitations under the License. | ||
| */ | ||
|
|
||
| #pragma once | ||
|
|
||
| #include <cuvs/core/c_api.h> | ||
| #include <cuvs/distance/distance_types.h> | ||
| #include <dlpack/dlpack.h> | ||
| #include <stdint.h> | ||
|
|
||
| #ifdef __cplusplus | ||
| extern "C" { | ||
| #endif | ||
|
|
||
| /** | ||
| * @defgroup bruteforce_c_index Bruteforce index | ||
| * @{ | ||
| */ | ||
| /** | ||
| * @brief Struct to hold address of cuvs::neighbors::brute_force::index and its active trained dtype | ||
| * | ||
| */ | ||
| typedef struct { | ||
| uintptr_t addr; | ||
| DLDataType dtype; | ||
| } bruteForceIndex; | ||
|
|
||
| typedef bruteForceIndex* cuvsBruteForceIndex_t; | ||
|
|
||
| /** | ||
| * @brief Allocate BRUTEFORCE index | ||
| * | ||
| * @param[in] index cuvsBruteForceIndex_t to allocate | ||
| * @return cuvsError_t | ||
| */ | ||
| cuvsError_t bruteForceIndexCreate(cuvsBruteForceIndex_t* index); | ||
|
|
||
| /** | ||
| * @brief De-allocate BRUTEFORCE index | ||
| * | ||
| * @param[in] index cuvsBruteForceIndex_t to de-allocate | ||
| */ | ||
| cuvsError_t bruteForceIndexDestroy(cuvsBruteForceIndex_t index); | ||
| /** | ||
| * @} | ||
| */ | ||
|
|
||
| /** | ||
| * @defgroup bruteforce_c_index_build Bruteforce index build | ||
| * @{ | ||
| */ | ||
| /** | ||
| * @brief Build a BRUTEFORCE index with a `DLManagedTensor` which has underlying | ||
| * `DLDeviceType` equal to `kDLCUDA`, `kDLCUDAHost`, `kDLCUDAManaged`, | ||
| * or `kDLCPU`. Also, acceptable underlying types are: | ||
| * 1. `kDLDataType.code == kDLFloat` and `kDLDataType.bits = 32` | ||
| * 2. `kDLDataType.code == kDLInt` and `kDLDataType.bits = 8` | ||
| * 3. `kDLDataType.code == kDLUInt` and `kDLDataType.bits = 8` | ||
| * | ||
| * @code {.c} | ||
| * #include <cuvs/core/c_api.h> | ||
| * #include <cuvs/neighbors/brute_force.h> | ||
| * | ||
| * // Create cuvsResources_t | ||
| * cuvsResources_t res; | ||
| * cuvsError_t res_create_status = cuvsResourcesCreate(&res); | ||
| * | ||
| * // Assume a populated `DLManagedTensor` type here | ||
| * DLManagedTensor dataset; | ||
| * | ||
| * // Create BRUTEFORCE index | ||
| * cuvsBruteForceIndex_t index; | ||
| * cuvsError_t index_create_status = bruteForceIndexCreate(&index); | ||
| * | ||
| * // Build the BRUTEFORCE Index | ||
| * cuvsError_t build_status = bruteForceBuild(res, &dataset_tensor, L2Expanded, 0.f, index); | ||
| * | ||
| * // de-allocate `index` and `res` | ||
| * cuvsError_t index_destroy_status = bruteForceIndexDestroy(index); | ||
| * cuvsError_t res_destroy_status = cuvsResourcesDestroy(res); | ||
| * @endcode | ||
| * | ||
| * @param[in] res cuvsResources_t opaque C handle | ||
| * @param[in] dataset DLManagedTensor* training dataset | ||
| * @param[in] metric metric | ||
| * @param[in] metric_arg metric_arg | ||
| * @param[out] index cuvsBruteForceIndex_t Newly built BRUTEFORCE index | ||
| * @return cuvsError_t | ||
| */ | ||
| cuvsError_t bruteForceBuild(cuvsResources_t res, | ||
| DLManagedTensor* dataset, | ||
| enum DistanceType metric, | ||
| float metric_arg, | ||
| cuvsBruteForceIndex_t index); | ||
| /** | ||
| * @} | ||
| */ | ||
|
|
||
| /** | ||
| * @defgroup bruteforce_c_index_search Bruteforce index search | ||
| * @{ | ||
| */ | ||
| /** | ||
| * @brief Search a BRUTEFORCE index with a `DLManagedTensor` which has underlying | ||
| * `DLDeviceType` equal to `kDLCUDA`, `kDLCUDAHost`, `kDLCUDAManaged`. | ||
| * It is also important to note that the BRUTEFORCE index must have been built | ||
| * with the same type of `queries`, such that `index.dtype.code == | ||
| * queries.dl_tensor.dtype.code` Types for input are: | ||
| * 1. `queries`: `kDLDataType.code == kDLFloat` and `kDLDataType.bits = 32` | ||
| * 2. `neighbors`: `kDLDataType.code == kDLUInt` and `kDLDataType.bits = 32` | ||
| * 3. `distances`: `kDLDataType.code == kDLFloat` and `kDLDataType.bits = 32` | ||
| * | ||
| * @code {.c} | ||
| * #include <cuvs/core/c_api.h> | ||
| * #include <cuvs/neighbors/brute_force.h> | ||
| * | ||
| * // Create cuvsResources_t | ||
| * cuvsResources_t res; | ||
| * cuvsError_t res_create_status = cuvsResourcesCreate(&res); | ||
| * | ||
| * // Assume a populated `DLManagedTensor` type here | ||
| * DLManagedTensor dataset; | ||
| * DLManagedTensor queries; | ||
| * DLManagedTensor neighbors; | ||
| * | ||
| * // Search the `index` built using `bruteForceBuild` | ||
| * cuvsError_t search_status = bruteForceSearch(res, index, &queries, &neighbors, &distances); | ||
| * | ||
| * // de-allocate `res` | ||
| * cuvsError_t res_destroy_status = cuvsResourcesDestroy(res); | ||
| * @endcode | ||
| * | ||
| * @param[in] res cuvsResources_t opaque C handle | ||
| * @param[in] index bruteForceIndex which has been returned by `bruteForceBuild` | ||
| * @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 bruteForceSearch(cuvsResources_t res, | ||
| cuvsBruteForceIndex_t index, | ||
| DLManagedTensor* queries, | ||
| DLManagedTensor* neighbors, | ||
| DLManagedTensor* distances); | ||
| /** | ||
| * @} | ||
| */ | ||
|
|
||
| #ifdef __cplusplus | ||
| } | ||
| #endif |
Oops, something went wrong.