diff --git a/cpp/include/cuvs/distance/distance_types.h b/cpp/include/cuvs/distance/distance_types.h index 8e9a4149c..6cc2a993b 100644 --- a/cpp/include/cuvs/distance/distance_types.h +++ b/cpp/include/cuvs/distance/distance_types.h @@ -13,6 +13,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ +#pragma once #ifdef __cplusplus extern "C" { @@ -67,4 +68,4 @@ enum DistanceType { #ifdef __cplusplus } -#endif \ No newline at end of file +#endif diff --git a/cpp/include/cuvs/neighbors/brute_force.h b/cpp/include/cuvs/neighbors/brute_force.h index 0bb4d6bdb..cb7ee2a30 100644 --- a/cpp/include/cuvs/neighbors/brute_force.h +++ b/cpp/include/cuvs/neighbors/brute_force.h @@ -36,9 +36,9 @@ extern "C" { typedef struct { uintptr_t addr; DLDataType dtype; -} bruteForceIndex; +} cuvsBruteForceIndex; -typedef bruteForceIndex* cuvsBruteForceIndex_t; +typedef cuvsBruteForceIndex* cuvsBruteForceIndex_t; /** * @brief Allocate BRUTEFORCE index @@ -46,14 +46,14 @@ typedef bruteForceIndex* cuvsBruteForceIndex_t; * @param[in] index cuvsBruteForceIndex_t to allocate * @return cuvsError_t */ -cuvsError_t bruteForceIndexCreate(cuvsBruteForceIndex_t* index); +cuvsError_t cuvsBruteForceIndexCreate(cuvsBruteForceIndex_t* index); /** * @brief De-allocate BRUTEFORCE index * * @param[in] index cuvsBruteForceIndex_t to de-allocate */ -cuvsError_t bruteForceIndexDestroy(cuvsBruteForceIndex_t index); +cuvsError_t cuvsBruteForceIndexDestroy(cuvsBruteForceIndex_t index); /** * @} */ @@ -83,13 +83,13 @@ cuvsError_t bruteForceIndexDestroy(cuvsBruteForceIndex_t index); * * // Create BRUTEFORCE index * cuvsBruteForceIndex_t index; - * cuvsError_t index_create_status = bruteForceIndexCreate(&index); + * cuvsError_t index_create_status = cuvsBruteForceIndexCreate(&index); * * // Build the BRUTEFORCE Index - * cuvsError_t build_status = bruteForceBuild(res, &dataset_tensor, L2Expanded, 0.f, index); + * cuvsError_t build_status = cuvsBruteForceBuild(res, &dataset_tensor, L2Expanded, 0.f, index); * * // de-allocate `index` and `res` - * cuvsError_t index_destroy_status = bruteForceIndexDestroy(index); + * cuvsError_t index_destroy_status = cuvsBruteForceIndexDestroy(index); * cuvsError_t res_destroy_status = cuvsResourcesDestroy(res); * @endcode * @@ -100,11 +100,11 @@ cuvsError_t bruteForceIndexDestroy(cuvsBruteForceIndex_t index); * @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); +cuvsError_t cuvsBruteForceBuild(cuvsResources_t res, + DLManagedTensor* dataset, + enum DistanceType metric, + float metric_arg, + cuvsBruteForceIndex_t index); /** * @} */ @@ -136,24 +136,24 @@ cuvsError_t bruteForceBuild(cuvsResources_t res, * DLManagedTensor queries; * DLManagedTensor neighbors; * - * // Search the `index` built using `bruteForceBuild` - * cuvsError_t search_status = bruteForceSearch(res, index, &queries, &neighbors, &distances); + * // Search the `index` built using `cuvsBruteForceBuild` + * cuvsError_t search_status = cuvsBruteForceSearch(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] index cuvsBruteForceIndex which has been returned by `cuvsBruteForceBuild` * @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); +cuvsError_t cuvsBruteForceSearch(cuvsResources_t res, + cuvsBruteForceIndex_t index, + DLManagedTensor* queries, + DLManagedTensor* neighbors, + DLManagedTensor* distances); /** * @} */ diff --git a/cpp/src/neighbors/brute_force_c.cpp b/cpp/src/neighbors/brute_force_c.cpp index 531fc9a57..88349e089 100644 --- a/cpp/src/neighbors/brute_force_c.cpp +++ b/cpp/src/neighbors/brute_force_c.cpp @@ -23,6 +23,7 @@ #include #include +#include #include #include #include @@ -49,7 +50,7 @@ void* _build(cuvsResources_t res, template void _search(cuvsResources_t res, - bruteForceIndex index, + cuvsBruteForceIndex index, DLManagedTensor* queries_tensor, DLManagedTensor* neighbors_tensor, DLManagedTensor* distances_tensor) @@ -70,19 +71,14 @@ void _search(cuvsResources_t res, } // namespace -extern "C" cuvsError_t bruteForceIndexCreate(cuvsBruteForceIndex_t* index) +extern "C" cuvsError_t cuvsBruteForceIndexCreate(cuvsBruteForceIndex_t* index) { - try { - *index = new bruteForceIndex{}; - return CUVS_SUCCESS; - } catch (...) { - return CUVS_ERROR; - } + return cuvs::core::translate_exceptions([=] { *index = new cuvsBruteForceIndex{}; }); } -extern "C" cuvsError_t bruteForceIndexDestroy(cuvsBruteForceIndex_t index_c_ptr) +extern "C" cuvsError_t cuvsBruteForceIndexDestroy(cuvsBruteForceIndex_t index_c_ptr) { - try { + return cuvs::core::translate_exceptions([=] { auto index = *index_c_ptr; if (index.dtype.code == kDLFloat) { @@ -96,19 +92,16 @@ extern "C" cuvsError_t bruteForceIndexDestroy(cuvsBruteForceIndex_t index_c_ptr) delete index_ptr; } delete index_c_ptr; - return CUVS_SUCCESS; - } catch (...) { - return CUVS_ERROR; - } + }); } -extern "C" cuvsError_t bruteForceBuild(cuvsResources_t res, - DLManagedTensor* dataset_tensor, - enum DistanceType metric, - float metric_arg, - cuvsBruteForceIndex_t index) +extern "C" cuvsError_t cuvsBruteForceBuild(cuvsResources_t res, + DLManagedTensor* dataset_tensor, + enum DistanceType metric, + float metric_arg, + cuvsBruteForceIndex_t index) { - try { + return cuvs::core::translate_exceptions([=] { auto dataset = dataset_tensor->dl_tensor; if (dataset.dtype.code == kDLFloat && dataset.dtype.bits == 32) { @@ -120,19 +113,16 @@ extern "C" cuvsError_t bruteForceBuild(cuvsResources_t res, dataset.dtype.code, dataset.dtype.bits); } - return CUVS_SUCCESS; - } catch (...) { - return CUVS_ERROR; - } + }); } -extern "C" cuvsError_t bruteForceSearch(cuvsResources_t res, - cuvsBruteForceIndex_t index_c_ptr, - DLManagedTensor* queries_tensor, - DLManagedTensor* neighbors_tensor, - DLManagedTensor* distances_tensor) +extern "C" cuvsError_t cuvsBruteForceSearch(cuvsResources_t res, + cuvsBruteForceIndex_t index_c_ptr, + DLManagedTensor* queries_tensor, + DLManagedTensor* neighbors_tensor, + DLManagedTensor* distances_tensor) { - try { + return cuvs::core::translate_exceptions([=] { auto queries = queries_tensor->dl_tensor; auto neighbors = neighbors_tensor->dl_tensor; auto distances = distances_tensor->dl_tensor; @@ -159,9 +149,5 @@ extern "C" cuvsError_t bruteForceSearch(cuvsResources_t res, queries.dtype.code, queries.dtype.bits); } - - return CUVS_SUCCESS; - } catch (...) { - return CUVS_ERROR; - } + }); } diff --git a/cpp/src/neighbors/ivf_flat_c.cpp b/cpp/src/neighbors/ivf_flat_c.cpp index b9488ec36..d8185c4f7 100644 --- a/cpp/src/neighbors/ivf_flat_c.cpp +++ b/cpp/src/neighbors/ivf_flat_c.cpp @@ -23,6 +23,7 @@ #include #include +#include #include #include #include @@ -86,17 +87,12 @@ void _search(cuvsResources_t res, extern "C" cuvsError_t ivfFlatIndexCreate(cuvsIvfFlatIndex_t* index) { - try { - *index = new ivfFlatIndex{}; - return CUVS_SUCCESS; - } catch (...) { - return CUVS_ERROR; - } + return cuvs::core::translate_exceptions([=] { *index = new ivfFlatIndex{}; }); } extern "C" cuvsError_t ivfFlatIndexDestroy(cuvsIvfFlatIndex_t index_c_ptr) { - try { + return cuvs::core::translate_exceptions([=] { auto index = *index_c_ptr; if (index.dtype.code == kDLFloat) { @@ -113,10 +109,7 @@ extern "C" cuvsError_t ivfFlatIndexDestroy(cuvsIvfFlatIndex_t index_c_ptr) delete index_ptr; } delete index_c_ptr; - return CUVS_SUCCESS; - } catch (...) { - return CUVS_ERROR; - } + }); } extern "C" cuvsError_t ivfFlatBuild(cuvsResources_t res, @@ -124,7 +117,7 @@ extern "C" cuvsError_t ivfFlatBuild(cuvsResources_t res, DLManagedTensor* dataset_tensor, cuvsIvfFlatIndex_t index) { - try { + return cuvs::core::translate_exceptions([=] { auto dataset = dataset_tensor->dl_tensor; if (dataset.dtype.code == kDLFloat && dataset.dtype.bits == 32) { @@ -144,10 +137,7 @@ extern "C" cuvsError_t ivfFlatBuild(cuvsResources_t res, dataset.dtype.code, dataset.dtype.bits); } - return CUVS_SUCCESS; - } catch (...) { - return CUVS_ERROR; - } + }); } extern "C" cuvsError_t ivfFlatSearch(cuvsResources_t res, @@ -157,7 +147,7 @@ extern "C" cuvsError_t ivfFlatSearch(cuvsResources_t res, DLManagedTensor* neighbors_tensor, DLManagedTensor* distances_tensor) { - try { + return cuvs::core::translate_exceptions([=] { auto queries = queries_tensor->dl_tensor; auto neighbors = neighbors_tensor->dl_tensor; auto distances = distances_tensor->dl_tensor; @@ -191,16 +181,12 @@ extern "C" cuvsError_t ivfFlatSearch(cuvsResources_t res, queries.dtype.code, queries.dtype.bits); } - - return CUVS_SUCCESS; - } catch (...) { - return CUVS_ERROR; - } + }); } extern "C" cuvsError_t cuvsIvfFlatIndexParamsCreate(cuvsIvfFlatIndexParams_t* params) { - try { + return cuvs::core::translate_exceptions([=] { *params = new ivfFlatIndexParams{.metric = L2Expanded, .metric_arg = 2.0f, .add_data_on_build = true, @@ -209,38 +195,21 @@ extern "C" cuvsError_t cuvsIvfFlatIndexParamsCreate(cuvsIvfFlatIndexParams_t* pa .kmeans_trainset_fraction = 0.5, .adaptive_centers = false, .conservative_memory_allocation = false}; - return CUVS_SUCCESS; - } catch (...) { - return CUVS_ERROR; - } + }); } extern "C" cuvsError_t cuvsIvfFlatIndexParamsDestroy(cuvsIvfFlatIndexParams_t params) { - try { - delete params; - return CUVS_SUCCESS; - } catch (...) { - return CUVS_ERROR; - } + return cuvs::core::translate_exceptions([=] { delete params; }); } extern "C" cuvsError_t cuvsIvfFlatSearchParamsCreate(cuvsIvfFlatSearchParams_t* params) { - try { - *params = new ivfFlatSearchParams{.n_probes = 20}; - return CUVS_SUCCESS; - } catch (...) { - return CUVS_ERROR; - } + return cuvs::core::translate_exceptions( + [=] { *params = new ivfFlatSearchParams{.n_probes = 20}; }); } extern "C" cuvsError_t cuvsIvfFlatSearchParamsDestroy(cuvsIvfFlatSearchParams_t params) { - try { - delete params; - return CUVS_SUCCESS; - } catch (...) { - return CUVS_ERROR; - } + return cuvs::core::translate_exceptions([=] { delete params; }); } diff --git a/cpp/test/neighbors/run_brute_force_c.c b/cpp/test/neighbors/run_brute_force_c.c index 9c7af13a6..ed775a2d6 100644 --- a/cpp/test/neighbors/run_brute_force_c.c +++ b/cpp/test/neighbors/run_brute_force_c.c @@ -44,10 +44,10 @@ void run_brute_force(int64_t n_rows, // create index cuvsBruteForceIndex_t index; - bruteForceIndexCreate(&index); + cuvsBruteForceIndexCreate(&index); // build index - bruteForceBuild(res, &dataset_tensor, metric, 0.0f, index); + cuvsBruteForceBuild(res, &dataset_tensor, metric, 0.0f, index); // create queries DLTensor DLManagedTensor queries_tensor; @@ -86,9 +86,9 @@ void run_brute_force(int64_t n_rows, distances_tensor.dl_tensor.strides = NULL; // search index - bruteForceSearch(res, index, &queries_tensor, &neighbors_tensor, &distances_tensor); + cuvsBruteForceSearch(res, index, &queries_tensor, &neighbors_tensor, &distances_tensor); // de-allocate index and res - bruteForceIndexDestroy(index); + cuvsBruteForceIndexDestroy(index); cuvsResourcesDestroy(res); } diff --git a/python/cuvs/cuvs/distance_type.pxd b/python/cuvs/cuvs/distance_type.pxd new file mode 100644 index 000000000..a1f0366a5 --- /dev/null +++ b/python/cuvs/cuvs/distance_type.pxd @@ -0,0 +1,41 @@ +# +# 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. +# +# cython: language_level=3 + + +cdef extern from "cuvs/distance/distance_types.h" nogil: + ctypedef enum DistanceType: + L2Expanded + L2SqrtExpanded + CosineExpanded + L1 + L2Unexpanded + L2SqrtUnexpanded + InnerProduct + Linf + Canberra + LpUnexpanded + CorrelationExpanded + JaccardExpanded + HellingerExpanded + Haversine + BrayCurtis + JensenShannon + HammingUnexpanded + KLDivergence + RusselRaoExpanded + DiceExpanded + Precomputed diff --git a/python/cuvs/cuvs/neighbors/CMakeLists.txt b/python/cuvs/cuvs/neighbors/CMakeLists.txt index eaf418c60..481ec9915 100644 --- a/python/cuvs/cuvs/neighbors/CMakeLists.txt +++ b/python/cuvs/cuvs/neighbors/CMakeLists.txt @@ -12,4 +12,5 @@ # the License. # ============================================================================= +add_subdirectory(brute_force) add_subdirectory(cagra) diff --git a/python/cuvs/cuvs/neighbors/brute_force/CMakeLists.txt b/python/cuvs/cuvs/neighbors/brute_force/CMakeLists.txt new file mode 100644 index 000000000..8534370ae --- /dev/null +++ b/python/cuvs/cuvs/neighbors/brute_force/CMakeLists.txt @@ -0,0 +1,25 @@ +# ============================================================================= +# Copyright (c) 2023-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. +# ============================================================================= + +# Set the list of Cython files to build +set(cython_sources brute_force.pyx) +set(linked_libraries cuvs::cuvs cuvs_c) + +# Build all of the Cython targets +rapids_cython_create_modules( + CXX + SOURCE_FILES "${cython_sources}" + LINKED_LIBRARIES "${linked_libraries}" ASSOCIATED_TARGETS cuvs MODULE_PREFIX + neighbors_brute_force_ +) diff --git a/python/cuvs/cuvs/neighbors/brute_force/__init__.pxd b/python/cuvs/cuvs/neighbors/brute_force/__init__.pxd new file mode 100644 index 000000000..e69de29bb diff --git a/python/cuvs/cuvs/neighbors/brute_force/__init__.py b/python/cuvs/cuvs/neighbors/brute_force/__init__.py new file mode 100644 index 000000000..b88c4b464 --- /dev/null +++ b/python/cuvs/cuvs/neighbors/brute_force/__init__.py @@ -0,0 +1,18 @@ +# 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. + + +from .brute_force import Index, build, search + +__all__ = ["Index", "build", "search"] diff --git a/python/cuvs/cuvs/neighbors/brute_force/brute_force.pxd b/python/cuvs/cuvs/neighbors/brute_force/brute_force.pxd new file mode 100644 index 000000000..77e484fba --- /dev/null +++ b/python/cuvs/cuvs/neighbors/brute_force/brute_force.pxd @@ -0,0 +1,47 @@ +# +# 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. +# +# cython: language_level=3 + +from libc.stdint cimport uintptr_t + +from cuvs.common.c_api cimport cuvsError_t, cuvsResources_t +from cuvs.common.cydlpack cimport DLDataType, DLManagedTensor +from cuvs.distance_type cimport DistanceType + + +cdef extern from "cuvs/neighbors/brute_force.h" nogil: + + ctypedef struct cuvsBruteForceIndex: + uintptr_t addr + DLDataType dtype + + ctypedef cuvsBruteForceIndex* cuvsBruteForceIndex_t + + cuvsError_t cuvsBruteForceIndexCreate(cuvsBruteForceIndex_t* index) + + cuvsError_t cuvsBruteForceIndexDestroy(cuvsBruteForceIndex_t index) + + cuvsError_t cuvsBruteForceBuild(cuvsResources_t res, + DLManagedTensor* dataset, + DistanceType metric, + float metric_arg, + cuvsBruteForceIndex_t index) except + + + cuvsError_t cuvsBruteForceSearch(cuvsResources_t res, + cuvsBruteForceIndex_t index, + DLManagedTensor* queries, + DLManagedTensor* neighbors, + DLManagedTensor* distances) except + diff --git a/python/cuvs/cuvs/neighbors/brute_force/brute_force.pyx b/python/cuvs/cuvs/neighbors/brute_force/brute_force.pyx new file mode 100644 index 000000000..ccb10e305 --- /dev/null +++ b/python/cuvs/cuvs/neighbors/brute_force/brute_force.pyx @@ -0,0 +1,216 @@ +# +# 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. +# +# cython: language_level=3 + +import numpy as np + +cimport cuvs.common.cydlpack + +from cuvs.common.resources import auto_sync_resources + +from cython.operator cimport dereference as deref +from libc.stdint cimport uint32_t +from libcpp cimport bool + +from cuvs.common cimport cydlpack +from cuvs.distance_type cimport DistanceType + +from pylibraft.common import auto_convert_output, cai_wrapper, device_ndarray +from pylibraft.common.cai_wrapper import wrap_array +from pylibraft.common.interruptible import cuda_interruptible +from pylibraft.distance.pairwise_distance import DISTANCE_TYPES +from pylibraft.neighbors.common import _check_input_array + +from cuvs.common.c_api cimport cuvsResources_t + +from cuvs.common.exceptions import check_cuvs + + +cdef class Index: + """ + Brute Force index object. This object stores the trained Brute Force + which can be used to perform nearest neighbors searches. + """ + + cdef cuvsBruteForceIndex_t index + cdef bool trained + + def __cinit__(self): + self.trained = False + check_cuvs(cuvsBruteForceIndexCreate(&self.index)) + + def __dealloc__(self): + if self.index is not NULL: + check_cuvs(cuvsBruteForceIndexDestroy(self.index)) + + @property + def trained(self): + return self.trained + + def __repr__(self): + return "Index(type=BruteForce)" + + +@auto_sync_resources +def build(dataset, metric="sqeuclidean", metric_arg=2.0, resources=None): + """ + Build the Brute Force index from the dataset for efficient search. + The following distance metrics are supported: + - L2 + + Parameters + ---------- + dataset : CUDA array interface compliant matrix shape (n_samples, dim) + Supported dtype [float, int8, uint8] + metric : Distance metric to use. Default is sqeuclidean + metric_arg : value of 'p' for Minkowski distances + {resources_docstring} + + Returns + ------- + index: cuvs.neighbors.brute_force.Index + + Examples + -------- + + >>> import cupy as cp + >>> from cuvs.neighbors import brute_force + >>> n_samples = 50000 + >>> n_features = 50 + >>> n_queries = 1000 + >>> k = 10 + >>> dataset = cp.random.random_sample((n_samples, n_features), + ... dtype=cp.float32) + >>> index = brute_force.build(dataset, metric="cosine") + >>> distances, neighbors = brute_force.search(index, dataset, k) + >>> distances = cp.asarray(distances) + >>> neighbors = cp.asarray(neighbors) + """ + + dataset_ai = wrap_array(dataset) + _check_input_array(dataset_ai, [np.dtype('float32')]) + + cdef cuvsResources_t res = resources.get_c_obj() + + cdef DistanceType c_metric = DISTANCE_TYPES[metric] + cdef Index idx = Index() + cdef cydlpack.DLManagedTensor* dataset_dlpack = \ + cydlpack.dlpack_c(dataset_ai) + + with cuda_interruptible(): + check_cuvs(cuvsBruteForceBuild( + res, + dataset_dlpack, + c_metric, + metric_arg, + idx.index + )) + idx.trained = True + + return idx + + +@auto_sync_resources +@auto_convert_output +def search(Index index, + queries, + k, + neighbors=None, + distances=None, + resources=None): + """ + Find the k nearest neighbors for each query. + + Parameters + ---------- + index : Index + Trained Brute Force index. + queries : CUDA array interface compliant matrix shape (n_samples, dim) + Supported dtype [float, int8, uint8] + k : int + The number of neighbors. + neighbors : Optional CUDA array interface compliant matrix shape + (n_queries, k), dtype int64_t. If supplied, neighbor + indices will be written here in-place. (default None) + distances : Optional CUDA array interface compliant matrix shape + (n_queries, k) If supplied, the distances to the + neighbors will be written here in-place. (default None) + {resources_docstring} + + Examples + -------- + >>> import cupy as cp + >>> from cuvs.neighbors import brute_force + >>> n_samples = 50000 + >>> n_features = 50 + >>> n_queries = 1000 + >>> dataset = cp.random.random_sample((n_samples, n_features), + ... dtype=cp.float32) + >>> # Build index + >>> index = brute_force.build(dataset, metric="sqeuclidean") + >>> # Search using the built index + >>> queries = cp.random.random_sample((n_queries, n_features), + ... dtype=cp.float32) + >>> k = 10 + >>> # Using a pooling allocator reduces overhead of temporary array + >>> # creation during search. This is useful if multiple searches + >>> # are performad with same query size. + >>> distances, neighbors = brute_force.search(index, queries, k) + >>> neighbors = cp.asarray(neighbors) + >>> distances = cp.asarray(distances) + """ + if not index.trained: + raise ValueError("Index needs to be built before calling search.") + + cdef cuvsResources_t res = resources.get_c_obj() + + queries_cai = wrap_array(queries) + _check_input_array(queries_cai, [np.dtype('float32')]) + + cdef uint32_t n_queries = queries_cai.shape[0] + + if neighbors is None: + neighbors = device_ndarray.empty((n_queries, k), dtype='int64') + + neighbors_cai = wrap_array(neighbors) + _check_input_array(neighbors_cai, [np.dtype('int64')], + exp_rows=n_queries, exp_cols=k) + + if distances is None: + distances = device_ndarray.empty((n_queries, k), dtype='float32') + + distances_cai = wrap_array(distances) + _check_input_array(distances_cai, [np.dtype('float32')], + exp_rows=n_queries, exp_cols=k) + + cdef cuvsError_t search_status + cdef cydlpack.DLManagedTensor* queries_dlpack = \ + cydlpack.dlpack_c(queries_cai) + cdef cydlpack.DLManagedTensor* neighbors_dlpack = \ + cydlpack.dlpack_c(neighbors_cai) + cdef cydlpack.DLManagedTensor* distances_dlpack = \ + cydlpack.dlpack_c(distances_cai) + + with cuda_interruptible(): + check_cuvs(cuvsBruteForceSearch( + res, + index.index, + queries_dlpack, + neighbors_dlpack, + distances_dlpack + )) + + return (distances, neighbors) diff --git a/python/cuvs/cuvs/test/test_brute_force.py b/python/cuvs/cuvs/test/test_brute_force.py new file mode 100644 index 000000000..146ab59f0 --- /dev/null +++ b/python/cuvs/cuvs/test/test_brute_force.py @@ -0,0 +1,90 @@ +# Copyright (c) 2022-2023, 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. +# + +import numpy as np +import pytest +from pylibraft.common import device_ndarray +from scipy.spatial.distance import cdist + +from cuvs.neighbors import brute_force + + +@pytest.mark.parametrize("n_index_rows", [32, 100]) +@pytest.mark.parametrize("n_query_rows", [32, 100]) +@pytest.mark.parametrize("n_cols", [40, 100]) +@pytest.mark.parametrize("k", [1, 5, 32]) +@pytest.mark.parametrize( + "metric", + [ + "euclidean", + "cityblock", + "chebyshev", + "canberra", + "correlation", + "russellrao", + "cosine", + "sqeuclidean", + # "inner_product", + ], +) +@pytest.mark.parametrize("inplace", [True, False]) +@pytest.mark.parametrize("dtype", [np.float32]) +def test_brute_force_knn( + n_index_rows, n_query_rows, n_cols, k, inplace, metric, dtype +): + index = np.random.random_sample((n_index_rows, n_cols)).astype(dtype) + queries = np.random.random_sample((n_query_rows, n_cols)).astype(dtype) + + # RussellRao expects boolean arrays + if metric == "russellrao": + index[index < 0.5] = 0.0 + index[index >= 0.5] = 1.0 + queries[queries < 0.5] = 0.0 + queries[queries >= 0.5] = 1.0 + + indices = np.zeros((n_query_rows, k), dtype="int64") + distances = np.zeros((n_query_rows, k), dtype=dtype) + + index_device = device_ndarray(index) + queries_device = device_ndarray(queries) + indices_device = device_ndarray(indices) + distances_device = device_ndarray(distances) + + brute_force_index = brute_force.build(index_device, metric) + ret_distances, ret_indices = brute_force.search( + brute_force_index, + queries_device, + k, + neighbors=indices_device, + distances=distances_device, + ) + + pw_dists = cdist(queries, index, metric=metric) + + distances_device = ret_distances if not inplace else distances_device + + actual_distances = distances_device.copy_to_host() + + actual_distances[actual_distances <= 1e-5] = 0.0 + argsort = np.argsort(pw_dists, axis=1) + + for i in range(pw_dists.shape[0]): + expected_indices = argsort[i] + gpu_dists = actual_distances[i] + + cpu_ordered = pw_dists[i, expected_indices] + np.testing.assert_allclose( + cpu_ordered[:k], gpu_dists, atol=1e-3, rtol=1e-3 + ) diff --git a/rust/cuvs-sys/build.rs b/rust/cuvs-sys/build.rs index ec9672569..f521f8af3 100644 --- a/rust/cuvs-sys/build.rs +++ b/rust/cuvs-sys/build.rs @@ -98,9 +98,9 @@ fn main() { // generated) .must_use_type("cuvsError_t") // Only generate bindings for cuvs/cagra types and functions - .allowlist_type("(cuvs|cagra|DL).*") - .allowlist_function("(cuvs|cagra).*") - .rustified_enum("(cuvs|cagra|DL).*") + .allowlist_type("(cuvs|bruteForce|cagra|DL).*") + .allowlist_function("(cuvs|bruteForce|cagra).*") + .rustified_enum("(cuvs|cagra|DL|DistanceType).*") // also need some basic cuda mem functions for copying data .allowlist_function("(cudaMemcpyAsync|cudaMemcpy)") .rustified_enum("cudaError") diff --git a/rust/cuvs-sys/cuvs_c_wrapper.h b/rust/cuvs-sys/cuvs_c_wrapper.h index ccca82632..e2c35f12d 100644 --- a/rust/cuvs-sys/cuvs_c_wrapper.h +++ b/rust/cuvs-sys/cuvs_c_wrapper.h @@ -17,4 +17,5 @@ // wrapper file containing all the C-API's we should automatically be creating rust // bindings for #include +#include #include diff --git a/rust/cuvs/src/brute_force.rs b/rust/cuvs/src/brute_force.rs new file mode 100644 index 000000000..e2e03c236 --- /dev/null +++ b/rust/cuvs/src/brute_force.rs @@ -0,0 +1,180 @@ +/* + * 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. + */ + +use std::io::{stderr, Write}; + +use crate::distance_type::DistanceType; +use crate::dlpack::ManagedTensor; +use crate::error::{check_cuvs, Result}; +use crate::resources::Resources; + +/// Brute Force KNN Index +#[derive(Debug)] +pub struct Index(ffi::cuvsBruteForceIndex_t); + +impl Index { + /// Builds a new Brute Force KNN Index from the dataset for efficient search. + /// + /// # Arguments + /// + /// * `res` - Resources to use + /// * `metric` - DistanceType to use for building the index + /// * `metric_arg` - Optional value of `p` for Minkowski distances + /// * `dataset` - A row-major matrix on either the host or device to index + pub fn build>( + res: &Resources, + metric: DistanceType, + metric_arg: Option, + dataset: T, + ) -> Result { + let dataset: ManagedTensor = dataset.into(); + let index = Index::new()?; + unsafe { + check_cuvs(ffi::cuvsBruteForceBuild( + res.0, + dataset.as_ptr(), + metric, + metric_arg.unwrap_or(2.0), + index.0, + ))?; + } + Ok(index) + } + + /// Creates a new empty index + pub fn new() -> Result { + unsafe { + let mut index = std::mem::MaybeUninit::::uninit(); + check_cuvs(ffi::cuvsBruteForceIndexCreate(index.as_mut_ptr()))?; + Ok(Index(index.assume_init())) + } + } + + /// Perform a Nearest Neighbors search on the Index + /// + /// # Arguments + /// + /// * `res` - Resources to use + /// * `queries` - A matrix in device memory to query for + /// * `neighbors` - Matrix in device memory that receives the indices of the nearest neighbors + /// * `distances` - Matrix in device memory that receives the distances of the nearest neighbors + pub fn search( + self, + res: &Resources, + queries: &ManagedTensor, + neighbors: &ManagedTensor, + distances: &ManagedTensor, + ) -> Result<()> { + unsafe { + check_cuvs(ffi::cuvsBruteForceSearch( + res.0, + self.0, + queries.as_ptr(), + neighbors.as_ptr(), + distances.as_ptr(), + )) + } + } +} + +impl Drop for Index { + fn drop(&mut self) { + if let Err(e) = check_cuvs(unsafe { ffi::cuvsBruteForceIndexDestroy(self.0) }) { + write!(stderr(), "failed to call cagraIndexDestroy {:?}", e) + .expect("failed to write to stderr"); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use ndarray::s; + use ndarray_rand::rand_distr::Uniform; + use ndarray_rand::RandomExt; + + fn test_bfknn(metric: DistanceType) { + let res = Resources::new().unwrap(); + + // Create a new random dataset to index + let n_datapoints = 16; + let n_features = 8; + let dataset_host = + ndarray::Array::::random((n_datapoints, n_features), Uniform::new(0., 1.0)); + + let dataset = ManagedTensor::from(&dataset_host) + .to_device(&res) + .unwrap(); + + println!("dataset {:#?}", dataset_host); + + // build the brute force index + let index = + Index::build(&res, metric, None, dataset).expect("failed to create brute force index"); + + res.sync_stream().unwrap(); + + // use the first 4 points from the dataset as queries : will test that we get them back + // as their own nearest neighbor + let n_queries = 4; + let queries = dataset_host.slice(s![0..n_queries, ..]); + + let k = 4; + + println!("queries! {:#?}", queries); + let queries = ManagedTensor::from(&queries).to_device(&res).unwrap(); + let mut neighbors_host = ndarray::Array::::zeros((n_queries, k)); + let neighbors = ManagedTensor::from(&neighbors_host) + .to_device(&res) + .unwrap(); + + let mut distances_host = ndarray::Array::::zeros((n_queries, k)); + let distances = ManagedTensor::from(&distances_host) + .to_device(&res) + .unwrap(); + + index + .search(&res, &queries, &neighbors, &distances) + .unwrap(); + + // Copy back to host memory + distances.to_host(&res, &mut distances_host).unwrap(); + neighbors.to_host(&res, &mut neighbors_host).unwrap(); + res.sync_stream().unwrap(); + + println!("distances {:#?}", distances_host); + println!("neighbors {:#?}", neighbors_host); + + // nearest neighbors should be themselves, since queries are from the + // dataset + assert_eq!(neighbors_host[[0, 0]], 0); + assert_eq!(neighbors_host[[1, 0]], 1); + assert_eq!(neighbors_host[[2, 0]], 2); + assert_eq!(neighbors_host[[3, 0]], 3); + } + +/* + #[test] + fn test_cosine() { + test_bfknn(DistanceType::CosineExpanded); + } +*/ + + #[test] + fn test_l2() { + test_bfknn(DistanceType::L2Expanded); + } +} diff --git a/rust/cuvs/src/distance_type.rs b/rust/cuvs/src/distance_type.rs new file mode 100644 index 000000000..4ac3e9164 --- /dev/null +++ b/rust/cuvs/src/distance_type.rs @@ -0,0 +1,17 @@ +/* + * 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. + */ + +pub type DistanceType = ffi::DistanceType; diff --git a/rust/cuvs/src/dlpack.rs b/rust/cuvs/src/dlpack.rs index 238caec9d..b74484a5c 100644 --- a/rust/cuvs/src/dlpack.rs +++ b/rust/cuvs/src/dlpack.rs @@ -67,16 +67,17 @@ impl ManagedTensor { D: ndarray::Dimension, >( &self, - _res: &Resources, + res: &Resources, arr: &mut ndarray::ArrayBase, ) -> Result<()> { unsafe { let bytes = dl_tensor_bytes(&self.0.dl_tensor); - check_cuda(ffi::cudaMemcpy( + check_cuda(ffi::cudaMemcpyAsync( arr.as_mut_ptr() as *mut std::ffi::c_void, self.0.dl_tensor.data, bytes, ffi::cudaMemcpyKind_cudaMemcpyDefault, + res.get_cuda_stream()?, ))?; Ok(()) } @@ -169,6 +170,16 @@ impl IntoDtype for i32 { } } +impl IntoDtype for i64 { + fn ffi_dtype() -> ffi::DLDataType { + ffi::DLDataType { + code: ffi::DLDataTypeCode::kDLInt as _, + bits: 64, + lanes: 1, + } + } +} + impl IntoDtype for u32 { fn ffi_dtype() -> ffi::DLDataType { ffi::DLDataType { diff --git a/rust/cuvs/src/lib.rs b/rust/cuvs/src/lib.rs index 273f04723..de974dad7 100644 --- a/rust/cuvs/src/lib.rs +++ b/rust/cuvs/src/lib.rs @@ -18,8 +18,9 @@ //! //! This crate provides Rust bindings for cuVS, allowing you to run //! approximate nearest neighbors search on the GPU. - +pub mod brute_force; pub mod cagra; +pub mod distance_type; mod dlpack; mod error; mod resources;