diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index ab883ef77..5060f4591 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -56,6 +56,7 @@ option(BUILD_TESTS "Build cuvs unit-tests" ON) option(BUILD_C_LIBRARY "Build raft C API library" OFF) option(BUILD_C_TESTS "Build raft C API tests" OFF) option(BUILD_ANN_BENCH "Build cuVS ann benchmarks" OFF) +option(BUILD_CAGRA_HNSWLIB "Build CAGRA+hnswlib interface" ON) option(CUDA_ENABLE_KERNELINFO "Enable kernel resource usage info" OFF) option(CUDA_ENABLE_LINEINFO "Enable the -lineinfo option for nvcc (useful for cuda-memcheck / profiler)" OFF @@ -192,6 +193,10 @@ if(BUILD_ANN_BENCH) rapids_cpm_gbench(BUILD_STATIC) endif() +if(BUILD_CAGRA_HNSWLIB) + include(cmake/thirdparty/get_hnswlib.cmake) +endif() + # ################################################################################################## # * cuvs --------------------------------------------------------------------- @@ -399,6 +404,7 @@ add_library( src/neighbors/detail/cagra/search_single_cta_half_uint64_dim256_t16.cu src/neighbors/detail/cagra/search_single_cta_half_uint64_dim512_t32.cu src/neighbors/detail/cagra/search_single_cta_half_uint64_dim1024_t32.cu + $<$:src/neighbors/hnsw.cpp> src/neighbors/ivf_flat_index.cpp src/neighbors/ivf_flat/ivf_flat_build_extend_float_int64_t.cu src/neighbors/ivf_flat/ivf_flat_build_extend_int8_t_int64_t.cu @@ -492,6 +498,11 @@ if(NOT BUILD_CPU_ONLY) ) endif() +if(BUILD_CAGRA_HNSWLIB) + target_link_libraries(cuvs PRIVATE hnswlib::hnswlib) + target_compile_definitions(cuvs PUBLIC CUVS_BUILD_CAGRA_HNSWLIB) +endif() + # Endian detection include(TestBigEndian) test_big_endian(BIG_ENDIAN) diff --git a/cpp/include/cuvs/neighbors/hnsw.hpp b/cpp/include/cuvs/neighbors/hnsw.hpp new file mode 100644 index 000000000..86f321564 --- /dev/null +++ b/cpp/include/cuvs/neighbors/hnsw.hpp @@ -0,0 +1,386 @@ +/* + * Copyright (c) 2024, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#ifdef CUVS_BUILD_CAGRA_HNSWLIB + +#pragma once + +#include "common.hpp" + +#include + +#include "cagra.hpp" +#include + +#include + +#include +#include +#include + +namespace cuvs::neighbors::hnsw { + +/** + * @defgroup hnsw Build CAGRA index and search with hnswlib + * @{ + */ + +struct search_params : cuvs::neighbors::search_params { + int ef; // size of the candidate list + int num_threads = 0; // number of host threads to use for concurrent searches. Value of 0 + // automatically maximizes parallelism +}; + +template +struct index : cuvs::neighbors::index { + public: + /** + * @brief load a base-layer-only hnswlib index originally saved from a built CAGRA index. + * This is a virtual class and it cannot be used directly. To create an index, use the factory + * function `cuvs::neighbors::hnsw::from_cagra` from the header + * `cuvs/neighbors/hnsw.hpp` + * + * @param[in] dim dimensions of the training dataset + * @param[in] metric distance metric to search. Supported metrics ("L2Expanded", "InnerProduct") + */ + index(int dim, cuvs::distance::DistanceType metric) : dim_{dim}, metric_{metric} {} + + /** + @brief Get underlying index + */ + virtual auto get_index() const -> void const* = 0; + + auto dim() const -> int const { return dim_; } + + auto metric() const -> cuvs::distance::DistanceType { return metric_; } + + /** + @brief Set ef for search + */ + virtual void set_ef(int ef) const; + + private: + int dim_; + cuvs::distance::DistanceType metric_; +}; + +/** + * @brief Construct an hnswlib base-layer-only index from a CAGRA index + * NOTE: 1. 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. + * 2. This function is only offered as a compiled symbol in `libraft.so` + * + * @param[in] res raft resources + * @param[in] cagra_index cagra index + * + * Usage example: + * @code{.cpp} + * // Build a CAGRA index + * using namespace raft::neighbors; + * // use default index parameters + * cagra::index_params index_params; + * // create and fill the index from a [N, D] dataset + * auto index = cagra::build(res, index_params, dataset); + * + * // Load CAGRA index as base-layer-only hnswlib index + * auto hnsw_index = hnsw::from_cagra(res, index); + * @endcode + */ +std::unique_ptr> from_cagra( + raft::resources const& res, const cuvs::neighbors::cagra::index& cagra_index); + +/** + * @brief Construct an hnswlib base-layer-only index from a CAGRA index + * NOTE: 1. 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. + * 2. This function is only offered as a compiled symbol in `libraft.so` + * + * @param[in] res raft resources + * @param[in] cagra_index cagra index + * + * Usage example: + * @code{.cpp} + * // Build a CAGRA index + * using namespace raft::neighbors; + * // use default index parameters + * cagra::index_params index_params; + * // create and fill the index from a [N, D] dataset + * auto index = cagra::build(res, index_params, dataset); + * + * // Load CAGRA index as base-layer-only hnswlib index + * auto hnsw_index = hnsw::from_cagra(res, index); + * @endcode + */ +std::unique_ptr> from_cagra( + raft::resources const& res, const cuvs::neighbors::cagra::index& cagra_index); + +/** + * @brief Construct an hnswlib base-layer-only index from a CAGRA index + * NOTE: 1. 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. + * 2. This function is only offered as a compiled symbol in `libraft.so` + * + * @param[in] res raft resources + * @param[in] cagra_index cagra index + * + * Usage example: + * @code{.cpp} + * // Build a CAGRA index + * using namespace raft::neighbors; + * // use default index parameters + * cagra::index_params index_params; + * // create and fill the index from a [N, D] dataset + * auto index = cagra::build(res, index_params, dataset); + * + * // Load CAGRA index as base-layer-only hnswlib index + * auto hnsw_index = hnsw::from_cagra(res, index); + * @endcode + */ +std::unique_ptr> from_cagra( + raft::resources const& res, const cuvs::neighbors::cagra::index& cagra_index); + +/** + * @brief Search hnswlib base-layer-only index constructed from a CAGRA index + * + * @param[in] res raft resources + * @param[in] params configure the search + * @param[in] idx cagra index + * @param[in] queries a host matrix view to a row-major matrix [n_queries, index->dim()] + * @param[out] neighbors a host matrix view to the indices of the neighbors in the source dataset + * [n_queries, k] + * @param[out] distances a host matrix view to the distances to the selected neighbors [n_queries, + * k] + * + * Usage example: + * @code{.cpp} + * // Build a CAGRA index + * using namespace cuvs::neighbors; + * // use default index parameters + * cagra::index_params index_params; + * // create and fill the index from a [N, D] dataset + * auto index = cagra::build(res, index_params, dataset); + * + * // Load CAGRA index as a base-layer HNSW index using the filesystem + * auto hnsw_index = hnsw::from_cagra(res, index); + * + * // Search K nearest neighbors as an hnswlib index + * // using host threads for concurrency + * hnsw::search_params search_params; + * search_params.ef = 50 // ef >= K; + * search_params.num_threads = 10; + * auto neighbors = raft::make_host_matrix(res, n_queries, k); + * auto distances = raft::make_host_matrix(res, n_queries, k); + * hnsw::search(res, search_params, *index.get(), queries, neighbors, distances); + * @endcode + */ +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); + +/** + * @brief Search hnswlib base-layer-only index constructed from a CAGRA index + * + * @param[in] res raft resources + * @param[in] params configure the search + * @param[in] idx cagra index + * @param[in] queries a host matrix view to a row-major matrix [n_queries, index->dim()] + * @param[out] neighbors a host matrix view to the indices of the neighbors in the source dataset + * [n_queries, k] + * @param[out] distances a host matrix view to the distances to the selected neighbors [n_queries, + * k] + * + * Usage example: + * @code{.cpp} + * // Build a CAGRA index + * using namespace cuvs::neighbors; + * // use default index parameters + * cagra::index_params index_params; + * // create and fill the index from a [N, D] dataset + * auto index = cagra::build(res, index_params, dataset); + * + * // Load CAGRA index as a base-layer HNSW index using the filesystem + * auto hnsw_index = hnsw::from_cagra(res, index); + * + * // Search K nearest neighbors as an hnswlib index + * // using host threads for concurrency + * hnsw::search_params search_params; + * search_params.ef = 50 // ef >= K; + * search_params.num_threads = 10; + * auto neighbors = raft::make_host_matrix(res, n_queries, k); + * auto distances = raft::make_host_matrix(res, n_queries, k); + * hnsw::search(res, search_params, *index.get(), queries, neighbors, distances); + * @endcode + */ +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); + +/** + * @brief Search hnswlib base-layer-only index constructed from a CAGRA index + * + * @param[in] res raft resources + * @param[in] params configure the search + * @param[in] idx cagra index + * @param[in] queries a host matrix view to a row-major matrix [n_queries, index->dim()] + * @param[out] neighbors a host matrix view to the indices of the neighbors in the source dataset + * [n_queries, k] + * @param[out] distances a host matrix view to the distances to the selected neighbors [n_queries, + * k] + * + * Usage example: + * @code{.cpp} + * // Build a CAGRA index + * using namespace cuvs::neighbors; + * // use default index parameters + * cagra::index_params index_params; + * // create and fill the index from a [N, D] dataset + * auto index = cagra::build(res, index_params, dataset); + * + * // Load CAGRA index as a base-layer HNSW index using the filesystem + * auto hnsw_index = hnsw::from_cagra(res, index); + * + * // Search K nearest neighbors as an hnswlib index + * // using host threads for concurrency + * hnsw::search_params search_params; + * search_params.ef = 50 // ef >= K; + * search_params.num_threads = 10; + * auto neighbors = raft::make_host_matrix(res, n_queries, k); + * auto distances = raft::make_host_matrix(res, n_queries, k); + * hnsw::search(res, search_params, *index.get(), queries, neighbors, distances); + * @endcode + */ +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); + +/** + * @brief De-serialize a CAGRA index saved to a file as an hnsw index + * + * @param[in] res raft resources + * @param[in] filename path to the file containing the serialized CAGRA index + * @param[in] dim dimensions of the training dataset + * @param[in] metric distance metric to search. Supported metrics ("L2Expanded", "InnerProduct") + * @param[out] index hnsw index + * + * Usage example: + * @code{.cpp} + * // Build a CAGRA index + * using namespace cuvs::neighbors; + * // use default index parameters + * cagra::index_params index_params; + * // create and fill the index from a [N, D] dataset + * auto index = cagra::build(res, index_params, dataset); + * + * // save a CAGRA index to a file + * cagra::serialize(res, index, "index.bin"); + * // De-serialize a CAGRA index as a base-layer HNSW index using the filesystem + * index* hnsw_index = nullptr; + * hnsw::deserialize(res, "index.bin", index->dim(), index->metric(), &hnsw_index); + * + * // Delete index after use + * delete hnsw_index; + * @endcode + */ +void deserialize(raft::resources const& res, + const std::string& filename, + int dim, + cuvs::distance::DistanceType metric, + index** index); + +/** + * @brief De-serialize a CAGRA index saved to a file as an hnsw index + * + * @param[in] res raft resources + * @param[in] filename path to the file containing the serialized CAGRA index + * @param[in] dim dimensions of the training dataset + * @param[in] metric distance metric to search. Supported metrics ("L2Expanded", "InnerProduct") + * @param[out] index hnsw index + * + * Usage example: + * @code{.cpp} + * // Build a CAGRA index + * using namespace cuvs::neighbors; + * // use default index parameters + * cagra::index_params index_params; + * // create and fill the index from a [N, D] dataset + * auto index = cagra::build(res, index_params, dataset); + * + * // save a CAGRA index to a file + * cagra::serialize(res, index, "index.bin"); + * // De-serialize a CAGRA index as a base-layer HNSW index using the filesystem + * index* hnsw_index = nullptr; + * hnsw::deserialize(res, "index.bin", index->dim(), index->metric(), &hnsw_index); + * + * // Delete index after use + * delete hnsw_index; + * @endcode + */ +void deserialize(raft::resources const& res, + const std::string& filename, + int dim, + cuvs::distance::DistanceType metric, + index** index); + +/** + * @brief De-serialize a CAGRA index saved to a file as an hnsw index + * + * @param[in] res raft resources + * @param[in] filename path to the file containing the serialized CAGRA index + * @param[in] dim dimensions of the training dataset + * @param[in] metric distance metric to search. Supported metrics ("L2Expanded", "InnerProduct") + * @param[out] index hnsw index + * + * Usage example: + * @code{.cpp} + * // Build a CAGRA index + * using namespace cuvs::neighbors; + * // use default index parameters + * cagra::index_params index_params; + * // create and fill the index from a [N, D] dataset + * auto index = cagra::build(res, index_params, dataset); + * + * // save a CAGRA index to a file + * cagra::serialize(res, index, "index.bin"); + * // De-serialize a CAGRA index as a base-layer HNSW index using the filesystem + * index* hnsw_index = nullptr; + * hnsw::deserialize(res, "index.bin", index->dim(), index->metric(), &hnsw_index); + * + * // Delete index after use + * delete hnsw_index; + * @endcode + */ +void deserialize(raft::resources const& res, + const std::string& filename, + int dim, + cuvs::distance::DistanceType metric, + index** index); + +/**@}*/ + +} // namespace cuvs::neighbors::hnsw + +#else +#error "This header is only available if cuVS CMake option `BUILD_CAGRA_HNSWLIB=ON" +#endif \ No newline at end of file diff --git a/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh b/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh index 41329975b..24cc2a22f 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 * sizeof(data_t) + sizeof(labeltype) - auto size_data_per_element = static_cast(index_.graph_degree() * sizeof(IdxT) + 4 + - index_.dim() * sizeof(T) + 8); + // dim * 4 + sizeof(labeltype) + auto size_data_per_element = + static_cast(index_.graph_degree() * sizeof(IdxT) + 4 + 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,17 +184,17 @@ 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 = host_dataset(i, j); - os.write(reinterpret_cast(&data_elem), sizeof(T)); + 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 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)); - // } - // } 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 new file mode 100644 index 000000000..0d1ae4ec9 --- /dev/null +++ b/cpp/src/neighbors/detail/hnsw.hpp @@ -0,0 +1,183 @@ +/* + * 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 +#include + +namespace cuvs::neighbors::hnsw::detail { + +template +struct hnsw_dist_t { + using type = void; +}; + +template <> +struct hnsw_dist_t { + using type = float; +}; + +template <> +struct hnsw_dist_t { + using type = int; +}; + +template <> +struct hnsw_dist_t { + using type = int; +}; + +template +struct index_impl : index { + public: + /** + * @brief load a base-layer-only hnswlib index originally saved from a built CAGRA index + * + * @param[in] filepath path to the index + * @param[in] dim dimensions of the training dataset + * @param[in] metric distance metric to search. Supported metrics ("L2Expanded", "InnerProduct") + */ + index_impl(const std::string& filepath, int dim, cuvs::distance::DistanceType metric) + : index{dim, metric} + { + if constexpr (std::is_same_v) { + if (metric == cuvs::distance::DistanceType::L2Expanded) { + space_ = std::make_unique(dim); + } else if (metric == cuvs::distance::DistanceType::InnerProduct) { + space_ = std::make_unique(dim); + } + } else if constexpr (std::is_same_v or std::is_same_v) { + if (metric == cuvs::distance::DistanceType::L2Expanded) { + space_ = std::make_unique>(dim); + } + } + + RAFT_EXPECTS(space_ != nullptr, "Unsupported metric type was used"); + + appr_alg_ = std::make_unique::type>>( + space_.get(), filepath); + + appr_alg_->base_layer_only = true; + } + + /** + @brief Get hnswlib index + */ + auto get_index() const -> void const* override { return appr_alg_.get(); } + + /** + @brief Set ef for search + */ + void set_ef(int ef) const override { appr_alg_->ef_ = ef; } + + private: + std::unique_ptr::type>> appr_alg_; + std::unique_ptr::type>> space_; +}; + +template +std::unique_ptr> from_cagra(raft::resources const& res, + const cuvs::neighbors::cagra::index& cagra_index) +{ + std::random_device dev; + std::mt19937 rng(dev()); + std::uniform_int_distribution dist(0); + auto uuid = std::to_string(dist(rng)); + std::string filepath = "/tmp/" + uuid + ".bin"; + cuvs::neighbors::cagra::serialize_to_hnswlib(res, filepath, cagra_index); + index* hnsw_index = nullptr; + cuvs::neighbors::hnsw::deserialize( + res, filepath, cagra_index.dim(), cagra_index.metric(), &hnsw_index); + std::filesystem::remove(filepath); + return std::unique_ptr>(hnsw_index); +} + +template +void get_search_knn_results(hnswlib::HierarchicalNSW const* idx, + const QueriesT* query, + int k, + uint64_t* indices, + float* distances) +{ + auto result = idx->searchKnn(query, k); + assert(result.size() >= static_cast(k)); + + for (int i = k - 1; i >= 0; --i) { + indices[i] = result.top().second; + distances[i] = result.top().first; + result.pop(); + } +} + +template +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) +{ + RAFT_EXPECTS( + queries.extent(0) == neighbors.extent(0) && queries.extent(0) == distances.extent(0), + "Number of rows in output neighbors and distances matrices must equal the number of queries."); + + RAFT_EXPECTS(neighbors.extent(1) == distances.extent(1), + "Number of columns in output neighbors and distances matrices must equal k"); + RAFT_EXPECTS(queries.extent(1) == idx.dim(), + "Number of query dimensions should equal number of dimensions in the index."); + + idx.set_ef(params.ef); + auto const* hnswlib_index = + reinterpret_cast const*>(idx.get_index()); + + // when num_threads == 0, automatically maximize parallelism + if (params.num_threads) { +#pragma omp parallel for num_threads(params.num_threads) + for (int64_t i = 0; i < queries.extent(0); ++i) { + get_search_knn_results(hnswlib_index, + queries.data_handle() + i * queries.extent(1), + neighbors.extent(1), + neighbors.data_handle() + i * neighbors.extent(1), + distances.data_handle() + i * distances.extent(1)); + } + } else { +#pragma omp parallel for + for (int64_t i = 0; i < queries.extent(0); ++i) { + get_search_knn_results(hnswlib_index, + queries.data_handle() + i * queries.extent(1), + neighbors.extent(1), + neighbors.data_handle() + i * neighbors.extent(1), + distances.data_handle() + i * distances.extent(1)); + } + } +} + +template +void deserialize(raft::resources const& res, + const std::string& filename, + int dim, + cuvs::distance::DistanceType metric, + index** idx) +{ + *idx = new detail::index_impl(filename, dim, metric); +} + +} // namespace cuvs::neighbors::hnsw::detail diff --git a/cpp/src/neighbors/hnsw.cpp b/cpp/src/neighbors/hnsw.cpp new file mode 100644 index 000000000..36cbb16c9 --- /dev/null +++ b/cpp/src/neighbors/hnsw.cpp @@ -0,0 +1,70 @@ +/* + * Copyright (c) 2024, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "detail/hnsw.hpp" +#include +#include +#include + +namespace cuvs::neighbors::hnsw { + +#define CUVS_INST_HNSW_FROM_CAGRA(T) \ + std::unique_ptr> from_cagra( \ + raft::resources const& res, const cuvs::neighbors::cagra::index& cagra_index) \ + { \ + return detail::from_cagra(res, cagra_index); \ + } + +CUVS_INST_HNSW_FROM_CAGRA(float); +CUVS_INST_HNSW_FROM_CAGRA(uint8_t); +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); \ + } + +CUVS_INST_HNSW_SEARCH(float, float); +CUVS_INST_HNSW_SEARCH(uint8_t, int); +CUVS_INST_HNSW_SEARCH(int8_t, int); + +#undef CUVS_INST_HNSW_SEARCH + +#define CUVS_INST_HNSW_DESERIALIZE(T) \ + void deserialize(raft::resources const& res, \ + const std::string& filename, \ + int dim, \ + cuvs::distance::DistanceType metric, \ + index** idx) \ + { \ + detail::deserialize(res, filename, dim, metric, idx); \ + } + +CUVS_INST_HNSW_DESERIALIZE(float); +CUVS_INST_HNSW_DESERIALIZE(uint8_t); +CUVS_INST_HNSW_DESERIALIZE(int8_t); + +#undef CUVS_INST_HNSW_DESERIALIZE + +} // namespace cuvs::neighbors::hnsw diff --git a/cpp/test/CMakeLists.txt b/cpp/test/CMakeLists.txt index 1fae2f70b..7921fffd3 100644 --- a/cpp/test/CMakeLists.txt +++ b/cpp/test/CMakeLists.txt @@ -91,16 +91,8 @@ endfunction() if(BUILD_TESTS) ConfigureTest( - NAME - NEIGHBORS_TEST - PATH - test/neighbors/brute_force.cu - test/neighbors/brute_force_prefiltered.cu - test/neighbors/refine.cu - GPUS - 1 - PERCENT - 100 + NAME NEIGHBORS_TEST PATH test/neighbors/brute_force.cu + test/neighbors/brute_force_prefiltered.cu test/neighbors/refine.cu GPUS 1 PERCENT 100 ) ConfigureTest( @@ -160,6 +152,10 @@ if(BUILD_TESTS) 100 ) + if(BUILD_CAGRA_HNSWLIB) + ConfigureTest(NAME NEIGHBORS_HNSW_TEST PATH test/neighbors/hnsw.cu GPUS 1 PERCENT 100) + endif() + ConfigureTest( NAME DISTANCE_TEST diff --git a/cpp/test/neighbors/hnsw.cu b/cpp/test/neighbors/hnsw.cu new file mode 100644 index 000000000..9fb88be05 --- /dev/null +++ b/cpp/test/neighbors/hnsw.cu @@ -0,0 +1,187 @@ +/* + * 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 "ann_utils.cuh" + +#include +#include +#include +#include +#include + +#include "naive_knn.cuh" + +#include + +#include +#include +#include +#include + +namespace cuvs::neighbors::hnsw { + +struct AnnHNSWInputs { + int n_rows; + int dim; + int graph_degree; + cuvs::distance::DistanceType metric; + int k; + int n_queries; + int ef; + double min_recall; +}; + +inline ::std::ostream& operator<<(::std::ostream& os, const AnnHNSWInputs& p) +{ + os << "dataset shape=" << p.n_rows << "x" << p.dim << ", graph_degree=" << p.graph_degree + << ", metric=" << static_cast(p.metric) << ", ef=" << (p.ef) << std::endl; + return os; +} + +template +class AnnHNSWTest : public ::testing::TestWithParam { + public: + AnnHNSWTest() + : stream_(raft::resource::get_cuda_stream(handle_)), + ps(::testing::TestWithParam::GetParam()), + database(0, stream_), + queries(0, stream_) + { + } + + protected: + void testHNSW() + { + std::vector indices_HNSW(ps.n_queries * ps.k); + std::vector distances_HNSW(ps.n_queries * ps.k); + std::vector indices_naive(ps.n_queries * ps.k); + std::vector distances_naive(ps.n_queries * ps.k); + + std::vector queries_h(ps.n_queries * ps.dim); + raft::update_host(queries_h.data(), queries.data(), ps.n_queries * ps.dim, stream_); + + { + rmm::device_uvector distances_naive_dev(ps.n_queries * ps.k, stream_); + rmm::device_uvector indices_naive_dev(ps.n_queries * ps.k, stream_); + naive_knn(handle_, + distances_naive_dev.data(), + indices_naive_dev.data(), + queries.data(), + database.data(), + ps.n_queries, + ps.n_rows, + ps.dim, + ps.k, + ps.metric); + raft::update_host( + indices_naive.data(), indices_naive_dev.data(), ps.n_queries * ps.k, stream_); + raft::update_host( + distances_naive.data(), distances_naive_dev.data(), ps.n_queries * ps.k, stream_); + raft::resource::sync_stream(handle_); + } + + { + cuvs::neighbors::cagra::index_params index_params; + index_params.metric = ps.metric; + index_params.graph_degree = ps.graph_degree; + index_params.intermediate_graph_degree = 2 * ps.graph_degree; + + auto database_view = raft::make_device_matrix_view( + (const DataT*)database.data(), ps.n_rows, ps.dim); + + auto index = cuvs::neighbors::cagra::build(handle_, index_params, database_view); + raft::resource::sync_stream(handle_); + + cuvs::neighbors::hnsw::search_params search_params; + search_params.ef = ps.ef; + auto hnsw_index = cuvs::neighbors::hnsw::from_cagra(handle_, index); + auto queries_HNSW_view = + raft::make_host_matrix_view(queries_h.data(), ps.n_queries, ps.dim); + auto indices_HNSW_view = + raft::make_host_matrix_view(indices_HNSW.data(), ps.n_queries, ps.k); + auto distances_HNSW_view = + raft::make_host_matrix_view(distances_HNSW.data(), ps.n_queries, ps.k); + cuvs::neighbors::hnsw::search(handle_, + search_params, + *hnsw_index.get(), + queries_HNSW_view, + indices_HNSW_view, + distances_HNSW_view); + } + + double min_recall = ps.min_recall; + EXPECT_TRUE(eval_neighbours(indices_naive, + indices_HNSW, + distances_naive, + distances_HNSW, + ps.n_queries, + ps.k, + 0.006, + min_recall)); + } + + void SetUp() override + { + database.resize(((size_t)ps.n_rows) * ps.dim, stream_); + queries.resize(((size_t)ps.n_queries) * ps.dim, stream_); + raft::random::RngState r(1234ULL); + if constexpr (std::is_same{}) { + raft::random::normal(handle_, r, database.data(), ps.n_rows * ps.dim, DataT(0.1), DataT(2.0)); + raft::random::normal( + handle_, r, queries.data(), ps.n_queries * ps.dim, DataT(0.1), DataT(2.0)); + + } else { + raft::random::uniformInt( + handle_, r, database.data(), ps.n_rows * ps.dim, DataT(1), DataT(20)); + raft::random::uniformInt( + handle_, r, queries.data(), ps.n_queries * ps.dim, DataT(1), DataT(20)); + } + raft::resource::sync_stream(handle_); + } + + void TearDown() override + { + raft::resource::sync_stream(handle_); + database.resize(0, stream_); + queries.resize(0, stream_); + } + + private: + raft::resources handle_; + rmm::cuda_stream_view stream_; + AnnHNSWInputs ps; + rmm::device_uvector database; + rmm::device_uvector queries; +}; + +const std::vector inputs = raft::util::itertools::product( + {1000, 2000}, // n_rows + {5, 10, 25, 50, 100, 250, 500, 1000}, // dim + {32, 64}, // graph_degree + {cuvs::distance::DistanceType::L2Expanded}, // metric + {50}, // k + {500}, // n_queries + {200}, // ef + {0.98} // min_recall +); + +typedef AnnHNSWTest AnnHNSW_F; +TEST_P(AnnHNSW_F, AnnHNSW) { this->testHNSW(); } + +INSTANTIATE_TEST_CASE_P(AnnHNSWTest, AnnHNSW_F, ::testing::ValuesIn(inputs)); + +} // namespace cuvs::neighbors::hnsw