From 09eb3ab16e907e635428f2657177be5fc31f5828 Mon Sep 17 00:00:00 2001 From: divyegala Date: Mon, 22 Jul 2024 15:08:15 -0700 Subject: [PATCH 01/18] c api and tests --- cpp/CMakeLists.txt | 1 + cpp/include/cuvs/neighbors/cagra.h | 33 ++++- cpp/include/cuvs/neighbors/hnsw.h | 202 ++++++++++++++++++++++++++++ cpp/include/cuvs/neighbors/hnsw.hpp | 2 + cpp/src/neighbors/cagra_c.cpp | 25 ++++ cpp/src/neighbors/hnsw_c.cpp | 166 +++++++++++++++++++++++ cpp/test/CMakeLists.txt | 4 + cpp/test/neighbors/ann_hnsw_c.cu | 134 ++++++++++++++++++ 8 files changed, 565 insertions(+), 2 deletions(-) create mode 100644 cpp/include/cuvs/neighbors/hnsw.h create mode 100644 cpp/src/neighbors/hnsw_c.cpp create mode 100644 cpp/test/neighbors/ann_hnsw_c.cu diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 5060f4591..412504c93 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -582,6 +582,7 @@ if(BUILD_C_LIBRARY) src/neighbors/ivf_flat_c.cpp src/neighbors/ivf_pq_c.cpp src/neighbors/cagra_c.cpp + src/neighbors/hnsw_c.cpp src/neighbors/refine/refine_c.cpp src/distance/pairwise_distance_c.cpp ) diff --git a/cpp/include/cuvs/neighbors/cagra.h b/cpp/include/cuvs/neighbors/cagra.h index 87541f7f0..cd492d1f4 100644 --- a/cpp/include/cuvs/neighbors/cagra.h +++ b/cpp/include/cuvs/neighbors/cagra.h @@ -337,7 +337,10 @@ cuvsError_t cuvsCagraBuild(cuvsResources_t res, * 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 == * queries.dl_tensor.dtype.code` Types for input are: - * 1. `queries`: `kDLDataType.code == kDLFloat` and `kDLDataType.bits = 32` + * 1. `queries`: + *` a. kDLDataType.code == kDLFloat` and `kDLDataType.bits = 32` + * b. `kDLDataType.code == kDLInt` and `kDLDataType.bits = 8` + * c. `kDLDataType.code == kDLUInt` and `kDLDataType.bits = 8` * 2. `neighbors`: `kDLDataType.code == kDLUInt` and `kDLDataType.bits = 32` * 3. `distances`: `kDLDataType.code == kDLFloat` and `kDLDataType.bits = 32` * @@ -394,7 +397,7 @@ cuvsError_t cuvsCagraSearch(cuvsResources_t res, * * Experimental, both the API and the serialization format are subject to change. * - * @code{.cpp} + * @code{.c} * #include * * // Create cuvsResources_t @@ -416,6 +419,32 @@ cuvsError_t cuvsCagraSerialize(cuvsResources_t res, cuvsCagraIndex_t index, bool include_dataset); +/** + * Save the CAGRA index to file in hnswlib format. + * + * Experimental, both the API and the serialization format are subject to change. + * + * @code{.c} + * #include + * #include + * + * // Create cuvsResources_t + * cuvsResources_t res; + * cuvsError_t res_create_status = cuvsResourcesCreate(&res); + * + * // create an index with `cuvsCagraBuild` + * cuvsCagraSerializeHnswlib(res, "/path/to/index", index); + * @endcode + * + * @param[in] res cuvsResources_t opaque C handle + * @param[in] filename the file name for saving the index + * @param[in] index CAGRA index + * + */ +cuvsError_t cuvsCagraSerializeToHnswlib(cuvsResources_t res, + const char* filename, + cuvsCagraIndex_t index); + /** * Load index from file. * diff --git a/cpp/include/cuvs/neighbors/hnsw.h b/cpp/include/cuvs/neighbors/hnsw.h new file mode 100644 index 000000000..e85d0085a --- /dev/null +++ b/cpp/include/cuvs/neighbors/hnsw.h @@ -0,0 +1,202 @@ +/* + * 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 +#include +#include +#include +#include + +#ifdef __cplusplus +extern "C" { +#endif + +/** + * @defgroup hnsw_c_search_params C API for hnswlib wrapper search params + * @{ + */ + +struct cuvsHnswSearchParams { + int32_t ef; + int32_t numThreads; +}; + +typedef struct cuvsHnswSearchParams* cuvsHnswSearchParams_t; + +/** + * @brief Allocate HNSW search params, and populate with default values + * + * @param[in] params cuvsHnswSearchParams_t to allocate + * @return cuvsError_t + */ +cuvsError_t cuvsHnswSearchParamsCreate(cuvsHnswSearchParams_t* params); + +/** + * @brief De-allocate HNSW search params + * + * @param[in] params cuvsHnswSearchParams_t to de-allocate + * @return cuvsError_t + */ +cuvsError_t cuvsHnswSearchParamsDestroy(cuvsHnswSearchParams_t params); + +/** + * @} + */ + +/** + * @defgroup hnsw_c_index C API for hnswlib wrapper index + * @{ + */ + +/** + * @brief Struct to hold address of cuvs::neighbors::Hnsw::index and its active trained dtype + * + */ +typedef struct { + uintptr_t addr; + DLDataType dtype; + +} cuvsHnswIndex; + +typedef cuvsHnswIndex* cuvsHnswIndex_t; + +/** + * @brief Allocate HNSW index + * + * @param[in] index cuvsHnswIndex_t to allocate + * @return HnswError_t + */ +cuvsError_t cuvsHnswIndexCreate(cuvsHnswIndex_t* index); + +/** + * @brief De-allocate HNSW index + * + * @param[in] index cuvsHnswIndex_t to de-allocate + */ +cuvsError_t cuvsHnswIndexDestroy(cuvsHnswIndex_t index); + +/** + * @} + */ + +/** + * @defgroup hnsw_c_index_search C API for CUDA ANN Graph-based nearest neighbor search + * @{ + */ +/** + * @brief Search a HNSW index with a `DLManagedTensor` which has underlying + * `DLDeviceType` equal to `kDLCPU`, `kDLCUDAHost`, or `kDLCUDAManaged`. + * It is also important to note that the HNSW Index must have been built + * with the same type of `queries`, such that `index.dtype.code == + * queries.dl_tensor.dtype.code` + * Supported types for input are: + * 1. `queries`: `kDLDataType.code == kDLFloat` or `kDLDataType.code == kDLInt` and + * `kDLDataType.bits = 32` + * 2. `neighbors`: `kDLDataType.code == kDLUInt` and `kDLDataType.bits = 64` + * 3. `distances`: `kDLDataType.code == kDLFloat` and `kDLDataType.bits = 32` + * + * @code {.c} + * #include + * #include + * + * // 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; + * + * // Create default search params + * cuvsHnswSearchParams_t params; + * cuvsError_t params_create_status = cuvsHnswSearchParamsCreate(¶ms); + * + * // Search the `index` built using `cuvsHnswBuild` + * cuvsError_t search_status = cuvsHnswSearch(res, params, index, &queries, &neighbors, + * &distances); + * + * // de-allocate `params` and `res` + * cuvsError_t params_destroy_status = cuvsHnswSearchParamsDestroy(params); + * cuvsError_t res_destroy_status = cuvsResourcesDestroy(res); + * @endcode + * + * @param[in] res cuvsResources_t opaque C handle + * @param[in] params cuvsHnswSearchParams_t used to search Hnsw index + * @param[in] index cuvsHnswIndex which has been returned by `cuvsHnswBuild` + * @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 cuvsHnswSearch(cuvsResources_t res, + cuvsHnswSearchParams_t params, + cuvsHnswIndex_t index, + DLManagedTensor* queries, + DLManagedTensor* neighbors, + DLManagedTensor* distances); + +/** + * @} + */ + +/** + * @defgroup hnsw_c_serialize HNSW C-API serialize functions + * @{ + */ + +/** + * Load hnswlib index from file which was serialized from a HNSW index. + * + * Experimental, both the API and the serialization format are subject to change. + * + * @code{.c} + * #include + * #include + * #include + * + * // Create cuvsResources_t + * cuvsResources_t res; + * cuvsError_t res_create_status = cuvsResourcesCreate(&res); + * + * // create an index with `cuvsCagraBuild` + * cuvsCagraSerializeHnswlib(res, "/path/to/index", index); + * + * // Load the serialized CAGRA index from file as an hnswlib index + * // The index should have the same dtype as the one used to build CAGRA the index + * cuvsHnswIndex_t hnsw_index; + * cuvsHnswIndexCreate(&hnsw_index); + * hnsw_index->dtype = index->dtype; + * cuvsCagraDeserialize(res, "/path/to/index", hnsw_index); + * @endcode + * + * @param[in] res cuvsResources_t opaque C handle + * @param[in] filename the name of the file that stores the index + * @param[out] index HNSW index loaded disk + */ +cuvsError_t cuvsHnswDeserialize(cuvsResources_t res, + const char* filename, + int dim, + cuvsDistanceType metric, + cuvsHnswIndex_t index); +/** + * @} + */ + +#ifdef __cplusplus +} +#endif diff --git a/cpp/include/cuvs/neighbors/hnsw.hpp b/cpp/include/cuvs/neighbors/hnsw.hpp index 86f321564..276e1a844 100644 --- a/cpp/include/cuvs/neighbors/hnsw.hpp +++ b/cpp/include/cuvs/neighbors/hnsw.hpp @@ -58,6 +58,8 @@ struct index : cuvs::neighbors::index { */ index(int dim, cuvs::distance::DistanceType metric) : dim_{dim}, metric_{metric} {} + virtual ~index() {} + /** @brief Get underlying index */ diff --git a/cpp/src/neighbors/cagra_c.cpp b/cpp/src/neighbors/cagra_c.cpp index 868b3dec0..164448f2c 100644 --- a/cpp/src/neighbors/cagra_c.cpp +++ b/cpp/src/neighbors/cagra_c.cpp @@ -130,6 +130,14 @@ void _serialize(cuvsResources_t res, cuvs::neighbors::cagra::serialize(*res_ptr, std::string(filename), *index_ptr, include_dataset); } +template +void _serialize_to_hnswlib(cuvsResources_t res, const char* filename, cuvsCagraIndex_t index) +{ + auto res_ptr = reinterpret_cast(res); + auto index_ptr = reinterpret_cast*>(index->addr); + cuvs::neighbors::cagra::serialize_to_hnswlib(*res_ptr, std::string(filename), *index_ptr); +} + template void* _deserialize(cuvsResources_t res, const char* filename) { @@ -326,3 +334,20 @@ extern "C" cuvsError_t cuvsCagraSerialize(cuvsResources_t res, } }); } + +extern "C" cuvsError_t cuvsCagraSerializeToHnswlib(cuvsResources_t res, + const char* filename, + cuvsCagraIndex_t index) +{ + return cuvs::core::translate_exceptions([=] { + if (index->dtype.code == kDLFloat && index->dtype.bits == 32) { + _serialize_to_hnswlib(res, filename, index); + } else if (index->dtype.code == kDLInt && index->dtype.bits == 8) { + _serialize_to_hnswlib(res, filename, index); + } else if (index->dtype.code == kDLUInt && index->dtype.bits == 8) { + _serialize_to_hnswlib(res, filename, index); + } else { + RAFT_FAIL("Unsupported index dtype: %d and bits: %d", index->dtype.code, index->dtype.bits); + } + }); +} diff --git a/cpp/src/neighbors/hnsw_c.cpp b/cpp/src/neighbors/hnsw_c.cpp new file mode 100644 index 000000000..ab5268a6d --- /dev/null +++ b/cpp/src/neighbors/hnsw_c.cpp @@ -0,0 +1,166 @@ + +/* + * 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. + */ + +#include "cuvs/distance/distance.h" +#include +#include + +#include +#include +#include +#include + +#include +#include +#include +#include +#include + +namespace { +template +void _search(cuvsResources_t res, + cuvsHnswSearchParams params, + cuvsHnswIndex index, + DLManagedTensor* queries_tensor, + DLManagedTensor* neighbors_tensor, + DLManagedTensor* distances_tensor) +{ + auto res_ptr = reinterpret_cast(res); + auto index_ptr = reinterpret_cast*>(index.addr); + + auto search_params = cuvs::neighbors::hnsw::search_params(); + search_params.ef = params.ef; + search_params.num_threads = params.numThreads; + + using queries_mdspan_type = raft::host_matrix_view; + using neighbors_mdspan_type = raft::host_matrix_view; + using distances_mdspan_type = raft::host_matrix_view; + auto queries_mds = cuvs::core::from_dlpack(queries_tensor); + auto neighbors_mds = cuvs::core::from_dlpack(neighbors_tensor); + auto distances_mds = cuvs::core::from_dlpack(distances_tensor); + cuvs::neighbors::hnsw::search( + *res_ptr, search_params, *index_ptr, queries_mds, neighbors_mds, distances_mds); +} + +template +void* _deserialize(cuvsResources_t res, const char* filename, int dim, cuvsDistanceType metric) +{ + auto res_ptr = reinterpret_cast(res); + cuvs::neighbors::hnsw::index* index = nullptr; + cuvs::neighbors::hnsw::deserialize(*res_ptr, std::string(filename), dim, metric, &index); + return index; +} +} // namespace + +extern "C" cuvsError_t cuvsHnswSearchParamsCreate(cuvsHnswSearchParams_t* params) +{ + return cuvs::core::translate_exceptions( + [=] { *params = new cuvsHnswSearchParams{.ef = 200, .numThreads = 0}; }); +} + +extern "C" cuvsError_t cuvsHnswSearchParamsDestroy(cuvsHnswSearchParams_t params) +{ + return cuvs::core::translate_exceptions([=] { delete params; }); +} + +extern "C" cuvsError_t cuvsHnswIndexCreate(cuvsHnswIndex_t* index) +{ + return cuvs::core::translate_exceptions([=] { *index = new cuvsHnswIndex{}; }); +} + +extern "C" cuvsError_t cuvsHnswIndexDestroy(cuvsHnswIndex_t index_c_ptr) +{ + return cuvs::core::translate_exceptions([=] { + auto index = *index_c_ptr; + + if (index.dtype.code == kDLFloat) { + auto index_ptr = reinterpret_cast*>(index.addr); + delete index_ptr; + } else if (index.dtype.code == kDLInt) { + auto index_ptr = reinterpret_cast*>(index.addr); + delete index_ptr; + } else if (index.dtype.code == kDLUInt) { + auto index_ptr = reinterpret_cast*>(index.addr); + delete index_ptr; + } + delete index_c_ptr; + }); +} + +extern "C" cuvsError_t cuvsHnswSearch(cuvsResources_t res, + cuvsHnswSearchParams_t params, + cuvsHnswIndex_t index_c_ptr, + DLManagedTensor* queries_tensor, + DLManagedTensor* neighbors_tensor, + DLManagedTensor* distances_tensor) +{ + return cuvs::core::translate_exceptions([=] { + auto queries = queries_tensor->dl_tensor; + auto neighbors = neighbors_tensor->dl_tensor; + auto distances = distances_tensor->dl_tensor; + + RAFT_EXPECTS(cuvs::core::is_dlpack_host_compatible(queries), + "queries should have host compatible memory"); + RAFT_EXPECTS(cuvs::core::is_dlpack_host_compatible(neighbors), + "neighbors should have host compatible memory"); + RAFT_EXPECTS(cuvs::core::is_dlpack_host_compatible(distances), + "distances should have host compatible memory"); + + RAFT_EXPECTS(neighbors.dtype.code == kDLUInt && neighbors.dtype.bits == 64, + "neighbors should be of type uint64_t"); + RAFT_EXPECTS(distances.dtype.code == kDLFloat && distances.dtype.bits == 32, + "distances should be of type float32"); + + auto index = *index_c_ptr; + RAFT_EXPECTS(queries.dtype.code == index.dtype.code, "type mismatch between index and queries"); + RAFT_EXPECTS(queries.dtype.bits == 32, "number of bits in queries dtype should be 32"); + + if (index.dtype.code == kDLFloat) { + _search( + res, *params, index, queries_tensor, neighbors_tensor, distances_tensor); + } else if (index.dtype.code == kDLUInt) { + _search( + res, *params, index, queries_tensor, neighbors_tensor, distances_tensor); + } else if (index.dtype.code == kDLInt) { + _search(res, *params, index, queries_tensor, neighbors_tensor, distances_tensor); + } else { + RAFT_FAIL("Unsupported index dtype: %d and bits: %d", queries.dtype.code, queries.dtype.bits); + } + }); +} + +extern "C" cuvsError_t cuvsHnswDeserialize(cuvsResources_t res, + const char* filename, + int dim, + cuvsDistanceType metric, + cuvsHnswIndex_t index) +{ + return cuvs::core::translate_exceptions([=] { + if (index->dtype.code == kDLFloat && index->dtype.bits == 32) { + index->addr = reinterpret_cast(_deserialize(res, filename, dim, metric)); + index->dtype.code = kDLFloat; + } else if (index->dtype.code == kDLUInt && index->dtype.bits == 8) { + index->addr = reinterpret_cast(_deserialize(res, filename, dim, metric)); + index->dtype.code = kDLInt; + } else if (index->dtype.code == kDLInt && index->dtype.bits == 8) { + index->addr = reinterpret_cast(_deserialize(res, filename, dim, metric)); + index->dtype.code = kDLUInt; + } else { + RAFT_FAIL("Unsupported dtype in file %s", filename); + } + }); +} diff --git a/cpp/test/CMakeLists.txt b/cpp/test/CMakeLists.txt index 7921fffd3..f3cc7b22f 100644 --- a/cpp/test/CMakeLists.txt +++ b/cpp/test/CMakeLists.txt @@ -205,6 +205,10 @@ if(BUILD_C_TESTS) ) ConfigureTest(NAME CAGRA_C_TEST PATH test/neighbors/ann_cagra_c.cu C_LIB) + + if (BUILD_CAGRA_HNSWLIB) + ConfigureTest(NAME HNSW_C_TEST PATH test/neighbors/ann_hnsw_c.cu C_LIB) + endif() endif() # ################################################################################################## diff --git a/cpp/test/neighbors/ann_hnsw_c.cu b/cpp/test/neighbors/ann_hnsw_c.cu new file mode 100644 index 000000000..345653783 --- /dev/null +++ b/cpp/test/neighbors/ann_hnsw_c.cu @@ -0,0 +1,134 @@ +/* + * 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. + */ + +#include "../test_utils.cuh" +#include "cuvs/distance/distance.h" +#include +#include + +#include +#include +#include + +#include +#include +#include +#include + +float dataset[4][2] = {{0.74021935, 0.9209938}, + {0.03902049, 0.9689629}, + {0.92514056, 0.4463501}, + {0.6673192, 0.10993068}}; +float queries[4][2] = {{0.48216683, 0.0428398}, + {0.5084142, 0.6545497}, + {0.51260436, 0.2643005}, + {0.05198065, 0.5789965}}; + +std::vector neighbors_exp = {3, 0, 3, 1}; +std::vector distances_exp = {0.03878258, 0.12472608, 0.04776672, 0.15224178}; + +TEST(CagraC, BuildSearch) +{ + // create cuvsResources_t + cuvsResources_t res; + cuvsResourcesCreate(&res); + + // create dataset DLTensor + DLManagedTensor dataset_tensor; + dataset_tensor.dl_tensor.data = dataset; + dataset_tensor.dl_tensor.device.device_type = kDLCPU; + dataset_tensor.dl_tensor.ndim = 2; + dataset_tensor.dl_tensor.dtype.code = kDLFloat; + dataset_tensor.dl_tensor.dtype.bits = 32; + dataset_tensor.dl_tensor.dtype.lanes = 1; + int64_t dataset_shape[2] = {4, 2}; + dataset_tensor.dl_tensor.shape = dataset_shape; + dataset_tensor.dl_tensor.strides = nullptr; + + // create index + cuvsCagraIndex_t index; + cuvsCagraIndexCreate(&index); + + // build index + cuvsCagraIndexParams_t build_params; + cuvsCagraIndexParamsCreate(&build_params); + cuvsCagraBuild(res, build_params, &dataset_tensor, index); + cuvsCagraSerializeToHnswlib(res, "/tmp/cagra_hnswlib.index", index); + + DLManagedTensor queries_tensor; + queries_tensor.dl_tensor.data = queries; + queries_tensor.dl_tensor.device.device_type = kDLCPU; + queries_tensor.dl_tensor.ndim = 2; + queries_tensor.dl_tensor.dtype.code = kDLFloat; + queries_tensor.dl_tensor.dtype.bits = 32; + queries_tensor.dl_tensor.dtype.lanes = 1; + int64_t queries_shape[2] = {4, 2}; + queries_tensor.dl_tensor.shape = queries_shape; + queries_tensor.dl_tensor.strides = nullptr; + + // create neighbors DLTensor + std::vector neighbors(4); + + DLManagedTensor neighbors_tensor; + neighbors_tensor.dl_tensor.data = neighbors.data(); + neighbors_tensor.dl_tensor.device.device_type = kDLCPU; + neighbors_tensor.dl_tensor.ndim = 2; + neighbors_tensor.dl_tensor.dtype.code = kDLUInt; + neighbors_tensor.dl_tensor.dtype.bits = 64; + neighbors_tensor.dl_tensor.dtype.lanes = 1; + int64_t neighbors_shape[2] = {4, 1}; + neighbors_tensor.dl_tensor.shape = neighbors_shape; + neighbors_tensor.dl_tensor.strides = nullptr; + + // create distances DLTensor + std::vector distances(4); + + DLManagedTensor distances_tensor; + distances_tensor.dl_tensor.data = distances.data(); + distances_tensor.dl_tensor.device.device_type = kDLCPU; + distances_tensor.dl_tensor.ndim = 2; + distances_tensor.dl_tensor.dtype.code = kDLFloat; + distances_tensor.dl_tensor.dtype.bits = 32; + distances_tensor.dl_tensor.dtype.lanes = 1; + int64_t distances_shape[2] = {4, 1}; + distances_tensor.dl_tensor.shape = distances_shape; + distances_tensor.dl_tensor.strides = nullptr; + + // create hnsw index + cuvsHnswIndex_t hnsw_index; + cuvsHnswIndexCreate(&hnsw_index); + hnsw_index->dtype = index->dtype; + cuvsHnswDeserialize(res, "/tmp/cagra_hnswlib.index", 2, L2Expanded, hnsw_index); + + // search index + cuvsHnswSearchParams_t search_params; + cuvsHnswSearchParamsCreate(&search_params); + cuvsHnswSearch( + res, search_params, hnsw_index, &queries_tensor, &neighbors_tensor, &distances_tensor); + + // verify output + ASSERT_TRUE(cuvs::hostVecMatch(neighbors_exp, neighbors, cuvs::Compare())); + ASSERT_TRUE(cuvs::hostVecMatch(distances_exp, distances, cuvs::CompareApprox(0.001f))); + + // delete device memory + + // de-allocate index and res + cuvsCagraIndexParamsDestroy(build_params); + cuvsCagraIndexDestroy(index); + cuvsHnswSearchParamsDestroy(search_params); + cuvsHnswIndexDestroy(hnsw_index); + cuvsResourcesDestroy(res); +} From 8bc035e7a9fa56f20424e80479a7c1d4cf5f1642 Mon Sep 17 00:00:00 2001 From: divyegala Date: Mon, 22 Jul 2024 15:14:24 -0700 Subject: [PATCH 02/18] remove unneeded comment --- cpp/test/neighbors/ann_hnsw_c.cu | 2 -- 1 file changed, 2 deletions(-) diff --git a/cpp/test/neighbors/ann_hnsw_c.cu b/cpp/test/neighbors/ann_hnsw_c.cu index 345653783..7c54cd46c 100644 --- a/cpp/test/neighbors/ann_hnsw_c.cu +++ b/cpp/test/neighbors/ann_hnsw_c.cu @@ -123,8 +123,6 @@ TEST(CagraC, BuildSearch) ASSERT_TRUE(cuvs::hostVecMatch(neighbors_exp, neighbors, cuvs::Compare())); ASSERT_TRUE(cuvs::hostVecMatch(distances_exp, distances, cuvs::CompareApprox(0.001f))); - // delete device memory - // de-allocate index and res cuvsCagraIndexParamsDestroy(build_params); cuvsCagraIndexDestroy(index); From f8327f5c5e9874d3de64317fafcd5fc8bdbbdae4 Mon Sep 17 00:00:00 2001 From: Divye Gala Date: Mon, 22 Jul 2024 21:59:50 -0400 Subject: [PATCH 03/18] Update ann_hnsw_c.cu --- cpp/test/neighbors/ann_hnsw_c.cu | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/test/neighbors/ann_hnsw_c.cu b/cpp/test/neighbors/ann_hnsw_c.cu index 7c54cd46c..df8cd5f47 100644 --- a/cpp/test/neighbors/ann_hnsw_c.cu +++ b/cpp/test/neighbors/ann_hnsw_c.cu @@ -15,7 +15,7 @@ */ #include "../test_utils.cuh" -#include "cuvs/distance/distance.h" +#include #include #include From 6c0df11877c81993762931fc90049e640b3dce26 Mon Sep 17 00:00:00 2001 From: Divye Gala Date: Mon, 22 Jul 2024 22:03:44 -0400 Subject: [PATCH 04/18] Update ann_hnsw_c.cu --- cpp/test/neighbors/ann_hnsw_c.cu | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/test/neighbors/ann_hnsw_c.cu b/cpp/test/neighbors/ann_hnsw_c.cu index df8cd5f47..c4947cc10 100644 --- a/cpp/test/neighbors/ann_hnsw_c.cu +++ b/cpp/test/neighbors/ann_hnsw_c.cu @@ -15,8 +15,8 @@ */ #include "../test_utils.cuh" -#include #include +#include #include #include From 8860a099952cf81524f1b455bc8e312a008e5aa7 Mon Sep 17 00:00:00 2001 From: divyegala Date: Tue, 23 Jul 2024 11:11:51 -0700 Subject: [PATCH 05/18] rename test --- cpp/test/neighbors/ann_hnsw_c.cu | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/test/neighbors/ann_hnsw_c.cu b/cpp/test/neighbors/ann_hnsw_c.cu index c4947cc10..8ccef24fb 100644 --- a/cpp/test/neighbors/ann_hnsw_c.cu +++ b/cpp/test/neighbors/ann_hnsw_c.cu @@ -40,7 +40,7 @@ float queries[4][2] = {{0.48216683, 0.0428398}, std::vector neighbors_exp = {3, 0, 3, 1}; std::vector distances_exp = {0.03878258, 0.12472608, 0.04776672, 0.15224178}; -TEST(CagraC, BuildSearch) +TEST(CagraHnswC, BuildSearch) { // create cuvsResources_t cuvsResources_t res; From 081eba5248f124edd22eb0ba3bd99a5d91f849ab Mon Sep 17 00:00:00 2001 From: divyegala Date: Tue, 23 Jul 2024 18:11:31 -0700 Subject: [PATCH 06/18] passing python tests --- cpp/include/cuvs/neighbors/cagra.h | 11 +- cpp/include/cuvs/neighbors/hnsw.h | 8 +- cpp/include/cuvs/neighbors/hnsw.hpp | 4 +- cpp/src/neighbors/cagra_c.cpp | 8 + .../detail/cagra/cagra_serialize.cuh | 19 +- cpp/src/neighbors/detail/hnsw.hpp | 13 +- cpp/src/neighbors/hnsw.cpp | 24 +- cpp/src/neighbors/hnsw_c.cpp | 22 +- python/cuvs/cuvs/neighbors/CMakeLists.txt | 1 + python/cuvs/cuvs/neighbors/cagra/cagra.pxd | 17 + python/cuvs/cuvs/neighbors/cagra/cagra.pyx | 17 +- .../cuvs/cuvs/neighbors/hnsw/CMakeLists.txt | 24 ++ python/cuvs/cuvs/neighbors/hnsw/__init__.pxd | 0 python/cuvs/cuvs/neighbors/hnsw/__init__.py | 25 ++ python/cuvs/cuvs/neighbors/hnsw/hnsw.pxd | 53 +++ python/cuvs/cuvs/neighbors/hnsw/hnsw.pyx | 364 ++++++++++++++++++ python/cuvs/cuvs/test/test_hnsw.py | 97 +++++ 17 files changed, 648 insertions(+), 59 deletions(-) create mode 100644 python/cuvs/cuvs/neighbors/hnsw/CMakeLists.txt create mode 100644 python/cuvs/cuvs/neighbors/hnsw/__init__.pxd create mode 100644 python/cuvs/cuvs/neighbors/hnsw/__init__.py create mode 100644 python/cuvs/cuvs/neighbors/hnsw/hnsw.pxd create mode 100644 python/cuvs/cuvs/neighbors/hnsw/hnsw.pyx create mode 100644 python/cuvs/cuvs/test/test_hnsw.py diff --git a/cpp/include/cuvs/neighbors/cagra.h b/cpp/include/cuvs/neighbors/cagra.h index cd492d1f4..00977983c 100644 --- a/cpp/include/cuvs/neighbors/cagra.h +++ b/cpp/include/cuvs/neighbors/cagra.h @@ -267,6 +267,15 @@ cuvsError_t cuvsCagraIndexCreate(cuvsCagraIndex_t* index); */ cuvsError_t cuvsCagraIndexDestroy(cuvsCagraIndex_t index); +/** + * @brief Get dimension of the CAGRA index + * + * @param[in] CAGRA index + * @param[out] dim + * @return cuvsError_t + */ +cuvsError_t cuvsCagraIndexDim(cuvsCagraIndex_t index, int* dim); + /** * @} */ @@ -338,7 +347,7 @@ cuvsError_t cuvsCagraBuild(cuvsResources_t res, * with the same type of `queries`, such that `index.dtype.code == * queries.dl_tensor.dtype.code` Types for input are: * 1. `queries`: - *` a. kDLDataType.code == kDLFloat` and `kDLDataType.bits = 32` + * a. kDLDataType.code == kDLFloat` and `kDLDataType.bits = 32` * b. `kDLDataType.code == kDLInt` and `kDLDataType.bits = 8` * c. `kDLDataType.code == kDLUInt` and `kDLDataType.bits = 8` * 2. `neighbors`: `kDLDataType.code == kDLUInt` and `kDLDataType.bits = 32` diff --git a/cpp/include/cuvs/neighbors/hnsw.h b/cpp/include/cuvs/neighbors/hnsw.h index e85d0085a..55d44e36c 100644 --- a/cpp/include/cuvs/neighbors/hnsw.h +++ b/cpp/include/cuvs/neighbors/hnsw.h @@ -33,7 +33,7 @@ extern "C" { struct cuvsHnswSearchParams { int32_t ef; - int32_t numThreads; + int32_t num_threads; }; typedef struct cuvsHnswSearchParams* cuvsHnswSearchParams_t; @@ -105,8 +105,10 @@ cuvsError_t cuvsHnswIndexDestroy(cuvsHnswIndex_t index); * with the same type of `queries`, such that `index.dtype.code == * queries.dl_tensor.dtype.code` * Supported types for input are: - * 1. `queries`: `kDLDataType.code == kDLFloat` or `kDLDataType.code == kDLInt` and - * `kDLDataType.bits = 32` + * 1. `queries`: + * a. kDLDataType.code == kDLFloat` and `kDLDataType.bits = 32` + * b. `kDLDataType.code == kDLInt` and `kDLDataType.bits = 8` + * c. `kDLDataType.code == kDLUInt` and `kDLDataType.bits = 8` * 2. `neighbors`: `kDLDataType.code == kDLUInt` and `kDLDataType.bits = 64` * 3. `distances`: `kDLDataType.code == kDLFloat` and `kDLDataType.bits = 32` * diff --git a/cpp/include/cuvs/neighbors/hnsw.hpp b/cpp/include/cuvs/neighbors/hnsw.hpp index 276e1a844..d15b502e3 100644 --- a/cpp/include/cuvs/neighbors/hnsw.hpp +++ b/cpp/include/cuvs/neighbors/hnsw.hpp @@ -232,7 +232,7 @@ void search(raft::resources const& res, void search(raft::resources const& res, const search_params& params, const index& idx, - raft::host_matrix_view queries, + raft::host_matrix_view queries, raft::host_matrix_view neighbors, raft::host_matrix_view distances); @@ -273,7 +273,7 @@ void search(raft::resources const& res, void search(raft::resources const& res, const search_params& params, const index& idx, - raft::host_matrix_view queries, + raft::host_matrix_view queries, raft::host_matrix_view neighbors, raft::host_matrix_view distances); diff --git a/cpp/src/neighbors/cagra_c.cpp b/cpp/src/neighbors/cagra_c.cpp index 164448f2c..229fa48f1 100644 --- a/cpp/src/neighbors/cagra_c.cpp +++ b/cpp/src/neighbors/cagra_c.cpp @@ -176,6 +176,14 @@ extern "C" cuvsError_t cuvsCagraIndexDestroy(cuvsCagraIndex_t index_c_ptr) }); } +extern "C" cuvsError_t cuvsCagraIndexDim(cuvsCagraIndex_t index, int* dim) +{ + return cuvs::core::translate_exceptions([=] { + auto index_ptr = reinterpret_cast*>(index->addr); + *dim = index_ptr->dim(); + }); +} + extern "C" cuvsError_t cuvsCagraBuild(cuvsResources_t res, cuvsCagraIndexParams_t params, DLManagedTensor* dataset_tensor, diff --git a/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh b/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh index 24cc2a22f..bb71d915b 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh @@ -119,9 +119,9 @@ void serialize_to_hnswlib(raft::resources const& res, os.write(reinterpret_cast(&curr_element_count), sizeof(std::size_t)); // Example:M: 16, dim = 128, data_t = float, index_t = uint32_t, list_size_type = uint32_t, // labeltype: size_t size_data_per_element_ = M * 2 * sizeof(index_t) + sizeof(list_size_type) + - // dim * 4 + sizeof(labeltype) - auto size_data_per_element = - static_cast(index_.graph_degree() * sizeof(IdxT) + 4 + index_.dim() * 4 + 8); + // dim * sizeof(T) + sizeof(labeltype) + auto size_data_per_element = static_cast(index_.graph_degree() * sizeof(IdxT) + 4 + + index_.dim() * sizeof(T) + 8); os.write(reinterpret_cast(&size_data_per_element), sizeof(std::size_t)); // label_offset std::size_t label_offset = size_data_per_element - 8; @@ -184,16 +184,9 @@ void serialize_to_hnswlib(raft::resources const& res, } auto data_row = host_dataset.data_handle() + (index_.dim() * i); - if constexpr (std::is_same_v) { - for (std::size_t j = 0; j < index_.dim(); ++j) { - auto data_elem = static_cast(host_dataset(i, j)); - os.write(reinterpret_cast(&data_elem), sizeof(float)); - } - } else if constexpr (std::is_same_v or std::is_same_v) { - for (std::size_t j = 0; j < index_.dim(); ++j) { - auto data_elem = static_cast(host_dataset(i, j)); - os.write(reinterpret_cast(&data_elem), sizeof(int)); - } + for (std::size_t j = 0; j < index_.dim(); ++j) { + auto data_elem = static_cast(host_dataset(i, j)); + os.write(reinterpret_cast(&data_elem), sizeof(T)); } os.write(reinterpret_cast(&i), sizeof(std::size_t)); diff --git a/cpp/src/neighbors/detail/hnsw.hpp b/cpp/src/neighbors/detail/hnsw.hpp index 0d1ae4ec9..ce1e03264 100644 --- a/cpp/src/neighbors/detail/hnsw.hpp +++ b/cpp/src/neighbors/detail/hnsw.hpp @@ -110,9 +110,9 @@ std::unique_ptr> from_cagra(raft::resources const& res, return std::unique_ptr>(hnsw_index); } -template -void get_search_knn_results(hnswlib::HierarchicalNSW const* idx, - const QueriesT* query, +template +void get_search_knn_results(hnswlib::HierarchicalNSW::type> const* idx, + const T* query, int k, uint64_t* indices, float* distances) @@ -127,11 +127,11 @@ void get_search_knn_results(hnswlib::HierarchicalNSW const* idx, } } -template +template void search(raft::resources const& res, const search_params& params, const index& idx, - raft::host_matrix_view queries, + raft::host_matrix_view queries, raft::host_matrix_view neighbors, raft::host_matrix_view distances) { @@ -146,7 +146,8 @@ void search(raft::resources const& res, idx.set_ef(params.ef); auto const* hnswlib_index = - reinterpret_cast const*>(idx.get_index()); + reinterpret_cast::type> const*>( + idx.get_index()); // when num_threads == 0, automatically maximize parallelism if (params.num_threads) { diff --git a/cpp/src/neighbors/hnsw.cpp b/cpp/src/neighbors/hnsw.cpp index 36cbb16c9..e6f3fbcc7 100644 --- a/cpp/src/neighbors/hnsw.cpp +++ b/cpp/src/neighbors/hnsw.cpp @@ -34,20 +34,20 @@ CUVS_INST_HNSW_FROM_CAGRA(int8_t); #undef CUVS_INST_HNSW_FROM_CAGRA -#define CUVS_INST_HNSW_SEARCH(T, QueriesT) \ - void search(raft::resources const& res, \ - const search_params& params, \ - const index& idx, \ - raft::host_matrix_view queries, \ - raft::host_matrix_view neighbors, \ - raft::host_matrix_view distances) \ - { \ - detail::search(res, params, idx, queries, neighbors, distances); \ +#define CUVS_INST_HNSW_SEARCH(T) \ + void search(raft::resources const& res, \ + const search_params& params, \ + const index& idx, \ + raft::host_matrix_view queries, \ + raft::host_matrix_view neighbors, \ + raft::host_matrix_view distances) \ + { \ + detail::search(res, params, idx, queries, neighbors, distances); \ } -CUVS_INST_HNSW_SEARCH(float, float); -CUVS_INST_HNSW_SEARCH(uint8_t, int); -CUVS_INST_HNSW_SEARCH(int8_t, int); +CUVS_INST_HNSW_SEARCH(float); +CUVS_INST_HNSW_SEARCH(uint8_t); +CUVS_INST_HNSW_SEARCH(int8_t); #undef CUVS_INST_HNSW_SEARCH diff --git a/cpp/src/neighbors/hnsw_c.cpp b/cpp/src/neighbors/hnsw_c.cpp index ab5268a6d..ea2a518c9 100644 --- a/cpp/src/neighbors/hnsw_c.cpp +++ b/cpp/src/neighbors/hnsw_c.cpp @@ -31,7 +31,7 @@ #include namespace { -template +template void _search(cuvsResources_t res, cuvsHnswSearchParams params, cuvsHnswIndex index, @@ -44,9 +44,9 @@ void _search(cuvsResources_t res, auto search_params = cuvs::neighbors::hnsw::search_params(); search_params.ef = params.ef; - search_params.num_threads = params.numThreads; + search_params.num_threads = params.num_threads; - using queries_mdspan_type = raft::host_matrix_view; + using queries_mdspan_type = raft::host_matrix_view; using neighbors_mdspan_type = raft::host_matrix_view; using distances_mdspan_type = raft::host_matrix_view; auto queries_mds = cuvs::core::from_dlpack(queries_tensor); @@ -69,7 +69,7 @@ void* _deserialize(cuvsResources_t res, const char* filename, int dim, cuvsDista extern "C" cuvsError_t cuvsHnswSearchParamsCreate(cuvsHnswSearchParams_t* params) { return cuvs::core::translate_exceptions( - [=] { *params = new cuvsHnswSearchParams{.ef = 200, .numThreads = 0}; }); + [=] { *params = new cuvsHnswSearchParams{.ef = 200, .num_threads = 0}; }); } extern "C" cuvsError_t cuvsHnswSearchParamsDestroy(cuvsHnswSearchParams_t params) @@ -126,17 +126,14 @@ extern "C" cuvsError_t cuvsHnswSearch(cuvsResources_t res, "distances should be of type float32"); auto index = *index_c_ptr; - RAFT_EXPECTS(queries.dtype.code == index.dtype.code, "type mismatch between index and queries"); - RAFT_EXPECTS(queries.dtype.bits == 32, "number of bits in queries dtype should be 32"); + RAFT_EXPECTS(queries.dtype.code == index.dtype.code, "type mismatch between queries and index"); if (index.dtype.code == kDLFloat) { - _search( - res, *params, index, queries_tensor, neighbors_tensor, distances_tensor); + _search(res, *params, index, queries_tensor, neighbors_tensor, distances_tensor); } else if (index.dtype.code == kDLUInt) { - _search( - res, *params, index, queries_tensor, neighbors_tensor, distances_tensor); + _search(res, *params, index, queries_tensor, neighbors_tensor, distances_tensor); } else if (index.dtype.code == kDLInt) { - _search(res, *params, index, queries_tensor, neighbors_tensor, distances_tensor); + _search(res, *params, index, queries_tensor, neighbors_tensor, distances_tensor); } else { RAFT_FAIL("Unsupported index dtype: %d and bits: %d", queries.dtype.code, queries.dtype.bits); } @@ -152,13 +149,10 @@ extern "C" cuvsError_t cuvsHnswDeserialize(cuvsResources_t res, return cuvs::core::translate_exceptions([=] { if (index->dtype.code == kDLFloat && index->dtype.bits == 32) { index->addr = reinterpret_cast(_deserialize(res, filename, dim, metric)); - index->dtype.code = kDLFloat; } else if (index->dtype.code == kDLUInt && index->dtype.bits == 8) { index->addr = reinterpret_cast(_deserialize(res, filename, dim, metric)); - index->dtype.code = kDLInt; } else if (index->dtype.code == kDLInt && index->dtype.bits == 8) { index->addr = reinterpret_cast(_deserialize(res, filename, dim, metric)); - index->dtype.code = kDLUInt; } else { RAFT_FAIL("Unsupported dtype in file %s", filename); } diff --git a/python/cuvs/cuvs/neighbors/CMakeLists.txt b/python/cuvs/cuvs/neighbors/CMakeLists.txt index 3579215fd..c5d17e34c 100644 --- a/python/cuvs/cuvs/neighbors/CMakeLists.txt +++ b/python/cuvs/cuvs/neighbors/CMakeLists.txt @@ -14,6 +14,7 @@ add_subdirectory(brute_force) add_subdirectory(cagra) +add_subdirectory(hnsw) add_subdirectory(ivf_flat) add_subdirectory(ivf_pq) diff --git a/python/cuvs/cuvs/neighbors/cagra/cagra.pxd b/python/cuvs/cuvs/neighbors/cagra/cagra.pxd index b23c2a4b3..04dc239fd 100644 --- a/python/cuvs/cuvs/neighbors/cagra/cagra.pxd +++ b/python/cuvs/cuvs/neighbors/cagra/cagra.pxd @@ -17,6 +17,7 @@ from libc.stdint cimport ( int8_t, + int32_t, int64_t, uint8_t, uint32_t, @@ -100,6 +101,8 @@ cdef extern from "cuvs/neighbors/cagra.h" nogil: cuvsError_t cuvsCagraIndexDestroy(cuvsCagraIndex_t index) + cuvsError_t cuvsCagraIndexDim(cuvsCagraIndex_t index, int32_t* dim) + cuvsError_t cuvsCagraBuild(cuvsResources_t res, cuvsCagraIndexParams* params, DLManagedTensor* dataset, @@ -117,6 +120,20 @@ cdef extern from "cuvs/neighbors/cagra.h" nogil: cuvsCagraIndex_t index, bool include_dataset) except + + cuvsError_t cuvsCagraSerializeToHnswlib(cuvsResources_t res, + const char * filename, + cuvsCagraIndex_t index) except + + cuvsError_t cuvsCagraDeserialize(cuvsResources_t res, const char * filename, cuvsCagraIndex_t index) except + + +cdef class Index: + """ + CAGRA index object. This object stores the trained CAGRA index state + which can be used to perform nearest neighbors searches. + """ + + cdef cuvsCagraIndex_t index + cdef bool trained + cdef str active_index_type diff --git a/python/cuvs/cuvs/neighbors/cagra/cagra.pyx b/python/cuvs/cuvs/neighbors/cagra/cagra.pyx index f940ab8bf..1473b063c 100644 --- a/python/cuvs/cuvs/neighbors/cagra/cagra.pyx +++ b/python/cuvs/cuvs/neighbors/cagra/cagra.pyx @@ -36,6 +36,7 @@ from pylibraft.neighbors.common import _check_input_array from libc.stdint cimport ( int8_t, + int32_t, int64_t, uint8_t, uint32_t, @@ -206,16 +207,9 @@ cdef class IndexParams: cdef class Index: - """ - CAGRA index object. This object stores the trained CAGRA index state - which can be used to perform nearest neighbors searches. - """ - - cdef cuvsCagraIndex_t index - cdef bool trained - def __cinit__(self): self.trained = False + self.active_index_type = None check_cuvs(cuvsCagraIndexCreate(&self.index)) def __dealloc__(self): @@ -226,6 +220,12 @@ cdef class Index: def trained(self): return self.trained + @property + def dim(self): + cdef int32_t dim + check_cuvs(cuvsCagraIndexDim(self.index, &dim)) + return dim + def __repr__(self): # todo(dgd): update repr as we expose data through C API attr_str = [] @@ -299,6 +299,7 @@ def build(IndexParams index_params, dataset, resources=None): idx.index )) idx.trained = True + idx.active_index_type = dataset_ai.dtype.name return idx diff --git a/python/cuvs/cuvs/neighbors/hnsw/CMakeLists.txt b/python/cuvs/cuvs/neighbors/hnsw/CMakeLists.txt new file mode 100644 index 000000000..1f9c422ca --- /dev/null +++ b/python/cuvs/cuvs/neighbors/hnsw/CMakeLists.txt @@ -0,0 +1,24 @@ +# ============================================================================= +# 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. +# ============================================================================= + +# Set the list of Cython files to build +set(cython_sources hnsw.pyx) +set(linked_libraries cuvs::cuvs cuvs::c_api) + +# 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_hnsw_ +) diff --git a/python/cuvs/cuvs/neighbors/hnsw/__init__.pxd b/python/cuvs/cuvs/neighbors/hnsw/__init__.pxd new file mode 100644 index 000000000..e69de29bb diff --git a/python/cuvs/cuvs/neighbors/hnsw/__init__.py b/python/cuvs/cuvs/neighbors/hnsw/__init__.py new file mode 100644 index 000000000..5efcdf68b --- /dev/null +++ b/python/cuvs/cuvs/neighbors/hnsw/__init__.py @@ -0,0 +1,25 @@ +# 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 .hnsw import Index, SearchParams, from_cagra, load, save, search + +__all__ = [ + "Index", + "SearchParams", + "load", + "save", + "search", + "from_cagra", +] diff --git a/python/cuvs/cuvs/neighbors/hnsw/hnsw.pxd b/python/cuvs/cuvs/neighbors/hnsw/hnsw.pxd new file mode 100644 index 000000000..fd251c004 --- /dev/null +++ b/python/cuvs/cuvs/neighbors/hnsw/hnsw.pxd @@ -0,0 +1,53 @@ +# +# 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 int32_t, uintptr_t + +from cuvs.common.c_api cimport cuvsError_t, cuvsResources_t +from cuvs.common.cydlpack cimport DLDataType, DLManagedTensor +from cuvs.distance_type cimport cuvsDistanceType + + +cdef extern from "cuvs/neighbors/hnsw.h" nogil: + ctypedef struct cuvsHnswSearchParams: + int32_t ef + int32_t num_threads + + ctypedef cuvsHnswSearchParams* cuvsHnswSearchParams_t + + ctypedef struct cuvsHnswIndex: + uintptr_t addr + DLDataType dtype + + ctypedef cuvsHnswIndex* cuvsHnswIndex_t + + cuvsError_t cuvsHnswIndexCreate(cuvsHnswIndex_t* index) + + cuvsError_t cuvsHnswIndexDestroy(cuvsHnswIndex_t index) + + cuvsError_t cuvsHnswSearch(cuvsResources_t res, + cuvsHnswSearchParams* params, + cuvsHnswIndex_t index, + DLManagedTensor* queries, + DLManagedTensor* neighbors, + DLManagedTensor* distances) except + + + cuvsError_t cuvsHnswDeserialize(cuvsResources_t res, + const char * filename, + int32_t dim, + cuvsDistanceType metric, + cuvsHnswIndex_t index) except + diff --git a/python/cuvs/cuvs/neighbors/hnsw/hnsw.pyx b/python/cuvs/cuvs/neighbors/hnsw/hnsw.pyx new file mode 100644 index 000000000..c780382bb --- /dev/null +++ b/python/cuvs/cuvs/neighbors/hnsw/hnsw.pyx @@ -0,0 +1,364 @@ +# +# 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 uint32_t +from libcpp cimport bool +from libcpp.string cimport string + +from cuvs.common.exceptions import check_cuvs +from cuvs.common.resources import auto_sync_resources + +from cuvs.common cimport cydlpack + +import numpy as np + +from cuvs.distance import DISTANCE_TYPES + +from cuvs.neighbors.cagra cimport cagra + +import os +import uuid + +from pylibraft.common import auto_convert_output +from pylibraft.common.cai_wrapper import wrap_array +from pylibraft.common.interruptible import cuda_interruptible +from pylibraft.neighbors.common import _check_input_array + + +cdef class SearchParams: + """ + HNSW search parameters + + Parameters + ---------- + ef: int, default = 200 + Maximum number of candidate list size used during search. + num_threads: int, default = 0 + Number of CPU threads used to increase search parallelism. + When set to 0, the number of threads is automatically determined. + """ + + cdef cuvsHnswSearchParams params + + def __init__(self, *, + ef=200, + num_threads=0): + self.params.ef = ef + self.params.num_threads = num_threads + + def __repr__(self): + attr_str = [attr + "=" + str(getattr(self, attr)) + for attr in [ + "ef", "num_threads"]] + return "SearchParams(type=HNSW, " + (", ".join(attr_str)) + ")" + + @property + def ef(self): + return self.params.ef + + @property + def num_threads(self): + return self.params.num_threads + + +cdef class Index: + """ + HNSW index object. This object stores the trained HNSW index state + which can be used to perform nearest neighbors searches. + """ + + cdef cuvsHnswIndex_t index + cdef bool trained + + def __cinit__(self): + self.trained = False + check_cuvs(cuvsHnswIndexCreate(&self.index)) + + def __dealloc__(self): + if self.index is not NULL: + check_cuvs(cuvsHnswIndexDestroy(self.index)) + + @property + def trained(self): + return self.trained + + def __repr__(self): + # todo(dgd): update repr as we expose data through C API + attr_str = [] + return "Index(type=HNSW, metric=L2" + (", ".join(attr_str)) + ")" + + +@auto_sync_resources +def save(filename, cagra.Index index, resources=None): + """ + Saves the CAGRA index to a file as an hnswlib index. + + Saving / loading the index is experimental. The serialization format is + subject to change. + + Parameters + ---------- + filename : string + Name of the file. + index : Index + Trained CAGRA index. + {resources_docstring} + + Examples + -------- + >>> import cupy as cp + >>> from cuvs.neighbors import cagra + >>> n_samples = 50000 + >>> n_features = 50 + >>> dataset = cp.random.random_sample((n_samples, n_features), + ... dtype=cp.float32) + >>> # Build index + >>> index = cagra.build(cagra.IndexParams(), dataset) + >>> # Serialize and deserialize the cagra index built + >>> hnsw.save("my_index.bin", index) + """ + cdef string c_filename = filename.encode('utf-8') + cdef cuvsResources_t res = resources.get_c_obj() + check_cuvs(cagra.cuvsCagraSerializeToHnswlib(res, + c_filename.c_str(), + index.index)) + + +@auto_sync_resources +def load(filename, dim, dtype, metric="sqeuclidean", resources=None): + """ + Loads base-layer-only hnswlib index from file, which was originally + saved as a built CAGRA index. + + Saving / loading the index is experimental. The serialization format is + subject to change, therefore loading an index saved with a previous + version of raft is not guaranteed to work. + + Parameters + ---------- + filename : string + Name of the file. + dim : int + Dimensions of the training dataest + dtype : np.dtype of the saved index + Valid values for dtype: [np.float32, np.byte, np.ubyte] + metric : string denoting the metric type, default="sqeuclidean" + Valid values for metric: ["sqeuclidean"], where + - sqeuclidean is the euclidean distance without the square root + operation, i.e.: distance(a,b) = \\sum_i (a_i - b_i)^2, + - inner product distance is defined as + distance(a, b) = \\sum_i a_i * b_i. + {resources_docstring} + + Returns + ------- + index : HnswIndex + + Examples + -------- + >>> import cupy as cp + >>> from cuvs.neighbors import cagra + >>> from cuvs.neighbors import hnsw + >>> n_samples = 50000 + >>> n_features = 50 + >>> dataset = cp.random.random_sample((n_samples, n_features), + ... dtype=cp.float32) + >>> # Build index + >>> index = cagra.build(cagra.IndexParams(), dataset) + >>> # Serialize the CAGRA index to hnswlib base layer only index format + >>> hnsw.save("my_index.bin", index) + >>> index = hnsw.load("my_index.bin", n_features, np.float32, + ... "sqeuclidean") + """ + cdef Index idx = Index() + cdef cuvsResources_t res = resources.get_c_obj() + cdef string c_filename = filename.encode('utf-8') + cdef cydlpack.DLDataType dl_dtype + if dtype == np.float32: + dl_dtype.code = cydlpack.kDLFloat + dl_dtype.bits = 32 + dl_dtype.lanes = 1 + elif dtype == np.ubyte: + dl_dtype.code = cydlpack.kDLUInt + dl_dtype.bits = 8 + dl_dtype.lanes = 1 + elif dtype == np.byte: + dl_dtype.code = cydlpack.kDLInt + dl_dtype.bits = 8 + dl_dtype.lanes = 1 + else: + raise ValueError("Only float32 is supported for dtype") + + idx.index.dtype = dl_dtype + cdef cuvsDistanceType distance_type = DISTANCE_TYPES[metric] + + check_cuvs(cuvsHnswDeserialize( + res, + c_filename.c_str(), + dim, + distance_type, + idx.index + )) + idx.trained = True + return idx + + +@auto_sync_resources +def from_cagra(cagra.Index index, resources=None): + """ + Returns an hnswlib base-layer-only index from a CAGRA index. + + NOTE: This method uses the filesystem to write the CAGRA index in + `/tmp/.bin` before reading it as an hnswlib index, + then deleting the temporary file. + + Saving / loading the index is experimental. The serialization format is + subject to change. + + Parameters + ---------- + index : Index + Trained CAGRA index. + {resources_docstring} + + Examples + -------- + >>> import cupy as cp + >>> from cuvs.neighbors import cagra + >>> from cuvs.neighbors import hnsw + >>> n_samples = 50000 + >>> n_features = 50 + >>> dataset = cp.random.random_sample((n_samples, n_features), + ... dtype=cp.float32) + >>> # Build index + >>> index = cagra.build(cagra.IndexParams(), dataset) + >>> # Serialize the CAGRA index to hnswlib base layer only index format + >>> hnsw_index = hnsw.from_cagra(index) + """ + uuid_num = uuid.uuid4() + filename = f"/tmp/{uuid_num}.bin" + save(filename, index, resources=resources) + hnsw_index = load(filename, index.dim, np.dtype(index.active_index_type), + "sqeuclidean", resources=resources) + os.remove(filename) + return hnsw_index + + +@auto_sync_resources +@auto_convert_output +def search(SearchParams search_params, + Index index, + queries, + k, + neighbors=None, + distances=None, + resources=None): + """ + Find the k nearest neighbors for each query. + + Parameters + ---------- + search_params : SearchParams + index : Index + Trained CAGRA index. + queries : CUDA array interface compliant matrix shape (n_samples, dim) + Supported dtype [float, int] + k : int + The number of neighbors. + neighbors : Optional CUDA array interface compliant matrix shape + (n_queries, k), dtype uint64_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 cagra, hnsw + >>> n_samples = 50000 + >>> n_features = 50 + >>> n_queries = 1000 + >>> dataset = cp.random.random_sample((n_samples, n_features), + ... dtype=cp.float32) + >>> # Build index + >>> index = cagra.build(cagra.IndexParams(), dataset) + >>> # Search using the built index + >>> queries = cp.random.random_sample((n_queries, n_features), + ... dtype=cp.float32) + >>> k = 10 + >>> search_params = hnsw.SearchParams( + ... ef=200, + ... num_threads=0 + ... ) + >>> # Convert CAGRA index to HNSW + >>> hnsw_index = hnsw.from_cagra(index) + >>> # Using a pooling allocator reduces overhead of temporary array + >>> # creation during search. This is useful if multiple searches + >>> # are performed with same query size. + >>> distances, neighbors = hnsw.search(search_params, 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.") + + # todo(dgd): we can make the check of dtype a parameter of wrap_array + # in RAFT to make this a single call + queries_ai = wrap_array(queries) + _check_input_array(queries_ai, [np.dtype('float32'), + np.dtype('uint8'), + np.dtype('int8')]) + + cdef uint32_t n_queries = queries_ai.shape[0] + + if neighbors is None: + neighbors = np.empty((n_queries, k), dtype='uint64') + + neighbors_ai = wrap_array(neighbors) + _check_input_array(neighbors_ai, [np.dtype('uint64')], + exp_rows=n_queries, exp_cols=k) + + if distances is None: + distances = np.empty((n_queries, k), dtype='float32') + + distances_ai = wrap_array(distances) + _check_input_array(distances_ai, [np.dtype('float32')], + exp_rows=n_queries, exp_cols=k) + + cdef cuvsHnswSearchParams* params = &search_params.params + cdef cydlpack.DLManagedTensor* queries_dlpack = \ + cydlpack.dlpack_c(queries_ai) + cdef cydlpack.DLManagedTensor* neighbors_dlpack = \ + cydlpack.dlpack_c(neighbors_ai) + cdef cydlpack.DLManagedTensor* distances_dlpack = \ + cydlpack.dlpack_c(distances_ai) + cdef cuvsResources_t res = resources.get_c_obj() + + with cuda_interruptible(): + check_cuvs(cuvsHnswSearch( + res, + params, + index.index, + queries_dlpack, + neighbors_dlpack, + distances_dlpack + )) + + return (distances, neighbors) diff --git a/python/cuvs/cuvs/test/test_hnsw.py b/python/cuvs/cuvs/test/test_hnsw.py new file mode 100644 index 000000000..8bd2e8b76 --- /dev/null +++ b/python/cuvs/cuvs/test/test_hnsw.py @@ -0,0 +1,97 @@ +# 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 +# +# h ttp://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 sklearn.neighbors import NearestNeighbors +from sklearn.preprocessing import normalize + +from cuvs.neighbors import cagra, hnsw +from cuvs.test.ann_utils import calc_recall, generate_data + + +def run_hnsw_build_search_test( + n_rows=10000, + n_cols=10, + n_queries=100, + k=10, + dtype=np.float32, + metric="sqeuclidean", + build_algo="ivf_pq", + intermediate_graph_degree=128, + graph_degree=64, + search_params={}, +): + dataset = generate_data((n_rows, n_cols), dtype) + if metric == "inner_product": + dataset = normalize(dataset, norm="l2", axis=1) + if dtype in [np.int8, np.uint8]: + pytest.skip( + "inner_product metric is not supported for int8/uint8 data" + ) + if build_algo == "nn_descent": + pytest.skip("inner_product metric is not supported for nn_descent") + + build_params = cagra.IndexParams( + metric=metric, + intermediate_graph_degree=intermediate_graph_degree, + graph_degree=graph_degree, + build_algo=build_algo, + ) + + index = cagra.build(build_params, dataset) + + assert index.trained + + hnsw_index = hnsw.from_cagra(index) + + queries = generate_data((n_queries, n_cols), dtype) + + search_params = hnsw.SearchParams(**search_params) + + out_dist, out_idx = hnsw.search(search_params, hnsw_index, queries, k) + + # Calculate reference values with sklearn + skl_metric = { + "sqeuclidean": "sqeuclidean", + "inner_product": "cosine", + "euclidean": "euclidean", + }[metric] + nn_skl = NearestNeighbors( + n_neighbors=k, algorithm="brute", metric=skl_metric + ) + nn_skl.fit(dataset) + skl_dist, skl_idx = nn_skl.kneighbors(queries, return_distance=True) + + recall = calc_recall(out_idx, skl_idx) + assert recall > 0.95 + + +@pytest.mark.parametrize("dtype", [np.float32, np.int8, np.uint8]) +@pytest.mark.parametrize("k", [10, 20]) +@pytest.mark.parametrize("ef", [30, 40]) +@pytest.mark.parametrize("num_threads", [2, 4]) +@pytest.mark.parametrize("metric", ["sqeuclidean"]) +@pytest.mark.parametrize("build_algo", ["ivf_pq", "nn_descent"]) +def test_hnsw(dtype, k, ef, num_threads, metric, build_algo): + # Note that inner_product tests use normalized input which we cannot + # represent in int8, therefore we test only sqeuclidean metric here. + run_hnsw_build_search_test( + dtype=dtype, + k=k, + metric=metric, + build_algo=build_algo, + search_params={"ef": ef, "num_threads": num_threads}, + ) From 5ba4fadca07763dcd5b070029a8b10ecb3f64b7c Mon Sep 17 00:00:00 2001 From: divyegala Date: Tue, 23 Jul 2024 18:13:14 -0700 Subject: [PATCH 07/18] documentation --- cpp/include/cuvs/neighbors/hnsw.h | 2 ++ 1 file changed, 2 insertions(+) diff --git a/cpp/include/cuvs/neighbors/hnsw.h b/cpp/include/cuvs/neighbors/hnsw.h index 55d44e36c..84db16866 100644 --- a/cpp/include/cuvs/neighbors/hnsw.h +++ b/cpp/include/cuvs/neighbors/hnsw.h @@ -188,6 +188,8 @@ cuvsError_t cuvsHnswSearch(cuvsResources_t res, * * @param[in] res cuvsResources_t opaque C handle * @param[in] filename the name of the file that stores the index + * @param[in] dim dimensions of the training dataset + * @param[in] metric distance metric to search. Supported metrics ("L2Expanded") * @param[out] index HNSW index loaded disk */ cuvsError_t cuvsHnswDeserialize(cuvsResources_t res, From 0c3d053569f2ca9d7f51a94444a40d5dd1a7369e Mon Sep 17 00:00:00 2001 From: divyegala Date: Wed, 24 Jul 2024 09:16:58 -0700 Subject: [PATCH 08/18] more docs --- cpp/include/cuvs/neighbors/cagra.h | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/cpp/include/cuvs/neighbors/cagra.h b/cpp/include/cuvs/neighbors/cagra.h index 00977983c..4306c94c4 100644 --- a/cpp/include/cuvs/neighbors/cagra.h +++ b/cpp/include/cuvs/neighbors/cagra.h @@ -270,8 +270,8 @@ cuvsError_t cuvsCagraIndexDestroy(cuvsCagraIndex_t index); /** * @brief Get dimension of the CAGRA index * - * @param[in] CAGRA index - * @param[out] dim + * @param[in] index CAGRA index + * @param[out] dim return dimension of the index * @return cuvsError_t */ cuvsError_t cuvsCagraIndexDim(cuvsCagraIndex_t index, int* dim); From 0c2d08292fa106eda375620e6a8be4b20c1825fa Mon Sep 17 00:00:00 2001 From: divyegala Date: Wed, 11 Sep 2024 09:35:51 -0700 Subject: [PATCH 09/18] passing tests --- cpp/src/neighbors/hnsw_c.cpp | 16 +++++----------- python/cuvs/cuvs/neighbors/hnsw/hnsw.pxd | 2 +- python/cuvs/cuvs/neighbors/hnsw/hnsw.pyx | 17 ++++++++++++----- 3 files changed, 18 insertions(+), 17 deletions(-) diff --git a/cpp/src/neighbors/hnsw_c.cpp b/cpp/src/neighbors/hnsw_c.cpp index ab5268a6d..a19875641 100644 --- a/cpp/src/neighbors/hnsw_c.cpp +++ b/cpp/src/neighbors/hnsw_c.cpp @@ -31,7 +31,7 @@ #include namespace { -template +template void _search(cuvsResources_t res, cuvsHnswSearchParams params, cuvsHnswIndex index, @@ -46,7 +46,7 @@ void _search(cuvsResources_t res, search_params.ef = params.ef; search_params.num_threads = params.numThreads; - using queries_mdspan_type = raft::host_matrix_view; + using queries_mdspan_type = raft::host_matrix_view; using neighbors_mdspan_type = raft::host_matrix_view; using distances_mdspan_type = raft::host_matrix_view; auto queries_mds = cuvs::core::from_dlpack(queries_tensor); @@ -127,16 +127,13 @@ extern "C" cuvsError_t cuvsHnswSearch(cuvsResources_t res, auto index = *index_c_ptr; RAFT_EXPECTS(queries.dtype.code == index.dtype.code, "type mismatch between index and queries"); - RAFT_EXPECTS(queries.dtype.bits == 32, "number of bits in queries dtype should be 32"); if (index.dtype.code == kDLFloat) { - _search( - res, *params, index, queries_tensor, neighbors_tensor, distances_tensor); + _search(res, *params, index, queries_tensor, neighbors_tensor, distances_tensor); } else if (index.dtype.code == kDLUInt) { - _search( - res, *params, index, queries_tensor, neighbors_tensor, distances_tensor); + _search(res, *params, index, queries_tensor, neighbors_tensor, distances_tensor); } else if (index.dtype.code == kDLInt) { - _search(res, *params, index, queries_tensor, neighbors_tensor, distances_tensor); + _search(res, *params, index, queries_tensor, neighbors_tensor, distances_tensor); } else { RAFT_FAIL("Unsupported index dtype: %d and bits: %d", queries.dtype.code, queries.dtype.bits); } @@ -152,13 +149,10 @@ extern "C" cuvsError_t cuvsHnswDeserialize(cuvsResources_t res, return cuvs::core::translate_exceptions([=] { if (index->dtype.code == kDLFloat && index->dtype.bits == 32) { index->addr = reinterpret_cast(_deserialize(res, filename, dim, metric)); - index->dtype.code = kDLFloat; } else if (index->dtype.code == kDLUInt && index->dtype.bits == 8) { index->addr = reinterpret_cast(_deserialize(res, filename, dim, metric)); - index->dtype.code = kDLInt; } else if (index->dtype.code == kDLInt && index->dtype.bits == 8) { index->addr = reinterpret_cast(_deserialize(res, filename, dim, metric)); - index->dtype.code = kDLUInt; } else { RAFT_FAIL("Unsupported dtype in file %s", filename); } diff --git a/python/cuvs/cuvs/neighbors/hnsw/hnsw.pxd b/python/cuvs/cuvs/neighbors/hnsw/hnsw.pxd index fd251c004..1cdc97406 100644 --- a/python/cuvs/cuvs/neighbors/hnsw/hnsw.pxd +++ b/python/cuvs/cuvs/neighbors/hnsw/hnsw.pxd @@ -25,7 +25,7 @@ from cuvs.distance_type cimport cuvsDistanceType cdef extern from "cuvs/neighbors/hnsw.h" nogil: ctypedef struct cuvsHnswSearchParams: int32_t ef - int32_t num_threads + int32_t numThreads ctypedef cuvsHnswSearchParams* cuvsHnswSearchParams_t diff --git a/python/cuvs/cuvs/neighbors/hnsw/hnsw.pyx b/python/cuvs/cuvs/neighbors/hnsw/hnsw.pyx index c780382bb..57b3b8ffe 100644 --- a/python/cuvs/cuvs/neighbors/hnsw/hnsw.pyx +++ b/python/cuvs/cuvs/neighbors/hnsw/hnsw.pyx @@ -58,7 +58,7 @@ cdef class SearchParams: ef=200, num_threads=0): self.params.ef = ef - self.params.num_threads = num_threads + self.params.numThreads = num_threads def __repr__(self): attr_str = [attr + "=" + str(getattr(self, attr)) @@ -72,7 +72,7 @@ cdef class SearchParams: @property def num_threads(self): - return self.params.num_threads + return self.params.numThreads cdef class Index: @@ -106,6 +106,9 @@ cdef class Index: def save(filename, cagra.Index index, resources=None): """ Saves the CAGRA index to a file as an hnswlib index. + The saved index is immutable and can only be searched by the hnswlib + wrapper in cuVS, as the format is not compatible with the original + hnswlib. Saving / loading the index is experimental. The serialization format is subject to change. @@ -142,11 +145,13 @@ def save(filename, cagra.Index index, resources=None): def load(filename, dim, dtype, metric="sqeuclidean", resources=None): """ Loads base-layer-only hnswlib index from file, which was originally - saved as a built CAGRA index. + saved as a built CAGRA index. The loaded index is immutable and can only + be searched by the hnswlib wrapper in cuVS, as the format is not + compatible with the original hnswlib. Saving / loading the index is experimental. The serialization format is subject to change, therefore loading an index saved with a previous - version of raft is not guaranteed to work. + version of cuVS is not guaranteed to work. Parameters ---------- @@ -224,7 +229,9 @@ def from_cagra(cagra.Index index, resources=None): NOTE: This method uses the filesystem to write the CAGRA index in `/tmp/.bin` before reading it as an hnswlib index, - then deleting the temporary file. + then deleting the temporary file. The returned index is immutable + and can only be searched by the hnswlib wrapper in cuVS, as the + format is not compatible with the original hnswlib. Saving / loading the index is experimental. The serialization format is subject to change. From 8c53e222538432967c3860f6f37f5fe64082430a Mon Sep 17 00:00:00 2001 From: divyegala Date: Mon, 16 Sep 2024 14:21:42 -0700 Subject: [PATCH 10/18] address review --- cpp/include/cuvs/neighbors/cagra.h | 2 +- cpp/src/neighbors/cagra_c.cpp | 2 +- python/cuvs/cuvs/neighbors/cagra/cagra.pxd | 2 +- python/cuvs/cuvs/neighbors/cagra/cagra.pyx | 2 +- python/cuvs/cuvs/neighbors/hnsw/hnsw.pyx | 18 +++++++++++------- 5 files changed, 15 insertions(+), 11 deletions(-) diff --git a/cpp/include/cuvs/neighbors/cagra.h b/cpp/include/cuvs/neighbors/cagra.h index d4c1a56e6..14331ebbc 100644 --- a/cpp/include/cuvs/neighbors/cagra.h +++ b/cpp/include/cuvs/neighbors/cagra.h @@ -274,7 +274,7 @@ cuvsError_t cuvsCagraIndexDestroy(cuvsCagraIndex_t index); * @param[out] dim return dimension of the index * @return cuvsError_t */ -cuvsError_t cuvsCagraIndexDim(cuvsCagraIndex_t index, int* dim); +cuvsError_t cuvsCagraIndexGetDims(cuvsCagraIndex_t index, int* dim); /** * @} diff --git a/cpp/src/neighbors/cagra_c.cpp b/cpp/src/neighbors/cagra_c.cpp index 229fa48f1..6985ff094 100644 --- a/cpp/src/neighbors/cagra_c.cpp +++ b/cpp/src/neighbors/cagra_c.cpp @@ -176,7 +176,7 @@ extern "C" cuvsError_t cuvsCagraIndexDestroy(cuvsCagraIndex_t index_c_ptr) }); } -extern "C" cuvsError_t cuvsCagraIndexDim(cuvsCagraIndex_t index, int* dim) +extern "C" cuvsError_t cuvsCagraIndexGetDims(cuvsCagraIndex_t index, int* dim) { return cuvs::core::translate_exceptions([=] { auto index_ptr = reinterpret_cast*>(index->addr); diff --git a/python/cuvs/cuvs/neighbors/cagra/cagra.pxd b/python/cuvs/cuvs/neighbors/cagra/cagra.pxd index 04dc239fd..bba5a91a8 100644 --- a/python/cuvs/cuvs/neighbors/cagra/cagra.pxd +++ b/python/cuvs/cuvs/neighbors/cagra/cagra.pxd @@ -101,7 +101,7 @@ cdef extern from "cuvs/neighbors/cagra.h" nogil: cuvsError_t cuvsCagraIndexDestroy(cuvsCagraIndex_t index) - cuvsError_t cuvsCagraIndexDim(cuvsCagraIndex_t index, int32_t* dim) + cuvsError_t cuvsCagraIndexGetDims(cuvsCagraIndex_t index, int32_t* dim) cuvsError_t cuvsCagraBuild(cuvsResources_t res, cuvsCagraIndexParams* params, diff --git a/python/cuvs/cuvs/neighbors/cagra/cagra.pyx b/python/cuvs/cuvs/neighbors/cagra/cagra.pyx index 1473b063c..95209dbeb 100644 --- a/python/cuvs/cuvs/neighbors/cagra/cagra.pyx +++ b/python/cuvs/cuvs/neighbors/cagra/cagra.pyx @@ -223,7 +223,7 @@ cdef class Index: @property def dim(self): cdef int32_t dim - check_cuvs(cuvsCagraIndexDim(self.index, &dim)) + check_cuvs(cuvsCagraIndexGetDims(self.index, &dim)) return dim def __repr__(self): diff --git a/python/cuvs/cuvs/neighbors/hnsw/hnsw.pyx b/python/cuvs/cuvs/neighbors/hnsw/hnsw.pyx index 57b3b8ffe..f87e6549b 100644 --- a/python/cuvs/cuvs/neighbors/hnsw/hnsw.pyx +++ b/python/cuvs/cuvs/neighbors/hnsw/hnsw.pyx @@ -49,7 +49,8 @@ cdef class SearchParams: Maximum number of candidate list size used during search. num_threads: int, default = 0 Number of CPU threads used to increase search parallelism. - When set to 0, the number of threads is automatically determined. + When set to 0, the number of threads is automatically determined + using OpenMP's `omp_get_max_threads()`. """ cdef cuvsHnswSearchParams params @@ -162,10 +163,10 @@ def load(filename, dim, dtype, metric="sqeuclidean", resources=None): dtype : np.dtype of the saved index Valid values for dtype: [np.float32, np.byte, np.ubyte] metric : string denoting the metric type, default="sqeuclidean" - Valid values for metric: ["sqeuclidean"], where + Valid values for metric: ["sqeuclidean", "inner_product"], where - sqeuclidean is the euclidean distance without the square root operation, i.e.: distance(a,b) = \\sum_i (a_i - b_i)^2, - - inner product distance is defined as + - inner_product distance is defined as distance(a, b) = \\sum_i a_i * b_i. {resources_docstring} @@ -225,13 +226,16 @@ def load(filename, dim, dtype, metric="sqeuclidean", resources=None): @auto_sync_resources def from_cagra(cagra.Index index, resources=None): """ - Returns an hnswlib base-layer-only index from a CAGRA index. + Returns an hnsw base-layer-only index from a CAGRA index. NOTE: This method uses the filesystem to write the CAGRA index in - `/tmp/.bin` before reading it as an hnswlib index, + `/tmp/.bin` before reading it as an hnsw index, then deleting the temporary file. The returned index is immutable - and can only be searched by the hnswlib wrapper in cuVS, as the - format is not compatible with the original hnswlib. + and can only be searched by the hnsw wrapper in cuVS, as the + format is not compatible with the original hnswlib library. + By `base_layer_only`, we mean that the hnsw index is created + without the additional layers that are used for the hierarchical + search in hnswlib. Instead, the base layer is used for the search. Saving / loading the index is experimental. The serialization format is subject to change. From ef98a4efd0b5e13a2397c50db7f262b0e0d6c524 Mon Sep 17 00:00:00 2001 From: divyegala Date: Fri, 27 Sep 2024 14:35:54 -0700 Subject: [PATCH 11/18] address review --- cpp/include/cuvs/neighbors/hnsw.hpp | 2 ++ docs/source/c_api/neighbors.rst | 1 + docs/source/cpp_api/neighbors.rst | 1 + docs/source/python_api/neighbors.rst | 1 + docs/source/python_api/neighbors_hnsw.rst | 30 +++++++++++++++++++++++ python/cuvs/cuvs/neighbors/hnsw/hnsw.pyx | 11 ++++++--- 6 files changed, 43 insertions(+), 3 deletions(-) create mode 100644 docs/source/python_api/neighbors_hnsw.rst diff --git a/cpp/include/cuvs/neighbors/hnsw.hpp b/cpp/include/cuvs/neighbors/hnsw.hpp index 5d8ba730f..d5abd6d55 100644 --- a/cpp/include/cuvs/neighbors/hnsw.hpp +++ b/cpp/include/cuvs/neighbors/hnsw.hpp @@ -173,6 +173,8 @@ std::unique_ptr> from_cagra( /**@}*/ +// TODO: Filtered Search APIs: https://github.com/rapidsai/cuvs/issues/363 + /** * @defgroup hnsw_cpp_index_search Search hnswlib index * @{ diff --git a/docs/source/c_api/neighbors.rst b/docs/source/c_api/neighbors.rst index dc55a74dc..9c3fce672 100644 --- a/docs/source/c_api/neighbors.rst +++ b/docs/source/c_api/neighbors.rst @@ -13,3 +13,4 @@ Nearest Neighbors neighbors_ivf_flat_c.rst neighbors_ivf_pq_c.rst neighbors_cagra_c.rst + neighbors_hnsw_c.rst diff --git a/docs/source/cpp_api/neighbors.rst b/docs/source/cpp_api/neighbors.rst index e5a9fc977..4f5dd6ff2 100644 --- a/docs/source/cpp_api/neighbors.rst +++ b/docs/source/cpp_api/neighbors.rst @@ -11,6 +11,7 @@ Nearest Neighbors neighbors_bruteforce.rst neighbors_cagra.rst + neighbors_hnsw.rst neighbors_ivf_flat.rst neighbors_ivf_pq.rst neighbors_nn_descent.rst diff --git a/docs/source/python_api/neighbors.rst b/docs/source/python_api/neighbors.rst index 022c50de3..cd4f2609c 100644 --- a/docs/source/python_api/neighbors.rst +++ b/docs/source/python_api/neighbors.rst @@ -11,5 +11,6 @@ Nearest Neighbors neighbors_brute_force.rst neighbors_cagra.rst + neighbors_hnsw.rst neighbors_ivf_flat.rst neighbors_ivf_pq.rst diff --git a/docs/source/python_api/neighbors_hnsw.rst b/docs/source/python_api/neighbors_hnsw.rst new file mode 100644 index 000000000..9922805b3 --- /dev/null +++ b/docs/source/python_api/neighbors_hnsw.rst @@ -0,0 +1,30 @@ +HNSW +==== + +This is a wrapper for hnswlib, to load a CAGRA index as an immutable HNSW index. The loaded HNSW index is only compatible in cuVS, and can be searched using wrapper functions. + +.. role:: py(code) + :language: python + :class: highlight + +Index search parameters +####################### + +.. autoclass:: cuvs.neighbors.hnsw.SearchParams + :members: + +Index +##### + +.. autoclass:: cuvs.neighbors.hnsw.Index + :members: + +Index Conversion +################ + +.. autofunction:: cuvs.neighbors.hnsw.from_cagra + +Index search +############ + +.. autofunction:: cuvs.neighbors.hnsw.search diff --git a/python/cuvs/cuvs/neighbors/hnsw/hnsw.pyx b/python/cuvs/cuvs/neighbors/hnsw/hnsw.pyx index f87e6549b..018fcfef9 100644 --- a/python/cuvs/cuvs/neighbors/hnsw/hnsw.pyx +++ b/python/cuvs/cuvs/neighbors/hnsw/hnsw.pyx @@ -224,12 +224,13 @@ def load(filename, dim, dtype, metric="sqeuclidean", resources=None): @auto_sync_resources -def from_cagra(cagra.Index index, resources=None): +def from_cagra(cagra.Index index, temporary_index_path=None, resources=None): """ Returns an hnsw base-layer-only index from a CAGRA index. NOTE: This method uses the filesystem to write the CAGRA index in - `/tmp/.bin` before reading it as an hnsw index, + `/tmp/.bin` or the parameter `temporary_index_path` + if not None before reading it as an hnsw index, then deleting the temporary file. The returned index is immutable and can only be searched by the hnsw wrapper in cuVS, as the format is not compatible with the original hnswlib library. @@ -244,6 +245,9 @@ def from_cagra(cagra.Index index, resources=None): ---------- index : Index Trained CAGRA index. + temporary_index_path : string, default = None + Path to save the temporary index file. If None, the temporary file + will be saved in `/tmp/.bin`. {resources_docstring} Examples @@ -261,7 +265,8 @@ def from_cagra(cagra.Index index, resources=None): >>> hnsw_index = hnsw.from_cagra(index) """ uuid_num = uuid.uuid4() - filename = f"/tmp/{uuid_num}.bin" + filename = temporary_index_path if temporary_index_path else \ + f"/tmp/{uuid_num}.bin" save(filename, index, resources=resources) hnsw_index = load(filename, index.dim, np.dtype(index.active_index_type), "sqeuclidean", resources=resources) From 97215f228e630dedf054383d41f84c5603d93bbd Mon Sep 17 00:00:00 2001 From: divyegala Date: Mon, 30 Sep 2024 15:46:00 -0700 Subject: [PATCH 12/18] revert some changes --- .../detail/cagra/cagra_serialize.cuh | 21 ++++++------------- 1 file changed, 6 insertions(+), 15 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh b/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh index f86ed9ef6..1b7ec137b 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh @@ -119,9 +119,9 @@ void serialize_to_hnswlib(raft::resources const& res, os.write(reinterpret_cast(&curr_element_count), sizeof(std::size_t)); // Example:M: 16, dim = 128, data_t = float, index_t = uint32_t, list_size_type = uint32_t, // labeltype: size_t size_data_per_element_ = M * 2 * sizeof(index_t) + sizeof(list_size_type) + - // dim * 4 + sizeof(labeltype) - auto size_data_per_element = - static_cast(index_.graph_degree() * sizeof(IdxT) + 4 + index_.dim() * 4 + 8); + // dim * sizeof(T) + sizeof(labeltype) + auto size_data_per_element = static_cast(index_.graph_degree() * sizeof(IdxT) + + sizeof(T) + index_.dim() * 4 + 8); os.write(reinterpret_cast(&size_data_per_element), sizeof(std::size_t)); // label_offset std::size_t label_offset = size_data_per_element - 8; @@ -184,18 +184,9 @@ void serialize_to_hnswlib(raft::resources const& res, } auto data_row = host_dataset.data_handle() + (index_.dim() * i); - if constexpr (std::is_same_v) { - for (std::size_t j = 0; j < index_.dim(); ++j) { - auto data_elem = static_cast(host_dataset(i, j)); - os.write(reinterpret_cast(&data_elem), sizeof(float)); - } - } else if constexpr (std::is_same_v or std::is_same_v) { - for (std::size_t j = 0; j < index_.dim(); ++j) { - auto data_elem = static_cast(host_dataset(i, j)); - os.write(reinterpret_cast(&data_elem), sizeof(int)); - } - } else { - RAFT_FAIL("Unsupported dataset type while saving CAGRA dataset to HNSWlib format"); + for (std::size_t j = 0; j < index_.dim(); ++j) { + auto data_elem = static_cast(host_dataset(i, j)); + os.write(reinterpret_cast(&data_elem), sizeof(T)); } os.write(reinterpret_cast(&i), sizeof(std::size_t)); From 4acd22b1e6708b15e606cddaf25a31ade0c9e9b1 Mon Sep 17 00:00:00 2001 From: divyegala Date: Tue, 1 Oct 2024 09:32:56 -0700 Subject: [PATCH 13/18] fix failing tests --- cpp/src/neighbors/detail/cagra/cagra_serialize.cuh | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh b/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh index 1b7ec137b..bb71d915b 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh @@ -120,8 +120,8 @@ void serialize_to_hnswlib(raft::resources const& res, // Example:M: 16, dim = 128, data_t = float, index_t = uint32_t, list_size_type = uint32_t, // labeltype: size_t size_data_per_element_ = M * 2 * sizeof(index_t) + sizeof(list_size_type) + // dim * sizeof(T) + sizeof(labeltype) - auto size_data_per_element = static_cast(index_.graph_degree() * sizeof(IdxT) + - sizeof(T) + index_.dim() * 4 + 8); + auto size_data_per_element = static_cast(index_.graph_degree() * sizeof(IdxT) + 4 + + index_.dim() * sizeof(T) + 8); os.write(reinterpret_cast(&size_data_per_element), sizeof(std::size_t)); // label_offset std::size_t label_offset = size_data_per_element - 8; From 006e77c828cd4d3b05b933812a62ef4f629a3c13 Mon Sep 17 00:00:00 2001 From: divyegala Date: Wed, 2 Oct 2024 22:13:32 -0700 Subject: [PATCH 14/18] add some stream syncs in nn_descent --- cpp/src/neighbors/detail/nn_descent.cuh | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/cpp/src/neighbors/detail/nn_descent.cuh b/cpp/src/neighbors/detail/nn_descent.cuh index 8c5767c50..d416bc686 100644 --- a/cpp/src/neighbors/detail/nn_descent.cuh +++ b/cpp/src/neighbors/detail/nn_descent.cuh @@ -1184,6 +1184,7 @@ void GNND::add_reverse_edges(Index_t* graph_ptr, graph_ptr, d_rev_graph_ptr, NUM_SAMPLES, list_sizes); raft::copy( h_rev_graph_ptr, d_rev_graph_ptr, nrow_ * NUM_SAMPLES, raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); } template @@ -1320,11 +1321,11 @@ void GNND::build(Data_t* data, const Index_t nrow, Index_t* out graph_buffer_.data_handle(), nrow_ * DEGREE_ON_DEVICE, raft::resource::get_cuda_stream(res)); - raft::resource::sync_stream(res); raft::copy(thrust::raw_pointer_cast(dists_host_buffer_.data()), dists_buffer_.data_handle(), nrow_ * DEGREE_ON_DEVICE, raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); graph_.sample_graph_new(thrust::raw_pointer_cast(graph_host_buffer_.data()), DEGREE_ON_DEVICE); } From 8d4d1a21eae37f8f54f5777a78845723bd5e07e4 Mon Sep 17 00:00:00 2001 From: divyegala Date: Wed, 2 Oct 2024 23:00:56 -0700 Subject: [PATCH 15/18] add more syncs, use thrust_policy --- cpp/src/neighbors/detail/nn_descent.cuh | 19 ++++++++++++++----- 1 file changed, 14 insertions(+), 5 deletions(-) diff --git a/cpp/src/neighbors/detail/nn_descent.cuh b/cpp/src/neighbors/detail/nn_descent.cuh index d416bc686..6a2e5b6a2 100644 --- a/cpp/src/neighbors/detail/nn_descent.cuh +++ b/cpp/src/neighbors/detail/nn_descent.cuh @@ -24,6 +24,7 @@ #include #include #include +#include #include #include // raft::util::arch::SM_* @@ -1162,15 +1163,19 @@ GNND::GNND(raft::resources const& res, const BuildConfig& build { static_assert(NUM_SAMPLES <= 32); - thrust::fill(thrust::device, + thrust::fill(raft::resource::get_thrust_policy(res), dists_buffer_.data_handle(), dists_buffer_.data_handle() + dists_buffer_.size(), std::numeric_limits::max()); - thrust::fill(thrust::device, + thrust::fill(raft::resource::get_thrust_policy(res), reinterpret_cast(graph_buffer_.data_handle()), reinterpret_cast(graph_buffer_.data_handle()) + graph_buffer_.size(), std::numeric_limits::max()); - thrust::fill(thrust::device, d_locks_.data_handle(), d_locks_.data_handle() + d_locks_.size(), 0); + thrust::fill(raft::resource::get_thrust_policy(res), + d_locks_.data_handle(), + d_locks_.data_handle() + d_locks_.size(), + 0); + raft::resource::sync_stream(res); }; template @@ -1190,7 +1195,7 @@ void GNND::add_reverse_edges(Index_t* graph_ptr, template void GNND::local_join(cudaStream_t stream) { - thrust::fill(thrust::device.on(stream), + thrust::fill(raft::resource::get_thrust_policy(res), dists_buffer_.data_handle(), dists_buffer_.data_handle() + dists_buffer_.size(), std::numeric_limits::max()); @@ -1209,6 +1214,7 @@ void GNND::local_join(cudaStream_t stream) DEGREE_ON_DEVICE, d_locks_.data_handle(), l2_norms_.data_handle()); + raft::resource::sync_stream(res); } template @@ -1240,10 +1246,11 @@ void GNND::build(Data_t* data, const Index_t nrow, Index_t* out batch.offset()); } - thrust::fill(thrust::device.on(stream), + thrust::fill(raft::resource::get_thrust_policy(res), (Index_t*)graph_buffer_.data_handle(), (Index_t*)graph_buffer_.data_handle() + graph_buffer_.size(), std::numeric_limits::max()); + raft::resource::sync_stream(res); graph_.clear(); graph_.init_random_graph(); @@ -1330,6 +1337,7 @@ void GNND::build(Data_t* data, const Index_t nrow, Index_t* out graph_.sample_graph_new(thrust::raw_pointer_cast(graph_host_buffer_.data()), DEGREE_ON_DEVICE); } + raft::resource::sync_stream(res); graph_.update_graph(thrust::raw_pointer_cast(graph_host_buffer_.data()), thrust::raw_pointer_cast(dists_host_buffer_.data()), DEGREE_ON_DEVICE, @@ -1415,6 +1423,7 @@ void build(raft::resources const& res, GNND nnd(res, build_config); nnd.build(dataset.data_handle(), dataset.extent(0), int_graph.data_handle()); + raft::resource::sync_stream(res); #pragma omp parallel for for (size_t i = 0; i < static_cast(dataset.extent(0)); i++) { From 0409d1262f6821779c165ae57d47d341af921149 Mon Sep 17 00:00:00 2001 From: divyegala Date: Thu, 3 Oct 2024 10:25:45 -0700 Subject: [PATCH 16/18] Revert "add some stream syncs in nn_descent" This reverts commit 006e77c828cd4d3b05b933812a62ef4f629a3c13. --- cpp/src/neighbors/detail/nn_descent.cuh | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/cpp/src/neighbors/detail/nn_descent.cuh b/cpp/src/neighbors/detail/nn_descent.cuh index 6a2e5b6a2..1b445488c 100644 --- a/cpp/src/neighbors/detail/nn_descent.cuh +++ b/cpp/src/neighbors/detail/nn_descent.cuh @@ -1189,7 +1189,6 @@ void GNND::add_reverse_edges(Index_t* graph_ptr, graph_ptr, d_rev_graph_ptr, NUM_SAMPLES, list_sizes); raft::copy( h_rev_graph_ptr, d_rev_graph_ptr, nrow_ * NUM_SAMPLES, raft::resource::get_cuda_stream(res)); - raft::resource::sync_stream(res); } template @@ -1328,11 +1327,11 @@ void GNND::build(Data_t* data, const Index_t nrow, Index_t* out graph_buffer_.data_handle(), nrow_ * DEGREE_ON_DEVICE, raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); raft::copy(thrust::raw_pointer_cast(dists_host_buffer_.data()), dists_buffer_.data_handle(), nrow_ * DEGREE_ON_DEVICE, raft::resource::get_cuda_stream(res)); - raft::resource::sync_stream(res); graph_.sample_graph_new(thrust::raw_pointer_cast(graph_host_buffer_.data()), DEGREE_ON_DEVICE); } From 366af068ff9b63e981bdf994bad842baba131a51 Mon Sep 17 00:00:00 2001 From: divyegala Date: Thu, 3 Oct 2024 10:25:56 -0700 Subject: [PATCH 17/18] Revert "add more syncs, use thrust_policy" This reverts commit 8d4d1a21eae37f8f54f5777a78845723bd5e07e4. --- cpp/src/neighbors/detail/nn_descent.cuh | 19 +++++-------------- 1 file changed, 5 insertions(+), 14 deletions(-) diff --git a/cpp/src/neighbors/detail/nn_descent.cuh b/cpp/src/neighbors/detail/nn_descent.cuh index 1b445488c..8c5767c50 100644 --- a/cpp/src/neighbors/detail/nn_descent.cuh +++ b/cpp/src/neighbors/detail/nn_descent.cuh @@ -24,7 +24,6 @@ #include #include #include -#include #include #include // raft::util::arch::SM_* @@ -1163,19 +1162,15 @@ GNND::GNND(raft::resources const& res, const BuildConfig& build { static_assert(NUM_SAMPLES <= 32); - thrust::fill(raft::resource::get_thrust_policy(res), + thrust::fill(thrust::device, dists_buffer_.data_handle(), dists_buffer_.data_handle() + dists_buffer_.size(), std::numeric_limits::max()); - thrust::fill(raft::resource::get_thrust_policy(res), + thrust::fill(thrust::device, reinterpret_cast(graph_buffer_.data_handle()), reinterpret_cast(graph_buffer_.data_handle()) + graph_buffer_.size(), std::numeric_limits::max()); - thrust::fill(raft::resource::get_thrust_policy(res), - d_locks_.data_handle(), - d_locks_.data_handle() + d_locks_.size(), - 0); - raft::resource::sync_stream(res); + thrust::fill(thrust::device, d_locks_.data_handle(), d_locks_.data_handle() + d_locks_.size(), 0); }; template @@ -1194,7 +1189,7 @@ void GNND::add_reverse_edges(Index_t* graph_ptr, template void GNND::local_join(cudaStream_t stream) { - thrust::fill(raft::resource::get_thrust_policy(res), + thrust::fill(thrust::device.on(stream), dists_buffer_.data_handle(), dists_buffer_.data_handle() + dists_buffer_.size(), std::numeric_limits::max()); @@ -1213,7 +1208,6 @@ void GNND::local_join(cudaStream_t stream) DEGREE_ON_DEVICE, d_locks_.data_handle(), l2_norms_.data_handle()); - raft::resource::sync_stream(res); } template @@ -1245,11 +1239,10 @@ void GNND::build(Data_t* data, const Index_t nrow, Index_t* out batch.offset()); } - thrust::fill(raft::resource::get_thrust_policy(res), + thrust::fill(thrust::device.on(stream), (Index_t*)graph_buffer_.data_handle(), (Index_t*)graph_buffer_.data_handle() + graph_buffer_.size(), std::numeric_limits::max()); - raft::resource::sync_stream(res); graph_.clear(); graph_.init_random_graph(); @@ -1336,7 +1329,6 @@ void GNND::build(Data_t* data, const Index_t nrow, Index_t* out graph_.sample_graph_new(thrust::raw_pointer_cast(graph_host_buffer_.data()), DEGREE_ON_DEVICE); } - raft::resource::sync_stream(res); graph_.update_graph(thrust::raw_pointer_cast(graph_host_buffer_.data()), thrust::raw_pointer_cast(dists_host_buffer_.data()), DEGREE_ON_DEVICE, @@ -1422,7 +1414,6 @@ void build(raft::resources const& res, GNND nnd(res, build_config); nnd.build(dataset.data_handle(), dataset.extent(0), int_graph.data_handle()); - raft::resource::sync_stream(res); #pragma omp parallel for for (size_t i = 0; i < static_cast(dataset.extent(0)); i++) { From ad40942906785f7891b7199caac5b6a76027e1a6 Mon Sep 17 00:00:00 2001 From: divyegala Date: Thu, 3 Oct 2024 10:26:21 -0700 Subject: [PATCH 18/18] 1000 rows in test --- python/cuvs/cuvs/test/test_hnsw.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/cuvs/cuvs/test/test_hnsw.py b/python/cuvs/cuvs/test/test_hnsw.py index 8bd2e8b76..0ae97266b 100644 --- a/python/cuvs/cuvs/test/test_hnsw.py +++ b/python/cuvs/cuvs/test/test_hnsw.py @@ -23,7 +23,7 @@ def run_hnsw_build_search_test( - n_rows=10000, + n_rows=1000, n_cols=10, n_queries=100, k=10,