diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index b72d7f165..d4821a273 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -202,6 +202,7 @@ endif() add_library( cuvs-cagra-search STATIC src/neighbors/cagra_search_float.cu + src/neighbors/cagra_search_half.cu src/neighbors/cagra_search_int8.cu src/neighbors/cagra_search_uint8.cu src/neighbors/detail/cagra/compute_distance.cu @@ -257,14 +258,10 @@ add_library( src/neighbors/detail/cagra/search_multi_cta_half_uint32.cu src/neighbors/detail/cagra/search_multi_cta_int8_uint32.cu src/neighbors/detail/cagra/search_multi_cta_uint8_uint32.cu - src/neighbors/detail/cagra/search_multi_cta_float_uint64.cu - src/neighbors/detail/cagra/search_multi_cta_half_uint64.cu src/neighbors/detail/cagra/search_single_cta_float_uint32.cu src/neighbors/detail/cagra/search_single_cta_half_uint32.cu src/neighbors/detail/cagra/search_single_cta_int8_uint32.cu src/neighbors/detail/cagra/search_single_cta_uint8_uint32.cu - src/neighbors/detail/cagra/search_single_cta_float_uint64.cu - src/neighbors/detail/cagra/search_single_cta_half_uint64.cu ) file(GLOB_RECURSE compute_distance_sources "src/neighbors/detail/cagra/compute_distance_*.cu") @@ -350,6 +347,7 @@ add_library( src/distance/pairwise_distance.cu src/neighbors/brute_force.cu src/neighbors/cagra_build_float.cu + src/neighbors/cagra_build_half.cu src/neighbors/cagra_build_int8.cu src/neighbors/cagra_build_uint8.cu src/neighbors/cagra_extend_float.cu @@ -357,6 +355,7 @@ add_library( src/neighbors/cagra_extend_uint8.cu src/neighbors/cagra_optimize.cu src/neighbors/cagra_serialize_float.cu + src/neighbors/cagra_serialize_half.cu src/neighbors/cagra_serialize_int8.cu src/neighbors/cagra_serialize_uint8.cu src/neighbors/detail/cagra/cagra_build.cpp @@ -378,6 +377,7 @@ add_library( src/neighbors/ivf_pq/ivf_pq_serialize.cu src/neighbors/ivf_pq/ivf_pq_deserialize.cu src/neighbors/ivf_pq/detail/ivf_pq_build_extend_float_int64_t.cu + src/neighbors/ivf_pq/detail/ivf_pq_build_extend_half_int64_t.cu src/neighbors/ivf_pq/detail/ivf_pq_build_extend_int8_t_int64_t.cu src/neighbors/ivf_pq/detail/ivf_pq_build_extend_uint8_t_int64_t.cu src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_half_fp8_false.cu @@ -395,13 +395,16 @@ add_library( src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_float_fp8_false_bitset64.cu src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_float_fp8_true_bitset64.cu src/neighbors/ivf_pq/detail/ivf_pq_search_float_int64_t.cu + src/neighbors/ivf_pq/detail/ivf_pq_search_half_int64_t.cu src/neighbors/ivf_pq/detail/ivf_pq_search_int8_t_int64_t.cu src/neighbors/ivf_pq/detail/ivf_pq_search_uint8_t_int64_t.cu src/neighbors/ivf_pq/detail/ivf_pq_search_with_filter_float_int64_t.cu + src/neighbors/ivf_pq/detail/ivf_pq_search_with_filter_half_int64_t.cu src/neighbors/ivf_pq/detail/ivf_pq_search_with_filter_int8_t_int64_t.cu src/neighbors/ivf_pq/detail/ivf_pq_search_with_filter_uint8_t_int64_t.cu src/neighbors/nn_descent.cu src/neighbors/nn_descent_float.cu + src/neighbors/nn_descent_half.cu src/neighbors/nn_descent_int8.cu src/neighbors/nn_descent_uint8.cu src/neighbors/refine/detail/refine_device_float_float.cu diff --git a/cpp/bench/ann/src/common/benchmark.hpp b/cpp/bench/ann/src/common/benchmark.hpp index 5d7b8934f..db3e533e0 100644 --- a/cpp/bench/ann/src/common/benchmark.hpp +++ b/cpp/bench/ann/src/common/benchmark.hpp @@ -687,17 +687,17 @@ inline auto run_main(int argc, char** argv) -> int override_kv, metric_objective, threads); - // } else if (dtype == "half") { - // dispatch_benchmark(cmdline - // conf, - // force_overwrite, - // build_mode, - // search_mode, - // data_prefix, - // index_prefix, - // override_kv, - // metric_objective, - // threads); + } else if (dtype == "half") { + dispatch_benchmark(cmdline, + conf, + force_overwrite, + build_mode, + search_mode, + data_prefix, + index_prefix, + override_kv, + metric_objective, + threads); } else if (dtype == "uint8") { dispatch_benchmark(cmdline, conf, diff --git a/cpp/bench/ann/src/cuvs/cuvs_benchmark.cu b/cpp/bench/ann/src/cuvs/cuvs_benchmark.cu index a7495c23a..a956ab139 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_benchmark.cu +++ b/cpp/bench/ann/src/cuvs/cuvs_benchmark.cu @@ -121,7 +121,7 @@ auto create_search_param(const std::string& algo_name, const nlohmann::json& con }; // namespace cuvs::bench REGISTER_ALGO_INSTANCE(float); -// REGISTER_ALGO_INSTANCE(half); +REGISTER_ALGO_INSTANCE(half); REGISTER_ALGO_INSTANCE(std::int8_t); REGISTER_ALGO_INSTANCE(std::uint8_t); diff --git a/cpp/bench/ann/src/cuvs/cuvs_cagra_half.cu b/cpp/bench/ann/src/cuvs/cuvs_cagra_half.cu index 6768034a2..b4a3235c4 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_cagra_half.cu +++ b/cpp/bench/ann/src/cuvs/cuvs_cagra_half.cu @@ -16,5 +16,5 @@ #include "cuvs_cagra_wrapper.h" namespace cuvs::bench { -// template class cuvs_cagra; +template class cuvs_cagra; } // namespace cuvs::bench diff --git a/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h index 9ca41cab0..ff854f890 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h +++ b/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h @@ -289,7 +289,11 @@ void cuvs_cagra::save(const std::string& file) const template void cuvs_cagra::save_to_hnswlib(const std::string& file) const { - cuvs::neighbors::cagra::serialize_to_hnswlib(handle_, file, *index_); + if constexpr (!std::is_same_v) { + cuvs::neighbors::cagra::serialize_to_hnswlib(handle_, file, *index_); + } else { + RAFT_FAIL("Cannot save fp16 index to hnswlib format"); + } } template diff --git a/cpp/bench/ann/src/cuvs/cuvs_ivf_pq.cu b/cpp/bench/ann/src/cuvs/cuvs_ivf_pq.cu index 3ffdd4a25..2df460966 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_ivf_pq.cu +++ b/cpp/bench/ann/src/cuvs/cuvs_ivf_pq.cu @@ -17,7 +17,7 @@ namespace cuvs::bench { template class cuvs_ivf_pq; -// template class cuvs_ivf_pq; +template class cuvs_ivf_pq; template class cuvs_ivf_pq; template class cuvs_ivf_pq; } // namespace cuvs::bench diff --git a/cpp/include/cuvs/neighbors/cagra.hpp b/cpp/include/cuvs/neighbors/cagra.hpp index fec95b563..20db7e8b7 100644 --- a/cpp/include/cuvs/neighbors/cagra.hpp +++ b/cpp/include/cuvs/neighbors/cagra.hpp @@ -613,6 +613,78 @@ auto build(raft::resources const& res, * * @return the constructed cagra index */ +auto build(raft::resources const& res, + const cuvs::neighbors::cagra::index_params& params, + raft::device_matrix_view dataset) + -> cuvs::neighbors::cagra::index; + +/** + * @brief Build the index from the dataset for efficient search. + * + * The build consist of two steps: build an intermediate knn-graph, and optimize it to + * create the final graph. The index_params struct controls the node degree of these + * graphs. + * + * The following distance metrics are supported: + * - L2 + * + * Usage example: + * @code{.cpp} + * 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); + * // use default search parameters + * cagra::search_params search_params; + * // search K nearest neighbours + * auto neighbors = raft::make_device_matrix(res, n_queries, k); + * auto distances = raft::make_device_matrix(res, n_queries, k); + * cagra::search(res, search_params, index, queries, neighbors, distances); + * @endcode + * + * @param[in] res + * @param[in] params parameters for building the index + * @param[in] dataset a matrix view (host) to a row-major matrix [n_rows, dim] + * + * @return the constructed cagra index + */ +auto build(raft::resources const& res, + const cuvs::neighbors::cagra::index_params& params, + raft::host_matrix_view dataset) + -> cuvs::neighbors::cagra::index; + +/** + * @brief Build the index from the dataset for efficient search. + * + * The build consist of two steps: build an intermediate knn-graph, and optimize it to + * create the final graph. The index_params struct controls the node degree of these + * graphs. + * + * The following distance metrics are supported: + * - L2 + * + * Usage example: + * @code{.cpp} + * 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); + * // use default search parameters + * cagra::search_params search_params; + * // search K nearest neighbours + * auto neighbors = raft::make_device_matrix(res, n_queries, k); + * auto distances = raft::make_device_matrix(res, n_queries, k); + * cagra::search(res, search_params, index, queries, neighbors, distances); + * @endcode + * + * @param[in] res + * @param[in] params parameters for building the index + * @param[in] dataset a matrix view (device) to a row-major matrix [n_rows, dim] + * + * @return the constructed cagra index + */ auto build(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, raft::device_matrix_view dataset) @@ -975,9 +1047,6 @@ void extend( * * See the [cagra::build](#cagra::build) documentation for a usage example. * - * @tparam T data element type - * @tparam IdxT type of the indices - * * @param[in] res raft resources * @param[in] params configure the search * @param[in] idx cagra index @@ -1000,8 +1069,26 @@ void search(raft::resources const& res, * * See the [cagra::build](#cagra::build) documentation for a usage example. * - * @tparam T data element type - * @tparam IdxT type of the indices + * @param[in] res raft resources + * @param[in] params configure the search + * @param[in] index cagra index + * @param[in] queries a device matrix view to a row-major matrix [n_queries, index->dim()] + * @param[out] neighbors a device matrix view to the indices of the neighbors in the source dataset + * [n_queries, k] + * @param[out] distances a device matrix view to the distances to the selected neighbors [n_queries, + * k] + */ +void search(raft::resources const& res, + cuvs::neighbors::cagra::search_params const& params, + const cuvs::neighbors::cagra::index& index, + raft::device_matrix_view queries, + raft::device_matrix_view neighbors, + raft::device_matrix_view distances); + +/** + * @brief Search ANN using the constructed index. + * + * See the [cagra::build](#cagra::build) documentation for a usage example. * * @param[in] res raft resources * @param[in] params configure the search @@ -1024,9 +1111,6 @@ void search(raft::resources const& res, * * See the [cagra::build](#cagra::build) documentation for a usage example. * - * @tparam T data element type - * @tparam IdxT type of the indices - * * @param[in] res raft resources * @param[in] params configure the search * @param[in] index cagra index @@ -1156,6 +1240,111 @@ void serialize(raft::resources const& handle, void deserialize(raft::resources const& handle, std::istream& is, cuvs::neighbors::cagra::index* index); +/** + * Save the index to file. + * + * Experimental, both the API and the serialization format are subject to change. + * + * @code{.cpp} + * #include + * #include + * + * raft::resources handle; + * + * // create a string with a filepath + * std::string filename("/path/to/index"); + * // create an index with `auto index = cuvs::neighbors::cagra::build(...);` + * cuvs::neighbors::cagra::serialize(handle, filename, index); + * @endcode + * + * @param[in] handle the raft handle + * @param[in] filename the file name for saving the index + * @param[in] index CAGRA index + * @param[in] include_dataset Whether or not to write out the dataset to the file. + * + */ +void serialize(raft::resources const& handle, + const std::string& filename, + const cuvs::neighbors::cagra::index& index, + bool include_dataset = true); + +/** + * Load index from file. + * + * Experimental, both the API and the serialization format are subject to change. + * + * @code{.cpp} + * #include + * #include + * + * raft::resources handle; + * + * // create a string with a filepath + * std::string filename("/path/to/index"); + + * cuvs::neighbors::cagra::index index; + * cuvs::neighbors::cagra::deserialize(handle, filename, &index); + * @endcode + * + * @param[in] handle the raft handle + * @param[in] filename the name of the file that stores the index + * @param[out] index the cagra index + */ +void deserialize(raft::resources const& handle, + const std::string& filename, + cuvs::neighbors::cagra::index* index); + +/** + * Write the index to an output stream + * + * Experimental, both the API and the serialization format are subject to change. + * + * @code{.cpp} + * #include + * #include + * + * raft::resources handle; + * + * // create an output stream + * std::ostream os(std::cout.rdbuf()); + * // create an index with `auto index = cuvs::neighbors::cagra::build(...);` + * cuvs::neighbors::cagra::serialize(handle, os, index); + * @endcode + * + * @param[in] handle the raft handle + * @param[in] os output stream + * @param[in] index CAGRA index + * @param[in] include_dataset Whether or not to write out the dataset to the file. + */ +void serialize(raft::resources const& handle, + std::ostream& os, + const cuvs::neighbors::cagra::index& index, + bool include_dataset = true); + +/** + * Load index from input stream + * + * Experimental, both the API and the serialization format are subject to change. + * + * @code{.cpp} + * #include + * #include + * + * raft::resources handle; + * + * // create an input stream + * std::istream is(std::cin.rdbuf()); + * cuvs::neighbors::cagra::index index; + * cuvs::neighbors::cagra::deserialize(handle, is, &index); + * @endcode + * + * @param[in] handle the raft handle + * @param[in] is input stream + * @param[out] index the cagra index + */ +void deserialize(raft::resources const& handle, + std::istream& is, + cuvs::neighbors::cagra::index* index); /** * Save the index to file. diff --git a/cpp/include/cuvs/neighbors/ivf_pq.hpp b/cpp/include/cuvs/neighbors/ivf_pq.hpp index b2db96686..8c378b1f0 100644 --- a/cpp/include/cuvs/neighbors/ivf_pq.hpp +++ b/cpp/include/cuvs/neighbors/ivf_pq.hpp @@ -16,6 +16,8 @@ #pragma once +#include + #include #include @@ -547,6 +549,52 @@ void build(raft::resources const& handle, raft::device_matrix_view dataset, cuvs::neighbors::ivf_pq::index* idx); +/** + * @brief Build the index from the dataset for efficient search. + * + * Usage example: + * @code{.cpp} + * using namespace cuvs::neighbors; + * // use default index parameters + * ivf_pq::index_params index_params; + * // create and fill the index from a [N, D] dataset + * auto index = ivf_pq::build(handle, index_params, dataset); + * @endcode + * + * @param[in] handle + * @param[in] index_params configure the index building + * @param[in] dataset a device matrix view to a row-major matrix [n_rows, dim] + * + * @return the constructed ivf-pq index + */ +auto build(raft::resources const& handle, + const cuvs::neighbors::ivf_pq::index_params& index_params, + raft::device_matrix_view dataset) + -> cuvs::neighbors::ivf_pq::index; + +/** + * @brief Build the index from the dataset for efficient search. + * + * Usage example: + * @code{.cpp} + * using namespace cuvs::neighbors; + * // use default index parameters + * ivf_pq::index_params index_params; + * // create and fill the index from a [N, D] dataset + * ivf_pq::index index; + * ivf_pq::build(handle, index_params, dataset, index); + * @endcode + * + * @param[in] handle + * @param[in] index_params configure the index building + * @param[in] dataset raft::device_matrix_view to a row-major matrix [n_rows, dim] + * @param[out] idx reference to ivf_pq::index + * + */ +void build(raft::resources const& handle, + const cuvs::neighbors::ivf_pq::index_params& index_params, + raft::device_matrix_view dataset, + cuvs::neighbors::ivf_pq::index* idx); /** * @brief Build the index from the dataset for efficient search. * @@ -726,6 +774,53 @@ void build(raft::resources const& handle, * * @return the constructed ivf-pq index */ +auto build(raft::resources const& handle, + const cuvs::neighbors::ivf_pq::index_params& index_params, + raft::host_matrix_view dataset) + -> cuvs::neighbors::ivf_pq::index; + +/** + * @brief Build the index from the dataset for efficient search. + * + * Usage example: + * @code{.cpp} + * using namespace cuvs::neighbors; + * // use default index parameters + * ivf_pq::index_params index_params; + * // create and fill the index from a [N, D] dataset + * ivf_pq::index index; + * ivf_pq::build(handle, index_params, dataset, index); + * @endcode + * + * @param[in] handle + * @param[in] index_params configure the index building + * @param[in] dataset raft::host_matrix_view to a row-major matrix [n_rows, dim] + * @param[out] idx reference to ivf_pq::index + * + */ +void build(raft::resources const& handle, + const cuvs::neighbors::ivf_pq::index_params& index_params, + raft::host_matrix_view dataset, + cuvs::neighbors::ivf_pq::index* idx); + +/** + * @brief Build the index from the dataset for efficient search. + * + * Usage example: + * @code{.cpp} + * using namespace cuvs::neighbors; + * // use default index parameters + * ivf_pq::index_params index_params; + * // create and fill the index from a [N, D] dataset + * auto index = ivf_pq::build(handle, index_params, dataset); + * @endcode + * + * @param[in] handle + * @param[in] index_params configure the index building + * @param[in] dataset a host_matrix_view to a row-major matrix [n_rows, dim] + * + * @return the constructed ivf-pq index + */ auto build(raft::resources const& handle, const cuvs::neighbors::ivf_pq::index_params& index_params, raft::host_matrix_view dataset) @@ -887,6 +982,62 @@ void extend(raft::resources const& handle, std::optional> new_indices, cuvs::neighbors::ivf_pq::index* idx); +/** + * @brief Extend the index with the new data. + * + * Usage example: + * @code{.cpp} + * using namespace cuvs::neighbors; + * ivf_pq::index_params index_params; + * index_params.add_data_on_build = false; // don't populate index on build + * index_params.kmeans_trainset_fraction = 1.0; // use whole dataset for kmeans training + * // train the index from a [N, D] dataset + * auto index_empty = ivf_pq::build(handle, index_params, dataset); + * // fill the index with the data + * std::optional> no_op = std::nullopt; + * auto index = ivf_pq::extend(handle, new_vectors, no_op, index_empty); + * @endcode + * + * @param[in] handle + * @param[in] new_vectors a device matrix view to a row-major matrix [n_rows, idx.dim()] + * @param[in] new_indices a device vector view to a vector of indices [n_rows]. + * If the original index is empty (`idx.size() == 0`), you can pass `std::nullopt` + * here to imply a continuous range `[0...n_rows)`. + * @param[inout] idx + */ +auto extend(raft::resources const& handle, + raft::device_matrix_view new_vectors, + std::optional> new_indices, + const cuvs::neighbors::ivf_pq::index& idx) + -> cuvs::neighbors::ivf_pq::index; + +/** + * @brief Extend the index with the new data. + * + * Usage example: + * @code{.cpp} + * using namespace cuvs::neighbors; + * ivf_pq::index_params index_params; + * index_params.add_data_on_build = false; // don't populate index on build + * index_params.kmeans_trainset_fraction = 1.0; // use whole dataset for kmeans training + * // train the index from a [N, D] dataset + * auto index_empty = ivf_pq::build(handle, index_params, dataset); + * // fill the index with the data + * std::optional> no_op = std::nullopt; + * ivf_pq::extend(handle, new_vectors, no_op, &index_empty); + * @endcode + * + * @param[in] handle + * @param[in] new_vectors a device matrix view to a row-major matrix [n_rows, idx.dim()] + * @param[in] new_indices a device vector view to a vector of indices [n_rows]. + * If the original index is empty (`idx.size() == 0`), you can pass `std::nullopt` + * here to imply a continuous range `[0...n_rows)`. + * @param[inout] idx + */ +void extend(raft::resources const& handle, + raft::device_matrix_view new_vectors, + std::optional> new_indices, + cuvs::neighbors::ivf_pq::index* idx); /** * @brief Extend the index with the new data. * @@ -1257,6 +1408,47 @@ void search(raft::resources const& handle, raft::device_matrix_view neighbors, raft::device_matrix_view distances); +/** + * @brief Search ANN using the constructed index. + * + * See the [ivf_pq::build](#ivf_pq::build) documentation for a usage example. + * + * Note, this function requires a temporary buffer to store intermediate results between cuda kernel + * calls, which may lead to undesirable allocations and slowdown. To alleviate the problem, you can + * pass a pool memory resource or a large enough pre-allocated memory resource to reduce or + * eliminate entirely allocations happening within `search`. + * The exact size of the temporary buffer depends on multiple factors and is an implementation + * detail. However, you can safely specify a small initial size for the memory pool, so that only a + * few allocations happen to grow it during the first invocations of the `search`. + * + * @code{.cpp} + * ... + * // use default search parameters + * ivf_pq::search_params search_params; + * // Use the same allocator across multiple searches to reduce the number of + * // cuda memory allocations + * ivf_pq::search(handle, search_params, index, queries1, out_inds1, out_dists1); + * ivf_pq::search(handle, search_params, index, queries2, out_inds2, out_dists2); + * ivf_pq::search(handle, search_params, index, queries3, out_inds3, out_dists3); + * ... + * @endcode + * + * @param[in] handle + * @param[in] search_params configure the search + * @param[in] index ivf-pq constructed index + * @param[in] queries a device matrix view to a row-major matrix [n_queries, index->dim()] + * @param[out] neighbors a device matrix view to the indices of the neighbors in the source dataset + * [n_queries, k] + * @param[out] distances a device matrix view to the distances to the selected neighbors [n_queries, + * k] + */ +void search(raft::resources const& handle, + const cuvs::neighbors::ivf_pq::search_params& search_params, + cuvs::neighbors::ivf_pq::index& index, + raft::device_matrix_view queries, + raft::device_matrix_view neighbors, + raft::device_matrix_view distances); + /** * @brief Search ANN using the constructed index. * @@ -1372,6 +1564,39 @@ void search_with_filtering( raft::device_matrix_view distances, cuvs::neighbors::filtering::bitset_filter sample_filter); +/** + * @brief Search ANN using the constructed index with the given filter. + * + * See the [ivf_pq::build](#ivf_pq::build) documentation for a usage example. + * + * Note, this function requires a temporary buffer to store intermediate results between cuda kernel + * calls, which may lead to undesirable allocations and slowdown. To alleviate the problem, you can + * pass a pool memory resource or a large enough pre-allocated memory resource to reduce or + * eliminate entirely allocations happening within `search`. + * The exact size of the temporary buffer depends on multiple factors and is an implementation + * detail. However, you can safely specify a small initial size for the memory pool, so that only a + * few allocations happen to grow it during the first invocations of the `search`. + * + * @param[in] handle + * @param[in] params configure the search + * @param[in] idx ivf-pq constructed index + * @param[in] queries a device matrix view to a row-major matrix [n_queries, index->dim()] + * @param[out] neighbors a device matrix view to the indices of the neighbors in the source dataset + * [n_queries, k] + * @param[out] distances a device matrix view to the distances to the selected neighbors [n_queries, + * k] + * @param[in] sample_filter a device bitset filter function that greenlights samples for a given + * query. + */ +void search_with_filtering( + raft::resources const& handle, + const search_params& params, + index& idx, + raft::device_matrix_view queries, + raft::device_matrix_view neighbors, + raft::device_matrix_view distances, + cuvs::neighbors::filtering::bitset_filter sample_filter); + /** * @brief Search ANN using the constructed index with the given filter. * @@ -1924,6 +2149,12 @@ void reconstruct_list_data(raft::resources const& res, uint32_t label, uint32_t offset); +void reconstruct_list_data(raft::resources const& res, + const index& index, + raft::device_matrix_view out_vectors, + uint32_t label, + uint32_t offset); + void reconstruct_list_data(raft::resources const& res, const index& index, raft::device_matrix_view out_vectors, @@ -1972,6 +2203,11 @@ void reconstruct_list_data(raft::resources const& res, raft::device_vector_view in_cluster_indices, raft::device_matrix_view out_vectors, uint32_t label); +void reconstruct_list_data(raft::resources const& res, + const index& index, + raft::device_vector_view in_cluster_indices, + raft::device_matrix_view out_vectors, + uint32_t label); void reconstruct_list_data(raft::resources const& res, const index& index, raft::device_vector_view in_cluster_indices, diff --git a/cpp/include/cuvs/neighbors/nn_descent.hpp b/cpp/include/cuvs/neighbors/nn_descent.hpp index 9f4300177..347ccf889 100644 --- a/cpp/include/cuvs/neighbors/nn_descent.hpp +++ b/cpp/include/cuvs/neighbors/nn_descent.hpp @@ -27,6 +27,8 @@ #include +#include + namespace cuvs::neighbors::nn_descent { /** * @defgroup nn_descent_cpp_index_params The nn-descent algorithm parameters. @@ -237,6 +239,68 @@ auto build(raft::resources const& res, raft::host_matrix_view dataset) -> cuvs::neighbors::nn_descent::index; +/** + * @brief Build nn-descent Index with dataset in device memory + * + * The following distance metrics are supported: + * - L2 + * + * Usage example: + * @code{.cpp} + * using namespace cuvs::neighbors; + * // use default index parameters + * nn_descent::index_params index_params; + * // create and fill the index from a [N, D] raft::device_matrix_view dataset + * auto index = nn_descent::build(res, index_params, dataset); + * // index.graph() provides a raft::host_matrix_view of an + * // all-neighbors knn graph of dimensions [N, k] of the input + * // dataset + * @endcode + * + * @param[in] res raft::resources is an object mangaging resources + * @param[in] params an instance of nn_descent::index_params that are parameters + * to run the nn-descent algorithm + * @param[in] dataset raft::device_matrix_view input dataset expected to be located + * in device memory + * @return index index containing all-neighbors knn graph in host memory + */ +auto build(raft::resources const& res, + index_params const& params, + raft::device_matrix_view dataset) + -> cuvs::neighbors::nn_descent::index; + +/** + * @brief Build nn-descent Index with dataset in host memory + * + * The following distance metrics are supported: + * - L2 + * + * Usage example: + * @code{.cpp} + * using namespace cuvs::neighbors::experimental; + * // use default index parameters + * nn_descent::index_params index_params; + * // create and fill the index from a [N, D] raft::host_matrix_view dataset + * auto index = cagra::build(res, index_params, dataset); + * // index.graph() provides a raft::host_matrix_view of an + * // all-neighbors knn graph of dimensions [N, k] of the input + * // dataset + * @endcode + * + * @tparam T data-type of the input dataset + * @tparam IdxT data-type for the output index + * @param res raft::resources is an object mangaging resources + * @param[in] params an instance of nn_descent::index_params that are parameters + * to run the nn-descent algorithm + * @param[in] dataset raft::host_matrix_view input dataset expected to be located + * in host memory + * @return index index containing all-neighbors knn graph in host memory + */ +auto build(raft::resources const& res, + index_params const& params, + raft::host_matrix_view dataset) + -> cuvs::neighbors::nn_descent::index; + /** * @brief Build nn-descent Index with dataset in device memory * diff --git a/cpp/src/neighbors/cagra_build_half.cu b/cpp/src/neighbors/cagra_build_half.cu new file mode 100644 index 000000000..2aba1dada --- /dev/null +++ b/cpp/src/neighbors/cagra_build_half.cu @@ -0,0 +1,39 @@ +/* + * 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 "cagra.cuh" +#include +#include + +namespace cuvs::neighbors::cagra { + +cuvs::neighbors::cagra::index build( + raft::resources const& handle, + const cuvs::neighbors::cagra::index_params& params, + raft::device_matrix_view dataset) +{ + return cuvs::neighbors::cagra::build(handle, params, dataset); +} + +cuvs::neighbors::cagra::index build( + raft::resources const& handle, + const cuvs::neighbors::cagra::index_params& params, + raft::host_matrix_view dataset) +{ + return cuvs::neighbors::cagra::build(handle, params, dataset); +} + +} // namespace cuvs::neighbors::cagra diff --git a/cpp/src/neighbors/cagra_search_half.cu b/cpp/src/neighbors/cagra_search_half.cu new file mode 100644 index 000000000..d80f2bc00 --- /dev/null +++ b/cpp/src/neighbors/cagra_search_half.cu @@ -0,0 +1,37 @@ +/* + * 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 "cagra.cuh" +#include + +namespace cuvs::neighbors::cagra { + +#define CUVS_INST_CAGRA_SEARCH(T, IdxT) \ + void search(raft::resources const& handle, \ + cuvs::neighbors::cagra::search_params const& params, \ + const cuvs::neighbors::cagra::index& index, \ + raft::device_matrix_view queries, \ + raft::device_matrix_view neighbors, \ + raft::device_matrix_view distances) \ + { \ + cuvs::neighbors::cagra::search(handle, params, index, queries, neighbors, distances); \ + } + +CUVS_INST_CAGRA_SEARCH(half, uint32_t); + +#undef CUVS_INST_CAGRA_SEARCH + +} // namespace cuvs::neighbors::cagra diff --git a/cpp/src/neighbors/detail/cagra/search_multi_cta_float_uint64.cu b/cpp/src/neighbors/cagra_serialize_half.cu similarity index 53% rename from cpp/src/neighbors/detail/cagra/search_multi_cta_float_uint64.cu rename to cpp/src/neighbors/cagra_serialize_half.cu index 88167b843..92ebd9b71 100644 --- a/cpp/src/neighbors/detail/cagra/search_multi_cta_float_uint64.cu +++ b/cpp/src/neighbors/cagra_serialize_half.cu @@ -14,21 +14,12 @@ * limitations under the License. */ -/* - * NOTE: this file is generated by search_multi_cta_00_generate.py - * - * Make changes there and run in this directory: - * - * > python search_multi_cta_00_generate.py - * - */ +#include "cagra_serialize.cuh" + +#include -#include "search_multi_cta_inst.cuh" +namespace cuvs::neighbors::cagra { -namespace cuvs::neighbors::cagra::detail::multi_cta_search { -instantiate_kernel_selection(float, - uint64_t, - float, - cuvs::neighbors::filtering::none_cagra_sample_filter); +CUVS_INST_CAGRA_SERIALIZE(half); -} // namespace cuvs::neighbors::cagra::detail::multi_cta_search +} // namespace cuvs::neighbors::cagra diff --git a/cpp/src/neighbors/detail/cagra/cagra_build.cuh b/cpp/src/neighbors/detail/cagra/cagra_build.cuh index 4a927add5..e5495dc3e 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_build.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_build.cuh @@ -33,8 +33,6 @@ #include // TODO: Fixme- this needs to be migrated -#include "../../ivf_pq/ivf_pq_build.cuh" -#include "../../ivf_pq/ivf_pq_search.cuh" #include "../../nn_descent.cuh" // TODO: This shouldn't be calling spatial/knn APIs @@ -156,8 +154,7 @@ void build_knn_graph( }(); RAFT_LOG_DEBUG("# Building IVF-PQ index %s", model_name.c_str()); - auto index = - cuvs::neighbors::ivf_pq::detail::build(res, pq.build_params, dataset); + auto index = cuvs::neighbors::ivf_pq::build(res, pq.build_params, dataset); // // search top (k + 1) neighbors @@ -169,7 +166,8 @@ void build_knn_graph( const auto num_queries = dataset.extent(0); // Use the same maximum batch size as the ivf_pq::search to avoid allocating more than needed. - using cuvs::neighbors::ivf_pq::detail::kMaxQueries; + constexpr uint32_t kMaxQueries = 4096; + // Heuristic: the build_knn_graph code should use only a fraction of the workspace memory; the // rest should be used by the ivf_pq::search. Here we say that the workspace size should be a good // multiple of what is required for the I/O batching below. diff --git a/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh b/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh index 24cc2a22f..f86ed9ef6 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh @@ -194,6 +194,8 @@ void serialize_to_hnswlib(raft::resources const& res, 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"); } os.write(reinterpret_cast(&i), sizeof(std::size_t)); diff --git a/cpp/src/neighbors/detail/cagra/compute_distance_00_generate.py b/cpp/src/neighbors/detail/cagra/compute_distance_00_generate.py index f8584c62e..aef31d161 100644 --- a/cpp/src/neighbors/detail/cagra/compute_distance_00_generate.py +++ b/cpp/src/neighbors/detail/cagra/compute_distance_00_generate.py @@ -66,8 +66,6 @@ half_uint32=("half", "uint32_t", "float"), int8_uint32=("int8_t", "uint32_t", "float"), uint8_uint32=("uint8_t", "uint32_t", "float"), - # float_uint64=("float", "uint64_t", "float"), - # half_uint64=("half", "uint64_t", "float"), ) metric_prefix = 'DistanceType::' diff --git a/cpp/src/neighbors/detail/cagra/search_multi_cta_00_generate.py b/cpp/src/neighbors/detail/cagra/search_multi_cta_00_generate.py index 3153a3a9f..4e3983e3f 100644 --- a/cpp/src/neighbors/detail/cagra/search_multi_cta_00_generate.py +++ b/cpp/src/neighbors/detail/cagra/search_multi_cta_00_generate.py @@ -58,8 +58,6 @@ half_uint32=("half", "uint32_t", "float"), int8_uint32=("int8_t", "uint32_t", "float"), uint8_uint32=("uint8_t", "uint32_t", "float"), - float_uint64=("float", "uint64_t", "float"), - half_uint64=("half", "uint64_t", "float"), ) # knn for type_path, (data_t, idx_t, distance_t) in search_types.items(): diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_00_generate.py b/cpp/src/neighbors/detail/cagra/search_single_cta_00_generate.py index e37ceb1fa..4693cd54d 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_00_generate.py +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_00_generate.py @@ -60,8 +60,6 @@ half_uint32=("half", "uint32_t", "float"), int8_uint32=("int8_t", "uint32_t", "float"), uint8_uint32=("uint8_t", "uint32_t", "float"), - float_uint64=("float", "uint64_t", "float"), - half_uint64=("half", "uint64_t", "float"), ) # knn diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_float_uint64.cu b/cpp/src/neighbors/detail/cagra/search_single_cta_float_uint64.cu deleted file mode 100644 index 0ef5c366f..000000000 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_float_uint64.cu +++ /dev/null @@ -1,34 +0,0 @@ -/* - * 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. - */ - -/* - * NOTE: this file is generated by search_single_cta_00_generate.py - * - * Make changes there and run in this directory: - * - * > python search_single_cta_00_generate.py - * - */ - -#include "search_single_cta_inst.cuh" - -namespace cuvs::neighbors::cagra::detail::single_cta_search { -instantiate_kernel_selection(float, - uint64_t, - float, - cuvs::neighbors::filtering::none_cagra_sample_filter); - -} // namespace cuvs::neighbors::cagra::detail::single_cta_search diff --git a/cpp/src/neighbors/ivf_pq/detail/generate_ivf_pq.py b/cpp/src/neighbors/ivf_pq/detail/generate_ivf_pq.py index 878c7ee21..9b3083c3b 100644 --- a/cpp/src/neighbors/ivf_pq/detail/generate_ivf_pq.py +++ b/cpp/src/neighbors/ivf_pq/detail/generate_ivf_pq.py @@ -57,6 +57,7 @@ types = dict( float_int64_t=("float", "int64_t"), + half_int64_t=("half", "int64_t"), int8_t_int64_t=("int8_t", "int64_t"), uint8_t_int64_t=("uint8_t", "int64_t"), ) diff --git a/cpp/src/neighbors/detail/cagra/search_multi_cta_half_uint64.cu b/cpp/src/neighbors/ivf_pq/detail/ivf_pq_build_extend_half_int64_t.cu similarity index 55% rename from cpp/src/neighbors/detail/cagra/search_multi_cta_half_uint64.cu rename to cpp/src/neighbors/ivf_pq/detail/ivf_pq_build_extend_half_int64_t.cu index dafb89cc3..2d7270957 100644 --- a/cpp/src/neighbors/detail/cagra/search_multi_cta_half_uint64.cu +++ b/cpp/src/neighbors/ivf_pq/detail/ivf_pq_build_extend_half_int64_t.cu @@ -1,5 +1,5 @@ /* - * Copyright (c) 2023-2024, NVIDIA CORPORATION. + * 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. @@ -15,20 +15,21 @@ */ /* - * NOTE: this file is generated by search_multi_cta_00_generate.py + * NOTE: this file is generated by generate_ivf_pq.py * * Make changes there and run in this directory: * - * > python search_multi_cta_00_generate.py + * > python generate_ivf_pq.py * */ -#include "search_multi_cta_inst.cuh" +#include -namespace cuvs::neighbors::cagra::detail::multi_cta_search { -instantiate_kernel_selection(half, - uint64_t, - float, - cuvs::neighbors::filtering::none_cagra_sample_filter); +#include "ivf_pq_build_extend_inst.cuh" -} // namespace cuvs::neighbors::cagra::detail::multi_cta_search +namespace cuvs::neighbors::ivf_pq { +CUVS_INST_IVF_PQ_BUILD_EXTEND(half, int64_t); + +#undef CUVS_INST_IVF_PQ_BUILD_EXTEND + +} // namespace cuvs::neighbors::ivf_pq diff --git a/cpp/src/neighbors/ivf_pq/detail/ivf_pq_search_half_int64_t.cu b/cpp/src/neighbors/ivf_pq/detail/ivf_pq_search_half_int64_t.cu new file mode 100644 index 000000000..e5556e593 --- /dev/null +++ b/cpp/src/neighbors/ivf_pq/detail/ivf_pq_search_half_int64_t.cu @@ -0,0 +1,46 @@ +/* + * 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. + */ + +/* + * NOTE: this file is generated by generate_ivf_pq.py + * + * Make changes there and run in this directory: + * + * > python generate_ivf_pq.py + * + */ + +#include + +#include "../ivf_pq_search.cuh" + +namespace cuvs::neighbors::ivf_pq { + +#define CUVS_INST_IVF_PQ_SEARCH(T, IdxT) \ + void search(raft::resources const& handle, \ + const cuvs::neighbors::ivf_pq::search_params& params, \ + cuvs::neighbors::ivf_pq::index& index, \ + raft::device_matrix_view queries, \ + raft::device_matrix_view neighbors, \ + raft::device_matrix_view distances) \ + { \ + cuvs::neighbors::ivf_pq::detail::search(handle, params, index, queries, neighbors, distances); \ + } +CUVS_INST_IVF_PQ_SEARCH(half, int64_t); + +#undef CUVS_INST_IVF_PQ_SEARCH + +} // namespace cuvs::neighbors::ivf_pq diff --git a/cpp/src/neighbors/ivf_pq/detail/ivf_pq_search_with_filter_half_int64_t.cu b/cpp/src/neighbors/ivf_pq/detail/ivf_pq_search_with_filter_half_int64_t.cu new file mode 100644 index 000000000..5874fba6c --- /dev/null +++ b/cpp/src/neighbors/ivf_pq/detail/ivf_pq_search_with_filter_half_int64_t.cu @@ -0,0 +1,49 @@ +/* + * 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. + */ + +/* + * NOTE: this file is generated by generate_ivf_pq.py + * + * Make changes there and run in this directory: + * + * > python generate_ivf_pq.py + * + */ + +#include + +#include "../ivf_pq_search.cuh" + +namespace cuvs::neighbors::ivf_pq { + +#define CUVS_INST_IVF_PQ_SEARCH_FILTER(T, IdxT) \ + void search_with_filtering( \ + raft::resources const& handle, \ + const cuvs::neighbors::ivf_pq::search_params& params, \ + cuvs::neighbors::ivf_pq::index& index, \ + raft::device_matrix_view queries, \ + raft::device_matrix_view neighbors, \ + raft::device_matrix_view distances, \ + cuvs::neighbors::filtering::bitset_filter sample_filter) \ + { \ + cuvs::neighbors::ivf_pq::detail::search_with_filtering( \ + handle, params, index, queries, neighbors, distances, sample_filter); \ + } +CUVS_INST_IVF_PQ_SEARCH_FILTER(half, int64_t); + +#undef CUVS_INST_IVF_PQ_SEARCH_FILTER + +} // namespace cuvs::neighbors::ivf_pq diff --git a/cpp/src/neighbors/nn_descent_half.cu b/cpp/src/neighbors/nn_descent_half.cu new file mode 100644 index 000000000..587993031 --- /dev/null +++ b/cpp/src/neighbors/nn_descent_half.cu @@ -0,0 +1,43 @@ +/* + * 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 "nn_descent.cuh" +#include + +namespace cuvs::neighbors::nn_descent { + +#define CUVS_INST_NN_DESCENT_BUILD(T, IdxT) \ + auto build(raft::resources const& handle, \ + const cuvs::neighbors::nn_descent::index_params& params, \ + raft::device_matrix_view dataset) \ + ->cuvs::neighbors::nn_descent::index \ + { \ + return cuvs::neighbors::nn_descent::build(handle, params, dataset); \ + }; \ + \ + auto build(raft::resources const& handle, \ + const cuvs::neighbors::nn_descent::index_params& params, \ + raft::host_matrix_view dataset) \ + ->cuvs::neighbors::nn_descent::index \ + { \ + return cuvs::neighbors::nn_descent::build(handle, params, dataset); \ + }; + +CUVS_INST_NN_DESCENT_BUILD(half, uint32_t); + +#undef CUVS_INST_NN_DESCENT_BUILD + +} // namespace cuvs::neighbors::nn_descent diff --git a/cpp/test/CMakeLists.txt b/cpp/test/CMakeLists.txt index e04c39318..b81ef6bfa 100644 --- a/cpp/test/CMakeLists.txt +++ b/cpp/test/CMakeLists.txt @@ -139,6 +139,7 @@ if(BUILD_TESTS) NEIGHBORS_ANN_CAGRA_TEST PATH neighbors/ann_cagra/test_float_uint32_t.cu + neighbors/ann_cagra/test_half_uint32_t.cu neighbors/ann_cagra/test_int8_t_uint32_t.cu neighbors/ann_cagra/test_uint8_t_uint32_t.cu GPUS diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_half_uint64.cu b/cpp/test/neighbors/ann_cagra/test_half_uint32_t.cu similarity index 53% rename from cpp/src/neighbors/detail/cagra/search_single_cta_half_uint64.cu rename to cpp/test/neighbors/ann_cagra/test_half_uint32_t.cu index b96ed0b22..f03de69d2 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_half_uint64.cu +++ b/cpp/test/neighbors/ann_cagra/test_half_uint32_t.cu @@ -14,21 +14,15 @@ * limitations under the License. */ -/* - * NOTE: this file is generated by search_single_cta_00_generate.py - * - * Make changes there and run in this directory: - * - * > python search_single_cta_00_generate.py - * - */ +#include + +#include "../ann_cagra.cuh" + +namespace cuvs::neighbors::cagra { -#include "search_single_cta_inst.cuh" +typedef AnnCagraTest AnnCagraTestF16_U32; +TEST_P(AnnCagraTestF16_U32, AnnCagra) { this->testCagra(); } -namespace cuvs::neighbors::cagra::detail::single_cta_search { -instantiate_kernel_selection(half, - uint64_t, - float, - cuvs::neighbors::filtering::none_cagra_sample_filter); +INSTANTIATE_TEST_CASE_P(AnnCagraTest, AnnCagraTestF16_U32, ::testing::ValuesIn(inputs)); -} // namespace cuvs::neighbors::cagra::detail::single_cta_search +} // namespace cuvs::neighbors::cagra